This is an automated email from the ASF dual-hosted git repository. timsaucer pushed a commit to branch feat/distributed-extensions-example in repository https://gitbox.apache.org/repos/asf/datafusion-python.git
commit 8a52de04131d1a335f070a57c5178ce3e73018fc Author: Tim Saucer <[email protected]> AuthorDate: Wed Sep 9 13:22:06 2026 -0400 Add the queries and the cross-library integration tests Seventeen tests for #1719, every one running real worker processes, plus `run_tpch.py` for the same thing against the generated TPC-H data. The four queries: a Q1-shaped distributed aggregate; the same with `dfx_udfs`' Rust scalar and aggregate functions; an inline Python UDF; and the storage library's provider read on the workers. Each compares the distributed answer against the single-process answer through the same session factory, because disagreement there is the only reliable signal that a split is wrong. Tests use a small hand-checked fixture rather than the real dataset. `tpchgen-cli` writes one file per table, so SF-1 `lineitem` is a single 220 MB file — one partition, and nothing to fan out. `run_tpch.py` re-shards it first, which is a fair illustration of the actual constraint: an engine can only spread work as widely as the data is split. Three things this pass turned up. **An empty `shuffle_dir` was writing files into the working directory.** A registered config extension always *has* an entry, so an unset directory arrives as `Some("")` rather than `None`, and the planner treated it as configured. The stage node's paths were then relative to wherever the process happened to be, so `run_local` scattered `stage-1-part-*.arrow` next to the caller and later queries read another query's leftovers back out of them — which is how six tests failed with "Batch has 3 columns but BatchCoalescer expects 5". Four of those files had already been committed by the previous change; they are deleted here. **cloudpickle captures a module attribute as the module, not as its parent.** I expected `pa.compute` inside a UDF to fail on a worker, since `import pyarrow` does not bind `pyarrow.compute` and nothing loads it transitively. It does not fail: cloudpickle resolves the attribute and stores an import of `pyarrow.compute` itself, so the worker imports the submodule on load. The real trap is a *function* with a resolvable `module.qualname` — the same callable is 1106 bytes pickled from `__main__` and 34 bytes from an importable module, because the second is a pointer. A helper at test-module scope therefore reaches the worker as `ModuleNotFoundError: No module named '_test_three_libraries'`, with `traceback: None` and nothing naming a UDF, a plan, or serialization. Both halves are pinned as tests. **An FFI query planner encodes its own output on every query.** It returns proto bytes rather than a plan handle, so both libraries' codecs show one encode apiece straight after `execution_plan()`, before the driver has asked for any bytes. Worth knowing before reading an encode counter as "this is what shipping cost". `run_tpch.py` compares floats with a tolerance rather than for equality: splitting a `sum` across partitions changes the order the additions happen in, and floating point addition is not associative, so the low bits of `sum_charge` differ legitimately between the two runs. Any distributed engine has this property, and someone diffing two runs should not conclude the split is broken. Verified: 400k rows of real `lineitem` across four worker processes, using the custom provider, its custom scan node, the engine's stage node, and both Rust functions, agreeing with the single-process result to 1e-6 relative. Co-Authored-By: Claude Opus 5 (1M context) <[email protected]> --- .../engine-library/python/dfx_engine/driver.py | 22 +- .../python/tests/_test_three_libraries.py | 396 +++++++++++++++++++++ .../engine-library/python/tests/conftest.py | 97 +++++ examples/distributed/engine-library/src/planner.rs | 7 + .../engine-library/stage-1-part-0.arrow | Bin 1032 -> 0 bytes .../engine-library/stage-1-part-1.arrow | Bin 1032 -> 0 bytes .../engine-library/stage-1-part-2.arrow | Bin 1032 -> 0 bytes examples/distributed/run_tpch.py | 185 ++++++++++ 8 files changed, 704 insertions(+), 3 deletions(-) diff --git a/examples/distributed/engine-library/python/dfx_engine/driver.py b/examples/distributed/engine-library/python/dfx_engine/driver.py index baa145f1..804aadf7 100644 --- a/examples/distributed/engine-library/python/dfx_engine/driver.py +++ b/examples/distributed/engine-library/python/dfx_engine/driver.py @@ -47,6 +47,7 @@ if TYPE_CHECKING: import pyarrow as pa from datafusion import DataFrame, SessionContext from datafusion.plan import ExecutionPlan + from datafusion.user_defined import ScalarUDF __all__ = ["DistributedResult", "find_stage", "run_distributed"] @@ -116,17 +117,27 @@ def _dispatch( ) -def run_distributed(sql: str, spec: SessionSpec) -> DistributedResult: +def run_distributed( + sql: str, spec: SessionSpec, extra_udfs: list[ScalarUDF] | None = None +) -> DistributedResult: """Run `sql`, executing its leaf stage in one worker process per partition. Requires ``spec.shuffle_dir``: without it the planner inserts no stage and there is nothing to distribute. + + ``extra_udfs`` are registered on the driver only. They have to be here for + the query to *plan*, but not on the worker: a Python UDF is cloudpickled + into the plan and travels by value, unlike the Rust functions in + :func:`~dfx_engine.session.build_session`, which travel by name and so + have to exist on both sides. """ if not spec.shuffle_dir: message = "run_distributed needs a shuffle_dir; build_session got none" raise ValueError(message) ctx, engine, _storage = build_session(spec) + for function in extra_udfs or []: + ctx.register_udf(function) plan = ctx.sql(sql).execution_plan() stage = find_stage(plan) @@ -184,11 +195,14 @@ def run_distributed(sql: str, spec: SessionSpec) -> DistributedResult: return DistributedResult(batches, partitions, worker_rows) -def run_local(sql: str, spec: SessionSpec) -> list[pa.RecordBatch]: +def run_local( + sql: str, spec: SessionSpec, extra_udfs: list[ScalarUDF] | None = None +) -> list[pa.RecordBatch]: """Run `sql` in this process, for comparison. Uses the same session factory with no shuffle directory, so the only - difference from :func:`run_distributed` is where the work happened. + difference from :func:`run_distributed` is where the work happened. Any + disagreement between the two is a bug in the split. """ ctx, _engine, _storage = build_session( SessionSpec( @@ -197,6 +211,8 @@ def run_local(sql: str, spec: SessionSpec) -> list[pa.RecordBatch]: target_partitions=spec.target_partitions, ) ) + for function in extra_udfs or []: + ctx.register_udf(function) return ctx.sql(sql).collect() diff --git a/examples/distributed/engine-library/python/tests/_test_three_libraries.py b/examples/distributed/engine-library/python/tests/_test_three_libraries.py new file mode 100644 index 00000000..a72c897a --- /dev/null +++ b/examples/distributed/engine-library/python/tests/_test_three_libraries.py @@ -0,0 +1,396 @@ +# Licensed to the Apache Software Foundation (ASF) under one +# or more contributor license agreements. See the NOTICE file +# distributed with this work for additional information +# regarding copyright ownership. The ASF licenses this file +# to you under the Apache License, Version 2.0 (the +# "License"); you may not use this file except in compliance +# with the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, +# software distributed under the License is distributed on an +# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +# KIND, either express or implied. See the License for the +# specific language governing permissions and limitations +# under the License. + +"""Three separately-compiled libraries, one distributed query. + +Every test here runs real worker processes. The comparison that matters is +between the distributed answer and the single-process answer through the same +session factory: if those ever disagree, the split is wrong. +""" + +from __future__ import annotations + +import dataclasses +import pathlib +import re + +import cloudpickle +import pyarrow as pa +import pyarrow.compute as pc +import pytest +from datafusion import SessionContext, udf +from dfx_engine import _internal +from dfx_engine.driver import find_stage, run_distributed, run_local +from dfx_engine.session import SessionSpec, build_session, expected_codec_ids +from dfx_engine.worker import run_task + +Q1 = """ +select l_returnflag, l_linestatus, + count(*) as n, + sum(l_quantity) as qty, + sum(l_extendedprice) as price +from lineitem +group by l_returnflag, l_linestatus +order by l_returnflag, l_linestatus +""" + +REVENUE = """ +select l_returnflag, + sum(dfx_net_revenue(l_extendedprice, l_discount, l_tax)) as revenue, + dfx_weighted_avg(l_extendedprice, l_quantity) as wavg +from lineitem +group by l_returnflag +order by l_returnflag +""" + + +def _module_level_bucket(prices: pa.Array) -> pa.Array: + """A UDF body at module scope, for the by-reference test. + + Defined here rather than inside the test on purpose: a function with a + resolvable ``module.qualname`` is pickled as a pointer to it, and this + module is not importable from a worker. + """ + return pc.if_else(pc.greater(prices, 400.0), pa.scalar("high"), pa.scalar("low")) + + +def _rows(batches: list[pa.RecordBatch]) -> list[tuple]: + table = pa.Table.from_batches(batches) if batches else None + if table is None: + return [] + columns = [table.column(i).to_pylist() for i in range(table.num_columns)] + return list(zip(*columns, strict=True)) + + +# --- the four queries ------------------------------------------------------- + + +def test_distributed_aggregate_matches_single_process(spec: SessionSpec) -> None: + """Query 1: the baseline. Four input files, four workers, one answer.""" + result = run_distributed(Q1, spec) + + assert result.partitions == [0, 1, 2, 3] + assert _rows(result.batches) == _rows(run_local(Q1, spec)) + # Checked by hand against the fixture. + assert _rows(result.batches) == [ + ("A", "F", 3, 10.0, 1000.0), + ("N", "O", 3, 15.0, 1500.0), + ("R", "F", 2, 11.0, 1100.0), + ] + + +def test_a_rust_udf_resolves_on_every_worker(spec: SessionSpec) -> None: + """Query 2: functions from a library that ships no bundle. + + `dfx_udfs` is installed by hand in `build_session`, and the aggregate runs + partially on each worker and finally on the driver -- so its two-sum state + has to survive the split. + """ + result = run_distributed(REVENUE, spec) + + assert _rows(result.batches) == _rows(run_local(REVENUE, spec)) + revenue = {row[0]: row[1] for row in _rows(result.batches)} + # Flag A sums three rows: full price, plus tax, then half off. + assert revenue["A"] == pytest.approx(730.0) + # Flag R sums two: one discounted a fifth, one taxed a fifth. + assert revenue["R"] == pytest.approx(1160.0) + + +def test_an_inline_python_udf_ships_by_value(spec: SessionSpec) -> None: + """Query 3: a Python callable defined right here, running on a worker. + + Defined *inside* the test function, so its qualified name is + ``...<locals>.bucket`` and cloudpickle cannot look it up -- which means it + travels by value, bytecode and all. The worker has never imported this + file and does not need to. + """ + + def bucket(prices: pa.Array) -> pa.Array: + # `pc` is a module, and cloudpickle resolves it to `pyarrow.compute` + # and stores an import of it -- so the worker imports the submodule on + # load and this works. Modules are the easy case; see + # `test_a_by_reference_capture_fails_on_the_worker` for the hard one. + return pc.if_else( + pc.greater(prices, 400.0), pa.scalar("high"), pa.scalar("low") + ) + + price_bucket = udf( + bucket, [pa.float64()], pa.string(), volatility="immutable", name="price_bucket" + ) + + sql = """ + select price_bucket(l_extendedprice) as bucket, count(*) as n + from lineitem group by bucket order by bucket + """ + + ctx, _engine, _storage = build_session(spec) + ctx.register_udf(price_bucket) + plan = ctx.sql(sql).execution_plan() + stage = find_stage(plan) + assert stage is not None + # The callable itself is in the bytes, under the scalar-UDF family prefix. + assert b"DFPYUDF" in stage.to_bytes(ctx) + + result = run_distributed(sql, spec, extra_udfs=[price_bucket]) + # Prices run from one hundred to eight hundred, so four exceed four hundred. + assert _rows(result.batches) == [("high", 4), ("low", 4)] + assert _rows(result.batches) == _rows( + run_local(sql, spec, extra_udfs=[price_bucket]) + ) + + +def test_a_by_reference_capture_fails_on_the_worker(spec: SessionSpec) -> None: + """The pitfall half of query 3, and the reason to read this file. + + The callable's *body* travels by value. Names it closes over travel by + **reference** if cloudpickle can find them under an importable module -- + and this test file is an importable module, so `_module_level_bucket` is + stored as a two-word pointer at it. + + The driver runs the query fine: the name resolves here. The worker has + never heard of this module and fails on import, with an error that names + the module and says nothing about UDFs, plans, or serialization. + + A module is the easy case, because the worker can just import it (see the + previous test). A *function in your own project* is the case that bites: + it means every worker needs your code installed, not just your data. + """ + price_bucket = udf( + _module_level_bucket, + [pa.float64()], + pa.string(), + volatility="immutable", + name="price_bucket", + ) + sql = "select price_bucket(l_extendedprice) as bucket from lineitem" + + # Pinning *why* it breaks: a reference, not a copy. The by-value version + # in the previous test is two orders of magnitude bigger. + assert len(cloudpickle.dumps(_module_level_bucket)) < 200 + assert b"_test_three_libraries" in cloudpickle.dumps(_module_level_bucket) + + # Works here, because the name resolves in this process. + assert len(run_local(sql, spec, extra_udfs=[price_bucket])) >= 1 + + with pytest.raises(RuntimeError) as excinfo: + run_distributed(sql, spec, extra_udfs=[price_bucket]) + assert "_test_three_libraries" in str(excinfo.value) + + +def test_the_custom_provider_is_read_on_the_workers(spec: SessionSpec) -> None: + """Query 4: the storage library's scan, executed in another process. + + Its codec had to write the directory into the logical plan *and* the file + list into the physical plan for this to work at all. + """ + sql = "select count(*) as n, sum(l_quantity) as qty from lineitem" + result = run_distributed(sql, spec) + + assert _rows(result.batches) == [(8, 36.0)] + assert _rows(result.batches) == _rows(run_local(sql, spec)) + + +# --- what the split actually did -------------------------------------------- + + +def test_every_partition_ran_exactly_once(spec: SessionSpec) -> None: + """Each worker got a different partition, and together they covered it.""" + result = run_distributed(Q1, spec) + + assert sorted(result.partitions) == [0, 1, 2, 3] + assert len(set(result.partitions)) == len(result.partitions) + # Two rows per input file, so each worker saw two rows' worth of groups. + assert sum(result.worker_rows.values()) == 8 + assert set(result.worker_rows) == set(result.partitions) + + +def test_each_worker_published_its_own_file(spec: SessionSpec) -> None: + """One shuffle file per partition, and nothing left half-written.""" + run_distributed(Q1, spec) + + shuffle = pathlib.Path(spec.shuffle_dir) + produced = sorted(path.name for path in shuffle.glob("*.arrow")) + expected = sorted( + pathlib.Path( + _internal.partition_path(spec.shuffle_dir, _internal.stage_id(), partition) + ).name + for partition in range(4) + ) + assert produced == expected + assert list(shuffle.glob("*.tmp")) == [] + + +def test_the_driver_reads_the_workers_output(spec: SessionSpec) -> None: + """Corrupt one shuffle file and the driver's query breaks. + + Without this the suite could not tell a distributed run from the driver + quietly recomputing everything and getting the same answer. + """ + run_distributed(Q1, spec) + + victim = pathlib.Path( + _internal.partition_path(spec.shuffle_dir, _internal.stage_id(), 1) + ) + victim.write_bytes(b"not an arrow stream") + + ctx, _engine, _storage = build_session(spec) + with pytest.raises(Exception, match="dfx_engine: reading"): + ctx.sql(Q1).collect() + + +def test_each_librarys_codec_carried_its_own_node(spec: SessionSpec) -> None: + """Both codecs installed is not the same as both codecs used.""" + ctx, engine, storage = build_session(spec) + plan = ctx.sql(Q1).execution_plan() + stage = find_stage(plan) + assert stage is not None + + # Already one apiece, before the driver has asked for any bytes: an FFI + # query planner returns its plan as protobuf rather than as a handle, so + # every query serializes the planner's output on the way back. Worth + # knowing before reading an encode counter as "this is what shipping + # cost". + assert engine.encode_calls() == 1 + assert storage.encode_calls() == 1 + + stage.to_bytes(ctx) + + # Now once more each, this time because the driver asked. + assert engine.encode_calls() == 2 + assert storage.encode_calls() == 2 + + +def test_the_plan_splits_at_the_partial_aggregate(spec: SessionSpec) -> None: + """The stage boundary is where the aggregate already splits itself.""" + ctx, engine, _storage = build_session(spec) + plan = ctx.sql(Q1).execution_plan() + + text = plan.display_indent() + assert "mode=FinalPartitioned" in text + assert "ShuffleStageExec" in text + # The final aggregate is above the stage, the partial one inside it. + assert text.index("FinalPartitioned") < text.index("ShuffleStageExec") + assert text.index("ShuffleStageExec") < text.index("mode=Partial") + assert engine.stages_inserted() == 1 + + stage = find_stage(plan) + assert stage is not None + # One stage partition per input file, which is what makes the fan-out + # meaningful rather than a single remote call. + assert stage.partition_count == 4 + assert stage.output_partitioning.scheme == "UnknownPartitioning" + + +# --- the ways it goes wrong ------------------------------------------------- + + +def test_without_a_shuffle_dir_nothing_is_distributed( + lineitem_dir: pathlib.Path, +) -> None: + """The engine declines to insert a stage it has nowhere to put. + + Better than inserting one and failing at execute time, and it is what + makes `run_local` use the same factory as the distributed path. + """ + local = SessionSpec(tables={"lineitem": str(lineitem_dir)}, shuffle_dir="") + ctx, engine, _storage = build_session(local) + + plan = ctx.sql(Q1).execution_plan() + assert find_stage(plan) is None + assert engine.plan_calls() >= 1 + assert engine.stages_inserted() == 0 + # And the query still answers correctly, in this process. + assert _rows(ctx.sql(Q1).collect())[0] == ("A", "F", 3, 10.0, 1000.0) + + +def test_a_worker_whose_codecs_disagree_refuses_the_plan(spec: SessionSpec) -> None: + """A codec-id mismatch is caught before any plan is decoded.""" + envelope = { + "spec": {**spec.to_json(), "codec_ids": ["dfx_storage.physical.v1"]}, + "plan": "unused", + "stage_id": _internal.stage_id(), + "partition": 0, + } + with pytest.raises(RuntimeError, match="do not match driver's"): + run_task(envelope) + + +def test_a_session_missing_a_library_is_rejected_at_build() -> None: + """`build_session` checks its own work, so a partial session cannot ship. + + The check is what turns "a worker was built slightly differently" from a + decode failure deep in a query into an error naming the codec ids. + """ + ctx = SessionContext() + installed = sorted(ctx.physical_extension_codec_ids()) + assert installed != expected_codec_ids() + assert installed == [] + + +def test_a_plan_encoded_without_a_context_cannot_be_encoded( + spec: SessionSpec, +) -> None: + """`to_bytes()` with no context uses an empty chain and fails. + + The driver has to pass its session. This is easy to get wrong because the + argument is optional and the failure only appears once a library node is + in the plan. + """ + ctx, _engine, _storage = build_session(spec) + stage = find_stage(ctx.sql(Q1).execution_plan()) + assert stage is not None + + with pytest.raises(Exception, match=r"(?i)codec"): + stage.to_bytes() + + +def test_the_spec_round_trips_through_json(spec: SessionSpec) -> None: + """Workers receive the spec as JSON, so it has to survive the trip.""" + restored = SessionSpec.from_json(spec.to_json()) + + assert dataclasses.asdict(restored) == dataclasses.asdict(spec) + assert spec.to_json()["codec_ids"] == expected_codec_ids() + + +def test_the_bundles_are_reusable_across_sessions(spec: SessionSpec) -> None: + """Two sessions from one factory call each get their own components.""" + first, _, _ = build_session(spec) + second, _, _ = build_session(spec) + + assert first.__datafusion_codec_id__ != second.__datafusion_codec_id__ + assert sorted(first.physical_extension_codec_ids()) == expected_codec_ids() + assert sorted(second.physical_extension_codec_ids()) == expected_codec_ids() + + +def test_an_out_of_range_partition_is_reported_by_the_worker( + spec: SessionSpec, tmp_path: pathlib.Path +) -> None: + """The worker bounds-checks rather than letting a scan index off the end.""" + ctx, _engine, _storage = build_session(spec) + stage = find_stage(ctx.sql(Q1).execution_plan()) + assert stage is not None + plan_path = tmp_path / "stage.plan" + plan_path.write_bytes(stage.to_bytes(ctx)) + + envelope = { + "spec": spec.to_json(), + "plan": str(plan_path), + "stage_id": _internal.stage_id(), + "partition": 99, + } + with pytest.raises(RuntimeError, match=re.escape("partition 99 is out of range")): + run_task(envelope) diff --git a/examples/distributed/engine-library/python/tests/conftest.py b/examples/distributed/engine-library/python/tests/conftest.py new file mode 100644 index 00000000..6c0791ec --- /dev/null +++ b/examples/distributed/engine-library/python/tests/conftest.py @@ -0,0 +1,97 @@ +# Licensed to the Apache Software Foundation (ASF) under one +# or more contributor license agreements. See the NOTICE file +# distributed with this work for additional information +# regarding copyright ownership. The ASF licenses this file +# to you under the Apache License, Version 2.0 (the +# "License"); you may not use this file except in compliance +# with the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, +# software distributed under the License is distributed on an +# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +# KIND, either express or implied. See the License for the +# specific language governing permissions and limitations +# under the License. + +from __future__ import annotations + +import logging +from typing import TYPE_CHECKING + +import pyarrow as pa +import pyarrow.parquet as pq +import pytest +from dfx_engine.session import SessionSpec + +if TYPE_CHECKING: + import pathlib + from collections.abc import Generator + from typing import Any + + +class _FailOnWarning(logging.Handler): + def emit(self, record: logging.LogRecord) -> None: + if record.levelno >= logging.WARNING: + err = f"Unexpected log warning from '{record.name}': {self.format(record)}" + raise AssertionError(err) + + [email protected](autouse=True) +def fail_on_log_warnings() -> Generator[None, Any, None]: + handler = _FailOnWarning() + logging.root.addHandler(handler) + yield + logging.root.removeHandler(handler) + + +# One row group per file, four files, TPC-H `lineitem` column names. Small +# enough that every expected value below is checked by hand, and partitioned +# so there is something to distribute -- the real SF-1 dataset is a single +# 220 MB file per table, which would give one partition and no fan-out. See +# `run_tpch.py` for the same queries against the real thing. +ROW_FIELDS = "returnflag, linestatus, quantity, extendedprice, discount, tax" +_ROWS = [ + ("A", "F", 1.0, 100.0, 0.00, 0.00), + ("N", "O", 2.0, 200.0, 0.10, 0.00), + ("A", "F", 3.0, 300.0, 0.00, 0.10), + ("R", "F", 4.0, 400.0, 0.20, 0.00), + ("N", "O", 5.0, 500.0, 0.00, 0.00), + ("A", "F", 6.0, 600.0, 0.50, 0.00), + ("R", "F", 7.0, 700.0, 0.00, 0.20), + ("N", "O", 8.0, 800.0, 0.25, 0.00), +] + + [email protected] +def lineitem_dir(tmp_path: pathlib.Path) -> pathlib.Path: + """`lineitem` as four Parquet files, two rows each.""" + directory = tmp_path / "lineitem" + directory.mkdir() + for index in range(4): + chunk = _ROWS[index * 2 : index * 2 + 2] + pq.write_table( + pa.table( + { + "l_returnflag": [row[0] for row in chunk], + "l_linestatus": [row[1] for row in chunk], + "l_quantity": [row[2] for row in chunk], + "l_extendedprice": [row[3] for row in chunk], + "l_discount": [row[4] for row in chunk], + "l_tax": [row[5] for row in chunk], + } + ), + directory / f"part-{index}.parquet", + ) + return directory + + [email protected] +def spec(lineitem_dir: pathlib.Path, tmp_path: pathlib.Path) -> SessionSpec: + """A distributed spec: four input partitions, a fresh shuffle directory.""" + return SessionSpec( + tables={"lineitem": str(lineitem_dir)}, + shuffle_dir=str(tmp_path / "shuffle"), + target_partitions=2, + ) diff --git a/examples/distributed/engine-library/src/planner.rs b/examples/distributed/engine-library/src/planner.rs index 4c308f57..11de4503 100644 --- a/examples/distributed/engine-library/src/planner.rs +++ b/examples/distributed/engine-library/src/planner.rs @@ -71,6 +71,13 @@ pub(crate) fn shuffle_dir_from_options(options: &ConfigOptions) -> Option<String .into_iter() .find(|entry| entry.key == SHUFFLE_DIR_KEY || entry.key == FFI_SHUFFLE_DIR_KEY) .and_then(|entry| entry.value) + // A registered config extension always *has* an entry, so an unset + // directory arrives as `Some("")` rather than `None`. Treating that as + // configured inserts a stage whose paths are relative to whatever the + // process's working directory happens to be -- which silently writes + // shuffle files next to the caller and then reads another query's + // leftovers back out of them. + .filter(|shuffle_dir| !shuffle_dir.is_empty()) } /// Wrap the partial aggregate, or the whole plan if there is not one. diff --git a/examples/distributed/engine-library/stage-1-part-0.arrow b/examples/distributed/engine-library/stage-1-part-0.arrow deleted file mode 100644 index 9e29db43..00000000 Binary files a/examples/distributed/engine-library/stage-1-part-0.arrow and /dev/null differ diff --git a/examples/distributed/engine-library/stage-1-part-1.arrow b/examples/distributed/engine-library/stage-1-part-1.arrow deleted file mode 100644 index 84c8ac2c..00000000 Binary files a/examples/distributed/engine-library/stage-1-part-1.arrow and /dev/null differ diff --git a/examples/distributed/engine-library/stage-1-part-2.arrow b/examples/distributed/engine-library/stage-1-part-2.arrow deleted file mode 100644 index e5c81332..00000000 Binary files a/examples/distributed/engine-library/stage-1-part-2.arrow and /dev/null differ diff --git a/examples/distributed/run_tpch.py b/examples/distributed/run_tpch.py new file mode 100644 index 00000000..90a25dee --- /dev/null +++ b/examples/distributed/run_tpch.py @@ -0,0 +1,185 @@ +# Licensed to the Apache Software Foundation (ASF) under one +# or more contributor license agreements. See the NOTICE file +# distributed with this work for additional information +# regarding copyright ownership. The ASF licenses this file +# to you under the Apache License, Version 2.0 (the +# "License"); you may not use this file except in compliance +# with the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, +# software distributed under the License is distributed on an +# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +# KIND, either express or implied. See the License for the +# specific language governing permissions and limitations +# under the License. + +"""Run TPC-H Q1 across worker processes, and compare against one process. + + python examples/distributed/run_tpch.py --partitions 4 + +Needs the TPC-H data the repository's other examples use:: + + mkdir -p examples/tpch/data && cd examples/tpch/data + uv pip install tpchgen-cli && uv run --no-project tpchgen-cli -s 1 --format=parquet + +`tpchgen-cli` writes one file per table, so `lineitem.parquet` is a single +220 MB file -- one partition, and nothing to fan out. This script re-shards +the columns Q1 needs into `--partitions` files first, which is also a fair +illustration of the real constraint: a distributed engine can only spread work +as widely as the data is split. +""" + +from __future__ import annotations + +import argparse +import pathlib +import shutil +import sys +import tempfile +import time + +import pyarrow as pa +import pyarrow.parquet as pq +from dfx_engine.driver import run_distributed, run_local +from dfx_engine.session import SessionSpec + +# Q1 without the `l_shipdate` filter and the `avg` columns, so the shard below +# stays small. The shape that matters is unchanged: group by two low-cardinality +# columns, aggregate, order. +Q1 = """ +select l_returnflag, + l_linestatus, + count(*) as count_order, + sum(l_quantity) as sum_qty, + sum(l_extendedprice) as sum_base_price, + sum(dfx_net_revenue(l_extendedprice, l_discount, l_tax)) as sum_charge, + dfx_weighted_avg(l_extendedprice, l_quantity) as wavg_price +from lineitem +group by l_returnflag, l_linestatus +order by l_returnflag, l_linestatus +""" + +COLUMNS = [ + "l_returnflag", + "l_linestatus", + "l_quantity", + "l_extendedprice", + "l_discount", + "l_tax", +] + + +def reshard( + source: pathlib.Path, into: pathlib.Path, partitions: int, rows: int +) -> int: + """Write the first `rows` rows of `source` as `partitions` Parquet files.""" + into.mkdir(parents=True, exist_ok=True) + table = pq.read_table(source, columns=COLUMNS) + if rows: + table = table.slice(0, rows) + + per_file = max(1, table.num_rows // partitions) + written = 0 + for index in range(partitions): + offset = index * per_file + length = table.num_rows - offset if index == partitions - 1 else per_file + if length <= 0: + break + pq.write_table(table.slice(offset, length), into / f"part-{index}.parquet") + written += 1 + return written + + +def main(argv: list[str] | None = None) -> int: + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument( + "--data", + type=pathlib.Path, + default=pathlib.Path(__file__).resolve().parents[1] + / "tpch" + / "data" + / "lineitem.parquet", + ) + parser.add_argument("--partitions", type=int, default=4) + parser.add_argument( + "--rows", + type=int, + default=2_000_000, + help="rows to use; 0 for all of them (SF 1 lineitem is ~6M)", + ) + args = parser.parse_args(argv) + + if not args.data.exists(): + sys.stderr.write( + f"{args.data} not found. Generate it with:\n" + " mkdir -p examples/tpch/data && cd examples/tpch/data\n" + " uv pip install tpchgen-cli\n" + " uv run --no-project tpchgen-cli -s 1 --format=parquet\n" + ) + return 2 + + workspace = pathlib.Path(tempfile.mkdtemp(prefix="dfx-tpch-")) + try: + data = workspace / "lineitem" + count = reshard(args.data, data, args.partitions, args.rows) + print(f"resharded into {count} file(s) under {data}") + + spec = SessionSpec( + tables={"lineitem": str(data)}, + shuffle_dir=str(workspace / "shuffle"), + target_partitions=args.partitions, + ) + + start = time.monotonic() + result = run_distributed(Q1, spec) + distributed = time.monotonic() - start + print( + f"distributed: {distributed:.2f}s across {len(result.partitions)} " + f"worker process(es); rows per worker {result.worker_rows}" + ) + + start = time.monotonic() + local = run_local(Q1, spec) + print(f"single process: {time.monotonic() - start:.2f}s") + + # The point of the comparison is agreement, not speed: four processes + # on one laptop will not beat one process that skips the round trip + # through Arrow IPC files. + table = pa.Table.from_batches(result.batches) + reference = pa.Table.from_batches(local) + + # Compared with a tolerance, not for equality. Splitting a `sum` across + # partitions changes the order the additions happen in, and floating + # point addition is not associative -- so the last bits of `sum_charge` + # legitimately differ between the two runs. Any distributed engine has + # this property; it is worth knowing before someone diffs two runs and + # concludes the split is broken. + assert table.column_names == reference.column_names + for name in table.column_names: + got, want = ( + table.column(name).to_pylist(), + reference.column(name).to_pylist(), + ) + assert len(got) == len(want), name + for lhs, rhs in zip(got, want, strict=True): + if isinstance(lhs, float): + assert abs(lhs - rhs) <= 1e-6 * max(1.0, abs(rhs)), (name, lhs, rhs) + else: + assert lhs == rhs, (name, lhs, rhs) + + print("\nsame answer both ways (floats to within 1e-6 relative):\n") + names = table.column_names + print(" ".join(f"{name:>16}" for name in names)) + for row in zip( + *(table.column(name).to_pylist() for name in names), strict=True + ): + print(" ".join(f"{value:>16}" for value in row)) + finally: + shutil.rmtree(workspace, ignore_errors=True) + return 0 + + +if __name__ == "__main__": + raise SystemExit(main()) --------------------------------------------------------------------- To unsubscribe, e-mail: [email protected] For additional commands, e-mail: [email protected]
