DeepSeek/smallpond
1
0
Fork 0
mirror of https://github.com/deepseek-ai/smallpond synced 2025-06-26 18:27:45 +00:00
Code Issues Actions Packages Projects Releases Wiki Activity

init

Browse Source
This commit is contained in:
Runji Wang
2025-02-25 18:16:31 +08:00
commit 770aa417d5
77 changed files with 18785 additions and 0 deletions
Download Patch File Download Diff File Expand all files Collapse all files

0
tests/__init__.py Normal file
View File

30
tests/conftest.py Normal file
View File

@@ -0,0 +1,30 @@
import os
import pytest
import ray
import smallpond
@pytest.fixture(scope="session")
def ray_address():
"""A global Ray instance for all tests"""
ray_address = ray.init(
address="local",
# disable dashboard in unit tests
include_dashboard=False,
).address_info["gcs_address"]
yield ray_address
ray.shutdown()
@pytest.fixture
def sp(ray_address: str, request):
"""A smallpond session for each test"""
runtime_root = os.getenv("TEST_RUNTIME_ROOT") or f"tests/runtime"
sp = smallpond.init(
data_root=os.path.join(runtime_root, request.node.name),
ray_address=ray_address,
)
yield sp
sp.shutdown()

186
tests/datagen.py Normal file
View File

@@ -0,0 +1,186 @@
import base64
import glob
import os
import random
import string
from concurrent.futures import ProcessPoolExecutor
from datetime import datetime, timedelta, timezone
from typing import Tuple
import pandas as pd
import pyarrow as pa
import pyarrow.parquet as pq
from filelock import FileLock
def generate_url_and_domain() -> Tuple[str, str]:
domain_part = "".join(
random.choices(string.ascii_lowercase, k=random.randint(5, 15))
)
tld = random.choice(["com", "net", "org", "cn", "edu", "gov", "co", "io"])
domain = f"www.{domain_part}.{tld}"
path_segments = []
for _ in range(random.randint(1, 3)):
segment = "".join(
random.choices(
string.ascii_lowercase + string.digits, k=random.randint(3, 10)
)
)
path_segments.append(segment)
path = "/" + "/".join(path_segments)
protocol = random.choice(["http", "https"])
if random.random() < 0.3:
path += random.choice([".html", ".php", ".htm", ".aspx"])
return f"{protocol}://{domain}{path}", domain
def generate_random_date() -> str:
start = datetime(2023, 1, 1, tzinfo=timezone.utc)
end = datetime(2023, 12, 31, tzinfo=timezone.utc)
delta = end - start
random_date = start + timedelta(
seconds=random.randint(0, int(delta.total_seconds()))
)
return random_date.strftime("%Y-%m-%dT%H:%M:%SZ")
def generate_content() -> bytes:
target_length = (
random.randint(1000, 100000)
if random.random() < 0.8
else random.randint(100000, 1000000)
)
before = b"<!DOCTYPE html><html><head><title>Random Page</title></head><body>"
after = b"</body></html>"
total_before_after = len(before) + len(after)
fill_length = max(target_length - total_before_after, 0)
filler = "".join(random.choices(string.printable, k=fill_length)).encode("ascii")[
:fill_length
]
return before + filler + after
def generate_arrow_parquet(path: str, num_rows=100):
data = []
for _ in range(num_rows):
url, domain = generate_url_and_domain()
date = generate_random_date()
content = generate_content()
data.append({"url": url, "domain": domain, "date": date, "content": content})
df = pd.DataFrame(data)
df.to_parquet(path, engine="pyarrow")
def generate_arrow_files(output_dir: str, num_files=10):
os.makedirs(output_dir, exist_ok=True)
with ProcessPoolExecutor(max_workers=10) as executor:
executor.map(
generate_arrow_parquet,
[f"{output_dir}/data{i}.parquet" for i in range(num_files)],
)
def concat_arrow_files(input_dir: str, output_dir: str, repeat: int = 10):
os.makedirs(output_dir, exist_ok=True)
files = glob.glob(os.path.join(input_dir, "*.parquet"))
table = pa.concat_tables([pa.parquet.read_table(file) for file in files] * repeat)
pq.write_table(table, os.path.join(output_dir, "large_array.parquet"))
def generate_random_string(length: int) -> str:
"""Generate a random string of a specified length"""
return "".join(random.choices(string.ascii_letters + string.digits, k=length))
def generate_random_url() -> str:
"""Generate a random URL"""
path = generate_random_string(random.randint(10, 20))
return (
f"com.{random.randint(10000, 999999)}.{random.randint(100, 9999)}/{path}.html"
)
def generate_random_data() -> str:
"""Generate random data"""
url = generate_random_url()
content = generate_random_string(random.randint(50, 100))
encoded_content = base64.b64encode(content.encode()).decode()
return f"{url}\t{encoded_content}"
def generate_url_parquet(path: str, num_rows=100):
"""Generate a parquet file with a specified number of random data lines"""
data = []
for _ in range(num_rows):
url = generate_random_url()
host = url.split("/")[0]
data.append({"host": host, "url": url})
df = pd.DataFrame(data)
df.to_parquet(path, engine="pyarrow")
def generate_url_parquet_files(output_dir: str, num_files: int = 10):
"""Generate multiple parquet files with a specified number of random data lines"""
os.makedirs(output_dir, exist_ok=True)
with ProcessPoolExecutor(max_workers=10) as executor:
executor.map(
generate_url_parquet,
[f"{output_dir}/urls{i}.parquet" for i in range(num_files)],
)
def generate_url_tsv_files(
output_dir: str, num_files: int = 10, lines_per_file: int = 100
):
"""Generate multiple files, each containing a specified number of random data lines"""
os.makedirs(output_dir, exist_ok=True)
for i in range(num_files):
with open(f"{output_dir}/urls{i}.tsv", "w") as f:
for _ in range(lines_per_file):
f.write(generate_random_data() + "\n")
def generate_long_path_list(path: str, num_lines: int = 1048576):
"""Generate a list of long paths"""
with open(path, "w", buffering=16 * 1024 * 1024) as f:
for i in range(num_lines):
path = os.path.abspath(f"tests/data/arrow/data{i % 10}.parquet")
f.write(f"{path}\n")
def generate_data(path: str = "tests/data"):
"""
Generate all data for testing.
"""
os.makedirs(path, exist_ok=True)
try:
with FileLock(path + "/data.lock"):
print("Generating data...")
if not os.path.exists(path + "/mock_urls"):
generate_url_tsv_files(
output_dir=path + "/mock_urls", num_files=10, lines_per_file=100
)
generate_url_parquet_files(output_dir=path + "/mock_urls", num_files=10)
if not os.path.exists(path + "/arrow"):
generate_arrow_files(output_dir=path + "/arrow", num_files=10)
if not os.path.exists(path + "/large_array"):
concat_arrow_files(
input_dir=path + "/arrow", output_dir=path + "/large_array"
)
if not os.path.exists(path + "/long_path_list.txt"):
generate_long_path_list(path=path + "/long_path_list.txt")
except Exception as e:
print(f"Error generating data: {e}")
if __name__ == "__main__":
generate_data()

187
tests/test_arrow.py Normal file
View File

@@ -0,0 +1,187 @@
import glob
import os.path
import tempfile
import unittest
import pyarrow.parquet as parquet
from loguru import logger
from smallpond.io.arrow import (
RowRange,
build_batch_reader_from_files,
cast_columns_to_large_string,
dump_to_parquet_files,
load_from_parquet_files,
)
from smallpond.utility import ConcurrentIter
from tests.test_fabric import TestFabric
class TestArrow(TestFabric, unittest.TestCase):
def test_load_from_parquet_files(self):
for dataset_path in (
"tests/data/arrow/*.parquet",
"tests/data/large_array/*.parquet",
):
with self.subTest(dataset_path=dataset_path):
parquet_files = glob.glob(dataset_path)
expected = self._load_parquet_files(parquet_files)
actual = load_from_parquet_files(parquet_files)
self._compare_arrow_tables(expected, actual)
def test_load_parquet_row_ranges(self):
for dataset_path in (
"tests/data/arrow/data0.parquet",
"tests/data/large_array/large_array.parquet",
):
with self.subTest(dataset_path=dataset_path):
metadata = parquet.read_metadata(dataset_path)
file_num_rows = metadata.num_rows
data_size = sum(
metadata.row_group(i).total_byte_size
for i in range(metadata.num_row_groups)
)
row_range = RowRange(
path=dataset_path,
begin=100,
end=200,
data_size=data_size,
file_num_rows=file_num_rows,
)
expected = self._load_parquet_files([dataset_path]).slice(
offset=100, length=100
)
actual = load_from_parquet_files([row_range])
self._compare_arrow_tables(expected, actual)
def test_dump_to_parquet_files(self):
for dataset_path in (
"tests/data/arrow/*.parquet",
"tests/data/large_array/*.parquet",
):
with self.subTest(dataset_path=dataset_path):
parquet_files = glob.glob(dataset_path)
expected = self._load_parquet_files(parquet_files)
with tempfile.TemporaryDirectory(
dir=self.output_root_abspath
) as output_dir:
ok = dump_to_parquet_files(expected, output_dir)
self.assertTrue(ok)
actual = self._load_parquet_files(
glob.glob(f"{output_dir}/*.parquet")
)
self._compare_arrow_tables(expected, actual)
def test_dump_load_empty_table(self):
# create empty table
empty_table = self._load_parquet_files(
["tests/data/arrow/data0.parquet"]
).slice(length=0)
self.assertEqual(empty_table.num_rows, 0)
# dump empty table
with tempfile.TemporaryDirectory(dir=self.output_root_abspath) as output_dir:
ok = dump_to_parquet_files(empty_table, output_dir)
self.assertTrue(ok)
parquet_files = glob.glob(f"{output_dir}/*.parquet")
# load empty table from file
actual_table = load_from_parquet_files(parquet_files)
self._compare_arrow_tables(empty_table, actual_table)
def test_parquet_batch_reader(self):
for dataset_path in (
"tests/data/arrow/*.parquet",
"tests/data/large_array/*.parquet",
):
with self.subTest(dataset_path=dataset_path):
parquet_files = glob.glob(dataset_path)
expected_num_rows = sum(
parquet.read_metadata(file).num_rows for file in parquet_files
)
with build_batch_reader_from_files(
parquet_files,
batch_size=expected_num_rows,
max_batch_byte_size=None,
) as batch_reader, ConcurrentIter(batch_reader) as concurrent_iter:
total_num_rows = 0
for batch in concurrent_iter:
print(
f"batch.num_rows {batch.num_rows}, max_batch_row_size {expected_num_rows}"
)
self.assertLessEqual(batch.num_rows, expected_num_rows)
total_num_rows += batch.num_rows
self.assertEqual(total_num_rows, expected_num_rows)
def test_table_to_batches(self):
for dataset_path in (
"tests/data/arrow/*.parquet",
"tests/data/large_array/*.parquet",
):
with self.subTest(dataset_path=dataset_path):
parquet_files = glob.glob(dataset_path)
table = self._load_parquet_files(parquet_files)
total_num_rows = 0
for batch in table.to_batches(max_chunksize=table.num_rows):
print(
f"batch.num_rows {batch.num_rows}, max_batch_row_size {table.num_rows}"
)
self.assertLessEqual(batch.num_rows, table.num_rows)
total_num_rows += batch.num_rows
self.assertEqual(total_num_rows, table.num_rows)
def test_arrow_schema_metadata(self):
table = self._load_parquet_files(glob.glob("tests/data/arrow/*.parquet"))
metadata = {b"a": b"1", b"b": b"2"}
table_with_meta = table.replace_schema_metadata(metadata)
print(f"table_with_meta.schema.metadata {table_with_meta.schema.metadata}")
with tempfile.TemporaryDirectory(dir=self.output_root_abspath) as output_dir:
self.assertTrue(
dump_to_parquet_files(
table_with_meta, output_dir, "arrow_schema_metadata", max_workers=2
)
)
parquet_files = glob.glob(
os.path.join(output_dir, "arrow_schema_metadata*.parquet")
)
loaded_table = load_from_parquet_files(
parquet_files, table.column_names[:1]
)
print(f"loaded_table.schema.metadata {loaded_table.schema.metadata}")
self.assertEqual(
table_with_meta.schema.metadata, loaded_table.schema.metadata
)
with parquet.ParquetFile(parquet_files[0]) as file:
print(f"file.schema_arrow.metadata {file.schema_arrow.metadata}")
self.assertEqual(
table_with_meta.schema.metadata, file.schema_arrow.metadata
)
def test_load_mixed_string_types(self):
parquet_paths = glob.glob("tests/data/arrow/*.parquet")
table = self._load_parquet_files(parquet_paths)
with tempfile.TemporaryDirectory(dir=self.output_root_abspath) as output_dir:
dump_to_parquet_files(cast_columns_to_large_string(table), output_dir)
parquet_paths += glob.glob(os.path.join(output_dir, "*.parquet"))
loaded_table = load_from_parquet_files(parquet_paths)
self.assertEqual(table.num_rows * 2, loaded_table.num_rows)
batch_reader = build_batch_reader_from_files(parquet_paths)
self.assertEqual(
table.num_rows * 2, sum(batch.num_rows for batch in batch_reader)
)
@logger.catch(reraise=True, message="failed to load parquet files")
def _load_from_parquet_files_with_log(self, paths, columns):
load_from_parquet_files(paths, columns)
def test_load_not_exist_column(self):
parquet_files = glob.glob("tests/data/arrow/*.parquet")
with self.assertRaises(AssertionError) as context:
self._load_from_parquet_files_with_log(parquet_files, ["not_exist_column"])
def test_change_ordering_of_columns(self):
parquet_files = glob.glob("tests/data/arrow/*.parquet")
loaded_table = load_from_parquet_files(parquet_files)
reversed_cols = list(reversed(loaded_table.column_names))
loaded_table = load_from_parquet_files(parquet_files, reversed_cols)
self.assertEqual(loaded_table.column_names, reversed_cols)

90
tests/test_bench.py Normal file
View File

@@ -0,0 +1,90 @@
import shutil
import unittest
from benchmarks.file_io_benchmark import file_io_benchmark
from benchmarks.gray_sort_benchmark import generate_random_records, gray_sort_benchmark
from benchmarks.hash_partition_benchmark import hash_partition_benchmark
from benchmarks.urls_sort_benchmark import urls_sort_benchmark
from smallpond.common import MB
from smallpond.logical.node import Context, LogicalPlan
from tests.test_fabric import TestFabric
class TestBench(TestFabric, unittest.TestCase):
fault_inject_prob = 0.05
def test_file_io_benchmark(self):
for io_engine in ("duckdb", "arrow", "stream"):
with self.subTest(io_engine=io_engine):
plan = file_io_benchmark(
["tests/data/mock_urls/*.parquet"],
npartitions=3,
io_engine=io_engine,
)
self.execute_plan(plan, enable_profiling=True)
def test_urls_sort_benchmark(self):
for engine_type in ("duckdb", "arrow"):
with self.subTest(engine_type=engine_type):
plan = urls_sort_benchmark(
["tests/data/mock_urls/*.tsv"],
num_data_partitions=3,
num_hash_partitions=3,
engine_type=engine_type,
)
self.execute_plan(plan, enable_profiling=True)
@unittest.skipIf(shutil.which("gensort") is None, "gensort not found")
def test_gray_sort_benchmark(self):
record_nbytes = 100
key_nbytes = 10
total_data_nbytes = 100 * MB
gensort_batch_nbytes = 10 * MB
num_data_partitions = 5
num_sort_partitions = 1 << 3
for shuffle_engine in ("duckdb", "arrow"):
for sort_engine in ("duckdb", "arrow", "polars"):
with self.subTest(
shuffle_engine=shuffle_engine, sort_engine=sort_engine
):
ctx = Context()
random_records = generate_random_records(
ctx,
record_nbytes,
key_nbytes,
total_data_nbytes,
gensort_batch_nbytes,
num_data_partitions,
num_sort_partitions,
)
plan = LogicalPlan(ctx, random_records)
exec_plan = self.execute_plan(plan, enable_profiling=True)
plan = gray_sort_benchmark(
record_nbytes,
key_nbytes,
total_data_nbytes,
gensort_batch_nbytes,
num_data_partitions,
num_sort_partitions,
input_paths=exec_plan.final_output.resolved_paths,
shuffle_engine=shuffle_engine,
sort_engine=sort_engine,
hive_partitioning=True,
validate_results=True,
)
self.execute_plan(plan, enable_profiling=True)
def test_hash_partition_benchmark(self):
for engine_type in ("duckdb", "arrow"):
with self.subTest(engine_type=engine_type):
plan = hash_partition_benchmark(
["tests/data/mock_urls/*.parquet"],
npartitions=5,
hash_columns=["url"],
engine_type=engine_type,
hive_partitioning=True,
partition_stats=True,
)
self.execute_plan(plan, enable_profiling=True)

73
tests/test_common.py Normal file
View File

@@ -0,0 +1,73 @@
import itertools
import unittest
import numpy as np
from hypothesis import given
from hypothesis import strategies as st
from smallpond.common import get_nth_partition, split_into_cols, split_into_rows
from tests.test_fabric import TestFabric
class TestCommon(TestFabric, unittest.TestCase):
def test_get_nth_partition(self):
items = [1, 2, 3]
# split into 1 partitions
self.assertListEqual([1, 2, 3], get_nth_partition(items, 0, 1))
# split into 2 partitions
self.assertListEqual([1, 2], get_nth_partition(items, 0, 2))
self.assertListEqual([3], get_nth_partition(items, 1, 2))
# split into 3 partitions
self.assertListEqual([1], get_nth_partition(items, 0, 3))
self.assertListEqual([2], get_nth_partition(items, 1, 3))
self.assertListEqual([3], get_nth_partition(items, 2, 3))
# split into 5 partitions
self.assertListEqual([1], get_nth_partition(items, 0, 5))
self.assertListEqual([2], get_nth_partition(items, 1, 5))
self.assertListEqual([3], get_nth_partition(items, 2, 5))
self.assertListEqual([], get_nth_partition(items, 3, 5))
self.assertListEqual([], get_nth_partition(items, 4, 5))
@given(st.data())
def test_split_into_rows(self, data: st.data):
nelements = data.draw(st.integers(1, 100))
npartitions = data.draw(st.integers(1, 2 * nelements))
items = list(range(nelements))
computed = split_into_rows(items, npartitions)
expected = [
get_nth_partition(items, n, npartitions) for n in range(npartitions)
]
self.assertEqual(expected, computed)
@given(st.data())
def test_split_into_cols(self, data: st.data):
nelements = data.draw(st.integers(1, 100))
npartitions = data.draw(st.integers(1, 2 * nelements))
items = list(range(nelements))
chunks = split_into_cols(items, npartitions)
self.assertEqual(npartitions, len(chunks))
self.assertListEqual(
items,
[x for row in itertools.zip_longest(*chunks) for x in row if x is not None],
)
chunk_sizes = set(len(chk) for chk in chunks)
if len(chunk_sizes) > 1:
small_size, large_size = sorted(chunk_sizes)
self.assertEqual(small_size + 1, large_size)
else:
(chunk_size,) = chunk_sizes
self.assertEqual(len(items), chunk_size * npartitions)
def test_split_into_rows_bench(self):
for nelements in [100000, 1000000]:
items = np.arange(nelements)
for npartitions in [1024, 4096, 10240, nelements, 2 * nelements]:
chunks = split_into_rows(items, npartitions)
self.assertEqual(npartitions, len(chunks))
def test_split_into_cols_bench(self):
for nelements in [100000, 1000000]:
items = np.arange(nelements)
for npartitions in [1024, 4096, 10240, nelements, 2 * nelements]:
chunks = split_into_cols(items, npartitions)
self.assertEqual(npartitions, len(chunks))

223
tests/test_dataframe.py Normal file
View File

@@ -0,0 +1,223 @@
from typing import List
import pandas as pd
import pyarrow as pa
import pytest
from smallpond.dataframe import Session
def test_pandas(sp: Session):
pandas_df = pd.DataFrame({"a": [1, 2, 3], "b": [4, 5, 6]})
df = sp.from_pandas(pandas_df)
assert df.to_pandas().equals(pandas_df)
def test_arrow(sp: Session):
arrow_table = pa.table({"a": [1, 2, 3], "b": [4, 5, 6]})
df = sp.from_arrow(arrow_table)
assert df.to_arrow() == arrow_table
def test_items(sp: Session):
df = sp.from_items([1, 2, 3])
assert df.take_all() == [{"item": 1}, {"item": 2}, {"item": 3}]
df = sp.from_items([{"a": 1, "b": 4}, {"a": 2, "b": 5}, {"a": 3, "b": 6}])
assert df.take_all() == [{"a": 1, "b": 4}, {"a": 2, "b": 5}, {"a": 3, "b": 6}]
def test_csv(sp: Session):
df = sp.read_csv(
"tests/data/mock_urls/*.tsv",
schema={"urlstr": "varchar", "valstr": "varchar"},
delim=r"\t",
)
assert df.count() == 1000
def test_parquet(sp: Session):
df = sp.read_parquet("tests/data/mock_urls/*.parquet")
assert df.count() == 1000
def test_take(sp: Session):
df = sp.from_pandas(pd.DataFrame({"a": [1, 2, 3], "b": [4, 5, 6]}))
assert df.take(2) == [{"a": 1, "b": 4}, {"a": 2, "b": 5}]
assert df.take_all() == [{"a": 1, "b": 4}, {"a": 2, "b": 5}, {"a": 3, "b": 6}]
def test_map(sp: Session):
df = sp.from_arrow(pa.table({"a": [1, 2, 3], "b": [4, 5, 6]}))
df1 = df.map("a + b as c")
assert df1.to_arrow() == pa.table({"c": [5, 7, 9]})
df2 = df.map(lambda r: {"c": r["a"] + r["b"]})
assert df2.to_arrow() == pa.table({"c": [5, 7, 9]})
# user need to specify the schema if can not be inferred from the mapping values
df3 = df.map(
lambda r: {"c": None if r["a"] == 1 else r["a"] + r["b"]},
schema=pa.schema([("c", pa.int64())]),
)
assert df3.to_arrow() == pa.table({"c": pa.array([None, 7, 9], type=pa.int64())})
def test_flat_map(sp: Session):
df = sp.from_arrow(pa.table({"a": [1, 2, 3], "b": [4, 5, 6]}))
df1 = df.flat_map(lambda r: [{"c": r["a"]}, {"c": r["b"]}])
assert df1.to_arrow() == pa.table({"c": [1, 4, 2, 5, 3, 6]})
df2 = df.flat_map("unnest(array[a, b]) as c")
assert df2.to_arrow() == pa.table({"c": [1, 4, 2, 5, 3, 6]})
# user need to specify the schema if can not be inferred from the mapping values
df3 = df.flat_map(lambda r: [{"c": None}], schema=pa.schema([("c", pa.int64())]))
assert df3.to_arrow() == pa.table(
{"c": pa.array([None, None, None], type=pa.int64())}
)
def test_map_batches(sp: Session):
df = sp.read_parquet("tests/data/mock_urls/*.parquet")
df = df.map_batches(
lambda batch: pa.table({"num_rows": [batch.num_rows]}),
batch_size=350,
)
assert df.take_all() == [{"num_rows": 350}, {"num_rows": 350}, {"num_rows": 300}]
def test_filter(sp: Session):
df = sp.from_arrow(pa.table({"a": [1, 2, 3], "b": [4, 5, 6]}))
df1 = df.filter("a > 1")
assert df1.to_arrow() == pa.table({"a": [2, 3], "b": [5, 6]})
df2 = df.filter(lambda r: r["a"] > 1)
assert df2.to_arrow() == pa.table({"a": [2, 3], "b": [5, 6]})
def test_random_shuffle(sp: Session):
df = sp.from_items(list(range(1000))).repartition(10, by_rows=True)
df = df.random_shuffle()
shuffled = [d["item"] for d in df.take_all()]
assert sorted(shuffled) == list(range(1000))
def count_inversions(arr: List[int]) -> int:
return sum(
sum(1 for j in range(i + 1, len(arr)) if arr[i] > arr[j])
for i in range(len(arr))
)
# check the shuffle is random enough
# the expected number of inversions is n*(n-1)/4 = 249750
assert 220000 <= count_inversions(shuffled) <= 280000
def test_partition_by(sp: Session):
df = sp.from_items(list(range(1000))).repartition(10, by="item % 10")
df = df.map("min(item % 10) as min, max(item % 10) as max")
assert df.take_all() == [{"min": i, "max": i} for i in range(10)]
def test_partition_by_key_out_of_range(sp: Session):
df = sp.from_items(list(range(1000))).repartition(10, by="item % 11")
try:
df.to_arrow()
except Exception as ex:
assert "partition key 10 is out of range 0-9" in str(ex)
else:
assert False, "expected exception"
def test_partition_by_hash(sp: Session):
df = sp.from_items(list(range(1000))).repartition(10, hash_by="item")
items = [d["item"] for d in df.take_all()]
assert sorted(items) == list(range(1000))
def test_count(sp: Session):
df = sp.from_items([1, 2, 3])
assert df.count() == 3
def test_limit(sp: Session):
df = sp.from_items(list(range(1000))).repartition(10, by_rows=True)
assert df.limit(2).count() == 2
@pytest.mark.skip(reason="limit can not be pushed down to sql node for now")
@pytest.mark.timeout(10)
def test_limit_large(sp: Session):
# limit will be fused with the previous select
# otherwise, it will be timeout
df = sp.partial_sql("select * from range(1000000000)")
assert df.limit(2).count() == 2
def test_partial_sql(sp: Session):
# no input deps
df = sp.partial_sql("select * from range(3)")
assert df.to_arrow() == pa.table({"range": [0, 1, 2]})
# join
df1 = sp.from_arrow(pa.table({"id1": [1, 2, 3], "val1": ["a", "b", "c"]}))
df2 = sp.from_arrow(pa.table({"id2": [1, 2, 3], "val2": ["d", "e", "f"]}))
joined = sp.partial_sql(
"select id1, val1, val2 from {0} join {1} on id1 = id2", df1, df2
)
assert joined.to_arrow() == pa.table(
{"id1": [1, 2, 3], "val1": ["a", "b", "c"], "val2": ["d", "e", "f"]},
schema=pa.schema(
[
("id1", pa.int64()),
("val1", pa.large_string()),
("val2", pa.large_string()),
]
),
)
def test_error_message(sp: Session):
df = sp.from_items([1, 2, 3])
df = sp.partial_sql("select a,, from {0}", df)
try:
df.to_arrow()
except Exception as ex:
# sql query should be in the exception message
assert "select a,, from" in str(ex)
else:
assert False, "expected exception"
def test_unpicklable_task_exception(sp: Session):
from loguru import logger
df = sp.from_items([1, 2, 3])
try:
df.map(lambda x: logger.info("use outside logger")).to_arrow()
except Exception as ex:
assert "Can't pickle task" in str(ex)
assert (
"HINT: DO NOT use externally imported loguru logger in your task. Please import it within the task."
in str(ex)
)
else:
assert False, "expected exception"
def test_log(sp: Session):
df = sp.from_items([1, 2, 3])
def log_record(x):
import logging
import sys
from loguru import logger
print("stdout")
print("stderr", file=sys.stderr)
logger.info("loguru")
logging.info("logging")
return x
df.map(log_record).to_arrow()
# TODO: check logs should be see in the log file
# FIXME: logs in unit test are not written to the log file
# because we share the same ray instance for all tests

174
tests/test_dataset.py Normal file
View File

@@ -0,0 +1,174 @@
import glob
import os.path
import unittest
from pathlib import PurePath
import duckdb
import pandas
import pyarrow as arrow
import pytest
from loguru import logger
from smallpond.common import DEFAULT_ROW_GROUP_SIZE, MB
from smallpond.logical.dataset import ParquetDataSet
from smallpond.utility import ConcurrentIter
from tests.test_fabric import TestFabric
class TestDataSet(TestFabric, unittest.TestCase):
def test_parquet_file_created_by_pandas(self):
num_urls = 0
for txt_file in glob.glob("tests/data/mock_urls/*.tsv"):
urls = pandas.read_csv(txt_file, delimiter="\t", names=["url"])
urls.to_parquet(
os.path.join(
self.output_root_abspath,
PurePath(os.path.basename(txt_file)).with_suffix(".parquet"),
)
)
num_urls += urls.size
dataset = ParquetDataSet([os.path.join(self.output_root_abspath, "*.parquet")])
self.assertEqual(num_urls, dataset.num_rows)
def _generate_parquet_dataset(
self, output_path, npartitions, num_rows, row_group_size
):
duckdb.sql(
f"""copy (
select range as i, range % {npartitions} as partition from range(0, {num_rows}) )
to '{output_path}'
(FORMAT PARQUET, ROW_GROUP_SIZE {row_group_size}, PARTITION_BY partition, OVERWRITE_OR_IGNORE true)"""
)
return ParquetDataSet([f"{output_path}/**/*.parquet"])
def _check_partition_datasets(
self, orig_dataset: ParquetDataSet, partition_func, npartition
):
# build partitioned datasets
partitioned_datasets = partition_func(npartition)
self.assertEqual(npartition, len(partitioned_datasets))
self.assertEqual(
orig_dataset.num_rows,
sum(dataset.num_rows for dataset in partitioned_datasets),
)
# load as arrow table
loaded_table = arrow.concat_tables(
[dataset.to_arrow_table(max_workers=1) for dataset in partitioned_datasets]
)
self.assertEqual(orig_dataset.num_rows, loaded_table.num_rows)
# compare arrow tables
orig_table = orig_dataset.to_arrow_table(max_workers=1)
self.assertEqual(orig_table.shape, loaded_table.shape)
self.assertTrue(orig_table.sort_by("i").equals(loaded_table.sort_by("i")))
# compare sql query results
join_query = f"""
select count(a.i) as num_rows
from {orig_dataset.sql_query_fragment()} as a
join ( {' union all '.join([dataset.sql_query_fragment() for dataset in partitioned_datasets])} ) as b on a.i = b.i"""
results = duckdb.sql(join_query).fetchall()
self.assertEqual(orig_dataset.num_rows, results[0][0])
def test_num_rows(self):
dataset = ParquetDataSet(["tests/data/arrow/*.parquet"])
self.assertEqual(dataset.num_rows, 1000)
def test_partition_by_files(self):
output_path = os.path.join(self.output_root_abspath, "test_partition_by_files")
orig_dataset = self._generate_parquet_dataset(
output_path, npartitions=11, num_rows=170 * 1000, row_group_size=10 * 1000
)
num_files = len(orig_dataset.resolved_paths)
for npartition in range(1, num_files + 1):
for random_shuffle in (False, True):
with self.subTest(npartition=npartition, random_shuffle=random_shuffle):
orig_dataset.reset(orig_dataset.paths, orig_dataset.root_dir)
self._check_partition_datasets(
orig_dataset,
lambda n: orig_dataset.partition_by_files(
n, random_shuffle=random_shuffle
),
npartition,
)
def test_partition_by_rows(self):
output_path = os.path.join(self.output_root_abspath, "test_partition_by_rows")
orig_dataset = self._generate_parquet_dataset(
output_path, npartitions=11, num_rows=170 * 1000, row_group_size=10 * 1000
)
num_files = len(orig_dataset.resolved_paths)
for npartition in range(1, 2 * num_files + 1):
for random_shuffle in (False, True):
with self.subTest(npartition=npartition, random_shuffle=random_shuffle):
orig_dataset.reset(orig_dataset.paths, orig_dataset.root_dir)
self._check_partition_datasets(
orig_dataset,
lambda n: orig_dataset.partition_by_rows(
n, random_shuffle=random_shuffle
),
npartition,
)
def test_resolved_many_paths(self):
with open("tests/data/long_path_list.txt", buffering=16 * MB) as fin:
filenames = list(map(os.path.basename, map(str.strip, fin.readlines())))
logger.info(f"loaded {len(filenames)} filenames")
dataset = ParquetDataSet(filenames)
self.assertEqual(len(dataset.resolved_paths), len(filenames))
def test_paths_with_char_ranges(self):
dataset_with_char_ranges = ParquetDataSet(
["tests/data/arrow/data[0-9].parquet"]
)
dataset_with_wildcards = ParquetDataSet(["tests/data/arrow/*.parquet"])
self.assertEqual(
len(dataset_with_char_ranges.resolved_paths),
len(dataset_with_wildcards.resolved_paths),
)
def test_to_arrow_table_batch_reader(self):
memdb = duckdb.connect(
database=":memory:", config={"arrow_large_buffer_size": "true"}
)
for dataset_path in (
"tests/data/arrow/*.parquet",
"tests/data/large_array/*.parquet",
):
for conn in (None, memdb):
print(f"dataset_path: {dataset_path}, conn: {conn}")
with self.subTest(dataset_path=dataset_path, conn=conn):
dataset = ParquetDataSet([dataset_path])
to_batches = dataset.to_arrow_table(
max_workers=1, conn=conn
).to_batches(max_chunksize=DEFAULT_ROW_GROUP_SIZE * 2)
batch_reader = dataset.to_batch_reader(
batch_size=DEFAULT_ROW_GROUP_SIZE * 2, conn=conn
)
with ConcurrentIter(
batch_reader, max_buffer_size=2
) as batch_reader:
for batch_iter in (to_batches, batch_reader):
total_num_rows = 0
for batch in batch_iter:
print(
f"batch.num_rows {batch.num_rows}, max_batch_row_size {DEFAULT_ROW_GROUP_SIZE*2}"
)
self.assertLessEqual(
batch.num_rows, DEFAULT_ROW_GROUP_SIZE * 2
)
total_num_rows += batch.num_rows
print(f"{dataset_path}: total_num_rows {total_num_rows}")
self.assertEqual(total_num_rows, dataset.num_rows)
@pytest.mark.parametrize("reader", ["arrow", "duckdb"])
@pytest.mark.parametrize("dataset_path", ["tests/data/arrow/*.parquet"])
# @pytest.mark.parametrize("dataset_path", ["tests/data/arrow/*.parquet", "tests/data/large_array/*.parquet"])
def test_arrow_reader(benchmark, reader: str, dataset_path: str):
dataset = ParquetDataSet([dataset_path])
conn = None
if reader == "duckdb":
conn = duckdb.connect(
database=":memory:", config={"arrow_large_buffer_size": "true"}
)
benchmark(dataset.to_arrow_table, conn=conn)
# result: arrow reader is 4x faster than duckdb reader in small dataset, 1.4x faster in large dataset

60
tests/test_deltalake.py Normal file
View File

@@ -0,0 +1,60 @@
import glob
import importlib
import tempfile
import unittest
from smallpond.io.arrow import cast_columns_to_large_string
from tests.test_fabric import TestFabric
@unittest.skipUnless(
importlib.util.find_spec("deltalake") is not None, "cannot find deltalake"
)
class TestDeltaLake(TestFabric, unittest.TestCase):
def test_read_write_deltalake(self):
from deltalake import DeltaTable, write_deltalake
for dataset_path in (
"tests/data/arrow/*.parquet",
"tests/data/large_array/*.parquet",
):
parquet_files = glob.glob(dataset_path)
expected = self._load_parquet_files(parquet_files)
with self.subTest(dataset_path=dataset_path), tempfile.TemporaryDirectory(
dir=self.output_root_abspath
) as output_dir:
write_deltalake(output_dir, expected, large_dtypes=True)
dt = DeltaTable(output_dir)
self._compare_arrow_tables(expected, dt.to_pyarrow_table())
def test_load_mixed_large_dtypes(self):
from deltalake import DeltaTable, write_deltalake
for dataset_path in (
"tests/data/arrow/*.parquet",
"tests/data/large_array/*.parquet",
):
parquet_files = glob.glob(dataset_path)
with self.subTest(dataset_path=dataset_path), tempfile.TemporaryDirectory(
dir=self.output_root_abspath
) as output_dir:
table = cast_columns_to_large_string(
self._load_parquet_files(parquet_files)
)
write_deltalake(output_dir, table, large_dtypes=True, mode="overwrite")
write_deltalake(output_dir, table, large_dtypes=False, mode="append")
loaded_table = DeltaTable(output_dir).to_pyarrow_table()
print("table:\n", table.schema)
print("loaded_table:\n", loaded_table.schema)
self.assertEqual(table.num_rows * 2, loaded_table.num_rows)
def test_delete_update(self):
import pandas as pd
from deltalake import DeltaTable, write_deltalake
with tempfile.TemporaryDirectory(dir=self.output_root_abspath) as output_dir:
df = pd.DataFrame({"num": [1, 2, 3], "animal": ["cat", "dog", "snake"]})
write_deltalake(output_dir, df, mode="overwrite")
dt = DeltaTable(output_dir)
dt.delete("animal = 'cat'")
dt.update(predicate="num = 3", new_values={"animal": "fish"})

46
tests/test_driver.py Normal file
View File

@@ -0,0 +1,46 @@
import os.path
import unittest
import uuid
from loguru import logger
from benchmarks.gray_sort_benchmark import gray_sort_benchmark
from examples.sort_mock_urls import sort_mock_urls
from smallpond.common import GB, MB
from smallpond.execution.driver import Driver
from tests.test_fabric import TestFabric
@unittest.skipUnless(os.getenv("ENABLE_DRIVER_TEST"), "unit test disabled")
class TestDriver(TestFabric, unittest.TestCase):
fault_inject_prob = 0.05
def create_driver(self, num_executors: int):
cmdline = f"scheduler --job_id {str(uuid.uuid4())} --job_name {self._testMethodName} --data_root {self.output_root_abspath} --num_executors {num_executors} --fault_inject_prob {self.fault_inject_prob}"
driver = Driver()
driver.parse_arguments(args=cmdline.split())
logger.info(f"{cmdline=} {driver.mode=} {driver.job_id=} {driver.data_root=}")
return driver
def test_standalone_mode(self):
plan = sort_mock_urls(["tests/data/mock_urls/*.tsv"], npartitions=3)
driver = self.create_driver(num_executors=0)
exec_plan = driver.run(plan, stop_process_on_done=False)
self.assertTrue(exec_plan.successful)
self.assertGreater(exec_plan.final_output.num_files, 0)
def test_run_on_remote_executors(self):
driver = self.create_driver(num_executors=2)
plan = gray_sort_benchmark(
record_nbytes=100,
key_nbytes=10,
total_data_nbytes=1 * GB,
gensort_batch_nbytes=100 * MB,
num_data_partitions=10,
num_sort_partitions=10,
validate_results=True,
)
exec_plan = driver.run(plan, stop_process_on_done=False)
self.assertTrue(exec_plan.successful)
self.assertGreater(exec_plan.final_output.num_files, 0)

886
tests/test_execution.py Normal file
View File

@@ -0,0 +1,886 @@
import functools
import os.path
import socket
import tempfile
import time
import unittest
from datetime import datetime
from typing import Iterable, List, Tuple
import pandas
import pyarrow as arrow
from loguru import logger
from pandas.core.api import DataFrame as DataFrame
from smallpond.common import GB, MB, split_into_rows
from smallpond.execution.task import (
DataSinkTask,
DataSourceTask,
JobId,
PartitionInfo,
PythonScriptTask,
RuntimeContext,
StreamOutput,
)
from smallpond.execution.workqueue import WorkStatus
from smallpond.logical.dataset import (
ArrowTableDataSet,
DataSet,
ParquetDataSet,
SqlQueryDataSet,
)
from smallpond.logical.node import (
ArrowBatchNode,
ArrowComputeNode,
ArrowStreamNode,
Context,
DataSetPartitionNode,
DataSinkNode,
DataSourceNode,
EvenlyDistributedPartitionNode,
HashPartitionNode,
LogicalPlan,
Node,
PandasBatchNode,
PandasComputeNode,
ProjectionNode,
PythonScriptNode,
RootNode,
SqlEngineNode,
)
from smallpond.logical.udf import UDFListType, UDFType
from tests.test_fabric import TestFabric
class OutputMsgPythonTask(PythonScriptTask):
def __init__(self, msg: str, *args, **kwargs) -> None:
super().__init__(*args, **kwargs)
self.msg = msg
def initialize(self):
pass
def finalize(self):
pass
def process(
self,
runtime_ctx: RuntimeContext,
input_datasets: List[DataSet],
output_path: str,
) -> bool:
logger.info(
f"msg: {self.msg}, num files: {input_datasets[0].num_files}, local gpu ranks: {self.local_gpu_ranks}"
)
self.inject_fault()
return True
# method1: inherit Task class and override spawn method
class OutputMsgPythonNode(PythonScriptNode):
def spawn(self, *args, **kwargs) -> OutputMsgPythonTask:
return OutputMsgPythonTask("python script", *args, **kwargs)
# method2: override process method
# this usage is not recommended and only for testing. use `process_func` instead.
class OutputMsgPythonNode2(PythonScriptNode):
def __init__(self, ctx: Context, input_deps: Tuple[Node, ...], msg: str) -> None:
super().__init__(ctx, input_deps)
self.msg = msg
def process(
self,
runtime_ctx: RuntimeContext,
input_datasets: List[DataSet],
output_path: str,
) -> bool:
logger.info(f"msg: {self.msg}, num files: {input_datasets[0].num_files}")
return True
# this usage is not recommended and only for testing. use `process_func` instead.
class CopyInputArrowNode(ArrowComputeNode):
def __init__(self, ctx: Context, input_deps: Tuple[Node, ...], msg: str) -> None:
super().__init__(ctx, input_deps)
self.msg = msg
def process(
self, runtime_ctx: RuntimeContext, input_tables: List[arrow.Table]
) -> arrow.Table:
return copy_input_arrow(runtime_ctx, input_tables, self.msg)
# this usage is not recommended and only for testing. use `process_func` instead.
class CopyInputStreamNode(ArrowStreamNode):
def __init__(self, ctx: Context, input_deps: Tuple[Node, ...], msg: str) -> None:
super().__init__(ctx, input_deps)
self.msg = msg
def process(
self, runtime_ctx: RuntimeContext, input_readers: List[arrow.RecordBatchReader]
) -> Iterable[arrow.Table]:
return copy_input_stream(runtime_ctx, input_readers, self.msg)
def copy_input_arrow(
runtime_ctx: RuntimeContext, input_tables: List[arrow.Table], msg: str
) -> arrow.Table:
logger.info(f"msg: {msg}, num rows: {input_tables[0].num_rows}")
time.sleep(runtime_ctx.secs_executor_probe_interval)
runtime_ctx.task.inject_fault()
return input_tables[0]
def copy_input_stream(
runtime_ctx: RuntimeContext, input_readers: List[arrow.RecordBatchReader], msg: str
) -> Iterable[arrow.Table]:
for index, batch in enumerate(input_readers[0]):
logger.info(f"msg: {msg}, batch index: {index}, num rows: {batch.num_rows}")
time.sleep(runtime_ctx.secs_executor_probe_interval)
yield StreamOutput(
arrow.Table.from_batches([batch]),
batch_indices=[index],
force_checkpoint=True,
)
runtime_ctx.task.inject_fault()
def copy_input_batch(
runtime_ctx: RuntimeContext, input_batches: List[arrow.Table], msg: str
) -> arrow.Table:
logger.info(f"msg: {msg}, num rows: {input_batches[0].num_rows}")
time.sleep(runtime_ctx.secs_executor_probe_interval)
runtime_ctx.task.inject_fault()
return input_batches[0]
def copy_input_data_frame(
runtime_ctx: RuntimeContext, input_dfs: List[DataFrame]
) -> DataFrame:
runtime_ctx.task.inject_fault()
return input_dfs[0]
def copy_input_data_frame_batch(
runtime_ctx: RuntimeContext, input_dfs: List[DataFrame]
) -> DataFrame:
runtime_ctx.task.inject_fault()
return input_dfs[0]
def merge_input_tables(
runtime_ctx: RuntimeContext, input_batches: List[arrow.Table]
) -> arrow.Table:
runtime_ctx.task.inject_fault()
output = arrow.concat_tables(input_batches)
logger.info(
f"input rows: {[len(batch) for batch in input_batches]}, output rows: {len(output)}"
)
return output
def merge_input_data_frames(
runtime_ctx: RuntimeContext, input_dfs: List[DataFrame]
) -> DataFrame:
runtime_ctx.task.inject_fault()
output = pandas.concat(input_dfs)
logger.info(
f"input rows: {[len(df) for df in input_dfs]}, output rows: {len(output)}"
)
return output
def parse_url(
runtime_ctx: RuntimeContext, input_tables: List[arrow.Table]
) -> arrow.Table:
urls = input_tables[0].columns[0]
hosts = [url.as_py().split("/", maxsplit=2)[0] for url in urls]
return input_tables[0].append_column("host", arrow.array(hosts))
def nonzero_exit_code(
runtime_ctx: RuntimeContext, input_datasets: List[DataSet], output_path: str
) -> bool:
import sys
if runtime_ctx.task._memory_boost == 1:
sys.exit(1)
return True
# create an empty file with a fixed name
def empty_file(
runtime_ctx: RuntimeContext, input_datasets: List[DataSet], output_path: str
) -> bool:
import os
with open(os.path.join(output_path, "file"), "w") as fout:
pass
return True
def return_fake_gpus(count: int = 8):
import GPUtil
return [GPUtil.GPU(i, *list(range(11))) for i in range(count)]
def split_url(urls: arrow.array) -> arrow.array:
url_parts = [url.as_py().split("/") for url in urls]
return arrow.array(url_parts, type=arrow.list_(arrow.string()))
def choose_random_urls(
runtime_ctx: RuntimeContext, input_tables: List[arrow.Table], k: int = 5
) -> arrow.Table:
# get the current running task
runtime_task = runtime_ctx.task
# access task-specific attributes
cpu_limit = runtime_task.cpu_limit
random_gen = runtime_task.python_random_gen
# input data
(url_table,) = input_tables
hosts, urls = url_table.columns
logger.info(f"{cpu_limit=} {len(urls)=}")
# generate ramdom samples
random_urls = random_gen.choices(urls.to_pylist(), k=k)
return arrow.Table.from_arrays([arrow.array(random_urls)], names=["random_urls"])
class TestExecution(TestFabric, unittest.TestCase):
fault_inject_prob = 0.05
def test_arrow_task(self):
for use_duckdb_reader in (False, True):
with self.subTest(use_duckdb_reader=use_duckdb_reader):
with tempfile.TemporaryDirectory(
dir=self.output_root_abspath
) as output_dir:
ctx = Context()
dataset = ParquetDataSet(["tests/data/mock_urls/*.parquet"])
data_table = dataset.to_arrow_table()
data_files = DataSourceNode(ctx, dataset)
data_partitions = DataSetPartitionNode(
ctx, (data_files,), npartitions=7
)
if use_duckdb_reader:
data_partitions = ProjectionNode(
ctx,
data_partitions,
columns=["*", "string_split(url, '/')[0] as parsed_host"],
)
arrow_compute = ArrowComputeNode(
ctx,
(data_partitions,),
process_func=functools.partial(
copy_input_arrow, msg="arrow compute"
),
use_duckdb_reader=use_duckdb_reader,
output_name="arrow_compute",
output_path=output_dir,
cpu_limit=2,
)
arrow_stream = ArrowStreamNode(
ctx,
(data_partitions,),
process_func=functools.partial(
copy_input_stream, msg="arrow stream"
),
streaming_batch_size=10,
secs_checkpoint_interval=0.5,
use_duckdb_reader=use_duckdb_reader,
output_name="arrow_stream",
output_path=output_dir,
cpu_limit=2,
)
arrow_batch = ArrowBatchNode(
ctx,
(data_partitions,),
process_func=functools.partial(
copy_input_batch, msg="arrow batch"
),
streaming_batch_size=10,
secs_checkpoint_interval=0.5,
use_duckdb_reader=use_duckdb_reader,
output_name="arrow_batch",
output_path=output_dir,
cpu_limit=2,
)
data_sink = DataSinkNode(
ctx,
(arrow_compute, arrow_stream, arrow_batch),
output_path=output_dir,
)
plan = LogicalPlan(ctx, data_sink)
exec_plan = self.execute_plan(
plan, fault_inject_prob=0.1, secs_executor_probe_interval=0.5
)
self.assertTrue(
all(map(os.path.exists, exec_plan.final_output.resolved_paths))
)
arrow_compute_output = ParquetDataSet(
[os.path.join(output_dir, "arrow_compute", "**/*.parquet")],
recursive=True,
)
arrow_stream_output = ParquetDataSet(
[os.path.join(output_dir, "arrow_stream", "**/*.parquet")],
recursive=True,
)
arrow_batch_output = ParquetDataSet(
[os.path.join(output_dir, "arrow_batch", "**/*.parquet")],
recursive=True,
)
self._compare_arrow_tables(
data_table,
arrow_compute_output.to_arrow_table().select(
data_table.column_names
),
)
self._compare_arrow_tables(
data_table,
arrow_stream_output.to_arrow_table().select(
data_table.column_names
),
)
self._compare_arrow_tables(
data_table,
arrow_batch_output.to_arrow_table().select(
data_table.column_names
),
)
def test_pandas_task(self):
with tempfile.TemporaryDirectory(dir=self.output_root_abspath) as output_dir:
ctx = Context()
dataset = ParquetDataSet(["tests/data/mock_urls/*.parquet"])
data_table = dataset.to_arrow_table()
data_files = DataSourceNode(ctx, dataset)
data_partitions = DataSetPartitionNode(ctx, (data_files,), npartitions=7)
pandas_compute = PandasComputeNode(
ctx,
(data_partitions,),
process_func=copy_input_data_frame,
output_name="pandas_compute",
output_path=output_dir,
cpu_limit=2,
)
pandas_batch = PandasBatchNode(
ctx,
(data_partitions,),
process_func=copy_input_data_frame_batch,
streaming_batch_size=10,
secs_checkpoint_interval=0.5,
output_name="pandas_batch",
output_path=output_dir,
cpu_limit=2,
)
data_sink = DataSinkNode(
ctx, (pandas_compute, pandas_batch), output_path=output_dir
)
plan = LogicalPlan(ctx, data_sink)
exec_plan = self.execute_plan(
plan, fault_inject_prob=0.1, secs_executor_probe_interval=0.5
)
self.assertTrue(
all(map(os.path.exists, exec_plan.final_output.resolved_paths))
)
pandas_compute_output = ParquetDataSet(
[os.path.join(output_dir, "pandas_compute", "**/*.parquet")],
recursive=True,
)
pandas_batch_output = ParquetDataSet(
[os.path.join(output_dir, "pandas_batch", "**/*.parquet")],
recursive=True,
)
self._compare_arrow_tables(
data_table, pandas_compute_output.to_arrow_table()
)
self._compare_arrow_tables(data_table, pandas_batch_output.to_arrow_table())
def test_variable_length_input_datasets(self):
ctx = Context()
small_dataset = ParquetDataSet(["tests/data/mock_urls/*.parquet"])
large_dataset = ParquetDataSet(["tests/data/mock_urls/*.parquet"] * 10)
small_partitions = DataSetPartitionNode(
ctx, (DataSourceNode(ctx, small_dataset),), npartitions=7
)
large_partitions = DataSetPartitionNode(
ctx, (DataSourceNode(ctx, large_dataset),), npartitions=7
)
arrow_batch = ArrowBatchNode(
ctx,
(small_partitions, large_partitions),
process_func=merge_input_tables,
streaming_batch_size=100,
secs_checkpoint_interval=0.5,
output_name="arrow_batch",
cpu_limit=2,
)
pandas_batch = PandasBatchNode(
ctx,
(small_partitions, large_partitions),
process_func=merge_input_data_frames,
streaming_batch_size=100,
secs_checkpoint_interval=0.5,
output_name="pandas_batch",
cpu_limit=2,
)
plan = LogicalPlan(ctx, RootNode(ctx, (arrow_batch, pandas_batch)))
exec_plan = self.execute_plan(
plan, fault_inject_prob=0.1, secs_executor_probe_interval=0.5
)
self.assertTrue(all(map(os.path.exists, exec_plan.final_output.resolved_paths)))
arrow_batch_output = ParquetDataSet(
[os.path.join(exec_plan.ctx.output_root, "arrow_batch", "**/*.parquet")],
recursive=True,
)
pandas_batch_output = ParquetDataSet(
[os.path.join(exec_plan.ctx.output_root, "pandas_batch", "**/*.parquet")],
recursive=True,
)
self.assertEqual(
small_dataset.num_rows + large_dataset.num_rows, arrow_batch_output.num_rows
)
self.assertEqual(
small_dataset.num_rows + large_dataset.num_rows,
pandas_batch_output.num_rows,
)
def test_projection_task(self):
ctx = Context()
# select columns when defining dataset
dataset = ParquetDataSet(["tests/data/mock_urls/*.parquet"], columns=["url"])
data_files = DataSourceNode(ctx, dataset)
data_partitions = DataSetPartitionNode(
ctx, (data_files,), npartitions=3, partition_by_rows=True
)
# projection as input of arrow node
generated_columns = ["filename", "file_row_number"]
urls_with_host = ArrowComputeNode(
ctx,
(ProjectionNode(ctx, data_partitions, ["url"], generated_columns),),
process_func=parse_url,
use_duckdb_reader=True,
)
# projection as input of sql node
distinct_urls_with_host = SqlEngineNode(
ctx,
(
ProjectionNode(
ctx,
data_partitions,
["url", "string_split(url, '/')[0] as host"],
generated_columns,
),
),
r"select distinct host, url, filename from {0}",
)
# unify different schemas
merged_diff_schemas = ProjectionNode(
ctx,
DataSetPartitionNode(
ctx, (distinct_urls_with_host, urls_with_host), npartitions=1
),
union_by_name=True,
)
host_partitions = HashPartitionNode(
ctx,
(merged_diff_schemas,),
npartitions=3,
hash_columns=["host"],
engine_type="duckdb",
output_name="host_partitions",
)
host_partitions.max_num_producer_tasks = 1
plan = LogicalPlan(ctx, host_partitions)
final_output = self.execute_plan(plan, fault_inject_prob=0.1).final_output
final_table = final_output.to_arrow_table()
self.assertEqual(
sorted(
[
"url",
"host",
*generated_columns,
HashPartitionNode.default_data_partition_column,
]
),
sorted(final_table.column_names),
)
def test_arrow_type_in_udfs(self):
ctx = Context()
dataset = ParquetDataSet(["tests/data/mock_urls/*.parquet"])
data_files = DataSourceNode(ctx, dataset)
data_partitions = DataSetPartitionNode(
ctx, (data_files,), npartitions=dataset.num_files
)
ctx.create_function(
"split_url",
split_url,
[UDFType.VARCHAR],
UDFListType(UDFType.VARCHAR),
use_arrow_type=True,
)
uniq_hosts = SqlEngineNode(
ctx,
(data_partitions,),
r"select split_url(url) as url_parts from {0}",
udfs=["split_url"],
)
plan = LogicalPlan(ctx, uniq_hosts)
self.execute_plan(plan)
def test_many_simple_tasks(self):
ctx = Context()
npartitions = 1000
dataset = ParquetDataSet(["tests/data/mock_urls/*.parquet"] * npartitions)
data_files = DataSourceNode(ctx, dataset)
data_partitions = EvenlyDistributedPartitionNode(
ctx, (data_files,), npartitions=npartitions
)
output_msg = OutputMsgPythonNode(ctx, (data_partitions,))
plan = LogicalPlan(ctx, output_msg)
self.execute_plan(
plan,
num_executors=10,
secs_executor_probe_interval=5,
enable_profiling=True,
)
def test_many_producers_and_partitions(self):
ctx = Context()
npartitions = 10000
dataset = ParquetDataSet(
["tests/data/mock_urls/*.parquet"] * (npartitions * 10)
)
data_files = DataSourceNode(ctx, dataset)
data_partitions = EvenlyDistributedPartitionNode(
ctx, (data_files,), npartitions=npartitions, cpu_limit=1
)
data_partitions.max_num_producer_tasks = 20
output_msg = OutputMsgPythonNode(ctx, (data_partitions,))
plan = LogicalPlan(ctx, output_msg)
self.execute_plan(
plan,
num_executors=10,
secs_executor_probe_interval=5,
enable_profiling=True,
)
def test_local_gpu_rank(self):
ctx = Context()
dataset = ParquetDataSet(["tests/data/mock_urls/*.parquet"])
data_files = DataSourceNode(ctx, dataset)
data_partitions = DataSetPartitionNode(
ctx, (data_files,), npartitions=dataset.num_files
)
output_msg = OutputMsgPythonNode(
ctx, (data_partitions,), cpu_limit=1, gpu_limit=0.5
)
plan = LogicalPlan(ctx, output_msg)
runtime_ctx = RuntimeContext(
JobId.new(),
datetime.now(),
self.output_root_abspath,
console_log_level="WARNING",
)
runtime_ctx.get_local_gpus = return_fake_gpus
runtime_ctx.initialize(socket.gethostname(), cleanup_root=True)
self.execute_plan(plan, runtime_ctx=runtime_ctx)
def test_python_node_with_process_method(self):
ctx = Context()
dataset = ParquetDataSet(["tests/data/mock_urls/*.parquet"])
data_files = DataSourceNode(ctx, dataset)
copy_input_arrow_node = CopyInputArrowNode(ctx, (data_files,), "hello")
copy_input_stream_node = CopyInputStreamNode(ctx, (data_files,), "hello")
output_msg = OutputMsgPythonNode2(
ctx, (copy_input_arrow_node, copy_input_stream_node), "hello"
)
plan = LogicalPlan(ctx, output_msg)
self.execute_plan(plan)
def test_sql_engine_oom(self):
ctx = Context()
dataset = ParquetDataSet(["tests/data/mock_urls/*.parquet"])
data_files = DataSourceNode(ctx, dataset)
uniq_urls = SqlEngineNode(
ctx, (data_files,), r"select distinct * from {0}", memory_limit=2 * MB
)
uniq_url_partitions = DataSetPartitionNode(ctx, (uniq_urls,), 2)
uniq_url_count = SqlEngineNode(
ctx,
(uniq_url_partitions,),
sql_query=r"select count(distinct columns(*)) from {0}",
memory_limit=2 * MB,
)
plan = LogicalPlan(ctx, uniq_url_count)
self.execute_plan(plan, max_fail_count=10)
@unittest.skip("flaky on CI")
def test_enforce_memory_limit(self):
ctx = Context()
dataset = ParquetDataSet(["tests/data/arrow/*.parquet"])
data_files = DataSourceNode(ctx, dataset)
arrow_compute = ArrowComputeNode(
ctx,
(data_files,),
process_func=functools.partial(copy_input_arrow, msg="arrow compute"),
memory_limit=1 * GB,
)
arrow_stream = ArrowStreamNode(
ctx,
(data_files,),
process_func=functools.partial(copy_input_stream, msg="arrow stream"),
memory_limit=1 * GB,
)
with tempfile.TemporaryDirectory(dir=self.output_root_abspath) as output_dir:
data_sink = DataSinkNode(
ctx, (arrow_compute, arrow_stream), output_path=output_dir
)
plan = LogicalPlan(ctx, data_sink)
self.execute_plan(
plan,
max_fail_count=10,
enforce_memory_limit=True,
nonzero_exitcode_as_oom=True,
)
def test_task_crash_as_oom(self):
ctx = Context()
dataset = ParquetDataSet(["tests/data/mock_urls/*.parquet"])
data_files = DataSourceNode(ctx, dataset)
nonzero_exitcode = PythonScriptNode(
ctx, (data_files,), process_func=nonzero_exit_code
)
plan = LogicalPlan(ctx, nonzero_exitcode)
exec_plan = self.execute_plan(
plan, num_executors=1, check_result=False, nonzero_exitcode_as_oom=False
)
self.assertFalse(exec_plan.successful)
exec_plan = self.execute_plan(
plan, num_executors=1, check_result=False, nonzero_exitcode_as_oom=True
)
self.assertTrue(exec_plan.successful)
def test_manifest_only_data_sink(self):
with open("tests/data/long_path_list.txt", buffering=16 * MB) as fin:
filenames = list(map(str.strip, fin.readlines()))
logger.info(f"loaded {len(filenames)} filenames")
with tempfile.TemporaryDirectory(dir=self.output_root_abspath) as output_dir:
ctx = Context()
dataset = ParquetDataSet(filenames)
data_files = DataSourceNode(ctx, dataset)
data_partitions = DataSetPartitionNode(ctx, (data_files,), npartitions=512)
data_sink = DataSinkNode(
ctx, (data_partitions,), output_path=output_dir, manifest_only=True
)
plan = LogicalPlan(ctx, data_sink)
self.execute_plan(plan)
with open(
os.path.join(output_dir, DataSinkTask.manifest_filename),
buffering=16 * MB,
) as fin:
num_lines = len(fin.readlines())
self.assertEqual(len(filenames), num_lines)
def test_sql_batched_processing(self):
for materialize_in_memory in (False, True):
with self.subTest(materialize_in_memory=materialize_in_memory):
ctx = Context()
dataset = ParquetDataSet(["tests/data/large_array/*.parquet"] * 2)
data_files = DataSourceNode(ctx, dataset)
content_length = SqlEngineNode(
ctx,
(data_files,),
r"select url, octet_length(content) as content_len from {0}",
materialize_in_memory=materialize_in_memory,
batched_processing=True,
cpu_limit=2,
memory_limit=2 * GB,
)
plan = LogicalPlan(ctx, content_length)
final_output: ParquetDataSet = self.execute_plan(plan).final_output
self.assertEqual(dataset.num_rows, final_output.num_rows)
def test_multiple_sql_queries(self):
for materialize_in_memory in (False, True):
with self.subTest(materialize_in_memory=materialize_in_memory):
ctx = Context()
dataset = ParquetDataSet(["tests/data/large_array/*.parquet"] * 2)
data_files = DataSourceNode(ctx, dataset)
content_length = SqlEngineNode(
ctx,
(data_files,),
[
r"create or replace temp table content_len_data as select url, octet_length(content) as content_len from {0}",
r"select * from content_len_data",
],
materialize_in_memory=materialize_in_memory,
batched_processing=True,
cpu_limit=2,
memory_limit=2 * GB,
)
plan = LogicalPlan(ctx, content_length)
final_output: ParquetDataSet = self.execute_plan(plan).final_output
self.assertEqual(dataset.num_rows, final_output.num_rows)
def test_temp_outputs_in_final_results(self):
ctx = Context()
dataset = ParquetDataSet(["tests/data/mock_urls/*.parquet"])
data_files = DataSourceNode(ctx, dataset)
data_partitions = DataSetPartitionNode(ctx, (data_files,), npartitions=10)
url_counts = SqlEngineNode(
ctx, (data_partitions,), r"select count(url) as cnt from {0}"
)
distinct_url_counts = SqlEngineNode(
ctx, (data_partitions,), r"select count(distinct url) as cnt from {0}"
)
merged_counts = DataSetPartitionNode(
ctx,
(
ProjectionNode(ctx, url_counts, ["cnt"]),
ProjectionNode(ctx, distinct_url_counts, ["cnt"]),
),
npartitions=1,
)
split_counts = DataSetPartitionNode(ctx, (merged_counts,), npartitions=10)
plan = LogicalPlan(ctx, split_counts)
final_output: ParquetDataSet = self.execute_plan(plan).final_output
self.assertEqual(data_partitions.npartitions * 2, final_output.num_rows)
def test_override_output_path(self):
ctx = Context()
dataset = ParquetDataSet(["tests/data/mock_urls/*.parquet"])
data_files = DataSourceNode(ctx, dataset)
data_partitions = DataSetPartitionNode(ctx, (data_files,), npartitions=10)
url_counts = SqlEngineNode(
ctx,
(data_partitions,),
r"select count(url) as cnt from {0}",
output_name="url_counts",
)
distinct_url_counts = SqlEngineNode(
ctx, (data_partitions,), r"select count(distinct url) as cnt from {0}"
)
merged_counts = DataSetPartitionNode(
ctx,
(
ProjectionNode(ctx, url_counts, ["cnt"]),
ProjectionNode(ctx, distinct_url_counts, ["cnt"]),
),
npartitions=1,
)
plan = LogicalPlan(ctx, merged_counts)
output_path = os.path.join(self.runtime_ctx.output_root, "final_output")
final_output = self.execute_plan(plan, output_path=output_path).final_output
self.assertTrue(os.path.exists(os.path.join(output_path, "url_counts")))
self.assertTrue(os.path.exists(os.path.join(output_path, "FinalResults")))
def test_data_sink_avoid_filename_conflicts(self):
ctx = Context()
dataset = ParquetDataSet(["tests/data/mock_urls/*.parquet"])
data_files = DataSourceNode(ctx, dataset)
data_partitions = DataSetPartitionNode(ctx, (data_files,), npartitions=10)
empty_files1 = PythonScriptNode(
ctx, (data_partitions,), process_func=empty_file
)
empty_files2 = PythonScriptNode(
ctx, (data_partitions,), process_func=empty_file
)
link_path = os.path.join(self.runtime_ctx.output_root, "link")
copy_path = os.path.join(self.runtime_ctx.output_root, "copy")
copy_input_path = os.path.join(self.runtime_ctx.output_root, "copy_input")
data_link = DataSinkNode(
ctx, (empty_files1, empty_files2), type="link", output_path=link_path
)
data_copy = DataSinkNode(
ctx, (empty_files1, empty_files2), type="copy", output_path=copy_path
)
data_copy_input = DataSinkNode(
ctx, (data_partitions,), type="copy", output_path=copy_input_path
)
plan = LogicalPlan(ctx, RootNode(ctx, (data_link, data_copy, data_copy_input)))
self.execute_plan(plan)
# there should be 21 files (20 input files + 1 manifest file) in the sink dir
self.assertEqual(21, len(os.listdir(link_path)))
self.assertEqual(21, len(os.listdir(copy_path)))
# file name should not be modified if no conflict
self.assertEqual(
set(
filename
for filename in os.listdir("tests/data/mock_urls")
if filename.endswith(".parquet")
),
set(
filename
for filename in os.listdir(copy_input_path)
if filename.endswith(".parquet")
),
)
def test_literal_datasets_as_data_sources(self):
ctx = Context()
num_rows = 10
query_dataset = SqlQueryDataSet(f"select i from range({num_rows}) as x(i)")
table_dataset = ArrowTableDataSet(
arrow.Table.from_arrays([list(range(num_rows))], names=["i"])
)
query_source = DataSourceNode(ctx, query_dataset)
table_source = DataSourceNode(ctx, table_dataset)
query_partitions = DataSetPartitionNode(
ctx, (query_source,), npartitions=num_rows, partition_by_rows=True
)
table_partitions = DataSetPartitionNode(
ctx, (table_source,), npartitions=num_rows, partition_by_rows=True
)
joined_rows = SqlEngineNode(
ctx,
(query_partitions, table_partitions),
r"select a.i as i, b.i as j from {0} as a join {1} as b on a.i = b.i",
)
plan = LogicalPlan(ctx, joined_rows)
final_output: ParquetDataSet = self.execute_plan(plan).final_output
self.assertEqual(num_rows, final_output.num_rows)
def test_partial_process_func(self):
ctx = Context()
dataset = ParquetDataSet(["tests/data/mock_urls/*.parquet"])
data_files = DataSourceNode(ctx, dataset)
data_partitions = DataSetPartitionNode(ctx, (data_files,), npartitions=3)
# use default value of k
random_urls_k5 = ArrowComputeNode(
ctx,
(data_partitions,),
process_func=choose_random_urls,
output_name="random_urls_k5",
)
# set value of k using functools.partial
random_urls_k10 = ArrowComputeNode(
ctx,
(data_partitions,),
process_func=functools.partial(choose_random_urls, k=10),
output_name="random_urls_k10",
)
random_urls_all = SqlEngineNode(
ctx,
(random_urls_k5, random_urls_k10),
r"select * from {0} union select * from {1}",
output_name="random_urls_all",
)
plan = LogicalPlan(ctx, random_urls_all)
exec_plan = self.execute_plan(plan)
self.assertEqual(
data_partitions.npartitions * 5,
exec_plan.get_output("random_urls_k5").to_arrow_table().num_rows,
)
self.assertEqual(
data_partitions.npartitions * 10,
exec_plan.get_output("random_urls_k10").to_arrow_table().num_rows,
)

294
tests/test_fabric.py Normal file
View File

@@ -0,0 +1,294 @@
import os.path
import queue
import sys
import unittest
from concurrent.futures import ThreadPoolExecutor
from datetime import datetime
from multiprocessing import Manager, Process
from typing import List, Optional
import fsspec
import numpy as np
import psutil
import pyarrow as arrow
import pyarrow.compute as pc
import pyarrow.parquet as parquet
from loguru import logger
from smallpond.common import DEFAULT_MAX_FAIL_COUNT, DEFAULT_MAX_RETRY_COUNT, GB, MB
from smallpond.execution.executor import Executor
from smallpond.execution.scheduler import Scheduler
from smallpond.execution.task import ExecutionPlan, JobId, RuntimeContext
from smallpond.io.arrow import cast_columns_to_large_string
from smallpond.logical.node import LogicalPlan
from smallpond.logical.planner import Planner
from tests.datagen import generate_data
generate_data()
def run_scheduler(
runtime_ctx: RuntimeContext, scheduler: Scheduler, queue: queue.Queue
):
runtime_ctx.initialize("scheduler")
scheduler.add_state_observer(Scheduler.StateObserver(SaveSchedState(queue)))
retval = scheduler.run()
print(f"scheduler exited with value {retval}", file=sys.stderr)
def run_executor(runtime_ctx: RuntimeContext, executor: Executor):
runtime_ctx.initialize(executor.id)
retval = executor.run()
print(f"{executor.id} exited with value {retval}", file=sys.stderr)
class SaveSchedState:
"""
A state observer that push the scheduler state into a queue when finished.
"""
def __init__(self, queue: queue.Queue):
self.queue = queue
def __call__(self, sched_state: Scheduler) -> bool:
if sched_state.num_local_running_works == 0:
self.queue.put(sched_state)
return True
class TestFabric(unittest.TestCase):
"""
A helper class that includes boilerplate code to test a logical plan.
"""
runtime_root = os.getenv("TEST_RUNTIME_ROOT") or f"tests/runtime"
runtime_ctx = None
fault_inject_prob = 0.00
queue_manager = None
sched_states: queue.Queue = None
latest_state: Scheduler = None
executors: List[Executor] = None
processes: List[Process] = None
@property
def output_dir(self):
return os.path.join(self.__class__.__name__, self._testMethodName)
@property
def output_root_abspath(self):
output_root = os.path.abspath(os.path.join(self.runtime_root, self.output_dir))
os.makedirs(output_root, exist_ok=True)
return output_root
def setUp(self) -> None:
try:
from pytest_cov.embed import cleanup_on_sigterm
except ImportError:
pass
else:
cleanup_on_sigterm()
self.runtime_ctx = RuntimeContext(
JobId.new(),
datetime.now(),
self.output_root_abspath,
console_log_level="WARNING",
)
self.runtime_ctx.initialize("setup")
return super().setUp()
def tearDown(self) -> None:
if self.sched_states is not None:
self.get_latest_sched_state()
assert self.sched_states.qsize() == 0
self.sched_states = None
if self.queue_manager is not None:
self.queue_manager.shutdown()
self.queue_manager = None
return super().tearDown()
def get_latest_sched_state(self) -> Scheduler:
while True:
try:
self.latest_state = self.sched_states.get(block=False)
except queue.Empty:
return self.latest_state
def join_running_procs(self, timeout=30):
for i, process in enumerate(self.processes):
if process.is_alive():
logger.info(f"join #{i} process: {process.name}")
process.join(timeout=None if i == 0 else timeout)
if process.exitcode is None:
logger.info(f"terminate #{i} process: {process.name}")
process.terminate()
process.join(timeout=timeout)
if process.exitcode is None:
logger.info(f"kill #{i} process: {process.name}")
process.kill()
process.join()
logger.info(
f"#{i} process {process.name} exited with code {process.exitcode}"
)
def start_execution(
self,
plan: LogicalPlan,
num_executors: int = 2,
secs_wq_poll_interval: float = 0.1,
secs_executor_probe_interval: float = 1,
max_num_missed_probes: int = 10,
max_retry_count: int = DEFAULT_MAX_RETRY_COUNT,
max_fail_count: int = DEFAULT_MAX_FAIL_COUNT,
prioritize_retry=False,
speculative_exec="enable",
stop_executor_on_failure=False,
enforce_memory_limit=False,
nonzero_exitcode_as_oom=False,
fault_inject_prob=None,
enable_profiling=False,
enable_diagnostic_metrics=False,
remove_empty_parquet=False,
skip_task_with_empty_input=False,
console_log_level="WARNING",
file_log_level="DEBUG",
output_path: Optional[str] = None,
runtime_ctx: Optional[RuntimeContext] = None,
):
"""
Start a scheduler and `num_executors` executors to execute `plan`.
When this function returns, the execution is mostly still running.
Parameters
----------
plan
A logical plan.
num_executors, optional
The number of executors
console_log_level, optional
Set to logger.INFO if more verbose loguru is needed for debug, by default "WARNING".
Returns
-------
A 3-tuple of type (Scheduler, List[Executor], List[Process]).
"""
if runtime_ctx is None:
runtime_ctx = RuntimeContext(
JobId.new(),
datetime.now(),
self.output_root_abspath,
num_executors=num_executors,
random_seed=123456,
enforce_memory_limit=enforce_memory_limit,
max_usable_cpu_count=min(64, psutil.cpu_count(logical=False)),
max_usable_gpu_count=0,
max_usable_memory_size=min(64 * GB, psutil.virtual_memory().total),
secs_wq_poll_interval=secs_wq_poll_interval,
secs_executor_probe_interval=secs_executor_probe_interval,
max_num_missed_probes=max_num_missed_probes,
fault_inject_prob=(
fault_inject_prob
if fault_inject_prob is not None
else self.fault_inject_prob
),
enable_profiling=enable_profiling,
enable_diagnostic_metrics=enable_diagnostic_metrics,
remove_empty_parquet=remove_empty_parquet,
skip_task_with_empty_input=skip_task_with_empty_input,
console_log_level=console_log_level,
file_log_level=file_log_level,
output_path=output_path,
)
self.queue_manager = Manager()
self.sched_states = self.queue_manager.Queue()
exec_plan = Planner(runtime_ctx).create_exec_plan(plan)
scheduler = Scheduler(
exec_plan,
max_retry_count=max_retry_count,
max_fail_count=max_fail_count,
prioritize_retry=prioritize_retry,
speculative_exec=speculative_exec,
stop_executor_on_failure=stop_executor_on_failure,
nonzero_exitcode_as_oom=nonzero_exitcode_as_oom,
)
self.latest_state = scheduler
self.executors = [
Executor.create(runtime_ctx, f"executor-{i}") for i in range(num_executors)
]
self.processes = [
Process(
target=run_scheduler,
# XXX: on macOS, scheduler state observer will be cleared when cross-process
# so we pass the queue and add the observer in the new process
args=(runtime_ctx, scheduler, self.sched_states),
name="scheduler",
)
]
self.processes += [
Process(target=run_executor, args=(runtime_ctx, executor), name=executor.id)
for executor in self.executors
]
for process in reversed(self.processes):
process.start()
return self.sched_states, self.executors, self.processes
def execute_plan(self, *args, check_result=True, **kvargs) -> ExecutionPlan:
"""
Start a scheduler and `num_executors` executors to execute `plan`,
and wait the execution completed, then assert if it succeeds.
Parameters
----------
plan
A logical plan.
num_executors, optional
The number of executors
console_log_level, optional
Set to logger.INFO if more verbose loguru is needed for debug, by default "WARNING".
Returns
-------
The completed ExecutionPlan instance.
"""
self.start_execution(*args, **kvargs)
self.join_running_procs()
latest_state = self.get_latest_sched_state()
if check_result:
self.assertTrue(latest_state.success)
return latest_state.exec_plan
def _load_parquet_files(
self, paths, filesystem: fsspec.AbstractFileSystem = None
) -> arrow.Table:
def read_parquet_file(path):
return arrow.Table.from_batches(
parquet.ParquetFile(
path, buffer_size=16 * MB, filesystem=filesystem
).iter_batches()
)
with ThreadPoolExecutor(16) as pool:
return arrow.concat_tables(pool.map(read_parquet_file, paths))
def _compare_arrow_tables(self, expected: arrow.Table, actual: arrow.Table):
def sorted_table(t: arrow.Table):
return t.sort_by([(col, "ascending") for col in t.column_names])
self.assertEqual(expected.shape, actual.shape)
self.assertEqual(expected.column_names, actual.column_names)
expected = sorted_table(cast_columns_to_large_string(expected))
actual = sorted_table(cast_columns_to_large_string(actual))
for col, x, y in zip(expected.column_names, expected.columns, actual.columns):
if not pc.equal(x, y):
x = x.to_numpy(zero_copy_only=False)
y = y.to_numpy(zero_copy_only=False)
logger.error(f" expect {col}: {x}")
logger.error(f" actual {col}: {y}")
np.testing.assert_array_equal(x, y, verbose=True)

25
tests/test_filesystem.py Normal file
View File

@@ -0,0 +1,25 @@
import os.path
import tempfile
import threading
import unittest
from smallpond.io.filesystem import dump, load
from tests.test_fabric import TestFabric
class TestFilesystem(TestFabric, unittest.TestCase):
def test_pickle_runtime_ctx(self):
with tempfile.TemporaryDirectory(dir=self.output_root_abspath) as output_dir:
pickle_path = os.path.join(output_dir, "runtime_ctx.pickle")
dump(self.runtime_ctx, pickle_path)
runtime_ctx = load(pickle_path)
self.assertEqual(self.runtime_ctx.job_id, runtime_ctx.job_id)
def test_pickle_trace(self):
with self.assertRaises(TypeError) as context:
with tempfile.TemporaryDirectory(
dir=self.output_root_abspath
) as output_dir:
thread = threading.Thread()
pickle_path = os.path.join(output_dir, "thread.pickle")
dump(thread, pickle_path)

103
tests/test_logical.py Normal file
View File

@@ -0,0 +1,103 @@
import unittest
from loguru import logger
from smallpond.logical.dataset import ParquetDataSet
from smallpond.logical.node import (
Context,
DataSetPartitionNode,
DataSourceNode,
EvenlyDistributedPartitionNode,
HashPartitionNode,
LogicalPlan,
SqlEngineNode,
)
from smallpond.logical.planner import Planner
from tests.test_fabric import TestFabric
class TestLogicalPlan(TestFabric, unittest.TestCase):
def test_join_chunkmeta_inodes(self):
ctx = Context()
chunkmeta_dump = DataSourceNode(
ctx, dataset=ParquetDataSet(["tests/data/chunkmeta*.parquet"])
)
chunkmeta_partitions = HashPartitionNode(
ctx, (chunkmeta_dump,), npartitions=2, hash_columns=["inodeId"]
)
inodes_dump = DataSourceNode(
ctx, dataset=ParquetDataSet(["tests/data/inodes*.parquet"])
)
inodes_partitions = HashPartitionNode(
ctx, (inodes_dump,), npartitions=2, hash_columns=["inode_id"]
)
num_gc_chunks = SqlEngineNode(
ctx,
(chunkmeta_partitions, inodes_partitions),
r"""
select count(chunkmeta_chunkId) from {0}
where chunkmeta.chunkmeta_chunkId NOT LIKE "F%" AND
chunkmeta.inodeId not in ( select distinct inode_id from {1} )""",
)
plan = LogicalPlan(ctx, num_gc_chunks)
logger.info(str(plan))
exec_plan = Planner(self.runtime_ctx).create_exec_plan(plan)
logger.info(str(exec_plan))
def test_partition_dims_not_compatible(self):
ctx = Context()
parquet_dataset = ParquetDataSet(["tests/data/mock_urls/*.parquet"])
data_source = DataSourceNode(ctx, parquet_dataset)
partition_dim_a = EvenlyDistributedPartitionNode(
ctx, (data_source,), npartitions=parquet_dataset.num_files, dimension="A"
)
partition_dim_b = EvenlyDistributedPartitionNode(
ctx, (data_source,), npartitions=parquet_dataset.num_files, dimension="B"
)
join_two_inputs = SqlEngineNode(
ctx,
(partition_dim_a, partition_dim_b),
r"select a.* from {0} as a join {1} as b on a.host = b.host",
)
plan = LogicalPlan(ctx, join_two_inputs)
logger.info(str(plan))
with self.assertRaises(AssertionError) as context:
Planner(self.runtime_ctx).create_exec_plan(plan)
def test_npartitions_not_compatible(self):
ctx = Context()
parquet_dataset = ParquetDataSet(["tests/data/mock_urls/*.parquet"])
data_source = DataSourceNode(ctx, parquet_dataset)
partition_dim_a = EvenlyDistributedPartitionNode(
ctx, (data_source,), npartitions=parquet_dataset.num_files, dimension="A"
)
partition_dim_a2 = EvenlyDistributedPartitionNode(
ctx,
(data_source,),
npartitions=parquet_dataset.num_files * 2,
dimension="A",
)
join_two_inputs1 = SqlEngineNode(
ctx,
(partition_dim_a, partition_dim_a2),
r"select a.* from {0} as a join {1} as b on a.host = b.host",
)
join_two_inputs2 = SqlEngineNode(
ctx,
(partition_dim_a2, partition_dim_a),
r"select a.* from {0} as a join {1} as b on a.host = b.host",
)
plan = LogicalPlan(
ctx,
DataSetPartitionNode(
ctx, (join_two_inputs1, join_two_inputs2), npartitions=1
),
)
logger.info(str(plan))
with self.assertRaises(AssertionError) as context:
Planner(self.runtime_ctx).create_exec_plan(plan)

659
tests/test_partition.py Normal file
View File

@@ -0,0 +1,659 @@
import os.path
import tempfile
import unittest
from typing import List
import pyarrow.compute as pc
from smallpond.common import DATA_PARTITION_COLUMN_NAME, GB
from smallpond.execution.task import RuntimeContext
from smallpond.logical.dataset import DataSet, ParquetDataSet
from smallpond.logical.node import (
ArrowComputeNode,
ConsolidateNode,
Context,
DataSetPartitionNode,
DataSinkNode,
DataSourceNode,
EvenlyDistributedPartitionNode,
HashPartitionNode,
LoadPartitionedDataSetNode,
LogicalPlan,
ProjectionNode,
SqlEngineNode,
UnionNode,
UserDefinedPartitionNode,
UserPartitionedDataSourceNode,
)
from tests.test_execution import parse_url
from tests.test_fabric import TestFabric
class CalculatePartitionFromFilename(UserDefinedPartitionNode):
def partition(self, runtime_ctx: RuntimeContext, dataset: DataSet) -> List[DataSet]:
partitioned_datasets: List[ParquetDataSet] = [
ParquetDataSet([]) for _ in range(self.npartitions)
]
for path in dataset.resolved_paths:
partition_idx = hash(path) % self.npartitions
partitioned_datasets[partition_idx].paths.append(path)
return partitioned_datasets
class TestPartition(TestFabric, unittest.TestCase):
def test_many_file_partitions(self):
ctx = Context()
dataset = ParquetDataSet(["tests/data/mock_urls/*.parquet"] * 10)
data_files = DataSourceNode(ctx, dataset)
data_partitions = DataSetPartitionNode(
ctx, (data_files,), npartitions=dataset.num_files
)
count_rows = SqlEngineNode(
ctx,
(data_partitions,),
"select count(*) from {0}",
cpu_limit=1,
memory_limit=1 * GB,
)
plan = LogicalPlan(ctx, count_rows)
self.execute_plan(plan)
def test_many_row_partitions(self):
ctx = Context()
dataset = ParquetDataSet(["tests/data/mock_urls/*.parquet"])
data_files = DataSourceNode(ctx, dataset)
data_partitions = DataSetPartitionNode(
ctx, (data_files,), npartitions=dataset.num_rows, partition_by_rows=True
)
count_rows = SqlEngineNode(
ctx,
(data_partitions,),
"select count(*) from {0}",
cpu_limit=1,
memory_limit=1 * GB,
)
plan = LogicalPlan(ctx, count_rows)
exec_plan = self.execute_plan(plan, num_executors=5)
self.assertEqual(
exec_plan.final_output.to_arrow_table().num_rows, dataset.num_rows
)
def test_empty_dataset_partition(self):
ctx = Context()
dataset = ParquetDataSet(["tests/data/mock_urls/*.parquet"])
data_files = DataSourceNode(ctx, dataset)
# create more partitions than files
data_partitions = EvenlyDistributedPartitionNode(
ctx, (data_files,), npartitions=dataset.num_files * 2
)
data_partitions.max_num_producer_tasks = 3
unique_urls = SqlEngineNode(
ctx,
(data_partitions,),
r"select distinct url from {0}",
cpu_limit=1,
memory_limit=1 * GB,
)
# nested partition
nested_partitioned_urls = EvenlyDistributedPartitionNode(
ctx, (unique_urls,), npartitions=3, dimension="nested", nested=True
)
parsed_urls = ArrowComputeNode(
ctx,
(nested_partitioned_urls,),
process_func=parse_url,
cpu_limit=1,
memory_limit=1 * GB,
)
plan = LogicalPlan(ctx, parsed_urls)
final_output = self.execute_plan(
plan, remove_empty_parquet=True, skip_task_with_empty_input=True
).final_output
self.assertTrue(isinstance(final_output, ParquetDataSet))
self.assertEqual(dataset.num_rows, final_output.num_rows)
def test_hash_partition(self):
for engine_type in ("duckdb", "arrow"):
for partition_by_rows in (False, True):
for hive_partitioning in (
(False, True) if engine_type == "duckdb" else (False,)
):
with self.subTest(
engine_type=engine_type,
partition_by_rows=partition_by_rows,
hive_partitioning=hive_partitioning,
):
ctx = Context()
dataset = ParquetDataSet(["tests/data/arrow/*.parquet"])
data_files = DataSourceNode(ctx, dataset)
npartitions = 3
data_partitions = DataSetPartitionNode(
ctx,
(data_files,),
npartitions=npartitions,
partition_by_rows=partition_by_rows,
)
hash_partitions = HashPartitionNode(
ctx,
(ProjectionNode(ctx, data_partitions, ["url"]),),
npartitions=npartitions,
hash_columns=["url"],
engine_type=engine_type,
hive_partitioning=hive_partitioning,
cpu_limit=2,
memory_limit=2 * GB,
output_name="hash_partitions",
)
row_count = SqlEngineNode(
ctx,
(hash_partitions,),
r"select count(*) as row_count from {0}",
cpu_limit=1,
memory_limit=1 * GB,
)
plan = LogicalPlan(ctx, row_count)
exec_plan = self.execute_plan(plan)
self.assertEqual(
dataset.num_rows,
pc.sum(
exec_plan.final_output.to_arrow_table().column(
"row_count"
)
).as_py(),
)
self.assertEqual(
npartitions,
len(
exec_plan.final_output.load_partitioned_datasets(
npartitions, DATA_PARTITION_COLUMN_NAME
)
),
)
self.assertEqual(
npartitions,
len(
exec_plan.get_output(
"hash_partitions"
).load_partitioned_datasets(
npartitions,
DATA_PARTITION_COLUMN_NAME,
hive_partitioning,
)
),
)
def test_empty_hash_partition(self):
for engine_type in ("duckdb", "arrow"):
for partition_by_rows in (False, True):
for hive_partitioning in (
(False, True) if engine_type == "duckdb" else (False,)
):
with self.subTest(
engine_type=engine_type,
partition_by_rows=partition_by_rows,
hive_partitioning=hive_partitioning,
):
ctx = Context()
dataset = ParquetDataSet(["tests/data/mock_urls/*.parquet"])
data_files = DataSourceNode(ctx, dataset)
npartitions = 3
npartitions_nested = 4
num_rows = 1
head_rows = SqlEngineNode(
ctx, (data_files,), f"select * from {{0}} limit {num_rows}"
)
data_partitions = DataSetPartitionNode(
ctx,
(head_rows,),
npartitions=npartitions,
partition_by_rows=partition_by_rows,
)
hash_partitions = HashPartitionNode(
ctx,
(data_partitions,),
npartitions=npartitions,
hash_columns=["url"],
data_partition_column="hash_partitions",
engine_type=engine_type,
hive_partitioning=hive_partitioning,
output_name="hash_partitions",
cpu_limit=2,
memory_limit=1 * GB,
)
nested_hash_partitions = HashPartitionNode(
ctx,
(hash_partitions,),
npartitions=npartitions_nested,
hash_columns=["url"],
data_partition_column="nested_hash_partitions",
nested=True,
engine_type=engine_type,
hive_partitioning=hive_partitioning,
output_name="nested_hash_partitions",
cpu_limit=2,
memory_limit=1 * GB,
)
select_every_row = SqlEngineNode(
ctx,
(nested_hash_partitions,),
r"select * from {0}",
cpu_limit=1,
memory_limit=1 * GB,
)
plan = LogicalPlan(ctx, select_every_row)
exec_plan = self.execute_plan(
plan, skip_task_with_empty_input=True
)
self.assertEqual(num_rows, exec_plan.final_output.num_rows)
self.assertEqual(
npartitions,
len(
exec_plan.final_output.load_partitioned_datasets(
npartitions, "hash_partitions"
)
),
)
self.assertEqual(
npartitions_nested,
len(
exec_plan.final_output.load_partitioned_datasets(
npartitions_nested, "nested_hash_partitions"
)
),
)
self.assertEqual(
npartitions,
len(
exec_plan.get_output(
"hash_partitions"
).load_partitioned_datasets(
npartitions, "hash_partitions"
)
),
)
self.assertEqual(
npartitions_nested,
len(
exec_plan.get_output(
"nested_hash_partitions"
).load_partitioned_datasets(
npartitions_nested, "nested_hash_partitions"
)
),
)
if hive_partitioning:
self.assertEqual(
npartitions,
len(
exec_plan.get_output(
"hash_partitions"
).load_partitioned_datasets(
npartitions,
"hash_partitions",
hive_partitioning=True,
)
),
)
self.assertEqual(
npartitions_nested,
len(
exec_plan.get_output(
"nested_hash_partitions"
).load_partitioned_datasets(
npartitions_nested,
"nested_hash_partitions",
hive_partitioning=True,
)
),
)
def test_load_partitioned_datasets(self):
def run_test_plan(
npartitions: int,
data_partition_column: str,
engine_type: str,
hive_partitioning: bool,
):
ctx = Context()
input_dataset = ParquetDataSet(["tests/data/mock_urls/*.parquet"])
input_data_files = DataSourceNode(ctx, input_dataset)
# create hash partitions
input_partitions = HashPartitionNode(
ctx,
(input_data_files,),
npartitions=npartitions,
hash_columns=["url"],
data_partition_column=data_partition_column,
engine_type=engine_type,
hive_partitioning=hive_partitioning,
output_name="input_partitions",
cpu_limit=1,
memory_limit=1 * GB,
)
split_urls = SqlEngineNode(
ctx,
(input_partitions,),
f"select url, string_split(url, '/')[0] as host from {{0}}",
cpu_limit=1,
memory_limit=1 * GB,
)
plan = LogicalPlan(ctx, split_urls)
exec_plan = self.execute_plan(plan)
self.assertEqual(
npartitions,
len(
exec_plan.final_output.load_partitioned_datasets(
npartitions, data_partition_column
)
),
)
self.assertEqual(
npartitions,
len(
exec_plan.get_output("input_partitions").load_partitioned_datasets(
npartitions, data_partition_column, hive_partitioning
)
),
)
return exec_plan
npartitions = 5
data_partition_column = "_human_readable_column_name_"
for engine_type in ("duckdb", "arrow"):
with self.subTest(engine_type=engine_type):
exec_plan1 = run_test_plan(
npartitions,
data_partition_column,
engine_type,
hive_partitioning=engine_type == "duckdb",
)
exec_plan2 = run_test_plan(
npartitions,
data_partition_column,
engine_type,
hive_partitioning=False,
)
ctx = Context()
output1 = DataSourceNode(
ctx, dataset=exec_plan1.get_output("input_partitions")
)
output2 = DataSourceNode(
ctx, dataset=exec_plan2.get_output("input_partitions")
)
split_urls1 = LoadPartitionedDataSetNode(
ctx,
(output1,),
npartitions=npartitions,
data_partition_column=data_partition_column,
hive_partitioning=engine_type == "duckdb",
)
split_urls2 = LoadPartitionedDataSetNode(
ctx,
(output2,),
npartitions=npartitions,
data_partition_column=data_partition_column,
hive_partitioning=False,
)
split_urls3 = SqlEngineNode(
ctx,
(split_urls1, split_urls2),
f"""
select split_urls1.url, string_split(split_urls2.url, '/')[0] as host
from {{0}} as split_urls1
join {{1}} as split_urls2
on split_urls1.url = split_urls2.url
""",
cpu_limit=1,
memory_limit=1 * GB,
)
plan = LogicalPlan(ctx, split_urls3)
exec_plan3 = self.execute_plan(plan)
# load each partition as arrow table and compare
final_output_partitions1 = (
exec_plan1.final_output.load_partitioned_datasets(
npartitions, data_partition_column
)
)
final_output_partitions3 = (
exec_plan3.final_output.load_partitioned_datasets(
npartitions, data_partition_column
)
)
self.assertEqual(npartitions, len(final_output_partitions3))
for x, y in zip(final_output_partitions1, final_output_partitions3):
self._compare_arrow_tables(x.to_arrow_table(), y.to_arrow_table())
def test_nested_partition(self):
ctx = Context()
parquet_files = ParquetDataSet(["tests/data/mock_urls/*.parquet"])
data_source = DataSourceNode(ctx, parquet_files)
SqlEngineNode.default_cpu_limit = 1
SqlEngineNode.default_memory_limit = 1 * GB
initial_reduce = r"select host, count(*) as cnt from {0} group by host"
combine_reduce_results = (
r"select host, cast(sum(cnt) as bigint) as cnt from {0} group by host"
)
join_query = r"select host, cnt from {0} where (exists (select * from {1} where {1}.host = {0}.host)) and (exists (select * from {2} where {2}.host = {0}.host))"
partition_by_hosts = HashPartitionNode(
ctx,
(data_source,),
npartitions=3,
hash_columns=["host"],
data_partition_column="host_partition",
)
partition_by_hosts_x_urls = HashPartitionNode(
ctx,
(partition_by_hosts,),
npartitions=5,
hash_columns=["url"],
data_partition_column="url_partition",
nested=True,
)
url_count_by_hosts_x_urls1 = SqlEngineNode(
ctx,
(partition_by_hosts_x_urls,),
initial_reduce,
output_name="url_count_by_hosts_x_urls1",
)
url_count_by_hosts1 = SqlEngineNode(
ctx,
(ConsolidateNode(ctx, url_count_by_hosts_x_urls1, ["host_partition"]),),
combine_reduce_results,
output_name="url_count_by_hosts1",
)
join_count_by_hosts_x_urls1 = SqlEngineNode(
ctx,
(url_count_by_hosts_x_urls1, url_count_by_hosts1, data_source),
join_query,
output_name="join_count_by_hosts_x_urls1",
)
partitioned_urls = LoadPartitionedDataSetNode(
ctx,
(partition_by_hosts_x_urls,),
data_partition_column="url_partition",
npartitions=5,
)
partitioned_hosts_x_urls = LoadPartitionedDataSetNode(
ctx,
(partitioned_urls,),
data_partition_column="host_partition",
npartitions=3,
nested=True,
)
partitioned_3dims = EvenlyDistributedPartitionNode(
ctx,
(partitioned_hosts_x_urls,),
npartitions=2,
dimension="inner_partition",
partition_by_rows=True,
nested=True,
)
url_count_by_3dims = SqlEngineNode(ctx, (partitioned_3dims,), initial_reduce)
url_count_by_hosts_x_urls2 = SqlEngineNode(
ctx,
(
ConsolidateNode(
ctx, url_count_by_3dims, ["host_partition", "url_partition"]
),
),
combine_reduce_results,
output_name="url_count_by_hosts_x_urls2",
)
url_count_by_hosts2 = SqlEngineNode(
ctx,
(ConsolidateNode(ctx, url_count_by_hosts_x_urls2, ["host_partition"]),),
combine_reduce_results,
output_name="url_count_by_hosts2",
)
url_count_by_hosts_expected = SqlEngineNode(
ctx,
(data_source,),
initial_reduce,
per_thread_output=False,
output_name="url_count_by_hosts_expected",
)
join_count_by_hosts_x_urls2 = SqlEngineNode(
ctx,
(url_count_by_hosts_x_urls2, url_count_by_hosts2, data_source),
join_query,
output_name="join_count_by_hosts_x_urls2",
)
union_url_count_by_hosts = UnionNode(
ctx, (url_count_by_hosts1, url_count_by_hosts2)
)
union_url_count_by_hosts_x_urls = UnionNode(
ctx,
(
url_count_by_hosts_x_urls1,
url_count_by_hosts_x_urls2,
join_count_by_hosts_x_urls1,
join_count_by_hosts_x_urls2,
),
)
with tempfile.TemporaryDirectory(dir=self.output_root_abspath) as output_dir:
data_sink = DataSinkNode(
ctx,
(
url_count_by_hosts_expected,
union_url_count_by_hosts,
union_url_count_by_hosts_x_urls,
),
output_path=output_dir,
manifest_only=True,
)
plan = LogicalPlan(ctx, data_sink)
exec_plan = self.execute_plan(plan, remove_empty_parquet=True)
# verify results
self._compare_arrow_tables(
exec_plan.get_output("url_count_by_hosts_x_urls1").to_arrow_table(),
exec_plan.get_output("url_count_by_hosts_x_urls2").to_arrow_table(),
)
self._compare_arrow_tables(
exec_plan.get_output("join_count_by_hosts_x_urls1").to_arrow_table(),
exec_plan.get_output("join_count_by_hosts_x_urls2").to_arrow_table(),
)
self._compare_arrow_tables(
exec_plan.get_output("url_count_by_hosts_x_urls1").to_arrow_table(),
exec_plan.get_output("join_count_by_hosts_x_urls1").to_arrow_table(),
)
self._compare_arrow_tables(
exec_plan.get_output("url_count_by_hosts1").to_arrow_table(),
exec_plan.get_output("url_count_by_hosts2").to_arrow_table(),
)
self._compare_arrow_tables(
exec_plan.get_output("url_count_by_hosts_expected").to_arrow_table(),
exec_plan.get_output("url_count_by_hosts1").to_arrow_table(),
)
def test_user_defined_partition(self):
ctx = Context()
parquet_files = ParquetDataSet(["tests/data/mock_urls/*.parquet"])
data_source = DataSourceNode(ctx, parquet_files)
file_partitions1 = CalculatePartitionFromFilename(
ctx, (data_source,), npartitions=3, dimension="by_filename_hash1"
)
url_count1 = SqlEngineNode(
ctx,
(file_partitions1,),
r"select host, count(*) as cnt from {0} group by host",
output_name="url_count1",
)
file_partitions2 = CalculatePartitionFromFilename(
ctx, (url_count1,), npartitions=3, dimension="by_filename_hash2"
)
url_count2 = SqlEngineNode(
ctx,
(file_partitions2,),
r"select host, cnt from {0}",
output_name="url_count2",
)
plan = LogicalPlan(ctx, url_count2)
exec_plan = self.execute_plan(plan, enable_diagnostic_metrics=True)
self._compare_arrow_tables(
exec_plan.get_output("url_count1").to_arrow_table(),
exec_plan.get_output("url_count2").to_arrow_table(),
)
def test_user_partitioned_data_source(self):
ctx = Context()
parquet_dataset = ParquetDataSet(["tests/data/mock_urls/*.parquet"])
data_source = DataSourceNode(ctx, parquet_dataset)
evenly_dist_data_source = EvenlyDistributedPartitionNode(
ctx, (data_source,), npartitions=parquet_dataset.num_files
)
parquet_datasets = [ParquetDataSet([p]) for p in parquet_dataset.resolved_paths]
partitioned_data_source = UserPartitionedDataSourceNode(ctx, parquet_datasets)
url_count_by_host1 = SqlEngineNode(
ctx,
(evenly_dist_data_source,),
r"select host, count(*) as cnt from {0} group by host",
output_name="url_count_by_host1",
cpu_limit=1,
memory_limit=1 * GB,
)
url_count_by_host2 = SqlEngineNode(
ctx,
(evenly_dist_data_source, partitioned_data_source),
r"select {1}.host, count(*) as cnt from {0} join {1} on {0}.host = {1}.host group by {1}.host",
output_name="url_count_by_host2",
cpu_limit=1,
memory_limit=1 * GB,
)
plan = LogicalPlan(
ctx, UnionNode(ctx, [url_count_by_host1, url_count_by_host2])
)
exec_plan = self.execute_plan(plan, enable_diagnostic_metrics=True)
self._compare_arrow_tables(
exec_plan.get_output("url_count_by_host1").to_arrow_table(),
exec_plan.get_output("url_count_by_host2").to_arrow_table(),
)
def test_partition_info_in_sql_query(self):
"""
User can refer to the partition info in the SQL query.
"""
ctx = Context()
parquet_dataset = ParquetDataSet(["tests/data/mock_urls/*.parquet"])
data_source = DataSourceNode(ctx, parquet_dataset)
evenly_dist_data_source = EvenlyDistributedPartitionNode(
ctx, (data_source,), npartitions=parquet_dataset.num_files
)
sql_query = SqlEngineNode(
ctx,
(evenly_dist_data_source,),
r"select host, {__data_partition__} as partition_info from {0}",
)
plan = LogicalPlan(ctx, sql_query)
exec_plan = self.execute_plan(plan)

70
tests/test_plan.py Normal file
View File

@@ -0,0 +1,70 @@
import os
import tempfile
import unittest
from examples.fstest import fstest
from examples.shuffle_data import shuffle_data
from examples.shuffle_mock_urls import shuffle_mock_urls
from examples.sort_mock_urls import sort_mock_urls
from examples.sort_mock_urls_v2 import sort_mock_urls_v2
from smallpond.dataframe import Session
from tests.test_fabric import TestFabric
class TestPlan(TestFabric, unittest.TestCase):
def test_sort_mock_urls(self):
for engine_type in ("duckdb", "arrow"):
with self.subTest(engine_type=engine_type):
plan = sort_mock_urls(
["tests/data/mock_urls/*.tsv"],
npartitions=3,
engine_type=engine_type,
)
self.execute_plan(plan)
def test_sort_mock_urls_external_output_path(self):
with tempfile.TemporaryDirectory(dir=self.output_root_abspath) as output_dir:
plan = sort_mock_urls(
["tests/data/mock_urls/*.tsv"],
npartitions=3,
external_output_path=output_dir,
)
self.execute_plan(plan)
def test_shuffle_mock_urls(self):
for engine_type in ("duckdb", "arrow"):
with self.subTest(engine_type=engine_type):
plan = shuffle_mock_urls(
["tests/data/mock_urls/*.parquet"],
npartitions=3,
sort_rand_keys=True,
)
self.execute_plan(plan)
def test_shuffle_data(self):
for engine_type in ("duckdb", "arrow"):
with self.subTest(engine_type=engine_type):
plan = shuffle_data(
["tests/data/mock_urls/*.parquet"],
num_data_partitions=3,
num_out_data_partitions=3,
engine_type=engine_type,
)
self.execute_plan(plan)
def test_fstest(sp: Session):
path = sp._runtime_ctx.output_root
fstest(
sp,
input_path=os.path.join(path, "*"),
output_path=path,
size="10M",
npartitions=3,
)
def test_sort_mock_urls_v2(sp: Session):
sort_mock_urls_v2(
sp, ["tests/data/mock_urls/*.tsv"], sp._runtime_ctx.output_root, npartitions=3
)

180
tests/test_scheduler.py Normal file
View File

@@ -0,0 +1,180 @@
import os.path
import random
import time
import unittest
from typing import List, Tuple
from loguru import logger
from smallpond.execution.scheduler import ExecutorState
from smallpond.execution.task import PythonScriptTask, RuntimeContext
from smallpond.logical.dataset import DataSet, ParquetDataSet
from smallpond.logical.node import (
Context,
DataSetPartitionNode,
DataSourceNode,
LogicalPlan,
Node,
PythonScriptNode,
)
from tests.test_fabric import TestFabric
class RandomSleepTask(PythonScriptTask):
def __init__(
self, *args, sleep_secs: float, fail_first_try: bool, **kwargs
) -> None:
super().__init__(*args, **kwargs)
self.sleep_secs = sleep_secs
self.fail_first_try = fail_first_try
def process(
self,
runtime_ctx: RuntimeContext,
input_datasets: List[DataSet],
output_path: str,
) -> bool:
logger.info(f"sleeping {self.sleep_secs} secs")
time.sleep(self.sleep_secs)
with open(os.path.join(output_path, self.output_filename), "w") as fout:
fout.write(f"{repr(self)}")
if self.fail_first_try and self.retry_count == 0:
return False
return True
class RandomSleepNode(PythonScriptNode):
def __init__(
self,
ctx: Context,
input_deps: Tuple[Node, ...],
*,
max_sleep_secs=5,
fail_first_try=False,
**kwargs,
):
super().__init__(ctx, input_deps, **kwargs)
self.max_sleep_secs = max_sleep_secs
self.fail_first_try = fail_first_try
def spawn(self, *args, **kwargs) -> RandomSleepTask:
sleep_secs = (
random.random() if len(self.generated_tasks) % 20 else self.max_sleep_secs
)
return RandomSleepTask(
*args, **kwargs, sleep_secs=sleep_secs, fail_first_try=self.fail_first_try
)
class TestScheduler(TestFabric, unittest.TestCase):
def create_random_sleep_plan(
self, npartitions, max_sleep_secs, fail_first_try=False
):
ctx = Context()
dataset = ParquetDataSet(["tests/data/mock_urls/*.parquet"])
data_files = DataSourceNode(ctx, dataset)
data_partitions = DataSetPartitionNode(
ctx, (data_files,), npartitions=npartitions, partition_by_rows=True
)
random_sleep = RandomSleepNode(
ctx,
(data_partitions,),
max_sleep_secs=max_sleep_secs,
fail_first_try=fail_first_try,
)
return LogicalPlan(ctx, random_sleep)
def check_executor_state(self, target_state: ExecutorState, nloops=200):
for _ in range(nloops):
latest_sched_state = self.get_latest_sched_state()
if any(
executor.state == target_state
for executor in latest_sched_state.remote_executors
):
logger.info(
f"found {target_state} executor in: {latest_sched_state.remote_executors}"
)
break
time.sleep(0.1)
else:
self.assertTrue(
False,
f"cannot find any executor in state {target_state}: {latest_sched_state.remote_executors}",
)
def test_standalone_mode(self):
plan = self.create_random_sleep_plan(npartitions=10, max_sleep_secs=1)
self.execute_plan(plan, num_executors=0)
def test_failed_executors(self):
num_exec = 6
num_fail = 4
plan = self.create_random_sleep_plan(npartitions=300, max_sleep_secs=10)
_, executors, processes = self.start_execution(
plan,
num_executors=num_exec,
secs_wq_poll_interval=0.1,
secs_executor_probe_interval=0.5,
console_log_level="WARNING",
)
latest_sched_state = self.get_latest_sched_state()
self.check_executor_state(ExecutorState.GOOD)
for i, (executor, process) in enumerate(
random.sample(list(zip(executors, processes[1:])), k=num_fail)
):
if i % 2 == 0:
logger.warning(f"kill executor: {executor}")
process.kill()
else:
logger.warning(f"skip probes: {executor}")
executor.skip_probes(latest_sched_state.ctx.max_num_missed_probes * 2)
self.join_running_procs()
latest_sched_state = self.get_latest_sched_state()
self.assertTrue(latest_sched_state.success)
self.assertGreater(len(latest_sched_state.abandoned_tasks), 0)
self.assertLessEqual(
1,
len(latest_sched_state.stopped_executors),
f"remote_executors: {latest_sched_state.remote_executors}",
)
self.assertLessEqual(
num_fail / 2,
len(latest_sched_state.failed_executors),
f"remote_executors: {latest_sched_state.remote_executors}",
)
def test_speculative_scheduling(self):
for speculative_exec in ("disable", "enable", "aggressive"):
with self.subTest(speculative_exec=speculative_exec):
plan = self.create_random_sleep_plan(npartitions=100, max_sleep_secs=10)
self.execute_plan(
plan,
num_executors=3,
secs_wq_poll_interval=0.1,
secs_executor_probe_interval=0.5,
prioritize_retry=(speculative_exec == "aggressive"),
speculative_exec=speculative_exec,
)
latest_sched_state = self.get_latest_sched_state()
if speculative_exec == "disable":
self.assertEqual(len(latest_sched_state.abandoned_tasks), 0)
else:
self.assertGreater(len(latest_sched_state.abandoned_tasks), 0)
def test_stop_executor_on_failure(self):
plan = self.create_random_sleep_plan(
npartitions=3, max_sleep_secs=5, fail_first_try=True
)
exec_plan = self.execute_plan(
plan,
num_executors=5,
secs_wq_poll_interval=0.1,
secs_executor_probe_interval=0.5,
check_result=False,
stop_executor_on_failure=True,
)
latest_sched_state = self.get_latest_sched_state()
self.assertGreater(len(latest_sched_state.abandoned_tasks), 0)

54
tests/test_session.py Normal file
View File

@@ -0,0 +1,54 @@
import os
from smallpond.dataframe import Session
def test_shutdown_cleanup(sp: Session):
assert os.path.exists(sp._runtime_ctx.queue_root), "queue directory should exist"
assert os.path.exists(
sp._runtime_ctx.staging_root
), "staging directory should exist"
assert os.path.exists(sp._runtime_ctx.temp_root), "temp directory should exist"
# create some tasks and complete them
df = sp.from_items([1, 2, 3])
df.write_parquet(sp._runtime_ctx.output_root)
sp.shutdown()
# shutdown should clean up directories
assert not os.path.exists(
sp._runtime_ctx.queue_root
), "queue directory should be cleared"
assert not os.path.exists(
sp._runtime_ctx.staging_root
), "staging directory should be cleared"
assert not os.path.exists(
sp._runtime_ctx.temp_root
), "temp directory should be cleared"
with open(sp._runtime_ctx.job_status_path) as fin:
assert "success" in fin.read(), "job status should be success"
def test_shutdown_no_cleanup_on_failure(sp: Session):
df = sp.from_items([1, 2, 3])
try:
# create a task that will fail
df.map(lambda x: x / 0).compute()
except Exception:
pass
else:
raise RuntimeError("task should fail")
sp.shutdown()
# shutdown should not clean up directories
assert os.path.exists(
sp._runtime_ctx.queue_root
), "queue directory should not be cleared"
assert os.path.exists(
sp._runtime_ctx.staging_root
), "staging directory should not be cleared"
assert os.path.exists(
sp._runtime_ctx.temp_root
), "temp directory should not be cleared"
with open(sp._runtime_ctx.job_status_path) as fin:
assert "failure" in fin.read(), "job status should be failure"

50
tests/test_utility.py Normal file
View File

@@ -0,0 +1,50 @@
import random
import subprocess
import time
import unittest
from typing import Iterable
from smallpond.utility import ConcurrentIter, execute_command
from tests.test_fabric import TestFabric
class TestUtility(TestFabric, unittest.TestCase):
def test_concurrent_iter_no_error(self):
def slow_iterator(iter: Iterable[int], sleep_ms: int):
for i in iter:
time.sleep(sleep_ms / 1000)
yield i
for n in [1, 5, 10, 50, 100]:
with ConcurrentIter(slow_iterator(range(n), 2)) as iter1:
with ConcurrentIter(slow_iterator(iter1, 5)) as iter2:
self.assertEqual(sum(slow_iterator(iter2, 1)), sum(range(n)))
def test_concurrent_iter_with_error(self):
def broken_iterator(iter: Iterable[int], sleep_ms: int):
for i in iter:
time.sleep(sleep_ms / 1000)
if random.randint(1, 10) == 1:
raise Exception("raised before yield")
yield i
if random.randint(1, 10) == 1:
raise Exception("raised after yield")
raise Exception("raised at the end")
for n in [1, 5, 10, 50, 100]:
with self.assertRaises(Exception):
with ConcurrentIter(range(n)) as iter:
print(sum(broken_iterator(iter, 1)))
with self.assertRaises(Exception):
with ConcurrentIter(broken_iterator(range(n), 2)) as iter1:
with ConcurrentIter(broken_iterator(iter1, 5)) as iter2:
print(sum(iter2))
def test_execute_command(self):
with self.assertRaises(subprocess.CalledProcessError):
for line in execute_command("ls non_existent_file"):
print(line)
for line in execute_command("echo hello"):
print(line)
for line in execute_command("cat /dev/null"):
print(line)

119
tests/test_workqueue.py Normal file
View File

@@ -0,0 +1,119 @@
import multiprocessing
import multiprocessing.dummy
import multiprocessing.queues
import queue
import tempfile
import time
import unittest
from loguru import logger
from smallpond.execution.workqueue import (
WorkItem,
WorkQueue,
WorkQueueInMemory,
WorkQueueOnFilesystem,
)
from tests.test_fabric import TestFabric
class PrintWork(WorkItem):
def __init__(self, name: str, message: str) -> None:
super().__init__(name, cpu_limit=1, gpu_limit=0, memory_limit=0)
self.message = message
def run(self) -> bool:
logger.debug(f"{self.key}: {self.message}")
return True
def producer(wq: WorkQueue, id: int, numItems: int, numConsumers: int) -> None:
print(f"wq.outbound_works: {wq.outbound_works}")
for i in range(numItems):
wq.push(PrintWork(f"item-{i}", message="hello"), buffering=(i % 3 == 1))
# wq.push(PrintWork(f"item-{i}", message="hello"))
if i % 5 == 0:
wq.flush()
for i in range(numConsumers):
wq.push(PrintWork(f"stop-{i}", message="stop"))
logger.success(f"producer {id} generated {numItems} items")
def consumer(wq: WorkQueue, id: int) -> int:
numItems = 0
numWaits = 0
running = True
while running:
items = wq.pop(count=1)
if not items:
numWaits += 1
time.sleep(0.01)
continue
for item in items:
assert isinstance(item, PrintWork)
if item.message == "stop":
running = False
break
item.exec()
numItems += 1
logger.success(f"consumer {id} collected {numItems} items, {numWaits} waits")
logger.complete()
return numItems
class WorkQueueTestBase(object):
wq: WorkQueue = None
pool: multiprocessing.Pool = None
def setUp(self) -> None:
logger.disable("smallpond.execution.workqueue")
return super().setUp()
def test_basics(self):
numItems = 200
for i in range(numItems):
self.wq.push(PrintWork(f"item-{i}", message="hello"))
numCollected = 0
for _ in range(numItems):
items = self.wq.pop()
logger.info(f"{len(items)} items")
numCollected += len(items)
if numItems == numCollected:
break
def test_multi_consumers(self):
numConsumers = 10
numItems = 200
result = self.pool.starmap_async(
consumer, [(self.wq, id) for id in range(numConsumers)]
)
producer(self.wq, 0, numItems, numConsumers)
logger.info("waiting for result")
numCollected = sum(result.get(timeout=20))
logger.info(f"expected vs collected: {numItems} vs {numCollected}")
self.assertEqual(numItems, numCollected)
logger.success("all done")
self.pool.terminate()
self.pool.join()
logger.success("workers stopped")
class TestWorkQueueInMemory(WorkQueueTestBase, TestFabric, unittest.TestCase):
def setUp(self) -> None:
super().setUp()
self.wq = WorkQueueInMemory(queue_type=queue.Queue)
self.pool = multiprocessing.dummy.Pool(10)
class TestWorkQueueOnFilesystem(WorkQueueTestBase, TestFabric, unittest.TestCase):
workq_root: str
def setUp(self) -> None:
super().setUp()
self.workq_root = tempfile.mkdtemp(dir=self.runtime_ctx.queue_root)
self.wq = WorkQueueOnFilesystem(self.workq_root, sort=True)
self.pool = multiprocessing.Pool(10)