From 6b18896dfcddd592e68a70a3525d557d8e218305 Mon Sep 17 00:00:00 2001 From: Tim Saucer Date: Wed, 9 Sep 2026 12:27:42 -0400 Subject: [PATCH 01/24] Add dfx_storage: a codec that writes durable metadata MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit First of three libraries for the multi-library distributed example in #1719, and the reference implementation the extension guide's `extension_codec_durable_metadata` section currently lacks. Every other example codec in this repository parks the live object in a process-global `HashMap` and encodes an integer token into it. The guide says plainly that this is a demonstration and not a pattern, then has nothing to point at that does it properly. This codec is that: it writes the file paths and sizes, the projection, the row limit, and the schema, so decoding needs nothing at all from the encoding process. The provider scans a directory of Parquet files and reports one output partition per file. That is the reason it exists rather than `register_parquet`: it fixes the mapping from partition index to file, so an engine can hand partition `i` to a worker and know which bytes that worker will read. Paths are sorted, because directory iteration order is unspecified and a worker that disagreed with the driver about which file is partition 3 would produce wrong answers silently rather than fail. `PartitionedParquetExec` is a leaf on purpose. A node with children hands them to the framework to encode with the host's codec, which is correct but means the interesting part of a codec — what it writes down — belongs to someone else. Everything this node needs to run is in the node, so that is what goes on the wire. It reuses `DataSourceExec` to do the actual reading; the point is to own the description of the scan across a process boundary, not to reimplement Parquet. Wire format is `DFXSTOR1 | json_len: u32 | json | arrow ipc schema`. JSON for the scalar fields because someone debugging a worker can read it, Arrow IPC for the schema because it is the only encoding that round-trips every Arrow type. The magic carries a version the codec refuses to guess at. The codec claims by downcasting to its own concrete type and hands anything else to the default codec, whose error is the chain's "not mine" signal. Tests pin one fact that makes the narrow claim obviously right: an extension codec is only ever consulted for nodes with no native encoding, so the only nodes that reach it are ones some library owns — a broad claim can only steal from a peer, never pick up slack. Ten tests, the load-bearing one being a genuinely separate interpreter spawned through `sys.executable` that builds its own session, checks the codec id it expects is installed, decodes a plan written by another process, and executes all three partitions. A token registry cannot pass that test, which is the point of writing it first. Co-Authored-By: Claude Opus 5 (1M context) --- .gitignore | 1 + Cargo.lock | 18 ++ Cargo.toml | 1 + .../distributed/storage-library/Cargo.toml | 51 ++++ examples/distributed/storage-library/build.rs | 20 ++ .../storage-library/pyproject.toml | 32 +++ .../python/tests/_test_portable_codec.py | 247 ++++++++++++++++++ .../storage-library/python/tests/conftest.py | 69 +++++ .../distributed/storage-library/src/codec.rs | 226 ++++++++++++++++ .../distributed/storage-library/src/exec.rs | 181 +++++++++++++ .../storage-library/src/extension.rs | 156 +++++++++++ .../distributed/storage-library/src/lib.rs | 40 +++ .../storage-library/src/table_provider.rs | 182 +++++++++++++ 13 files changed, 1224 insertions(+) create mode 100644 examples/distributed/storage-library/Cargo.toml create mode 100644 examples/distributed/storage-library/build.rs create mode 100644 examples/distributed/storage-library/pyproject.toml create mode 100644 examples/distributed/storage-library/python/tests/_test_portable_codec.py create mode 100644 examples/distributed/storage-library/python/tests/conftest.py create mode 100644 examples/distributed/storage-library/src/codec.rs create mode 100644 examples/distributed/storage-library/src/exec.rs create mode 100644 examples/distributed/storage-library/src/extension.rs create mode 100644 examples/distributed/storage-library/src/lib.rs create mode 100644 examples/distributed/storage-library/src/table_provider.rs diff --git a/.gitignore b/.gitignore index 614d82327..ef00c0fd6 100644 --- a/.gitignore +++ b/.gitignore @@ -36,3 +36,4 @@ docs/mdbook/book .pyo3_build_config +examples/distributed/*/.venv/ diff --git a/Cargo.lock b/Cargo.lock index 6a7f68438..13ac8b645 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -1666,6 +1666,24 @@ dependencies = [ "url", ] +[[package]] +name = "dfx-storage" +version = "54.0.0" +dependencies = [ + "arrow", + "async-trait", + "datafusion", + "datafusion-catalog", + "datafusion-common", + "datafusion-ffi", + "datafusion-proto", + "datafusion-python-util", + "pyo3", + "pyo3-build-config", + "pyo3-log", + "serde_json", +] + [[package]] name = "digest" version = "0.10.7" diff --git a/Cargo.toml b/Cargo.toml index 0fabd5437..929c3bbad 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -32,6 +32,7 @@ members = [ "crates/util", "examples/datafusion-ffi-example", "examples/datafusion-ffi-query-planner-example", + "examples/distributed/storage-library", ] resolver = "3" diff --git a/examples/distributed/storage-library/Cargo.toml b/examples/distributed/storage-library/Cargo.toml new file mode 100644 index 000000000..0e01e0523 --- /dev/null +++ b/examples/distributed/storage-library/Cargo.toml @@ -0,0 +1,51 @@ +# 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. + +[package] +name = "dfx-storage" +version.workspace = true +edition.workspace = true +rust-version.workspace = true +license.workspace = true +description = "Example extension library: a partitioned Parquet table provider whose codec encodes durable metadata" +homepage.workspace = true +repository.workspace = true +publish = false + +[dependencies] +arrow = { workspace = true } +async-trait = { workspace = true } +datafusion = { workspace = true } +datafusion-catalog = { workspace = true, default-features = false } +datafusion-common = { workspace = true, default-features = false } +datafusion-ffi = { workspace = true } +datafusion-proto = { workspace = true } +datafusion-python-util.workspace = true +pyo3 = { workspace = true, features = [ + "extension-module", + "abi3", + "abi3-py310", +] } +pyo3-log = { workspace = true } +serde_json = { workspace = true } + +[build-dependencies] +pyo3-build-config = { workspace = true } + +[lib] +name = "dfx_storage" +crate-type = ["cdylib", "rlib"] diff --git a/examples/distributed/storage-library/build.rs b/examples/distributed/storage-library/build.rs new file mode 100644 index 000000000..4878d8b0e --- /dev/null +++ b/examples/distributed/storage-library/build.rs @@ -0,0 +1,20 @@ +// 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. + +fn main() { + pyo3_build_config::add_extension_module_link_args(); +} diff --git a/examples/distributed/storage-library/pyproject.toml b/examples/distributed/storage-library/pyproject.toml new file mode 100644 index 000000000..24bc294fd --- /dev/null +++ b/examples/distributed/storage-library/pyproject.toml @@ -0,0 +1,32 @@ +# 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. + +[build-system] +requires = ["maturin>=1.6,<2.0"] +build-backend = "maturin" + +[project] +name = "dfx_storage" +requires-python = ">=3.10" +classifiers = [ + "Programming Language :: Rust", + "Programming Language :: Python :: Implementation :: CPython", +] +dynamic = ["version"] + +[tool.maturin] +features = ["pyo3/extension-module"] diff --git a/examples/distributed/storage-library/python/tests/_test_portable_codec.py b/examples/distributed/storage-library/python/tests/_test_portable_codec.py new file mode 100644 index 000000000..716d92fd6 --- /dev/null +++ b/examples/distributed/storage-library/python/tests/_test_portable_codec.py @@ -0,0 +1,247 @@ +# 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. + +"""The claim this library exists to make: its plans decode in another process. + +Every other example codec in this repository parks the live object in a +process-global map and encodes a token. These tests are written to fail if +this one ever does that -- the decoding side is a separate interpreter, so a +token would have nothing to look up. +""" + +from __future__ import annotations + +import itertools +import re +import subprocess +import sys +import textwrap +from typing import TYPE_CHECKING + +import pytest +from datafusion import SessionContext +from datafusion.plan import ExecutionPlan +from dfx_storage import DfxStorageExtension, PartitionedParquetTable + +if TYPE_CHECKING: + import pathlib + + +def _configured(directory: pathlib.Path) -> tuple[SessionContext, DfxStorageExtension]: + bundle = DfxStorageExtension() + ctx = SessionContext().with_extensions(bundle) + ctx.register_table("readings", PartitionedParquetTable(str(directory))) + return ctx, bundle + + +def test_provider_reports_one_partition_per_file(readings_dir: pathlib.Path) -> None: + """The file is the partition, which is the axis an engine splits along.""" + ctx, _ = _configured(readings_dir) + plan = ctx.sql("select sensor_id from readings").execution_plan() + + assert plan.partition_count == 3 + # Not `Hash`: rows are grouped by which file they landed in, which says + # nothing about their values. + assert plan.output_partitioning.scheme == "UnknownPartitioning" + assert plan.output_partitioning.hash_expressions is None + + +def test_each_partition_reads_exactly_one_file(readings_dir: pathlib.Path) -> None: + """Partition i reads file i, so two workers never read the same bytes.""" + ctx, _ = _configured(readings_dir) + plan = ctx.sql("select sensor_id from readings").execution_plan() + + per_partition = [ + sorted( + value + for batch in ctx.execute(plan, partition) + for value in batch.to_pyarrow().column("sensor_id").to_pylist() + ) + for partition in range(plan.partition_count) + ] + + assert per_partition == [[0, 1, 2], [100, 101, 102], [200, 201, 202]] + # Disjoint, and together the whole table. + everything = sorted(itertools.chain.from_iterable(per_partition)) + assert everything == [0, 1, 2, 100, 101, 102, 200, 201, 202] + + +def test_this_librarys_codec_carried_the_node(readings_dir: pathlib.Path) -> None: + """Assert *this* codec did the work, not merely that the query succeeded. + + Both codecs being installed does not mean this one saw the node; a codec + installed earlier that claims broadly would have taken it. + """ + ctx, bundle = _configured(readings_dir) + plan = ctx.sql("select sensor_id from readings").execution_plan() + assert bundle.encode_calls() == 0 + + blob = plan.to_bytes(ctx) + assert bundle.encode_calls() == 1 + + restored = ExecutionPlan.from_bytes(ctx, blob) + assert bundle.decode_calls() == 1 + assert "PartitionedParquetExec" in restored.display_indent() + + +def test_the_payload_is_metadata_not_a_token(readings_dir: pathlib.Path) -> None: + """The bytes name the files, so they mean something in another process.""" + ctx, _ = _configured(readings_dir) + blob = ctx.sql("select sensor_id from readings").execution_plan().to_bytes(ctx) + + assert b"DFXSTOR1" in blob + for index in range(3): + assert f"part-{index}.parquet".encode() in blob + + +def test_the_same_bytes_decode_twice(readings_dir: pathlib.Path) -> None: + """A token registry consumes its entry on decode. Durable metadata does not. + + This is what lets one encoded plan fan out to several workers. + """ + ctx, bundle = _configured(readings_dir) + blob = ctx.sql("select sensor_id from readings").execution_plan().to_bytes(ctx) + + first = ExecutionPlan.from_bytes(ctx, blob) + second = ExecutionPlan.from_bytes(ctx, blob) + + assert bundle.decode_calls() == 2 + assert first.partition_count == second.partition_count == 3 + + +def test_a_projection_survives_the_round_trip(readings_dir: pathlib.Path) -> None: + """The projection is part of the descriptor, not re-derived on decode.""" + ctx, _ = _configured(readings_dir) + plan = ctx.sql("select reading from readings").execution_plan() + restored = ExecutionPlan.from_bytes(ctx, plan.to_bytes(ctx)) + + rows = [ + value + for partition in range(restored.partition_count) + for batch in ctx.execute(restored, partition) + for value in batch.to_pyarrow().column("reading").to_pylist() + ] + assert sorted(rows) == [1.5, 1.5, 1.5, 2.5, 2.5, 2.5, 3.5, 3.5, 3.5] + + +def test_stock_nodes_never_reach_this_codec(readings_dir: pathlib.Path) -> None: + """An extension codec is only consulted for nodes with no native encoding. + + The aggregate and filter above the scan all have their own `try_to_proto`, + so the framework encodes them itself and this codec is never offered them. + That is why claiming a broad category is so damaging: the only nodes that + ever arrive here are ones *some* library owns, so a broad claim can only + ever steal from a peer, never pick up slack. + """ + ctx, bundle = _configured(readings_dir) + plan = ctx.sql("select count(*) from readings where reading > 2.0").execution_plan() + + plan.to_bytes(ctx) + + # Exactly one node in that plan is ours, and nothing else was offered. + assert bundle.encode_calls() == 1 + assert bundle.declined_calls() == 0 + + +WORKER = textwrap.dedent( + """ + import sys + from datafusion import SessionContext + from datafusion.plan import ExecutionPlan + from dfx_storage import DfxStorageExtension + + blob_path, expected_id = sys.argv[1], sys.argv[2] + + # A session built from scratch: this process has never registered the + # table, and shares nothing with the one that wrote the plan. + bundle = DfxStorageExtension() + ctx = SessionContext().with_extensions(bundle) + + installed = ctx.physical_extension_codec_ids() + assert expected_id in installed, f"codec id {expected_id} not in {installed}" + + with open(blob_path, "rb") as handle: + plan = ExecutionPlan.from_bytes(ctx, handle.read()) + + total = 0 + for partition in range(plan.partition_count): + for batch in ctx.execute(plan, partition): + total += batch.to_pyarrow().num_rows + print(f"partitions={plan.partition_count} rows={total} decoded={bundle.decode_calls()}") + """ +) + + +def test_a_separate_process_decodes_and_executes_the_plan( + readings_dir: pathlib.Path, tmp_path: pathlib.Path +) -> None: + """The whole point. No shared session, no shared registry, no token. + + Spawned through `sys.executable` rather than `multiprocessing`: the tokio + runtime backing this extension is a process-global, so `fork` is unsafe, + and a hardcoded `python` could differ in minor version from this one. + """ + ctx, _ = _configured(readings_dir) + blob = ctx.sql("select sensor_id from readings").execution_plan().to_bytes(ctx) + blob_path = tmp_path / "plan.bin" + blob_path.write_bytes(blob) + + worker = tmp_path / "worker.py" + worker.write_text(WORKER) + + result = subprocess.run( # noqa: S603 + [sys.executable, str(worker), str(blob_path), "dfx_storage.physical.v1"], + capture_output=True, + text=True, + check=False, + ) + + assert result.returncode == 0, result.stderr + assert "partitions=3 rows=9 decoded=1" in result.stdout + + +def test_a_worker_without_the_codec_says_which_one_is_missing( + readings_dir: pathlib.Path, +) -> None: + """The failure names the codec, which is the whole value of pinned ids.""" + ctx, _ = _configured(readings_dir) + blob = ctx.sql("select sensor_id from readings").execution_plan().to_bytes(ctx) + + # A session with no extension codecs at all, standing in for a worker + # whose bootstrap forgot to install the bundle. + bare = SessionContext() + with pytest.raises( + Exception, match=re.escape("dfx_storage.physical.v1") + ) as excinfo: + ExecutionPlan.from_bytes(bare, blob) + assert "not installed on this session" in str(excinfo.value) + + +def test_the_bundle_is_reusable_across_sessions(readings_dir: pathlib.Path) -> None: + """One bundle object, two sessions: components are built per install.""" + bundle = DfxStorageExtension() + first = SessionContext().with_extensions(bundle) + second = SessionContext().with_extensions(bundle) + + for ctx in (first, second): + ctx.register_table("readings", PartitionedParquetTable(str(readings_dir))) + assert ( + ctx.sql("select count(*) from readings").collect()[0].column(0)[0].as_py() + == 9 + ) + + assert first.__datafusion_codec_id__ != second.__datafusion_codec_id__ diff --git a/examples/distributed/storage-library/python/tests/conftest.py b/examples/distributed/storage-library/python/tests/conftest.py new file mode 100644 index 000000000..e37db8f82 --- /dev/null +++ b/examples/distributed/storage-library/python/tests/conftest.py @@ -0,0 +1,69 @@ +# 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 + +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) + + +@pytest.fixture(autouse=True) +def fail_on_log_warnings() -> Generator[None, Any, None]: + handler = _FailOnWarning() + logging.root.addHandler(handler) + yield + logging.root.removeHandler(handler) + + +@pytest.fixture +def readings_dir(tmp_path: pathlib.Path) -> pathlib.Path: + """Three Parquet files, so the provider reports three partitions. + + Written as `part-0/1/2` rather than in one file because the file *is* the + partition for this provider, and a single-file table would hide every + partition-routing mistake. + """ + directory = tmp_path / "readings" + directory.mkdir() + for index in range(3): + base = index * 100 + pq.write_table( + pa.table( + { + "sensor_id": [base, base + 1, base + 2], + "reading": [1.5, 2.5, 3.5], + } + ), + directory / f"part-{index}.parquet", + ) + return directory diff --git a/examples/distributed/storage-library/src/codec.rs b/examples/distributed/storage-library/src/codec.rs new file mode 100644 index 000000000..9868e0014 --- /dev/null +++ b/examples/distributed/storage-library/src/codec.rs @@ -0,0 +1,226 @@ +// 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. + +//! A physical codec that writes durable metadata. +//! +//! The other example crates in this repository park the live object in a +//! process-global `HashMap` and encode an integer token into it. That makes +//! Rust type identity observable in a test, and it is explicitly not a +//! pattern: the token means the same bytes cannot be decoded twice, one plan +//! cannot fan out to several readers, and a plan that never reaches a decoder +//! leaks. None of that is acceptable for a plan that leaves the process. +//! +//! This codec writes down what a fresh [`PartitionedParquetExec`] can be built +//! from -- the file paths and sizes, the projection, the row limit, and the +//! schema -- so decoding needs nothing from the encoding process. Sending the +//! same bytes to ten workers works, and so does sending them tomorrow. +//! +//! # Wire format +//! +//! ```text +//! DFXSTOR1 | json_len: u32 (LE) | json | arrow ipc schema +//! ``` +//! +//! The magic is checked before anything else is read, and the trailing `1` is +//! a version this codec refuses to guess at. JSON carries the small scalar +//! fields because a human debugging a worker can read it; the schema is Arrow +//! IPC because that is the only encoding guaranteed to round-trip every Arrow +//! type, including extension types and field metadata. + +use std::fmt; +use std::sync::Arc; +use std::sync::atomic::{AtomicUsize, Ordering}; + +use arrow::datatypes::Schema; +use arrow::ipc::reader::StreamReader; +use arrow::ipc::writer::StreamWriter; +use datafusion::common::{Result, internal_datafusion_err, internal_err}; +use datafusion::execution::TaskContext; +use datafusion::physical_plan::ExecutionPlan; +use datafusion_proto::physical_plan::{ + DefaultPhysicalExtensionCodec, PhysicalExtensionCodec, PhysicalProtoConverterExtension, +}; + +use crate::exec::{FileSlice, PartitionedParquetExec}; + +/// Framing magic. The trailing digit is the payload version. +const MAGIC: &[u8; 8] = b"DFXSTOR1"; + +/// How often this codec claimed one of its own nodes. +/// +/// Exposed to Python so a test can assert that *this* codec carried the node, +/// rather than inferring it from a query that merely succeeded. Both codecs +/// being installed does not mean yours saw the node -- see +/// `extension_codec_order`. +#[derive(Default, Debug)] +pub(crate) struct CodecCounters { + pub(crate) encoded: AtomicUsize, + pub(crate) decoded: AtomicUsize, + pub(crate) declined: AtomicUsize, +} + +pub(crate) struct DfxStoragePhysicalCodec { + /// Anything this library does not own is handed to the default codec, + /// whose error is the chain's "not mine" signal. + inner: DefaultPhysicalExtensionCodec, + pub(crate) counters: Arc, +} + +impl DfxStoragePhysicalCodec { + pub(crate) fn new(counters: Arc) -> Self { + Self { + inner: DefaultPhysicalExtensionCodec {}, + counters, + } + } +} + +impl fmt::Debug for DfxStoragePhysicalCodec { + fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result { + formatter + .debug_struct("DfxStoragePhysicalCodec") + .finish_non_exhaustive() + } +} + +fn schema_to_ipc_bytes(schema: &Schema) -> Result> { + let mut buf: Vec = Vec::new(); + { + let mut writer = StreamWriter::try_new(&mut buf, schema) + .map_err(|err| internal_datafusion_err!("dfx_storage: writing schema: {err}"))?; + writer + .finish() + .map_err(|err| internal_datafusion_err!("dfx_storage: writing schema: {err}"))?; + } + Ok(buf) +} + +fn schema_from_ipc_bytes(bytes: &[u8]) -> Result { + let reader = StreamReader::try_new(std::io::Cursor::new(bytes), None) + .map_err(|err| internal_datafusion_err!("dfx_storage: reading schema: {err}"))?; + Ok(reader.schema().as_ref().clone()) +} + +impl PhysicalExtensionCodec for DfxStoragePhysicalCodec { + fn try_encode( + &self, + node: Arc, + buf: &mut Vec, + proto_converter: &dyn PhysicalProtoConverterExtension, + ) -> Result<()> { + // Downcast to our own concrete type. Claiming a broad category -- + // `ForeignExecutionPlan`, say -- would take nodes from every library + // installed after this one, and the query would still succeed, so + // nothing would point at the codec that stole them. + let Some(exec) = node.downcast_ref::() else { + self.counters.declined.fetch_add(1, Ordering::SeqCst); + return self.inner.try_encode(node, buf, proto_converter); + }; + + let descriptor = serde_json::json!({ + "files": exec.files.iter().map(|file| { + serde_json::json!({ "path": file.path, "size": file.size }) + }).collect::>(), + "projection": exec.projection, + "limit": exec.limit, + }); + let json = serde_json::to_vec(&descriptor) + .map_err(|err| internal_datafusion_err!("dfx_storage: encoding descriptor: {err}"))?; + let schema = schema_to_ipc_bytes(&exec.table_schema)?; + + buf.extend_from_slice(MAGIC); + let json_len = u32::try_from(json.len()) + .map_err(|_| internal_datafusion_err!("dfx_storage: descriptor too large to encode"))?; + buf.extend_from_slice(&json_len.to_le_bytes()); + buf.extend_from_slice(&json); + buf.extend_from_slice(&schema); + + self.counters.encoded.fetch_add(1, Ordering::SeqCst); + Ok(()) + } + + fn try_decode( + &self, + buf: &[u8], + inputs: &[Arc], + ctx: &TaskContext, + proto_converter: &dyn PhysicalProtoConverterExtension, + ) -> Result> { + // The chain routes a framed payload by id, so reaching this codec + // already means the payload is ours. Checking the magic anyway is + // cheap and turns a version skew into a clear error instead of a + // misparse. + let Some(rest) = buf.strip_prefix(MAGIC) else { + self.counters.declined.fetch_add(1, Ordering::SeqCst); + return self.inner.try_decode(buf, inputs, ctx, proto_converter); + }; + if !inputs.is_empty() { + return internal_err!( + "PartitionedParquetExec is a leaf, got {} input(s)", + inputs.len() + ); + } + + let (len_bytes, rest) = rest.split_at_checked(4).ok_or_else(|| { + internal_datafusion_err!("dfx_storage: payload truncated before descriptor length") + })?; + let json_len = u32::from_le_bytes( + len_bytes + .try_into() + .map_err(|_| internal_datafusion_err!("dfx_storage: bad descriptor length"))?, + ) as usize; + let (json, schema_bytes) = rest.split_at_checked(json_len).ok_or_else(|| { + internal_datafusion_err!( + "dfx_storage: descriptor claims {json_len} bytes, {} remain", + rest.len() + ) + })?; + + let descriptor: serde_json::Value = serde_json::from_slice(json) + .map_err(|err| internal_datafusion_err!("dfx_storage: bad descriptor: {err}"))?; + let files = descriptor["files"] + .as_array() + .ok_or_else(|| internal_datafusion_err!("dfx_storage: descriptor has no file list"))? + .iter() + .map(|file| { + let path = file["path"].as_str().ok_or_else(|| { + internal_datafusion_err!("dfx_storage: file entry has no path") + })?; + let size = file["size"].as_u64().ok_or_else(|| { + internal_datafusion_err!("dfx_storage: file entry {path} has no size") + })?; + Ok(FileSlice { + path: path.to_string(), + size, + }) + }) + .collect::>>()?; + let projection = descriptor["projection"].as_array().map(|indices| { + indices + .iter() + .filter_map(|index| index.as_u64().map(|index| index as usize)) + .collect::>() + }); + let limit = descriptor["limit"].as_u64().map(|limit| limit as usize); + let schema = Arc::new(schema_from_ipc_bytes(schema_bytes)?); + + self.counters.decoded.fetch_add(1, Ordering::SeqCst); + Ok(Arc::new(PartitionedParquetExec::new( + files, schema, projection, limit, + )?)) + } +} diff --git a/examples/distributed/storage-library/src/exec.rs b/examples/distributed/storage-library/src/exec.rs new file mode 100644 index 000000000..3f09f0391 --- /dev/null +++ b/examples/distributed/storage-library/src/exec.rs @@ -0,0 +1,181 @@ +// 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. + +//! This library's own execution plan node. +//! +//! A leaf, and deliberately so. A node with children hands them to the +//! framework to encode with the *host's* codec, which is the right thing but +//! means the interesting part of a codec -- what it writes down -- is somebody +//! else's problem. Everything this node needs to run lives in the node +//! itself: which files, which columns, how many rows. That is what +//! [`crate::codec`] writes to the wire, and it is why a plan built here can be +//! decoded in a process that has never seen this table registered. +//! +//! One output partition per file. That is the axis a distributed engine +//! splits along: partition `i` reads file `i` and nothing else, so two workers +//! never touch the same bytes and no coordination is needed. + +use std::fmt; +use std::sync::Arc; + +use arrow::datatypes::SchemaRef; +use datafusion::common::Result; +use datafusion::common::tree_node::TreeNodeRecursion; +use datafusion::datasource::listing::PartitionedFile; +use datafusion::datasource::physical_plan::{FileScanConfigBuilder, ParquetSource}; +use datafusion::datasource::source::DataSourceExec; +use datafusion::execution::TaskContext; +use datafusion::execution::object_store::ObjectStoreUrl; +use datafusion::physical_expr::{EquivalenceProperties, Partitioning, PhysicalExpr}; +use datafusion::physical_plan::execution_plan::{Boundedness, EmissionType}; +use datafusion::physical_plan::{ + DisplayAs, DisplayFormatType, ExecutionPlan, PlanProperties, SendableRecordBatchStream, +}; + +/// One Parquet file, and the size the object store will report for it. +/// +/// The size travels with the path because `PartitionedFile` needs it up front +/// and a decoding process should not have to stat the file to rebuild a plan. +#[derive(Debug, Clone, PartialEq, Eq)] +pub(crate) struct FileSlice { + pub(crate) path: String, + pub(crate) size: u64, +} + +/// Scans a fixed list of Parquet files, one file per output partition. +#[derive(Debug)] +pub(crate) struct PartitionedParquetExec { + /// Output partition `i` reads `files[i]`. + pub(crate) files: Vec, + /// The table's full schema, before projection. + pub(crate) table_schema: SchemaRef, + /// Column indices into `table_schema`, or `None` for all of them. + pub(crate) projection: Option>, + pub(crate) limit: Option, + properties: Arc, +} + +impl PartitionedParquetExec { + pub(crate) fn new( + files: Vec, + table_schema: SchemaRef, + projection: Option>, + limit: Option, + ) -> Result { + let projected_schema = match projection.as_ref() { + Some(indices) => Arc::new(table_schema.project(indices)?), + None => Arc::clone(&table_schema), + }; + // `UnknownPartitioning`, not `Hash`: the rows are split by which file + // they happen to live in, which says nothing about their values. A + // plan that claimed a hash partitioning here would let the optimizer + // skip a repartition it actually needs. + let properties = Arc::new(PlanProperties::new( + EquivalenceProperties::new(projected_schema), + Partitioning::UnknownPartitioning(files.len()), + EmissionType::Incremental, + Boundedness::Bounded, + )); + Ok(Self { + files, + table_schema, + projection, + limit, + properties, + }) + } + + /// Build the stock scan for a single one of our files. + /// + /// Reusing `DataSourceExec` for the actual reading is the point: this node + /// exists to own the *description* of the scan across a process boundary, + /// not to reimplement Parquet. + fn scan_for(&self, partition: usize) -> Result> { + let slice = self.files.get(partition).ok_or_else(|| { + datafusion::common::internal_datafusion_err!( + "PartitionedParquetExec has {} partition(s), asked for {partition}", + self.files.len() + ) + })?; + let source = Arc::new(ParquetSource::new(Arc::clone(&self.table_schema))); + let config = FileScanConfigBuilder::new(ObjectStoreUrl::local_filesystem(), source) + .with_file(PartitionedFile::new(slice.path.clone(), slice.size)) + .with_projection_indices(self.projection.clone())? + .with_limit(self.limit) + .build(); + Ok(DataSourceExec::from_data_source(config)) + } +} + +impl DisplayAs for PartitionedParquetExec { + fn fmt_as(&self, _t: DisplayFormatType, f: &mut fmt::Formatter<'_>) -> fmt::Result { + write!(f, "PartitionedParquetExec: files={}", self.files.len())?; + if let Some(projection) = self.projection.as_ref() { + write!(f, ", projection={projection:?}")?; + } + if let Some(limit) = self.limit { + write!(f, ", limit={limit}")?; + } + Ok(()) + } +} + +impl ExecutionPlan for PartitionedParquetExec { + fn name(&self) -> &str { + Self::static_name() + } + + fn properties(&self) -> &Arc { + &self.properties + } + + fn children(&self) -> Vec<&Arc> { + vec![] + } + + fn apply_expressions( + &self, + _f: &mut dyn FnMut(&Arc) -> Result, + ) -> Result { + // The projection is column indices, not expressions, and any pushed + // down filter is held by the `DataSourceExec` this node builds. + Ok(TreeNodeRecursion::Continue) + } + + fn with_new_children( + self: Arc, + children: Vec>, + ) -> Result> { + if !children.is_empty() { + return datafusion::common::internal_err!( + "PartitionedParquetExec is a leaf, got {} children", + children.len() + ); + } + Ok(self) + } + + fn execute( + &self, + partition: usize, + context: Arc, + ) -> Result { + // Partition 0 of the single-file scan: each of our partitions is one + // whole file, so the inner scan only ever has one of its own. + self.scan_for(partition)?.execute(0, context) + } +} diff --git a/examples/distributed/storage-library/src/extension.rs b/examples/distributed/storage-library/src/extension.rs new file mode 100644 index 000000000..b42944b3f --- /dev/null +++ b/examples/distributed/storage-library/src/extension.rs @@ -0,0 +1,156 @@ +// 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. + +//! This library's extension bundle: codecs only, no planner. +//! +//! A provider library has no business installing a query planner, so this +//! bundle implements `__datafusion_session_components__` and stops there. +//! `with_extensions` accepts a bundle that implements only one of the two +//! hooks. + +use std::fmt; +use std::sync::Arc; +use std::sync::atomic::Ordering; + +use datafusion_ffi::proto::physical_extension_codec::FFI_PhysicalExtensionCodec; +use datafusion_proto::physical_plan::PhysicalExtensionCodec; +use datafusion_python_util::{ + create_physical_extension_capsule, ffi_task_context_provider_from_pycapsule, get_tokio_runtime, +}; +use pyo3::prelude::*; +use pyo3::types::{PyCapsule, PyDict}; + +use crate::codec::{CodecCounters, DfxStoragePhysicalCodec}; + +/// Wire id this codec's payloads carry. +/// +/// Pinned rather than left to default to the exporting class's import path, +/// because these payloads outlive the process that wrote them: a driver that +/// imports the class as `dfx_storage.BundledPhysicalCodec` and a worker that +/// imports it under any other name would otherwise disagree about the id and +/// every decode would fail. +const PHYSICAL_CODEC_ID: &str = "dfx_storage.physical.v1"; + +/// Carries this library's physical codec as an object rather than a capsule. +/// +/// `with_extensions` requires an object: a codec's wire id is read off the +/// thing it is handed over as, and a capsule has no type to read one from. +#[pyclass(name = "BundledPhysicalCodec", module = "dfx_storage")] +pub(crate) struct BundledPhysicalCodec { + codec: FFI_PhysicalExtensionCodec, +} + +#[pymethods] +impl BundledPhysicalCodec { + #[getter] + fn __datafusion_codec_id__(&self) -> &'static str { + PHYSICAL_CODEC_ID + } + + /// `session` is unused: the codec was bound to its task-context provider + /// when the bundle was installed, which is why the bundle receives the + /// context at all. + #[pyo3(signature = (session=None))] + fn __datafusion_physical_extension_codec__<'py>( + &self, + py: Python<'py>, + session: Option>, + ) -> PyResult> { + let _ = session; + create_physical_extension_capsule(py, &self.codec) + } +} + +/// Extension bundle for `SessionContext.with_extensions`. +/// +/// Reusable configuration, not bound state: every +/// `__datafusion_session_components__` call builds fresh components against +/// the context it is handed, so one bundle may be installed on several +/// sessions. +#[pyclass(from_py_object, name = "DfxStorageExtension", module = "dfx_storage")] +#[derive(Default, Clone)] +pub(crate) struct DfxStorageExtension { + counters: Arc, +} + +impl fmt::Debug for DfxStorageExtension { + fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result { + formatter + .debug_struct("DfxStorageExtension") + .field("counters", &self.counters) + .finish_non_exhaustive() + } +} + +#[pymethods] +impl DfxStorageExtension { + #[new] + fn new() -> Self { + Self::default() + } + + /// How often this codec encoded one of its own nodes. + fn encode_calls(&self) -> usize { + self.counters.encoded.load(Ordering::SeqCst) + } + + /// How often it rebuilt one, which is the half that happens on a worker. + fn decode_calls(&self) -> usize { + self.counters.decoded.load(Ordering::SeqCst) + } + + /// How often it was offered a node it does not own and passed it on. + /// + /// Non-zero is healthy: it means the chain is asking this codec about + /// other libraries' nodes and it is declining them. + fn declined_calls(&self) -> usize { + self.counters.declined.load(Ordering::SeqCst) + } + + /// The wire id, so a driver can put it in a worker's task envelope and + /// the worker can check it before decoding anything. + #[staticmethod] + fn physical_codec_id() -> &'static str { + PHYSICAL_CODEC_ID + } + + fn __datafusion_session_components__<'py>( + &self, + py: Python<'py>, + ctx: Bound<'py, PyAny>, + ) -> PyResult> { + // Take the provider off the context supplied by the host, so the + // codec's decode callbacks resolve against the session that will run + // the query. + let provider = ffi_task_context_provider_from_pycapsule(&ctx)?; + let runtime = get_tokio_runtime().handle().clone(); + + let codec: Arc = + Arc::new(DfxStoragePhysicalCodec::new(Arc::clone(&self.counters))); + let ffi = FFI_PhysicalExtensionCodec::new(codec, Some(runtime), provider); + let physical = Py::new(py, BundledPhysicalCodec { codec: ffi })?; + + // No logical codec: this library defines no logical extension node. + // Its table provider crosses FFI as a provider, not as a plan node. + let components = py + .import("datafusion")? + .getattr("SessionExtensionComponents")?; + let kwargs = PyDict::new(py); + kwargs.set_item("physical_extension_codecs", (physical,))?; + components.call((), Some(&kwargs)) + } +} diff --git a/examples/distributed/storage-library/src/lib.rs b/examples/distributed/storage-library/src/lib.rs new file mode 100644 index 000000000..bb01cf9c7 --- /dev/null +++ b/examples/distributed/storage-library/src/lib.rs @@ -0,0 +1,40 @@ +// 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. + +//! Example storage library: a partitioned Parquet table provider, its own +//! execution plan node, and a physical codec that writes durable metadata. +//! +//! One of three libraries in `examples/distributed`. This one owns tables. + +use pyo3::prelude::*; + +use crate::extension::{BundledPhysicalCodec, DfxStorageExtension}; +use crate::table_provider::PyPartitionedParquetTable; + +mod codec; +mod exec; +mod extension; +mod table_provider; + +#[pymodule] +fn dfx_storage(m: &Bound<'_, PyModule>) -> PyResult<()> { + pyo3_log::init(); + m.add_class::()?; + m.add_class::()?; + m.add_class::()?; + Ok(()) +} diff --git a/examples/distributed/storage-library/src/table_provider.rs b/examples/distributed/storage-library/src/table_provider.rs new file mode 100644 index 000000000..4b980c03f --- /dev/null +++ b/examples/distributed/storage-library/src/table_provider.rs @@ -0,0 +1,182 @@ +// 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. + +//! A table provider over a directory of Parquet files. +//! +//! One output partition per file, which is the whole reason this provider +//! exists rather than `SessionContext.register_parquet`: it fixes the mapping +//! from partition index to file, so a distributed engine can hand partition +//! `i` to a worker and know exactly which bytes that worker will read. + +use std::fs; +use std::path::Path; +use std::sync::Arc; + +use arrow::datatypes::{Schema, SchemaRef}; +use async_trait::async_trait; +use datafusion::catalog::{Session, TableProvider}; +use datafusion::common::{DataFusionError, Result, plan_err}; +use datafusion::datasource::TableType; +use datafusion::logical_expr::{Expr, TableProviderFilterPushDown}; +use datafusion::parquet::arrow::parquet_to_arrow_schema; +use datafusion::parquet::file::reader::{FileReader, SerializedFileReader}; +use datafusion::physical_plan::ExecutionPlan; +use datafusion_ffi::table_provider::FFI_TableProvider; +use datafusion_python_util::ffi_logical_codec_from_pycapsule; +use pyo3::prelude::*; +use pyo3::types::PyCapsule; + +use crate::exec::{FileSlice, PartitionedParquetExec}; + +/// Scans `*.parquet` under `directory`, one partition per file. +#[derive(Debug)] +pub(crate) struct PartitionedParquetTable { + files: Vec, + schema: SchemaRef, +} + +impl PartitionedParquetTable { + /// Read the directory listing and the first file's schema, once. + /// + /// Sorted by path so that partition `i` means the same file in every + /// process that opens the same directory. Directory iteration order is + /// not specified, and a worker that disagreed with the driver about which + /// file is partition 3 would silently produce wrong answers. + pub(crate) fn try_new(directory: &Path) -> Result { + let mut paths: Vec<_> = fs::read_dir(directory) + .map_err(|err| DataFusionError::External(Box::new(err)))? + .collect::>>() + .map_err(|err| DataFusionError::External(Box::new(err)))? + .into_iter() + .map(|entry| entry.path()) + .filter(|path| path.extension().is_some_and(|ext| ext == "parquet")) + .collect(); + paths.sort(); + + if paths.is_empty() { + return plan_err!("no .parquet files under {}", directory.display()); + } + + let mut files = Vec::with_capacity(paths.len()); + for path in &paths { + let metadata = + fs::metadata(path).map_err(|err| DataFusionError::External(Box::new(err)))?; + let path = path + .to_str() + .ok_or_else(|| DataFusionError::Plan(format!("non-UTF-8 path {path:?}")))?; + files.push(FileSlice { + path: path.to_string(), + size: metadata.len(), + }); + } + + let schema = Self::read_schema(&paths[0])?; + Ok(Self { + files, + schema: Arc::new(schema), + }) + } + + fn read_schema(path: &Path) -> Result { + let file = fs::File::open(path).map_err(|err| DataFusionError::External(Box::new(err)))?; + let reader = SerializedFileReader::new(file) + .map_err(|err| DataFusionError::ParquetError(Box::new(err)))?; + let metadata = reader.metadata().file_metadata(); + parquet_to_arrow_schema(metadata.schema_descr(), metadata.key_value_metadata()) + .map_err(|err| DataFusionError::ParquetError(Box::new(err))) + } +} + +#[async_trait] +impl TableProvider for PartitionedParquetTable { + fn schema(&self) -> SchemaRef { + Arc::clone(&self.schema) + } + + fn table_type(&self) -> TableType { + TableType::Base + } + + fn supports_filters_pushdown( + &self, + filters: &[&Expr], + ) -> Result> { + // Every filter is re-applied above the scan. Claiming `Exact` would + // tell the optimizer to drop the `FilterExec`, and this node does not + // pass predicates down to the Parquet reader. + Ok(vec![ + TableProviderFilterPushDown::Unsupported; + filters.len() + ]) + } + + async fn scan( + &self, + _state: &dyn Session, + projection: Option<&Vec>, + _filters: &[Expr], + limit: Option, + ) -> Result> { + Ok(Arc::new(PartitionedParquetExec::new( + self.files.clone(), + Arc::clone(&self.schema), + projection.cloned(), + limit, + )?)) + } +} + +/// Python handle for [`PartitionedParquetTable`]. +#[pyclass(name = "PartitionedParquetTable", module = "dfx_storage")] +pub(crate) struct PyPartitionedParquetTable { + directory: String, +} + +#[pymethods] +impl PyPartitionedParquetTable { + /// Open every `*.parquet` file under `directory` as one table. + #[new] + fn new(directory: String) -> Self { + Self { directory } + } + + /// Number of files, and so the number of output partitions. + fn partition_count(&self) -> PyResult { + Ok(self.build()?.files.len()) + } + + fn __datafusion_table_provider__<'py>( + &self, + py: Python<'py>, + session: Bound<'py, PyAny>, + ) -> PyResult> { + let provider = Arc::new(self.build()?); + // The codec comes off the session this provider is being installed + // on, never from a `SessionContext` built here: one built inline is + // already dropped by the time the capsule is used. + let codec = ffi_logical_codec_from_pycapsule(session, None)?; + let ffi = FFI_TableProvider::new_with_ffi_codec(provider, false, None, codec); + PyCapsule::new_with_value(py, ffi, cr"datafusion_table_provider") + } +} + +impl PyPartitionedParquetTable { + fn build(&self) -> PyResult { + PartitionedParquetTable::try_new(Path::new(&self.directory)) + .map_err(|err| pyo3::exceptions::PyValueError::new_err(err.to_string())) + } +} From 532995cb6688fe98a655170d932a19fcda4e42d3 Mon Sep 17 00:00:00 2001 From: Tim Saucer Date: Wed, 9 Sep 2026 12:41:13 -0400 Subject: [PATCH 02/24] Add dfx_udfs: the library that cannot be installed as a bundle MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Second of three libraries for #1719, and the one carrying the mixed-workflow case: it exposes no `__datafusion_session_components__`, so callers register its three functions and install its two codecs by hand. That is not an artificial handicap. `SessionExtensionComponents` carries codec fields only, so a function library has nowhere to put its functions — the rename in 5a1bfeba noted that UDF and provider fields will join later. Until they do, this is what a function library actually looks like, and the example should show what that costs rather than pretend every dependency has caught up. A test asserts the shape rather than describing it: `with_extensions` rejects this object, naming the hook it lacks. The functions are `dfx_net_revenue` (the TPC-H revenue expression), `dfx_weighted_avg`, and `dfx_revenue_rank`. The aggregate is written out rather than delegating to a built-in because its state is the point: two running sums, which is what lets DataFusion compute a partial aggregate per partition and merge the results. An aggregate that could only be evaluated over its whole input at once would give a different answer once split, which is exactly what a worker does to it. A test pins that by checking the plan really is `mode=Partial` and the answer is still right. Both codecs are name-only — `try_encode_*` writes nothing and `try_decode_*` rebuilds from `name`. Three worker tests, each a separate interpreter, pin what that buys, and they disagree with each other in the useful way: codec installed, nothing registered -> works, decode_calls == 1 functions registered, no codec -> works, decode_calls == 0 neither -> fails, naming dfx_net_revenue So installing the codec is an *alternative* to registering the functions, not an addition to it. The middle case is the trap worth knowing: on the driver, where the functions are registered, the registry is tried first and the codec is never consulted — so a codec that was broken or missing looks fine right up until a worker needs it. `try_decode_*` checks the name before the buffer, in that order. An empty encoding leaves `fun_definition` unset and carries no codec id, so it is the one path where a payload is offered to every installed codec in turn; a codec that trusted `buf` first would answer for names it does not own. There are two codecs because there are two plan layers and a library cannot know which one its callers will serialize — an engine shipping physical plans exercises only the physical one. Co-Authored-By: Claude Opus 5 (1M context) --- Cargo.lock | 18 ++ Cargo.toml | 1 + examples/distributed/udf-library/Cargo.toml | 51 +++ examples/distributed/udf-library/build.rs | 20 ++ .../distributed/udf-library/pyproject.toml | 32 ++ .../udf-library/python/tests/_test_udfs.py | 295 ++++++++++++++++++ .../udf-library/python/tests/conftest.py | 79 +++++ examples/distributed/udf-library/src/codec.rs | 243 +++++++++++++++ .../distributed/udf-library/src/functions.rs | 288 +++++++++++++++++ examples/distributed/udf-library/src/lib.rs | 46 +++ .../distributed/udf-library/src/python.rs | 231 ++++++++++++++ 11 files changed, 1304 insertions(+) create mode 100644 examples/distributed/udf-library/Cargo.toml create mode 100644 examples/distributed/udf-library/build.rs create mode 100644 examples/distributed/udf-library/pyproject.toml create mode 100644 examples/distributed/udf-library/python/tests/_test_udfs.py create mode 100644 examples/distributed/udf-library/python/tests/conftest.py create mode 100644 examples/distributed/udf-library/src/codec.rs create mode 100644 examples/distributed/udf-library/src/functions.rs create mode 100644 examples/distributed/udf-library/src/lib.rs create mode 100644 examples/distributed/udf-library/src/python.rs diff --git a/Cargo.lock b/Cargo.lock index 13ac8b645..6ae37a13f 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -1684,6 +1684,24 @@ dependencies = [ "serde_json", ] +[[package]] +name = "dfx-udfs" +version = "54.0.0" +dependencies = [ + "arrow", + "arrow-schema", + "datafusion", + "datafusion-common", + "datafusion-expr", + "datafusion-ffi", + "datafusion-functions-window", + "datafusion-proto", + "datafusion-python-util", + "pyo3", + "pyo3-build-config", + "pyo3-log", +] + [[package]] name = "digest" version = "0.10.7" diff --git a/Cargo.toml b/Cargo.toml index 929c3bbad..fb7db2744 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -33,6 +33,7 @@ members = [ "examples/datafusion-ffi-example", "examples/datafusion-ffi-query-planner-example", "examples/distributed/storage-library", + "examples/distributed/udf-library", ] resolver = "3" diff --git a/examples/distributed/udf-library/Cargo.toml b/examples/distributed/udf-library/Cargo.toml new file mode 100644 index 000000000..dac3eeff5 --- /dev/null +++ b/examples/distributed/udf-library/Cargo.toml @@ -0,0 +1,51 @@ +# 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. + +[package] +name = "dfx-udfs" +version.workspace = true +edition.workspace = true +rust-version.workspace = true +license.workspace = true +description = "Example extension library: user defined functions plus the codecs that make them portable" +homepage.workspace = true +repository.workspace = true +publish = false + +[dependencies] +arrow = { workspace = true } +arrow-schema = { workspace = true } +datafusion = { workspace = true } +datafusion-common = { workspace = true, default-features = false } +datafusion-expr = { workspace = true } +datafusion-ffi = { workspace = true } +datafusion-functions-window = { workspace = true } +datafusion-proto = { workspace = true } +datafusion-python-util.workspace = true +pyo3 = { workspace = true, features = [ + "extension-module", + "abi3", + "abi3-py310", +] } +pyo3-log = { workspace = true } + +[build-dependencies] +pyo3-build-config = { workspace = true } + +[lib] +name = "dfx_udfs" +crate-type = ["cdylib", "rlib"] diff --git a/examples/distributed/udf-library/build.rs b/examples/distributed/udf-library/build.rs new file mode 100644 index 000000000..4878d8b0e --- /dev/null +++ b/examples/distributed/udf-library/build.rs @@ -0,0 +1,20 @@ +// 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. + +fn main() { + pyo3_build_config::add_extension_module_link_args(); +} diff --git a/examples/distributed/udf-library/pyproject.toml b/examples/distributed/udf-library/pyproject.toml new file mode 100644 index 000000000..8e87abf47 --- /dev/null +++ b/examples/distributed/udf-library/pyproject.toml @@ -0,0 +1,32 @@ +# 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. + +[build-system] +requires = ["maturin>=1.6,<2.0"] +build-backend = "maturin" + +[project] +name = "dfx_udfs" +requires-python = ">=3.10" +classifiers = [ + "Programming Language :: Rust", + "Programming Language :: Python :: Implementation :: CPython", +] +dynamic = ["version"] + +[tool.maturin] +features = ["pyo3/extension-module"] diff --git a/examples/distributed/udf-library/python/tests/_test_udfs.py b/examples/distributed/udf-library/python/tests/_test_udfs.py new file mode 100644 index 000000000..54044578a --- /dev/null +++ b/examples/distributed/udf-library/python/tests/_test_udfs.py @@ -0,0 +1,295 @@ +# 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. + +"""What a function library owes a caller who is going to ship its plans.""" + +from __future__ import annotations + +import re +import subprocess +import sys +import textwrap +from typing import TYPE_CHECKING + +import pytest +from datafusion import SessionConfig, SessionContext, udaf, udf, udwf +from datafusion.plan import ExecutionPlan +from dfx_udfs import ( + CodecObservations, + NetRevenueUDF, + RevenueRankUDWF, + WeightedAvgUDAF, +) + +if TYPE_CHECKING: + import pathlib + + +def _session(directory: pathlib.Path, *, with_codecs: bool = True) -> tuple: + """Build a session the way this library requires: by hand. + + There is no `with_extensions(...)` here, and that is the point. Compare + with `dfx_storage`, which is one call. This library needs two codec + installs and three registrations, in an order the caller has to get right + on their own. + """ + observations = CodecObservations() + ctx = SessionContext(SessionConfig().with_target_partitions(2)) + if with_codecs: + ctx = ctx.with_logical_extension_codec(observations.logical_codec()) + ctx = ctx.with_physical_extension_codec(observations.physical_codec()) + ctx.register_udf(udf(NetRevenueUDF())) + ctx.register_udaf(udaf(WeightedAvgUDAF())) + ctx.register_udwf(udwf(RevenueRankUDWF())) + ctx.register_parquet("lineitem", str(directory)) + return ctx, observations + + +def test_the_scalar_function_computes_tpch_revenue(lineitem: pathlib.Path) -> None: + """`price * (1 - discount) * (1 + tax)`, checked by hand.""" + ctx, _ = _session(lineitem) + rows = ctx.sql( + "select dfx_net_revenue(l_extendedprice, l_discount, l_tax) as revenue " + "from lineitem order by revenue" + ).collect() + + revenue = [value for batch in rows for value in batch.column(0).to_pylist()] + # 100*1*1, 200*0.5*1, 400*0.75*1.1 + assert revenue == pytest.approx([100.0, 100.0, 330.0]) + + +def test_the_aggregate_is_correct_when_split_across_partitions( + lineitem: pathlib.Path, +) -> None: + """The two-sum state is what makes a partial/final split come out right. + + The input is two files and the session has two target partitions, so + DataFusion runs a partial aggregate per partition and merges. An aggregate + that could not be computed that way would give a different answer here + than over a single partition -- which is exactly what happens on a worker. + """ + ctx, _ = _session(lineitem) + plan = ctx.sql( + "select dfx_weighted_avg(l_extendedprice, l_quantity) from lineitem" + ).execution_plan() + assert "AggregateExec: mode=Partial" in plan.display_indent() + + result = ctx.sql( + "select dfx_weighted_avg(l_extendedprice, l_quantity) as wavg from lineitem" + ).collect()[0] + # (100*1 + 200*3 + 400*4) / (1 + 3 + 4) = 2300/8 + assert result.column(0)[0].as_py() == pytest.approx(287.5) + + +def test_the_window_function_runs(lineitem: pathlib.Path) -> None: + """A window function under a name this library owns.""" + ctx, _ = _session(lineitem) + rows = ctx.sql( + "select dfx_revenue_rank() over (order by l_extendedprice desc) as rnk " + "from lineitem order by rnk" + ).collect() + + ranks = [value for batch in rows for value in batch.column(0).to_pylist()] + assert ranks == [1, 2, 3] + + +def test_a_registered_session_never_reaches_the_codec(lineitem: pathlib.Path) -> None: + """The registry is tried first, so the codec is the fallback, not the path. + + This is the fact that makes a missing worker-side setup so easy to miss: + on the driver, where the functions are registered, the codec is never + consulted and so a codec that was broken or absent would look fine. + """ + ctx, observations = _session(lineitem) + plan = ctx.sql( + "select dfx_net_revenue(l_extendedprice, l_discount, l_tax) from lineitem" + ).execution_plan() + + ExecutionPlan.from_bytes(ctx, plan.to_bytes(ctx)) + + assert observations.decode_calls() == 0 + + +def test_the_payload_carries_no_bytes(lineitem: pathlib.Path) -> None: + """Encoded by name: no payload, so nothing to tag with a codec id. + + That is why `try_decode_udf` has to check the name before the buffer -- + with no id to route on, the chain offers the payload to every installed + codec in turn. + """ + ctx, _ = _session(lineitem) + blob = ( + ctx.sql( + "select dfx_net_revenue(l_extendedprice, l_discount, l_tax) from lineitem" + ) + .execution_plan() + .to_bytes(ctx) + ) + + assert b"dfx_net_revenue" in blob + # No chained envelope for this function: an empty encoding is not framed. + assert b"dfx_udfs.physical.v1" not in blob + + +WORKER = textwrap.dedent( + """ + import sys + from datafusion import SessionContext + from datafusion.plan import ExecutionPlan + from datafusion import udaf, udf, udwf + from dfx_udfs import ( + CodecObservations, NetRevenueUDF, RevenueRankUDWF, WeightedAvgUDAF, + ) + + blob_path, mode = sys.argv[1], sys.argv[2] + observations = CodecObservations() + ctx = SessionContext() + + if mode == "codec": + # Install the codec and register nothing. The functions are rebuilt + # from their names. + ctx = ctx.with_physical_extension_codec(observations.physical_codec()) + elif mode == "registry": + # The mirror image: register the functions, install no codec. + ctx.register_udf(udf(NetRevenueUDF())) + ctx.register_udaf(udaf(WeightedAvgUDAF())) + ctx.register_udwf(udwf(RevenueRankUDWF())) + elif mode == "neither": + pass + + ctx.register_parquet("lineitem", sys.argv[3]) + with open(blob_path, "rb") as handle: + plan = ExecutionPlan.from_bytes(ctx, handle.read()) + + total = 0.0 + for partition in range(plan.partition_count): + for batch in ctx.execute(plan, partition): + total += sum(batch.to_pyarrow().column(0).to_pylist()) + print(f"total={total:.1f} decoded={observations.decode_calls()}") + """ +) + + +def _run_worker( + tmp_path: pathlib.Path, blob: bytes, mode: str, data: pathlib.Path +) -> subprocess.CompletedProcess[str]: + blob_path = tmp_path / "plan.bin" + blob_path.write_bytes(blob) + worker = tmp_path / "worker.py" + worker.write_text(WORKER) + return subprocess.run( # noqa: S603 + [sys.executable, str(worker), str(blob_path), mode, str(data)], + capture_output=True, + text=True, + check=False, + ) + + +@pytest.fixture +def revenue_plan(lineitem: pathlib.Path) -> bytes: + ctx, _ = _session(lineitem) + return ( + ctx.sql( + "select dfx_net_revenue(l_extendedprice, l_discount, l_tax) from lineitem" + ) + .execution_plan() + .to_bytes(ctx) + ) + + +def test_a_worker_with_only_the_codec_can_run_the_plan( + revenue_plan: bytes, lineitem: pathlib.Path, tmp_path: pathlib.Path +) -> None: + """Installing the codec is an alternative to registering the functions. + + A separate process that has never registered `dfx_net_revenue` rebuilds it + from the name in the plan. + """ + result = _run_worker(tmp_path, revenue_plan, "codec", lineitem) + + assert result.returncode == 0, result.stderr + assert "total=530.0" in result.stdout + # The codec, not a registry hit, is what answered. + assert "decoded=1" in result.stdout + + +def test_a_worker_with_only_the_registrations_can_run_the_plan( + revenue_plan: bytes, lineitem: pathlib.Path, tmp_path: pathlib.Path +) -> None: + """And registering the functions is an alternative to the codec.""" + result = _run_worker(tmp_path, revenue_plan, "registry", lineitem) + + assert result.returncode == 0, result.stderr + assert "total=530.0" in result.stdout + assert "decoded=0" in result.stdout + + +def test_a_worker_with_neither_names_the_function_it_cannot_find( + revenue_plan: bytes, lineitem: pathlib.Path, tmp_path: pathlib.Path +) -> None: + """The failure is legible, and it arrives at decode rather than at execute. + + This is the whole cost of a function library that ships without a codec + and without documenting what a worker must register. + """ + result = _run_worker(tmp_path, revenue_plan, "neither", lineitem) + + assert result.returncode != 0 + assert "dfx_net_revenue" in result.stderr + + +def test_the_codec_declines_names_it_does_not_own(lineitem: pathlib.Path) -> None: + """A name-only payload reaches every codec, so declining matters. + + With no bytes there is no codec id to route on. A codec that answered for + any name it was handed would hijack another library's functions. + """ + ctx, observations = _session(lineitem, with_codecs=True) + # `abs` is a built-in, so the plan references a name this library does not + # own; decoding offers it around. + blob = ( + ctx.sql("select abs(l_discount) from lineitem").execution_plan().to_bytes(ctx) + ) + ExecutionPlan.from_bytes(ctx, blob) + + assert observations.decode_calls() == 0 + + +def test_the_codec_ids_are_pinned() -> None: + """Renaming the exporting class must not invalidate written plans.""" + observations = CodecObservations() + + assert observations.logical_codec().__datafusion_codec_id__ == "dfx_udfs.logical.v1" + assert ( + observations.physical_codec().__datafusion_codec_id__ == "dfx_udfs.physical.v1" + ) + + +def test_this_library_cannot_be_installed_as_a_bundle() -> None: + """The mixed-workflow case, asserted rather than described. + + `SessionExtensionComponents` carries codec fields only, so a function + library has nowhere to put its functions and this one does not pretend + otherwise. `with_extensions` rejects it by name. + """ + observations = CodecObservations() + + assert not hasattr(observations, "__datafusion_session_components__") + assert not hasattr(observations, "__datafusion_session_planner__") + + with pytest.raises(TypeError, match=re.escape("__datafusion_session_components__")): + SessionContext().with_extensions(observations) diff --git a/examples/distributed/udf-library/python/tests/conftest.py b/examples/distributed/udf-library/python/tests/conftest.py new file mode 100644 index 000000000..2d30c4358 --- /dev/null +++ b/examples/distributed/udf-library/python/tests/conftest.py @@ -0,0 +1,79 @@ +# 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 + +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) + + +@pytest.fixture(autouse=True) +def fail_on_log_warnings() -> Generator[None, Any, None]: + handler = _FailOnWarning() + logging.root.addHandler(handler) + yield + logging.root.removeHandler(handler) + + +@pytest.fixture +def lineitem(tmp_path: pathlib.Path) -> pathlib.Path: + """A TPC-H-shaped slice, small enough to check the arithmetic by hand. + + Two files, because the aggregate has to be correct when DataFusion splits + it into a partial pass per partition and merges the results. + """ + directory = tmp_path / "lineitem" + directory.mkdir() + pq.write_table( + pa.table( + { + "l_extendedprice": [100.0, 200.0], + "l_discount": [0.0, 0.5], + "l_tax": [0.0, 0.0], + "l_quantity": [1.0, 3.0], + } + ), + directory / "part-0.parquet", + ) + pq.write_table( + pa.table( + { + "l_extendedprice": [400.0], + "l_discount": [0.25], + "l_tax": [0.1], + "l_quantity": [4.0], + } + ), + directory / "part-1.parquet", + ) + return directory diff --git a/examples/distributed/udf-library/src/codec.rs b/examples/distributed/udf-library/src/codec.rs new file mode 100644 index 000000000..c19df4dbc --- /dev/null +++ b/examples/distributed/udf-library/src/codec.rs @@ -0,0 +1,243 @@ +// 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. + +//! Codecs that make this library's functions portable. +//! +//! A UDF library that stops at exporting functions is only usable in the +//! process that registered them. The moment a plan referencing +//! `dfx_net_revenue` is serialized and read somewhere else, *something* has to +//! turn that name back into a function, and there are exactly two candidates: +//! the receiving session's function registry, or a codec. +//! +//! Both codecs here are name-only: `try_encode_*` writes nothing, and +//! `try_decode_*` rebuilds from `name`. That shape is supported directly -- +//! an encoder that writes no bytes leaves `fun_definition` unset, and the +//! decoder then tries the registry first and the codec second. So installing +//! this library's codec on a worker is an *alternative* to registering the +//! three functions there, not an addition to it. Either is enough; neither is +//! a failure that shows up before the query runs. +//! +//! There are two codecs because there are two plan layers and this library +//! cannot know which one its callers will serialize. A distributed engine +//! shipping physical plans exercises only the physical one; `LogicalPlan. +//! to_bytes` exercises only the logical one. + +use std::fmt; +use std::sync::Arc; +use std::sync::atomic::{AtomicUsize, Ordering}; + +use datafusion::common::{Result, not_impl_err}; +use datafusion::logical_expr::{AggregateUDF, ScalarUDF, WindowUDF}; +use datafusion_proto::logical_plan::{DefaultLogicalExtensionCodec, LogicalExtensionCodec}; +use datafusion_proto::physical_plan::{DefaultPhysicalExtensionCodec, PhysicalExtensionCodec}; + +use crate::functions::{aggregate_by_name, scalar_by_name, window_by_name}; + +/// Counts, so a test can assert this codec did the work rather than infer it +/// from a query that merely succeeded. +#[derive(Default, Debug)] +pub(crate) struct CodecCounters { + pub(crate) decoded: AtomicUsize, + pub(crate) declined: AtomicUsize, +} + +/// Reject a payload for a function whose name is its whole encoding. +/// +/// Checking `name` before `buf` is the order that matters. An empty payload +/// carries no codec id, so it is the one path where a payload is offered to +/// every installed codec in turn -- meaning this hook can be called with +/// another library's function name. Trusting `buf` first would have this codec +/// answer for names it does not own. +fn reject_payload(name: &str, buf: &[u8]) -> Result<()> { + if buf.is_empty() { + return Ok(()); + } + not_impl_err!( + "{name} is encoded by name and carries no payload, but {} bytes were supplied", + buf.len() + ) +} + +macro_rules! decode_by_name { + ($self:ident, $name:expr, $buf:expr, $lookup:ident, $kind:literal) => {{ + let Some(function) = $lookup($name) else { + $self.counters.declined.fetch_add(1, Ordering::SeqCst); + return not_impl_err!("{} is not a dfx_udfs {}", $name, $kind); + }; + reject_payload($name, $buf)?; + $self.counters.decoded.fetch_add(1, Ordering::SeqCst); + Ok(function) + }}; +} + +/// Logical half. See the module docs for why there are two. +pub(crate) struct DfxUdfsLogicalCodec { + inner: DefaultLogicalExtensionCodec, + pub(crate) counters: Arc, +} + +impl DfxUdfsLogicalCodec { + pub(crate) fn new(counters: Arc) -> Self { + Self { + inner: DefaultLogicalExtensionCodec {}, + counters, + } + } +} + +impl fmt::Debug for DfxUdfsLogicalCodec { + fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result { + formatter + .debug_struct("DfxUdfsLogicalCodec") + .finish_non_exhaustive() + } +} + +impl LogicalExtensionCodec for DfxUdfsLogicalCodec { + fn try_decode( + &self, + buf: &[u8], + inputs: &[datafusion::logical_expr::LogicalPlan], + ctx: &datafusion::execution::TaskContext, + ) -> Result { + self.inner.try_decode(buf, inputs, ctx) + } + + fn try_encode( + &self, + node: &datafusion::logical_expr::Extension, + buf: &mut Vec, + ) -> Result<()> { + self.inner.try_encode(node, buf) + } + + fn try_decode_table_provider( + &self, + buf: &[u8], + table_ref: &datafusion::common::TableReference, + schema: arrow_schema::SchemaRef, + ctx: &datafusion::execution::TaskContext, + ) -> Result> { + self.inner + .try_decode_table_provider(buf, table_ref, schema, ctx) + } + + fn try_encode_table_provider( + &self, + table_ref: &datafusion::common::TableReference, + node: Arc, + buf: &mut Vec, + ) -> Result<()> { + self.inner.try_encode_table_provider(table_ref, node, buf) + } + + /// Writes nothing: returning `Ok` with an empty buffer is how a codec + /// says "encoded by name". + fn try_encode_udf(&self, _node: &ScalarUDF, _buf: &mut Vec) -> Result<()> { + Ok(()) + } + + fn try_decode_udf(&self, name: &str, buf: &[u8]) -> Result> { + decode_by_name!(self, name, buf, scalar_by_name, "scalar function") + } + + fn try_encode_udaf(&self, _node: &AggregateUDF, _buf: &mut Vec) -> Result<()> { + Ok(()) + } + + fn try_decode_udaf(&self, name: &str, buf: &[u8]) -> Result> { + decode_by_name!(self, name, buf, aggregate_by_name, "aggregate function") + } + + fn try_encode_udwf(&self, _node: &WindowUDF, _buf: &mut Vec) -> Result<()> { + Ok(()) + } + + fn try_decode_udwf(&self, name: &str, buf: &[u8]) -> Result> { + decode_by_name!(self, name, buf, window_by_name, "window function") + } +} + +/// Physical half. This is the one a distributed engine exercises, because it +/// ships physical plans. +pub(crate) struct DfxUdfsPhysicalCodec { + inner: DefaultPhysicalExtensionCodec, + pub(crate) counters: Arc, +} + +impl DfxUdfsPhysicalCodec { + pub(crate) fn new(counters: Arc) -> Self { + Self { + inner: DefaultPhysicalExtensionCodec {}, + counters, + } + } +} + +impl fmt::Debug for DfxUdfsPhysicalCodec { + fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result { + formatter + .debug_struct("DfxUdfsPhysicalCodec") + .finish_non_exhaustive() + } +} + +impl PhysicalExtensionCodec for DfxUdfsPhysicalCodec { + fn try_decode( + &self, + buf: &[u8], + inputs: &[Arc], + ctx: &datafusion::execution::TaskContext, + proto_converter: &dyn datafusion_proto::physical_plan::PhysicalProtoConverterExtension, + ) -> Result> { + // This library owns no execution plan nodes, only functions. + self.inner.try_decode(buf, inputs, ctx, proto_converter) + } + + fn try_encode( + &self, + node: Arc, + buf: &mut Vec, + proto_converter: &dyn datafusion_proto::physical_plan::PhysicalProtoConverterExtension, + ) -> Result<()> { + self.inner.try_encode(node, buf, proto_converter) + } + + fn try_encode_udf(&self, _node: &ScalarUDF, _buf: &mut Vec) -> Result<()> { + Ok(()) + } + + fn try_decode_udf(&self, name: &str, buf: &[u8]) -> Result> { + decode_by_name!(self, name, buf, scalar_by_name, "scalar function") + } + + fn try_encode_udaf(&self, _node: &AggregateUDF, _buf: &mut Vec) -> Result<()> { + Ok(()) + } + + fn try_decode_udaf(&self, name: &str, buf: &[u8]) -> Result> { + decode_by_name!(self, name, buf, aggregate_by_name, "aggregate function") + } + + fn try_encode_udwf(&self, _node: &WindowUDF, _buf: &mut Vec) -> Result<()> { + Ok(()) + } + + fn try_decode_udwf(&self, name: &str, buf: &[u8]) -> Result> { + decode_by_name!(self, name, buf, window_by_name, "window function") + } +} diff --git a/examples/distributed/udf-library/src/functions.rs b/examples/distributed/udf-library/src/functions.rs new file mode 100644 index 000000000..ec7e95eb5 --- /dev/null +++ b/examples/distributed/udf-library/src/functions.rs @@ -0,0 +1,288 @@ +// 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. + +//! The functions this library contributes. +//! +//! Deliberately arithmetic on TPC-H columns rather than anything clever: the +//! interesting part of this crate is that a function has to be *reachable* in +//! whichever process ends up evaluating it, not what the function computes. + +use std::sync::Arc; + +use arrow::array::{Array, ArrayRef, AsArray, Float64Array}; +use arrow::datatypes::{DataType, Field, FieldRef, Float64Type}; +use datafusion::common::{Result, ScalarValue, exec_err}; +use datafusion::logical_expr::function::{ + AccumulatorArgs, PartitionEvaluatorArgs, StateFieldsArgs, WindowUDFFieldArgs, +}; +use datafusion::logical_expr::{ + Accumulator, AggregateUDF, AggregateUDFImpl, ColumnarValue, PartitionEvaluator, + ScalarFunctionArgs, ScalarUDF, ScalarUDFImpl, Signature, Volatility, WindowUDF, WindowUDFImpl, +}; +use datafusion_functions_window::rank::rank_udwf; + +/// Name of the scalar function, used by the codec as the whole encoding. +pub(crate) const NET_REVENUE: &str = "dfx_net_revenue"; +/// Name of the aggregate function. +pub(crate) const WEIGHTED_AVG: &str = "dfx_weighted_avg"; +/// Name of the window function. +pub(crate) const REVENUE_RANK: &str = "dfx_revenue_rank"; + +/// `extendedprice * (1 - discount) * (1 + tax)`, the TPC-H revenue expression. +#[derive(Debug, Clone, PartialEq, Eq, Hash)] +pub(crate) struct NetRevenue { + signature: Signature, +} + +impl Default for NetRevenue { + fn default() -> Self { + Self { + signature: Signature::exact( + vec![DataType::Float64, DataType::Float64, DataType::Float64], + Volatility::Immutable, + ), + } + } +} + +impl ScalarUDFImpl for NetRevenue { + fn name(&self) -> &str { + NET_REVENUE + } + + fn signature(&self) -> &Signature { + &self.signature + } + + fn return_type(&self, _arg_types: &[DataType]) -> Result { + Ok(DataType::Float64) + } + + fn invoke_with_args(&self, args: ScalarFunctionArgs) -> Result { + let arrays = ColumnarValue::values_to_arrays(&args.args)?; + let [price, discount, tax] = arrays.as_slice() else { + return exec_err!("{NET_REVENUE} takes 3 arguments, got {}", arrays.len()); + }; + let price = price.as_primitive::(); + let discount = discount.as_primitive::(); + let tax = tax.as_primitive::(); + + let values: Float64Array = (0..price.len()) + .map(|row| { + if price.is_null(row) || discount.is_null(row) || tax.is_null(row) { + return None; + } + Some(price.value(row) * (1.0 - discount.value(row)) * (1.0 + tax.value(row))) + }) + .collect(); + Ok(ColumnarValue::Array(Arc::new(values))) + } +} + +/// `sum(value * weight) / sum(weight)`. +/// +/// Written out rather than delegating to a built-in because the state is the +/// point: two running sums, which is what lets DataFusion compute this in a +/// partial aggregate on each worker and merge the results on the driver. An +/// aggregate that could only be evaluated over the whole input at once would +/// not survive being split across processes. +#[derive(Debug, Clone, PartialEq, Eq, Hash)] +pub(crate) struct WeightedAvg { + signature: Signature, +} + +impl Default for WeightedAvg { + fn default() -> Self { + Self { + signature: Signature::exact( + vec![DataType::Float64, DataType::Float64], + Volatility::Immutable, + ), + } + } +} + +impl AggregateUDFImpl for WeightedAvg { + fn name(&self) -> &str { + WEIGHTED_AVG + } + + fn signature(&self) -> &Signature { + &self.signature + } + + fn return_type(&self, _arg_types: &[DataType]) -> Result { + Ok(DataType::Float64) + } + + fn accumulator(&self, _args: AccumulatorArgs) -> Result> { + Ok(Box::new(WeightedAvgAccumulator::default())) + } + + /// The two partial sums, in the order [`WeightedAvgAccumulator::state`] + /// returns them. + fn state_fields(&self, args: StateFieldsArgs) -> Result> { + Ok(vec![ + Arc::new(Field::new( + format!("{}[weighted_sum]", args.name), + DataType::Float64, + false, + )), + Arc::new(Field::new( + format!("{}[weight_sum]", args.name), + DataType::Float64, + false, + )), + ]) + } +} + +#[derive(Debug, Default)] +struct WeightedAvgAccumulator { + weighted_sum: f64, + weight_sum: f64, +} + +impl Accumulator for WeightedAvgAccumulator { + fn update_batch(&mut self, values: &[ArrayRef]) -> Result<()> { + let [value, weight] = values else { + return exec_err!("{WEIGHTED_AVG} takes 2 arguments, got {}", values.len()); + }; + let value = value.as_primitive::(); + let weight = weight.as_primitive::(); + for row in 0..value.len() { + // A null in either argument contributes to neither sum, so the + // result is the weighted average of the rows that had both. + if value.is_null(row) || weight.is_null(row) { + continue; + } + self.weighted_sum += value.value(row) * weight.value(row); + self.weight_sum += weight.value(row); + } + Ok(()) + } + + /// Merge partial states, which is the step that runs on the driver over + /// results computed on the workers. + fn merge_batch(&mut self, states: &[ArrayRef]) -> Result<()> { + let [weighted_sum, weight_sum] = states else { + return exec_err!("{WEIGHTED_AVG} has 2 state columns, got {}", states.len()); + }; + let weighted_sum = weighted_sum.as_primitive::(); + let weight_sum = weight_sum.as_primitive::(); + for row in 0..weighted_sum.len() { + self.weighted_sum += weighted_sum.value(row); + self.weight_sum += weight_sum.value(row); + } + Ok(()) + } + + fn state(&mut self) -> Result> { + Ok(vec![ + ScalarValue::Float64(Some(self.weighted_sum)), + ScalarValue::Float64(Some(self.weight_sum)), + ]) + } + + fn evaluate(&mut self) -> Result { + // No rows, or every weight zero: null rather than a division by zero, + // matching what `avg` does for an empty input. + if self.weight_sum == 0.0 { + return Ok(ScalarValue::Float64(None)); + } + Ok(ScalarValue::Float64(Some( + self.weighted_sum / self.weight_sum, + ))) + } + + fn size(&self) -> usize { + std::mem::size_of_val(self) + } +} + +/// Ranks rows within a window, under a name this library owns. +/// +/// Delegates to the built-in `rank`: the reason it is here is to give the +/// library a window function whose *name* has to resolve on whichever process +/// evaluates it, which is the same portability question the other two raise. +#[derive(Debug, Clone)] +pub(crate) struct RevenueRank { + inner: Arc, +} + +impl Default for RevenueRank { + fn default() -> Self { + Self { inner: rank_udwf() } + } +} + +impl PartialEq for RevenueRank { + fn eq(&self, other: &Self) -> bool { + self.inner == other.inner + } +} + +impl Eq for RevenueRank {} + +impl std::hash::Hash for RevenueRank { + fn hash(&self, state: &mut H) { + self.inner.hash(state); + } +} + +impl WindowUDFImpl for RevenueRank { + fn name(&self) -> &str { + REVENUE_RANK + } + + fn signature(&self) -> &Signature { + self.inner.signature() + } + + fn partition_evaluator( + &self, + args: PartitionEvaluatorArgs, + ) -> Result> { + self.inner.inner().partition_evaluator(args) + } + + fn field(&self, field_args: WindowUDFFieldArgs) -> Result { + self.inner.inner().field(field_args) + } + + fn coerce_types(&self, arg_types: &[DataType]) -> Result> { + self.inner.coerce_types(arg_types) + } +} + +/// Rebuild one of this library's functions from its name alone. +/// +/// This is the whole decode path: the names are the encoding, so a process +/// that has this library's codec installed can reconstruct any of them +/// without the driver having sent bytes and without the function having been +/// registered locally. +pub(crate) fn scalar_by_name(name: &str) -> Option> { + (name == NET_REVENUE).then(|| Arc::new(ScalarUDF::from(NetRevenue::default()))) +} + +pub(crate) fn aggregate_by_name(name: &str) -> Option> { + (name == WEIGHTED_AVG).then(|| Arc::new(AggregateUDF::from(WeightedAvg::default()))) +} + +pub(crate) fn window_by_name(name: &str) -> Option> { + (name == REVENUE_RANK).then(|| Arc::new(WindowUDF::from(RevenueRank::default()))) +} diff --git a/examples/distributed/udf-library/src/lib.rs b/examples/distributed/udf-library/src/lib.rs new file mode 100644 index 000000000..95fa87a35 --- /dev/null +++ b/examples/distributed/udf-library/src/lib.rs @@ -0,0 +1,46 @@ +// 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. + +//! Example function library: a scalar UDF, an aggregate, a window function, +//! and the two codecs that let plans referencing them be read elsewhere. +//! +//! One of three libraries in `examples/distributed`. This one owns functions, +//! and is the one that cannot be installed with `with_extensions` -- see +//! [`crate::python`]. + +use pyo3::prelude::*; + +use crate::python::{ + PyCodecObservations, PyLogicalCodec, PyNetRevenue, PyPhysicalCodec, PyRevenueRank, + PyWeightedAvg, +}; + +mod codec; +mod functions; +mod python; + +#[pymodule] +fn dfx_udfs(m: &Bound<'_, PyModule>) -> PyResult<()> { + pyo3_log::init(); + m.add_class::()?; + m.add_class::()?; + m.add_class::()?; + m.add_class::()?; + m.add_class::()?; + m.add_class::()?; + Ok(()) +} diff --git a/examples/distributed/udf-library/src/python.rs b/examples/distributed/udf-library/src/python.rs new file mode 100644 index 000000000..397aa6727 --- /dev/null +++ b/examples/distributed/udf-library/src/python.rs @@ -0,0 +1,231 @@ +// 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. + +//! The Python surface, and the one thing it deliberately lacks. +//! +//! This library exposes **no** `__datafusion_session_components__`, so it +//! cannot be installed with `SessionContext.with_extensions`. Callers register +//! the three functions and install the two codecs by hand, which is the older +//! and more error-prone path -- and currently the honest one for a function +//! library, because `SessionExtensionComponents` carries codec fields only. +//! There is nowhere for a UDF to go. +//! +//! Keeping one library on the manual path is the point: a real deployment +//! mixes libraries built against different versions of the protocol, and the +//! example should show what that costs rather than pretend every dependency +//! has caught up. + +use std::sync::Arc; +use std::sync::atomic::Ordering; + +use datafusion::logical_expr::{AggregateUDF, ScalarUDF, WindowUDF}; +use datafusion_ffi::proto::logical_extension_codec::FFI_LogicalExtensionCodec; +use datafusion_ffi::proto::physical_extension_codec::FFI_PhysicalExtensionCodec; +use datafusion_ffi::udaf::FFI_AggregateUDF; +use datafusion_ffi::udf::FFI_ScalarUDF; +use datafusion_ffi::udwf::FFI_WindowUDF; +use datafusion_proto::logical_plan::LogicalExtensionCodec; +use datafusion_proto::physical_plan::PhysicalExtensionCodec; +use datafusion_python_util::{ + create_logical_extension_capsule, create_physical_extension_capsule, + ffi_task_context_provider_from_pycapsule, get_tokio_runtime, +}; +use pyo3::prelude::*; +use pyo3::types::PyCapsule; + +use crate::codec::{CodecCounters, DfxUdfsLogicalCodec, DfxUdfsPhysicalCodec}; +use crate::functions::{NetRevenue, RevenueRank, WeightedAvg}; + +/// Wire ids, pinned so a rename of the exporting class cannot invalidate +/// plans already written. +const LOGICAL_CODEC_ID: &str = "dfx_udfs.logical.v1"; +const PHYSICAL_CODEC_ID: &str = "dfx_udfs.physical.v1"; + +/// `dfx_net_revenue(extendedprice, discount, tax)`. +#[pyclass(name = "NetRevenueUDF", module = "dfx_udfs")] +#[derive(Default)] +pub(crate) struct PyNetRevenue; + +#[pymethods] +impl PyNetRevenue { + #[new] + fn new() -> Self { + Self + } + + fn __datafusion_scalar_udf__<'py>(&self, py: Python<'py>) -> PyResult> { + let func = Arc::new(ScalarUDF::from(NetRevenue::default())); + PyCapsule::new_with_value(py, FFI_ScalarUDF::from(func), cr"datafusion_scalar_udf") + } +} + +/// `dfx_weighted_avg(value, weight)`. +#[pyclass(name = "WeightedAvgUDAF", module = "dfx_udfs")] +#[derive(Default)] +pub(crate) struct PyWeightedAvg; + +#[pymethods] +impl PyWeightedAvg { + #[new] + fn new() -> Self { + Self + } + + fn __datafusion_aggregate_udf__<'py>( + &self, + py: Python<'py>, + ) -> PyResult> { + let func = Arc::new(AggregateUDF::from(WeightedAvg::default())); + PyCapsule::new_with_value( + py, + FFI_AggregateUDF::from(func), + cr"datafusion_aggregate_udf", + ) + } +} + +/// `dfx_revenue_rank()`, as a window function. +#[pyclass(name = "RevenueRankUDWF", module = "dfx_udfs")] +#[derive(Default)] +pub(crate) struct PyRevenueRank; + +#[pymethods] +impl PyRevenueRank { + #[new] + fn new() -> Self { + Self + } + + fn __datafusion_window_udf__<'py>(&self, py: Python<'py>) -> PyResult> { + let func = Arc::new(WindowUDF::from(RevenueRank::default())); + PyCapsule::new_with_value(py, FFI_WindowUDF::from(func), cr"datafusion_window_udf") + } +} + +/// Shared decode counters, so a test can see which codec answered. +#[pyclass(from_py_object, name = "CodecObservations", module = "dfx_udfs")] +#[derive(Default, Clone)] +pub(crate) struct PyCodecObservations { + counters: Arc, +} + +#[pymethods] +impl PyCodecObservations { + #[new] + fn new() -> Self { + Self::default() + } + + /// How often a codec rebuilt one of this library's functions from its name. + /// + /// Zero after a successful query on a session that registered the + /// functions: the registry is tried first, so the codec is only reached + /// when the receiving session does *not* have them. + fn decode_calls(&self) -> usize { + self.counters.decoded.load(Ordering::SeqCst) + } + + /// How often a codec was asked about a name it does not own. + /// + /// Non-zero is expected. A name-only payload has no codec id to route on, + /// so it is offered to every installed codec in turn. + fn declined_calls(&self) -> usize { + self.counters.declined.load(Ordering::SeqCst) + } + + /// Build the logical codec, sharing these counters. + fn logical_codec(&self) -> PyLogicalCodec { + PyLogicalCodec { + counters: Arc::clone(&self.counters), + } + } + + /// Build the physical codec, sharing these counters. + fn physical_codec(&self) -> PyPhysicalCodec { + PyPhysicalCodec { + counters: Arc::clone(&self.counters), + } + } +} + +/// Install with `ctx.with_logical_extension_codec(...)`. +#[pyclass(name = "DfxUdfsLogicalCodec", module = "dfx_udfs")] +pub(crate) struct PyLogicalCodec { + counters: Arc, +} + +#[pymethods] +impl PyLogicalCodec { + #[new] + fn new() -> Self { + Self { + counters: Arc::default(), + } + } + + #[getter] + fn __datafusion_codec_id__(&self) -> &'static str { + LOGICAL_CODEC_ID + } + + fn __datafusion_logical_extension_codec__<'py>( + &self, + py: Python<'py>, + session: Bound<'py, PyAny>, + ) -> PyResult> { + let provider = ffi_task_context_provider_from_pycapsule(&session)?; + let runtime = get_tokio_runtime().handle().clone(); + let codec: Arc = + Arc::new(DfxUdfsLogicalCodec::new(Arc::clone(&self.counters))); + let ffi = FFI_LogicalExtensionCodec::new(codec, Some(runtime), provider); + create_logical_extension_capsule(py, &ffi) + } +} + +/// Install with `ctx.with_physical_extension_codec(...)`. +#[pyclass(name = "DfxUdfsPhysicalCodec", module = "dfx_udfs")] +pub(crate) struct PyPhysicalCodec { + counters: Arc, +} + +#[pymethods] +impl PyPhysicalCodec { + #[new] + fn new() -> Self { + Self { + counters: Arc::default(), + } + } + + #[getter] + fn __datafusion_codec_id__(&self) -> &'static str { + PHYSICAL_CODEC_ID + } + + fn __datafusion_physical_extension_codec__<'py>( + &self, + py: Python<'py>, + session: Bound<'py, PyAny>, + ) -> PyResult> { + let provider = ffi_task_context_provider_from_pycapsule(&session)?; + let runtime = get_tokio_runtime().handle().clone(); + let codec: Arc = + Arc::new(DfxUdfsPhysicalCodec::new(Arc::clone(&self.counters))); + let ffi = FFI_PhysicalExtensionCodec::new(codec, Some(runtime), provider); + create_physical_extension_capsule(py, &ffi) + } +} From dba062c5c81db4c2e97a81c1aa89ece81a2117f6 Mon Sep 17 00:00:00 2001 From: Tim Saucer Date: Wed, 9 Sep 2026 13:00:51 -0400 Subject: [PATCH 03/24] Add dfx_engine: a toy distributed engine in both halves MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Third library for #1719, and the one that makes the other two do something. Rust owns the query planner, the stage node, its codec, and a config extension; Python owns the session factory, the driver, and the worker entry point. A real engine needs both, so this crate is a mixed maturin package rather than a pure extension module. The split point is the partial aggregate. DataFusion already breaks a GROUP BY into a partial pass per input partition and a final pass that merges them, so the partial passes are independent by construction and only their output has to come back. Wrapping that subtree in a `ShuffleStageExec` is the whole rewrite. The planner plans against `LocalOptimizerSession`, which borrows the foreign session but owns the stock optimizer rule list. Without it the rules run back across FFI and hand the library `ForeignExecutionPlan` wrappers, which cannot be serialized (that is G1) and cannot be rewritten either — an engine cannot split a subtree it holds only an opaque handle to. This was validated as a spike before any of it was built. One node does both halves of the shuffle. `execute(i)` reads the file for partition `i` if it exists and otherwise computes its child and writes it on the way past, so the same node is the thing a worker runs and the thing the driver reads, and nothing has to rewrite the plan in between. A query with no workers still gets the right answer, having done the work itself. The shuffle directory travels inside the node and therefore inside its encoding, so a worker and a driver cannot disagree about where results go. `session.py` is the piece the whole example exists to motivate. There is no way to snapshot a SessionContext and restore it elsewhere — SessionConfig is write-only from Python, and `df_settings` is readable but lists `datafusion.runtime.*` keys that have no namespace to set them back into — so worker parity cannot be automated. It has to be built the same way twice from data small enough to put in a message, which is what `SessionSpec` is. Both sides call `build_session`; anything a query depends on that is not in the spec is a bug waiting for a worker to find it. Two findings this turned up, both now documented in the code: `dfx_storage` needed a *logical* codec, not just a physical one. Its scan node is physical, so a physical codec looks sufficient — but installing any FFI query planner means the session hands that planner the logical plan as protobuf, and a logical plan holds its tables as `Arc`. With no `try_encode_table_provider` the session fails at `execution_plan()` with "Error serializing custom table", before anything is distributed. The payload is the directory, since everything else the provider holds is read back from it. A foreign node does not print its own name. The host shows `FFI_ExecutionPlan: ShuffleStageExec, number_of_children=1`, so the driver's tree walk has to match on containment; an anchored match works in a single-library test and fails the moment a real extension is involved. Verified end to end: three Parquet files, three worker processes, each producing one partition of partial aggregate, driver merging them to the same answer the single-process path gives. Corrupting one shuffle file breaks the driver's query, which is how we know the workers did the work rather than the driver quietly recomputing it. Co-Authored-By: Claude Opus 5 (1M context) --- Cargo.lock | 18 ++ Cargo.toml | 1 + .../distributed/engine-library/Cargo.toml | 51 ++++ examples/distributed/engine-library/build.rs | 20 ++ .../distributed/engine-library/pyproject.toml | 34 +++ .../python/dfx_engine/__init__.py | 40 +++ .../python/dfx_engine/driver.py | 206 +++++++++++++++ .../python/dfx_engine/session.py | 170 +++++++++++++ .../python/dfx_engine/worker.py | 113 +++++++++ .../distributed/engine-library/src/codec.rs | 136 ++++++++++ .../distributed/engine-library/src/config.rs | 113 +++++++++ .../engine-library/src/extension.rs | 200 +++++++++++++++ .../distributed/engine-library/src/lib.rs | 70 ++++++ .../engine-library/src/local_session.rs | 161 ++++++++++++ .../distributed/engine-library/src/planner.rs | 170 +++++++++++++ .../distributed/engine-library/src/stage.rs | 236 ++++++++++++++++++ .../engine-library/stage-1-part-0.arrow | Bin 0 -> 1032 bytes .../engine-library/stage-1-part-1.arrow | Bin 0 -> 1032 bytes .../engine-library/stage-1-part-2.arrow | Bin 0 -> 1032 bytes .../distributed/storage-library/src/codec.rs | 110 +++++++- .../storage-library/src/extension.rs | 49 +++- .../storage-library/src/table_provider.rs | 5 + 22 files changed, 1897 insertions(+), 6 deletions(-) create mode 100644 examples/distributed/engine-library/Cargo.toml create mode 100644 examples/distributed/engine-library/build.rs create mode 100644 examples/distributed/engine-library/pyproject.toml create mode 100644 examples/distributed/engine-library/python/dfx_engine/__init__.py create mode 100644 examples/distributed/engine-library/python/dfx_engine/driver.py create mode 100644 examples/distributed/engine-library/python/dfx_engine/session.py create mode 100644 examples/distributed/engine-library/python/dfx_engine/worker.py create mode 100644 examples/distributed/engine-library/src/codec.rs create mode 100644 examples/distributed/engine-library/src/config.rs create mode 100644 examples/distributed/engine-library/src/extension.rs create mode 100644 examples/distributed/engine-library/src/lib.rs create mode 100644 examples/distributed/engine-library/src/local_session.rs create mode 100644 examples/distributed/engine-library/src/planner.rs create mode 100644 examples/distributed/engine-library/src/stage.rs create mode 100644 examples/distributed/engine-library/stage-1-part-0.arrow create mode 100644 examples/distributed/engine-library/stage-1-part-1.arrow create mode 100644 examples/distributed/engine-library/stage-1-part-2.arrow diff --git a/Cargo.lock b/Cargo.lock index 6ae37a13f..d946352c5 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -1666,6 +1666,24 @@ dependencies = [ "url", ] +[[package]] +name = "dfx-engine" +version = "54.0.0" +dependencies = [ + "arrow", + "async-trait", + "datafusion", + "datafusion-common", + "datafusion-ffi", + "datafusion-proto", + "datafusion-python-util", + "datafusion-session", + "futures", + "pyo3", + "pyo3-build-config", + "pyo3-log", +] + [[package]] name = "dfx-storage" version = "54.0.0" diff --git a/Cargo.toml b/Cargo.toml index fb7db2744..14f34e599 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -34,6 +34,7 @@ members = [ "examples/datafusion-ffi-query-planner-example", "examples/distributed/storage-library", "examples/distributed/udf-library", + "examples/distributed/engine-library", ] resolver = "3" diff --git a/examples/distributed/engine-library/Cargo.toml b/examples/distributed/engine-library/Cargo.toml new file mode 100644 index 000000000..c26d326db --- /dev/null +++ b/examples/distributed/engine-library/Cargo.toml @@ -0,0 +1,51 @@ +# 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. + +[package] +name = "dfx-engine" +version.workspace = true +edition.workspace = true +rust-version.workspace = true +license.workspace = true +description = "Example extension library: a toy distributed engine that splits plans into stages" +homepage.workspace = true +repository.workspace = true +publish = false + +[dependencies] +arrow = { workspace = true } +async-trait = { workspace = true } +datafusion = { workspace = true } +datafusion-common = { workspace = true, default-features = false } +datafusion-ffi = { workspace = true } +datafusion-proto = { workspace = true } +datafusion-python-util.workspace = true +datafusion-session = { workspace = true } +futures = { workspace = true } +pyo3 = { workspace = true, features = [ + "extension-module", + "abi3", + "abi3-py310", +] } +pyo3-log = { workspace = true } + +[build-dependencies] +pyo3-build-config = { workspace = true } + +[lib] +name = "_internal" +crate-type = ["cdylib", "rlib"] diff --git a/examples/distributed/engine-library/build.rs b/examples/distributed/engine-library/build.rs new file mode 100644 index 000000000..4878d8b0e --- /dev/null +++ b/examples/distributed/engine-library/build.rs @@ -0,0 +1,20 @@ +// 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. + +fn main() { + pyo3_build_config::add_extension_module_link_args(); +} diff --git a/examples/distributed/engine-library/pyproject.toml b/examples/distributed/engine-library/pyproject.toml new file mode 100644 index 000000000..ca9de9818 --- /dev/null +++ b/examples/distributed/engine-library/pyproject.toml @@ -0,0 +1,34 @@ +# 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. + +[build-system] +requires = ["maturin>=1.6,<2.0"] +build-backend = "maturin" + +[project] +name = "dfx_engine" +requires-python = ">=3.10" +classifiers = [ + "Programming Language :: Rust", + "Programming Language :: Python :: Implementation :: CPython", +] +dynamic = ["version"] + +[tool.maturin] +features = ["pyo3/extension-module"] +python-source = "python" +module-name = "dfx_engine._internal" diff --git a/examples/distributed/engine-library/python/dfx_engine/__init__.py b/examples/distributed/engine-library/python/dfx_engine/__init__.py new file mode 100644 index 000000000..d98cbf1d8 --- /dev/null +++ b/examples/distributed/engine-library/python/dfx_engine/__init__.py @@ -0,0 +1,40 @@ +# 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. + +"""A toy distributed engine, as an extension library. + +Two halves, because a real engine has two: the Rust side owns the query +planner, the stage node, the codec that carries it, and a config extension; +the Python side owns the session factory, the driver, and the worker entry +point. + +Start with :mod:`dfx_engine.session` -- ``build_session`` is the piece the +rest of the example exists to motivate. +""" + +from dfx_engine import _internal +from dfx_engine._internal import DfxEngineConfig, DfxEngineExtension +from dfx_engine.session import SessionSpec, build_session, expected_codec_ids + +__all__ = [ + "DfxEngineConfig", + "DfxEngineExtension", + "SessionSpec", + "_internal", + "build_session", + "expected_codec_ids", +] diff --git a/examples/distributed/engine-library/python/dfx_engine/driver.py b/examples/distributed/engine-library/python/dfx_engine/driver.py new file mode 100644 index 000000000..baa145f15 --- /dev/null +++ b/examples/distributed/engine-library/python/dfx_engine/driver.py @@ -0,0 +1,206 @@ +# 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. + +"""The driver: split a query into tasks, fan them out, collect the answer. + +The shape is deliberately boring, because the interesting part is not the +scheduling. What matters is the four things the driver has to get right, each +of which is a way a real deployment goes wrong: + +1. It serializes the stage **with** its session. ``to_bytes(None)`` uses an + empty codec chain and cannot encode any library's node. +2. It ships the codec ids it used, so a worker can refuse a plan it would + misread rather than decode it with the wrong codec. +3. It puts the shuffle directory in the session config, not in the message, + so the directory travels *inside* the encoded plan and the two sides + cannot disagree. +4. It waits for every worker before reading, because the stage node decides + whether to read or recompute by looking at the filesystem. +""" + +from __future__ import annotations + +import json +import pathlib +import subprocess +import sys +from typing import TYPE_CHECKING + +from dfx_engine import _internal +from dfx_engine.session import SessionSpec, build_session + +if TYPE_CHECKING: + import pyarrow as pa + from datafusion import DataFrame, SessionContext + from datafusion.plan import ExecutionPlan + +__all__ = ["DistributedResult", "find_stage", "run_distributed"] + + +class DistributedResult: + """What a distributed run produced, and how.""" + + def __init__( + self, + batches: list[pa.RecordBatch], + partitions: list[int], + worker_rows: dict[int, int], + ) -> None: + self.batches = batches + self.partitions = partitions + """Partition indices that were dispatched, one per worker.""" + self.worker_rows = worker_rows + """Rows each worker produced, keyed by partition index.""" + + +STAGE_NODE_NAME = "ShuffleStageExec" + + +def find_stage(plan: ExecutionPlan) -> ExecutionPlan | None: + """Locate the stage node the planner inserted. + + Matched on the display string because a Python caller has no way to + downcast a Rust plan node -- there is no ``isinstance`` across an FFI + boundary. + + Note the *containment* test. The node was built inside this library and + handed back to the host, so what the host prints is not + ``ShuffleStageExec: stage=1`` but:: + + FFI_ExecutionPlan: ShuffleStageExec, number_of_children=1 + + A foreign node reports its own name nested inside the wrapper's, which + makes anchored matches on plan text quietly wrong -- the kind of thing + that works in a single-library test and fails the moment a real extension + is involved. + """ + if STAGE_NODE_NAME in plan.display(): + return plan + for child in plan.children(): + found = find_stage(child) + if found is not None: + return found + return None + + +def _dispatch( + envelope: dict, envelope_dir: pathlib.Path, partition: int +) -> subprocess.Popen[str]: + """Start one worker for one partition. + + ``sys.executable``, not ``python``: a worker on a different Python minor + version cannot load a cloudpickled inline UDF, and that failure is far + from its cause. + """ + path = envelope_dir / f"task-{partition}.json" + path.write_text(json.dumps(envelope)) + return subprocess.Popen( # noqa: S603 + [sys.executable, "-m", "dfx_engine.worker", str(path)], + stdout=subprocess.PIPE, + stderr=subprocess.PIPE, + text=True, + ) + + +def run_distributed(sql: str, spec: SessionSpec) -> 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. + """ + 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) + plan = ctx.sql(sql).execution_plan() + + stage = find_stage(plan) + if stage is None: + message = ( + "no ShuffleStageExec in the plan; the engine's planner did not run, " + "or its config extension was not registered" + ) + raise RuntimeError(message) + + shuffle_dir = pathlib.Path(spec.shuffle_dir) + shuffle_dir.mkdir(parents=True, exist_ok=True) + + # Encode the stage subtree, through the session that owns the codecs. + plan_path = shuffle_dir / "stage.plan" + plan_path.write_bytes(stage.to_bytes(ctx)) + + stage_id = _internal.stage_id() + partitions = list(range(stage.partition_count)) + envelopes = [ + { + "spec": spec.to_json(), + "plan": str(plan_path), + "stage_id": stage_id, + "partition": partition, + } + for partition in partitions + ] + + # One process per partition, all in flight together. This is the claim the + # example is making: each worker reads a different file and writes a + # different result, so they need no coordination beyond the directory. + workers = [ + _dispatch(envelope, shuffle_dir, partition) + for envelope, partition in zip(envelopes, partitions, strict=True) + ] + + worker_rows: dict[int, int] = {} + failures = [] + for partition, worker in zip(partitions, workers, strict=True): + stdout, stderr = worker.communicate() + if worker.returncode != 0: + failures.append(f"partition {partition} failed:\n{stderr}") + continue + worker_rows[partition] = json.loads(stdout)["rows"] + + if failures: + raise RuntimeError("\n".join(failures)) + + # Now run the whole query here. Every stage partition has a file, so the + # stage node streams them instead of recomputing -- the driver does only + # the final merge. + batches = ctx.sql(sql).collect() + _ = engine + return DistributedResult(batches, partitions, worker_rows) + + +def run_local(sql: str, spec: SessionSpec) -> 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. + """ + ctx, _engine, _storage = build_session( + SessionSpec( + tables=spec.tables, + shuffle_dir="", + target_partitions=spec.target_partitions, + ) + ) + return ctx.sql(sql).collect() + + +def dataframe_for(sql: str, spec: SessionSpec) -> tuple[SessionContext, DataFrame]: + """Session and DataFrame for `sql`, for tests that want to inspect a plan.""" + ctx, _engine, _storage = build_session(spec) + return ctx, ctx.sql(sql) diff --git a/examples/distributed/engine-library/python/dfx_engine/session.py b/examples/distributed/engine-library/python/dfx_engine/session.py new file mode 100644 index 000000000..14390d01f --- /dev/null +++ b/examples/distributed/engine-library/python/dfx_engine/session.py @@ -0,0 +1,170 @@ +# 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. + +"""One session factory, used by the driver and by every worker. + +This module is the answer to the question the rest of the example exists to +raise: *what exactly does a worker have to reproduce?* + +There is no way to snapshot a :class:`~datafusion.SessionContext` and restore +it somewhere else. :class:`~datafusion.SessionConfig` is write-only from +Python, and ``information_schema.df_settings`` -- which can be read -- lists +``datafusion.runtime.*`` keys that have no config namespace to set them back +into. So worker parity cannot be automated away; it has to be *built the same +way twice*, from data small enough to put in a message. + +That is what :class:`SessionSpec` is, and why both sides call +:func:`build_session` rather than each assembling a context of their own. +Anything a query depends on that is not in the spec is a bug waiting for a +worker to find it. +""" + +from __future__ import annotations + +import dataclasses +from typing import TYPE_CHECKING + +import dfx_storage +import dfx_udfs +from datafusion import SessionConfig, SessionContext, udaf, udf, udwf + +from dfx_engine import _internal + +if TYPE_CHECKING: + from collections.abc import Mapping + +__all__ = ["SessionSpec", "build_session", "expected_codec_ids"] + + +@dataclasses.dataclass(frozen=True) +class SessionSpec: + """Everything needed to rebuild an equivalent session. + + Small and explicit on purpose: it travels to workers as JSON, so anything + that cannot be written down here cannot be relied on by a shipped plan. + """ + + tables: Mapping[str, str] + """Table name to the directory ``dfx_storage`` should scan for it.""" + + shuffle_dir: str + """Where stages exchange results. Empty means "run in this process".""" + + target_partitions: int = 2 + """Pinned rather than defaulted to the core count. + + Two machines with different core counts would otherwise disagree about + how many partitions a re-planned query has. + """ + + def to_json(self) -> dict: + """Render for a worker's task envelope.""" + return { + "tables": dict(self.tables), + "shuffle_dir": self.shuffle_dir, + "target_partitions": self.target_partitions, + "codec_ids": expected_codec_ids(), + } + + @staticmethod + def from_json(payload: Mapping) -> SessionSpec: + """Rebuild from a task envelope, ignoring the codec ids. + + The ids are checked against the session after it is built rather than + used to build it -- see :func:`build_session`. + """ + return SessionSpec( + tables=payload["tables"], + shuffle_dir=payload["shuffle_dir"], + target_partitions=payload["target_partitions"], + ) + + +def expected_codec_ids() -> list[str]: + """The physical codec ids a correctly-built session carries. + + Read from the libraries rather than written out here, so adding a library + to :func:`build_session` and forgetting this list is not possible. + """ + return sorted( + [ + dfx_storage.DfxStorageExtension.physical_codec_id(), + _internal.DfxEngineExtension.physical_codec_id(), + "dfx_udfs.physical.v1", + ] + ) + + +def build_session( + spec: SessionSpec, +) -> tuple[ + SessionContext, _internal.DfxEngineExtension, dfx_storage.DfxStorageExtension +]: + """Build the session both the driver and the workers run on. + + The order below is not arbitrary: + + 1. The engine's config extension is registered on the ``SessionConfig`` + *before* the context exists, because ``dfx_engine.shuffle_dir`` cannot + be set into a namespace that has not been declared. + 2. The two bundle libraries go in through a single + :meth:`~datafusion.SessionContext.with_extensions` call, so their + codecs are all installed before the engine's planner is bound. Passing + them in separate calls would bind the planner against a partial chain. + 3. ``dfx_udfs`` is installed by hand, because it ships no bundle hook. + Its codecs must go on before anything serializes a plan referencing its + functions. + 4. Tables are registered last. Registration order does not matter to + ``with_extensions``, but doing it after means the same code path builds + a driver and a worker. + + Returns the context plus the two bundles, whose counters let a test assert + which codec carried which node. + """ + engine = _internal.DfxEngineExtension() + storage = dfx_storage.DfxStorageExtension() + + config = SessionConfig().with_target_partitions(spec.target_partitions) + # Declares the `dfx_engine` namespace. Without this, setting + # `dfx_engine.shuffle_dir` raises rather than being ignored. + config = config.with_extension(_internal.DfxEngineConfig(spec.shuffle_dir)) + + ctx = SessionContext(config) + ctx = ctx.with_extensions(storage, engine) + + # The manual path, for the library that has no bundle. Two codec installs + # and three registrations, in place of one call. + observations = dfx_udfs.CodecObservations() + ctx = ctx.with_logical_extension_codec(observations.logical_codec()) + ctx = ctx.with_physical_extension_codec(observations.physical_codec()) + ctx.register_udf(udf(dfx_udfs.NetRevenueUDF())) + ctx.register_udaf(udaf(dfx_udfs.WeightedAvgUDAF())) + ctx.register_udwf(udwf(dfx_udfs.RevenueRankUDWF())) + + for name, directory in spec.tables.items(): + ctx.register_table(name, dfx_storage.PartitionedParquetTable(directory)) + + installed = sorted(ctx.physical_extension_codec_ids()) + expected = expected_codec_ids() + if installed != expected: + message = ( + f"session codec ids {installed} do not match the expected " + f"{expected}; a plan encoded elsewhere will fail to decode" + ) + raise RuntimeError(message) + + return ctx, engine, storage diff --git a/examples/distributed/engine-library/python/dfx_engine/worker.py b/examples/distributed/engine-library/python/dfx_engine/worker.py new file mode 100644 index 000000000..0250bf104 --- /dev/null +++ b/examples/distributed/engine-library/python/dfx_engine/worker.py @@ -0,0 +1,113 @@ +# 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. + +"""One worker: rebuild the session, decode one stage, run one partition. + +Run as ``python -m dfx_engine.worker ``. Started as a fresh +interpreter rather than a :mod:`multiprocessing` child on purpose: + +- The tokio runtime behind these extensions is a process-global, so ``fork`` + is unsafe. ``spawn`` would be fine but buys nothing here. +- Launching ``sys.executable`` makes the Python minor version match the + driver's by construction. Inline Python UDFs travel as cloudpickle payloads + stamped with the sender's ``(major, minor)``, and a mismatch is a hard + error -- so a hardcoded ``python`` on ``PATH`` would be a real bug. + +The order of operations in :func:`main` is the interesting part, and every +step is there because getting it wrong fails somewhere unhelpful. +""" + +from __future__ import annotations + +import json +import pathlib +import sys + +from datafusion.plan import ExecutionPlan + +from dfx_engine import _internal +from dfx_engine.session import SessionSpec, build_session + + +def run_task(envelope: dict) -> int: + """Execute one ``(stage, partition)`` and publish the result. + + Returns the number of rows written. + """ + spec = SessionSpec.from_json(envelope["spec"]) + partition = envelope["partition"] + + # 1. Build the session exactly as the driver did. Anything the driver + # relied on that is not in the spec is missing here. + ctx, _engine, _storage = build_session(spec) + + # 2. Check the codec ids *before* decoding. Without this the failure is a + # decode error naming a codec id, which is legible but arrives after + # the work of building a session; with it the mismatch is reported + # against the envelope that caused it. + expected = sorted(envelope["spec"]["codec_ids"]) + installed = sorted(ctx.physical_extension_codec_ids()) + if installed != expected: + message = f"worker codec ids {installed} do not match driver's {expected}" + raise RuntimeError(message) + + # 3. Decode. The stage node's shuffle directory travels inside the plan, + # so the worker cannot write somewhere the driver will not look. + plan = ExecutionPlan.from_bytes(ctx, pathlib.Path(envelope["plan"]).read_bytes()) + + # 4. Bounds-check before executing. A plan's partition count is a property + # of the plan, not of the spec, so a driver that miscounted is caught + # here rather than deep inside a scan. + if partition >= plan.partition_count: + message = ( + f"partition {partition} is out of range for a stage with " + f"{plan.partition_count} partition(s)" + ) + raise RuntimeError(message) + + # 5. Execute, and drain the stream. The stage node finds no result file + # for this partition -- this worker is the one producing it -- so it + # computes its child and writes the file as the batches go past. + # Draining is what makes that happen: the node does the work lazily, + # so a caller that dropped the stream would publish nothing. + rows = sum(batch.to_pyarrow().num_rows for batch in ctx.execute(plan, partition)) + + published = pathlib.Path( + _internal.partition_path(spec.shuffle_dir, envelope["stage_id"], partition) + ) + if not published.exists(): + message = f"stage partition {partition} produced no file at {published}" + raise RuntimeError(message) + + return rows + + +def main(argv: list[str] | None = None) -> int: + argv = sys.argv[1:] if argv is None else argv + if len(argv) != 1: + sys.stderr.write("usage: python -m dfx_engine.worker \n") + return 2 + + envelope = json.loads(pathlib.Path(argv[0]).read_text()) + rows = run_task(envelope) + # Read back by the driver, so it can report what each worker did. + print(json.dumps({"partition": envelope["partition"], "rows": rows})) + return 0 + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/examples/distributed/engine-library/src/codec.rs b/examples/distributed/engine-library/src/codec.rs new file mode 100644 index 000000000..f4188ab78 --- /dev/null +++ b/examples/distributed/engine-library/src/codec.rs @@ -0,0 +1,136 @@ +// 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. + +//! Carries this engine's own node. +//! +//! This is the codec half of the bundle, and the reason the bundle ships both +//! halves. The planner emits a [`ShuffleStageExec`]; nothing else in the +//! process knows that type, so without this codec the plans that planner +//! produces cannot be serialized at all -- and an engine whose whole job is +//! sending plans to workers would not get off the ground. +//! +//! The payload is the stage id and the shuffle directory, and nothing else. +//! The child is not encoded here: after `try_encode` returns, the framework +//! encodes `children()` itself using the *host's* chain, so the scan +//! underneath is claimed by whichever library owns it. Claiming a whole +//! subtree would cut those libraries out -- and could not work anyway, since +//! the FFI codec wrapper drops the caller's converter and substitutes a bare +//! default. + +use std::fmt; +use std::sync::Arc; +use std::sync::atomic::{AtomicUsize, Ordering}; + +use datafusion::common::{Result, internal_datafusion_err, internal_err}; +use datafusion::execution::TaskContext; +use datafusion::physical_plan::ExecutionPlan; +use datafusion_proto::physical_plan::{ + DefaultPhysicalExtensionCodec, PhysicalExtensionCodec, PhysicalProtoConverterExtension, +}; + +use crate::stage::ShuffleStageExec; + +/// Framing magic; the trailing digit is the payload version. +const MAGIC: &[u8; 8] = b"DFXENG01"; + +#[derive(Default, Debug)] +pub(crate) struct CodecCounters { + pub(crate) encoded: AtomicUsize, + pub(crate) decoded: AtomicUsize, +} + +pub(crate) struct DfxEnginePhysicalCodec { + inner: DefaultPhysicalExtensionCodec, + pub(crate) counters: Arc, +} + +impl DfxEnginePhysicalCodec { + pub(crate) fn new(counters: Arc) -> Self { + Self { + inner: DefaultPhysicalExtensionCodec {}, + counters, + } + } +} + +impl fmt::Debug for DfxEnginePhysicalCodec { + fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result { + formatter + .debug_struct("DfxEnginePhysicalCodec") + .finish_non_exhaustive() + } +} + +impl PhysicalExtensionCodec for DfxEnginePhysicalCodec { + fn try_encode( + &self, + node: Arc, + buf: &mut Vec, + proto_converter: &dyn PhysicalProtoConverterExtension, + ) -> Result<()> { + // This engine's own type only. Anything else goes to the default + // codec, whose error is the chain's "not mine" signal. + let Some(stage) = node.downcast_ref::() else { + return self.inner.try_encode(node, buf, proto_converter); + }; + + buf.extend_from_slice(MAGIC); + buf.extend_from_slice(&stage.stage_id.to_le_bytes()); + buf.extend_from_slice(stage.shuffle_dir.as_bytes()); + + self.counters.encoded.fetch_add(1, Ordering::SeqCst); + Ok(()) + } + + fn try_decode( + &self, + buf: &[u8], + inputs: &[Arc], + ctx: &TaskContext, + proto_converter: &dyn PhysicalProtoConverterExtension, + ) -> Result> { + let Some(rest) = buf.strip_prefix(MAGIC) else { + return self.inner.try_decode(buf, inputs, ctx, proto_converter); + }; + + let (stage_id, shuffle_dir) = rest.split_at_checked(4).ok_or_else(|| { + internal_datafusion_err!("dfx_engine: payload truncated before stage id") + })?; + let stage_id = u32::from_le_bytes( + stage_id + .try_into() + .map_err(|_| internal_datafusion_err!("dfx_engine: bad stage id"))?, + ); + let shuffle_dir = std::str::from_utf8(shuffle_dir) + .map_err(|err| internal_datafusion_err!("dfx_engine: bad shuffle dir: {err}"))?; + + // The child arrives already decoded, by the host's chain. + let [input] = inputs else { + return internal_err!( + "ShuffleStageExec expects exactly one input, got {}", + inputs.len() + ); + }; + + self.counters.decoded.fetch_add(1, Ordering::SeqCst); + Ok(Arc::new(ShuffleStageExec::new( + stage_id, + shuffle_dir.to_string(), + Arc::clone(input), + ))) + } +} diff --git a/examples/distributed/engine-library/src/config.rs b/examples/distributed/engine-library/src/config.rs new file mode 100644 index 000000000..051999056 --- /dev/null +++ b/examples/distributed/engine-library/src/config.rs @@ -0,0 +1,113 @@ +// 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. + +//! This engine's session config. +//! +//! A config extension rather than a constructor argument, because the driver +//! and every worker have to agree on the shuffle directory and the config is +//! the one thing that travels with the session. It also has to be *registered* +//! before anything can set `dfx_engine.shuffle_dir`: an unknown namespace is +//! an error, not a no-op, which is the first thing a worker bootstrap gets +//! wrong. + +use std::any::Any; + +use datafusion_common::config::{ + ConfigEntry, ConfigExtension, ConfigField, ExtensionOptions, Visit, +}; +use datafusion_common::{DataFusionError, config_err}; +use datafusion_ffi::config::extension_options::FFI_ExtensionOptions; +use pyo3::exceptions::PyRuntimeError; +use pyo3::prelude::*; +use pyo3::types::PyCapsule; + +/// Options under the `dfx_engine` prefix. +#[pyclass(from_py_object, name = "DfxEngineConfig", module = "dfx_engine")] +#[derive(Clone, Debug, Default)] +pub(crate) struct DfxEngineConfig { + /// Directory stage results are exchanged through. Empty means "do not + /// distribute": the planner leaves the plan alone and it runs in process. + pub(crate) shuffle_dir: String, +} + +#[pymethods] +impl DfxEngineConfig { + #[new] + #[pyo3(signature = (shuffle_dir=String::new()))] + fn new(shuffle_dir: String) -> Self { + Self { shuffle_dir } + } + + fn __datafusion_extension_options__<'py>( + &self, + py: Python<'py>, + ) -> PyResult> { + let mut config = FFI_ExtensionOptions::default(); + config + .add_config(self) + .map_err(|err| PyRuntimeError::new_err(err.to_string()))?; + PyCapsule::new_with_value(py, config, cr"datafusion_extension_options") + } +} + +impl ConfigExtension for DfxEngineConfig { + const PREFIX: &'static str = "dfx_engine"; +} + +impl ExtensionOptions for DfxEngineConfig { + fn as_any(&self) -> &dyn Any { + self + } + + fn as_any_mut(&mut self) -> &mut dyn Any { + self + } + + fn cloned(&self) -> Box { + Box::new(self.clone()) + } + + fn set(&mut self, key: &str, value: &str) -> datafusion_common::Result<()> { + ConfigField::set(self, key, value) + } + + fn entries(&self) -> Vec { + vec![ConfigEntry { + key: "shuffle_dir".to_owned(), + value: Some(self.shuffle_dir.clone()), + description: "directory stage results are exchanged through", + }] + } +} + +impl ConfigField for DfxEngineConfig { + fn visit(&self, v: &mut V, _key: &str, _description: &'static str) { + self.shuffle_dir.visit( + v, + "shuffle_dir", + "directory stage results are exchanged through", + ); + } + + fn set(&mut self, key: &str, value: &str) -> Result<(), DataFusionError> { + let (key, rem) = key.split_once('.').unwrap_or((key, "")); + match key { + "shuffle_dir" => self.shuffle_dir.set(rem, value), + _ => config_err!("Config value \"{key}\" not found on DfxEngineConfig"), + } + } +} diff --git a/examples/distributed/engine-library/src/extension.rs b/examples/distributed/engine-library/src/extension.rs new file mode 100644 index 000000000..a53541d00 --- /dev/null +++ b/examples/distributed/engine-library/src/extension.rs @@ -0,0 +1,200 @@ +// 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. + +//! This library's extension bundle: codecs *and* a planner. +//! +//! Both hooks, because the two halves are useless apart. The planner emits a +//! node only this library's codec can carry, so installing the planner without +//! the codec produces plans that cannot be serialized -- and installing the +//! codec without the planner produces nothing for it to carry. Shipping them +//! as one object is what makes that impossible to get wrong, and it is why +//! `with_extensions` installs every bundle's codecs before it binds any +//! planner. + +use std::fmt; +use std::sync::Arc; +use std::sync::atomic::Ordering; + +use datafusion_ffi::proto::physical_extension_codec::FFI_PhysicalExtensionCodec; +use datafusion_ffi::query_planner::FFI_QueryPlanner; +use datafusion_proto::physical_plan::PhysicalExtensionCodec; +use datafusion_python_util::{ + create_physical_extension_capsule, create_query_planner_capsule, + ffi_logical_codec_from_pycapsule, ffi_physical_codec_from_pycapsule, + ffi_query_planner_from_pycapsule, ffi_task_context_provider_from_pycapsule, get_tokio_runtime, +}; +use datafusion_session::QueryPlanner; +use pyo3::prelude::*; +use pyo3::types::{PyCapsule, PyDict}; + +use crate::codec::{CodecCounters, DfxEnginePhysicalCodec}; +use crate::planner::{DistributedQueryPlanner, PlannerObservations}; + +/// Wire id this codec's payloads carry, pinned because they cross processes. +const PHYSICAL_CODEC_ID: &str = "dfx_engine.physical.v1"; + +/// Carries this library's physical codec as an object rather than a capsule. +/// +/// `with_extensions` requires an object: a codec's wire id is read off the +/// thing it is handed over as, and a capsule has no type to read one from. +/// Wrapping also keeps the id *this library's* -- an id derived from the +/// contributing bundle would follow whichever object the caller passed, so an +/// application packaging this engine inside a bundle of its own would silently +/// re-tag these payloads and they would stop decoding on the workers. +#[pyclass(name = "BundledPhysicalCodec", module = "dfx_engine")] +pub(crate) struct BundledPhysicalCodec { + codec: FFI_PhysicalExtensionCodec, +} + +#[pymethods] +impl BundledPhysicalCodec { + #[getter] + fn __datafusion_codec_id__(&self) -> &'static str { + PHYSICAL_CODEC_ID + } + + /// `session` is unused: the codec was bound to its task-context provider + /// when the bundle was installed, which is why the bundle receives the + /// context at all. + #[pyo3(signature = (session=None))] + fn __datafusion_physical_extension_codec__<'py>( + &self, + py: Python<'py>, + session: Option>, + ) -> PyResult> { + let _ = session; + create_physical_extension_capsule(py, &self.codec) + } +} + +/// Extension bundle for `SessionContext.with_extensions`. +/// +/// Reusable configuration, not bound state: components are built fresh against +/// whichever context each install hands over, so one bundle works on several +/// sessions. +#[pyclass(from_py_object, name = "DfxEngineExtension", module = "dfx_engine")] +#[derive(Default, Clone)] +pub(crate) struct DfxEngineExtension { + observations: Arc, + counters: Arc, +} + +impl fmt::Debug for DfxEngineExtension { + fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result { + formatter + .debug_struct("DfxEngineExtension") + .field("observations", &self.observations) + .finish_non_exhaustive() + } +} + +#[pymethods] +impl DfxEngineExtension { + #[new] + fn new() -> Self { + Self::default() + } + + /// How often this engine's planner was asked for a physical plan. + fn plan_calls(&self) -> usize { + self.observations.plan_calls.load(Ordering::SeqCst) + } + + /// How often it inserted a stage, which is the rewrite it exists to do. + fn stages_inserted(&self) -> usize { + self.observations.stages_inserted.load(Ordering::SeqCst) + } + + /// How often this codec encoded one of its own stage nodes -- the step + /// that happens when the driver ships work. + fn encode_calls(&self) -> usize { + self.counters.encoded.load(Ordering::SeqCst) + } + + /// How often it rebuilt one, which happens on a worker. + fn decode_calls(&self) -> usize { + self.counters.decoded.load(Ordering::SeqCst) + } + + /// The wire id, so a driver can put it in a worker's task envelope and the + /// worker can check it before decoding anything. + #[staticmethod] + fn physical_codec_id() -> &'static str { + PHYSICAL_CODEC_ID + } + + fn __datafusion_session_components__<'py>( + &self, + py: Python<'py>, + ctx: Bound<'py, PyAny>, + ) -> PyResult> { + // Bind to the context the host supplied -- the session these + // components will run on -- and build fresh ones every call. The + // task-context provider comes off that context rather than from a + // `SessionContext` built here, so decode callbacks resolve names + // against the session that will actually run the query. + let provider = ffi_task_context_provider_from_pycapsule(&ctx)?; + let runtime = get_tokio_runtime().handle().clone(); + + let codec: Arc = + Arc::new(DfxEnginePhysicalCodec::new(Arc::clone(&self.counters))); + let ffi = FFI_PhysicalExtensionCodec::new(codec, Some(runtime), provider); + let physical = Py::new(py, BundledPhysicalCodec { codec: ffi })?; + + let components = py + .import("datafusion")? + .getattr("SessionExtensionComponents")?; + let kwargs = PyDict::new(py); + kwargs.set_item("physical_extension_codecs", (physical,))?; + components.call((), Some(&kwargs)) + } + + /// Contribute this engine's planner, nesting it on whatever came before. + /// + /// Runs after every bundle's codecs are installed, so `ctx` carries the + /// final chains and the planner is not left encoding through a partial + /// set. `fallback` is the planner assembled so far; delegating to it is + /// what makes several planner-shipping libraries composable, and + /// returning a planner that ignored it would discard every layer beneath. + fn __datafusion_session_planner__<'py>( + &self, + py: Python<'py>, + ctx: Bound<'py, PyAny>, + fallback: Bound<'py, PyAny>, + ) -> PyResult> { + let fallback = ffi_query_planner_from_pycapsule(&fallback, Some(&ctx))?; + let planner: Arc = Arc::new(DistributedQueryPlanner { + observations: Arc::clone(&self.observations), + // Deliberately not layered. Delegating would hand physical + // planning to the host and bring the plan back as opaque foreign + // nodes, which this engine cannot split -- so it plans for itself + // and the fallback goes unused. A planner that only rearranged + // stock nodes would keep it. + fallback: None, + }); + let _ = fallback; + + // The planner takes the *host's* codecs, not ones built here. By now + // those are the final chains, and this library has no business + // minting a task-context provider of its own. + let host_logical = ffi_logical_codec_from_pycapsule(ctx.clone(), None)?; + let host_physical = ffi_physical_codec_from_pycapsule(ctx, None)?; + let ffi_planner = + FFI_QueryPlanner::new_with_ffi_codecs(planner, host_logical, host_physical); + create_query_planner_capsule(py, &ffi_planner) + } +} diff --git a/examples/distributed/engine-library/src/lib.rs b/examples/distributed/engine-library/src/lib.rs new file mode 100644 index 000000000..2dd3b7e17 --- /dev/null +++ b/examples/distributed/engine-library/src/lib.rs @@ -0,0 +1,70 @@ +// 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. + +//! A toy distributed engine, in the two halves a real one has. +//! +//! The Rust half is here: a query planner that splits the plan into stages, +//! the node that marks a stage, the codec that carries it, and a config +//! extension so the driver and its workers agree on where results go. +//! +//! The Python half is in `python/dfx_engine`: the session factory both sides +//! build from, the worker entry point, and the driver that fans work out. An +//! engine needs both, which is why this crate is a mixed maturin package +//! rather than a pure extension module. +//! +//! One of three libraries in `examples/distributed`. This one owns execution. + +use pyo3::prelude::*; + +use crate::config::DfxEngineConfig; +use crate::extension::{BundledPhysicalCodec, DfxEngineExtension}; + +mod codec; +mod config; +mod extension; +mod local_session; +mod planner; +mod stage; + +/// Where the results of one stage partition live. +/// +/// Exported so the Python worker writes the path the Rust node will read, +/// rather than the convention being spelled out on both sides of the +/// boundary and drifting. +#[pyfunction] +fn partition_path(shuffle_dir: &str, stage_id: u32, partition: usize) -> String { + stage::partition_path(shuffle_dir, stage_id, partition) + .to_string_lossy() + .into_owned() +} + +/// The stage id this engine's planner produces. +#[pyfunction] +fn stage_id() -> u32 { + planner::STAGE_ID +} + +#[pymodule] +fn _internal(m: &Bound<'_, PyModule>) -> PyResult<()> { + pyo3_log::init(); + m.add_class::()?; + m.add_class::()?; + m.add_class::()?; + m.add_function(wrap_pyfunction!(partition_path, m)?)?; + m.add_function(wrap_pyfunction!(stage_id, m)?)?; + Ok(()) +} diff --git a/examples/distributed/engine-library/src/local_session.rs b/examples/distributed/engine-library/src/local_session.rs new file mode 100644 index 000000000..3660c5c84 --- /dev/null +++ b/examples/distributed/engine-library/src/local_session.rs @@ -0,0 +1,161 @@ +// 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. + +//! A `Session` that borrows another one but owns its optimizer rules. +//! +//! Physical planning applies `session.physical_optimizers()`. When the session +//! arrived over FFI those rules are the *host's*, so each one runs back across +//! the boundary and hands this library a `ForeignExecutionPlan`. A stock +//! `CooperativeExec` produced that way has no reachable `try_to_proto`, so a +//! planner that must serialize its result -- and `FFI_QueryPlanner` always +//! must, it returns proto bytes rather than a handle -- fails on a node that +//! is perfectly serializable in the process that made it. See the "Known gaps" +//! section of the extension guide. +//! +//! Wrapping the session with a locally-owned copy of the same rule set keeps +//! every rewrite inside this library, where the nodes stay concrete. That is +//! also what lets this engine split the plan: it cannot rewrite a subtree it +//! is only holding an opaque handle to. + +use std::any::Any; +use std::collections::HashMap; +use std::sync::Arc; + +use datafusion::catalog::CatalogProviderList; +use datafusion::common::config::{ConfigOptions, TableOptions}; +use datafusion::common::{DFSchema, Result}; +use datafusion::execution::TaskContext; +use datafusion::execution::config::SessionConfig; +use datafusion::execution::runtime_env::RuntimeEnv; +use datafusion::logical_expr::execution_props::ExecutionProps; +use datafusion::logical_expr::registry::ExtensionTypeRegistryRef; +use datafusion::logical_expr::{ + AggregateUDF, Expr, HigherOrderUDF, LogicalPlan, ScalarUDF, WindowUDF, +}; +use datafusion::physical_expr::PhysicalExpr; +use datafusion::physical_optimizer::PhysicalOptimizerRule; +use datafusion::physical_plan::ExecutionPlan; +use datafusion_session::Session; + +/// Borrows `inner` for everything except the physical optimizer rules. +pub(crate) struct LocalOptimizerSession<'a> { + inner: &'a dyn Session, + rules: Vec>, +} + +impl<'a> LocalOptimizerSession<'a> { + /// Wrap `inner` with the stock DataFusion rule set, owned here. + pub(crate) fn new(inner: &'a dyn Session) -> Self { + Self { + inner, + rules: datafusion::physical_optimizer::optimizer::PhysicalOptimizer::default().rules, + } + } +} + +#[async_trait::async_trait] +impl Session for LocalOptimizerSession<'_> { + /// The one override. Everything below delegates. + fn physical_optimizers(&self) -> &[Arc] { + &self.rules + } + + fn session_id(&self) -> &str { + self.inner.session_id() + } + + fn config(&self) -> &SessionConfig { + self.inner.config() + } + + fn catalog_list(&self) -> Arc { + self.inner.catalog_list() + } + + fn config_options(&self) -> &ConfigOptions { + self.inner.config_options() + } + + fn optimize(&self, plan: &LogicalPlan) -> Result { + self.inner.optimize(plan) + } + + async fn create_physical_plan( + &self, + logical_plan: &LogicalPlan, + ) -> Result> { + self.inner.create_physical_plan(logical_plan).await + } + + fn create_physical_expr( + &self, + expr: Expr, + df_schema: &DFSchema, + ) -> Result> { + self.inner.create_physical_expr(expr, df_schema) + } + + fn scalar_functions(&self) -> &HashMap> { + self.inner.scalar_functions() + } + + fn higher_order_functions(&self) -> &HashMap> { + self.inner.higher_order_functions() + } + + fn aggregate_functions(&self) -> &HashMap> { + self.inner.aggregate_functions() + } + + fn window_functions(&self) -> &HashMap> { + self.inner.window_functions() + } + + fn extension_type_registry(&self) -> &ExtensionTypeRegistryRef { + self.inner.extension_type_registry() + } + + fn runtime_env(&self) -> &Arc { + self.inner.runtime_env() + } + + fn execution_props(&self) -> &ExecutionProps { + self.inner.execution_props() + } + + fn as_any(&self) -> &dyn Any { + // Delegated, not `self`: the return type is implicitly `&dyn Any + + // 'static` and this wrapper only lives as long as its borrow. It also + // keeps `as_any().is::()` answering about the real + // session rather than the wrapper. + self.inner.as_any() + } + + fn table_options(&self) -> &TableOptions { + self.inner.table_options() + } + + fn table_options_mut(&mut self) -> &mut TableOptions { + // The wrapper only borrows `inner`, so it cannot hand out a mutable + // reference. Physical planning never calls this; verified in the spike. + unimplemented!("LocalOptimizerSession does not support table_options_mut") + } + + fn task_ctx(&self) -> Arc { + self.inner.task_ctx() + } +} diff --git a/examples/distributed/engine-library/src/planner.rs b/examples/distributed/engine-library/src/planner.rs new file mode 100644 index 000000000..4c308f57e --- /dev/null +++ b/examples/distributed/engine-library/src/planner.rs @@ -0,0 +1,170 @@ +// 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. + +//! Where the plan gets split into stages. +//! +//! The split point is the partial aggregate. DataFusion already breaks a +//! `GROUP BY` into a partial pass per input partition and a final pass that +//! merges them, which is exactly the shape a distributed engine wants: the +//! partial passes are independent, so they can run anywhere, and only their +//! output has to come back. Wrapping the partial aggregate in a +//! [`ShuffleStageExec`] is the whole rewrite. +//! +//! A query with no aggregate gets its whole plan wrapped instead, so there is +//! always exactly one stage and the orchestration in Python has one shape to +//! deal with. + +use std::sync::Arc; +use std::sync::atomic::{AtomicUsize, Ordering}; + +use async_trait::async_trait; +use datafusion::common::Result; +use datafusion::config::ConfigOptions; +use datafusion::logical_expr::LogicalPlan; +use datafusion::physical_plan::ExecutionPlan; +use datafusion::physical_plan::aggregates::{AggregateExec, AggregateMode}; +use datafusion::physical_plan::execution_plan::{ChildrenPropertiesMode, ReplaceChildrenOptions}; +use datafusion::physical_planner::{DefaultPhysicalPlanner, PhysicalPlanner}; +use datafusion_session::{QueryPlanner, Session}; + +use crate::local_session::LocalOptimizerSession; +use crate::stage::ShuffleStageExec; + +/// Config key naming the directory stages exchange results through. +/// +/// Read from the session rather than baked in, because the driver picks a +/// fresh directory per query and the workers have to be told the same one. +pub(crate) const SHUFFLE_DIR_KEY: &str = "dfx_engine.shuffle_dir"; + +/// The same setting once the session has crossed the FFI boundary, where +/// every foreign config extension is namespaced under `datafusion_ffi`. +const FFI_SHUFFLE_DIR_KEY: &str = "datafusion_ffi.dfx_engine.shuffle_dir"; + +/// The one stage id this engine produces. A real engine would number a chain +/// of them; one is enough to show the mechanism. +pub(crate) const STAGE_ID: u32 = 1; + +/// What the planner did, so a test can assert it rather than infer it. +#[derive(Default, Debug)] +pub(crate) struct PlannerObservations { + pub(crate) plan_calls: AtomicUsize, + pub(crate) stages_inserted: AtomicUsize, +} + +pub(crate) fn shuffle_dir_from_options(options: &ConfigOptions) -> Option { + options + .entries() + .into_iter() + .find(|entry| entry.key == SHUFFLE_DIR_KEY || entry.key == FFI_SHUFFLE_DIR_KEY) + .and_then(|entry| entry.value) +} + +/// Wrap the partial aggregate, or the whole plan if there is not one. +/// +/// Returns the rewritten plan and whether a stage was inserted. Only the +/// topmost partial aggregate is wrapped: an aggregate nested inside another +/// stage's subtree already travels with it. +fn insert_stage( + plan: Arc, + shuffle_dir: &str, +) -> Result<(Arc, bool)> { + if let Some(aggregate) = plan.downcast_ref::() + && matches!(aggregate.mode(), AggregateMode::Partial) + { + let stage = ShuffleStageExec::new(STAGE_ID, shuffle_dir.to_string(), Arc::clone(&plan)); + return Ok((Arc::new(stage), true)); + } + + let mut inserted = false; + let mut children = Vec::new(); + for child in plan.children() { + let (child, child_inserted) = insert_stage(Arc::clone(child), shuffle_dir)?; + inserted |= child_inserted; + children.push(child); + } + if !inserted { + return Ok((plan, false)); + } + // `Keep`: the replacement is a `ShuffleStageExec` wrapping the node it + // replaced, and that node takes its properties from its child, so the + // parent's view of its children is unchanged. + let options = ReplaceChildrenOptions::new(ChildrenPropertiesMode::Keep); + Ok((plan.replace_children(children, options)?, true)) +} + +#[derive(Debug)] +pub(crate) struct DistributedQueryPlanner { + pub(crate) observations: Arc, + /// Planner to layer on top of, if the session already had one. + /// + /// Held so several planner-shipping libraries compose. Note that + /// `Session::create_physical_plan` cannot be used for this: it dispatches + /// through the session's installed planner, so calling it from inside that + /// planner recurses until the stack overflows. + pub(crate) fallback: Option>, +} + +#[async_trait] +impl QueryPlanner for DistributedQueryPlanner { + async fn create_physical_plan( + &self, + logical_plan: &LogicalPlan, + session: &dyn Session, + ) -> Result> { + self.observations.plan_calls.fetch_add(1, Ordering::SeqCst); + + let plan = match self.fallback.as_ref() { + // Delegating hands physical planning to whoever is underneath, + // including the host. That is correct for composition, but it + // means the plan comes back as opaque foreign nodes this engine + // cannot split -- so a fallback and a split are exclusive, and + // the split is what this library is for. + Some(fallback) => return fallback.create_physical_plan(logical_plan, session).await, + None => { + // Plan against a session that owns the stock rule set locally + // instead of reaching back over FFI for the host's. Without + // this the plan contains `ForeignExecutionPlan` wrappers that + // cannot be serialized and cannot be rewritten. + let local = LocalOptimizerSession::new(session); + DefaultPhysicalPlanner::default() + .create_physical_plan(logical_plan, &local) + .await? + } + }; + + let Some(shuffle_dir) = shuffle_dir_from_options(session.config_options()) else { + // No shuffle directory configured: leave the plan alone and let it + // run in this process. An engine that inserted stages with nowhere + // to put their output would fail at execute time instead. + return Ok(plan); + }; + + let (plan, inserted) = insert_stage(plan, &shuffle_dir)?; + if inserted { + self.observations + .stages_inserted + .fetch_add(1, Ordering::SeqCst); + return Ok(plan); + } + + // Nothing to split at, so the whole plan is the stage. + self.observations + .stages_inserted + .fetch_add(1, Ordering::SeqCst); + Ok(Arc::new(ShuffleStageExec::new(STAGE_ID, shuffle_dir, plan))) + } +} diff --git a/examples/distributed/engine-library/src/stage.rs b/examples/distributed/engine-library/src/stage.rs new file mode 100644 index 000000000..16bf43c91 --- /dev/null +++ b/examples/distributed/engine-library/src/stage.rs @@ -0,0 +1,236 @@ +// 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. + +//! The node that makes a subtree a unit of remote work. +//! +//! One node does both halves of a shuffle, which is what keeps this example +//! small enough to read. `execute(i)` looks for the file a worker would have +//! written for partition `i`; if it is there it streams it, and if it is not +//! it runs the child instead. +//! +//! That is not a fallback bolted on -- it is what lets *the same node* be the +//! thing the worker runs and the thing the driver reads: +//! +//! - The driver ships this node to a worker. The worker's shuffle directory is +//! empty, so the node computes its child, and the worker writes the result +//! to the file for its partition. +//! - The driver then executes the very same plan. The files now exist, so the +//! node streams them instead of recomputing. +//! +//! Nothing has to rewrite the plan between those two steps, and a query run +//! with no workers at all still produces the right answer -- it just computes +//! everything locally. The shuffle directory travels inside the node, and so +//! inside its encoding, which is what stops a worker and a driver disagreeing +//! about where results go. + +use std::path::{Path, PathBuf}; +use std::sync::Arc; +use std::{fmt, fs}; + +use arrow::ipc::reader::StreamReader; +use arrow::ipc::writer::StreamWriter; +use datafusion::common::tree_node::TreeNodeRecursion; +use datafusion::common::{DataFusionError, Result, internal_datafusion_err, internal_err}; +use datafusion::execution::{SendableRecordBatchStream, TaskContext}; +use datafusion::physical_expr::PhysicalExpr; +use datafusion::physical_plan::memory::MemoryStream; +use datafusion::physical_plan::stream::RecordBatchStreamAdapter; +use datafusion::physical_plan::{DisplayAs, DisplayFormatType, ExecutionPlan, PlanProperties}; +use futures::StreamExt; + +/// Where the results of one stage partition live. +/// +/// Both halves of the exchange are in this file, so the convention has one +/// definition. It is also exported to Python, where the driver uses it to see +/// which partitions have been produced without having to know the layout. +pub(crate) fn partition_path(shuffle_dir: &str, stage_id: u32, partition: usize) -> PathBuf { + Path::new(shuffle_dir).join(format!("stage-{stage_id}-part-{partition}.arrow")) +} + +/// Marks a subtree as one stage of a distributed query. +#[derive(Debug)] +pub(crate) struct ShuffleStageExec { + pub(crate) stage_id: u32, + pub(crate) shuffle_dir: String, + pub(crate) input: Arc, + properties: Arc, +} + +impl ShuffleStageExec { + pub(crate) fn new(stage_id: u32, shuffle_dir: String, input: Arc) -> Self { + // Reading the child's results back yields the child's partitioning: + // one file per partition, in partition order. + let properties = Arc::clone(input.properties()); + Self { + stage_id, + shuffle_dir, + input, + properties, + } + } + + /// Compute this partition and write it where a reader will look. + /// + /// The batches are collected before anything is written, because an Arrow + /// IPC stream needs a schema up front and the file has to be complete + /// before it is published. A production engine would stream to the file + /// and track completion separately; holding one partition in memory is + /// the simplification this example makes. + /// + /// Published by rename, so a reader can never observe a half-written + /// file. The driver waits for workers to exit before reading, but relying + /// on that alone would break for anyone who overlapped the two. + fn write_partition( + &self, + partition: usize, + context: Arc, + ) -> Result { + let input = Arc::clone(&self.input); + let schema = input.schema(); + let final_path = partition_path(&self.shuffle_dir, self.stage_id, partition); + let shuffle_dir = self.shuffle_dir.clone(); + + let collected = async move { + let mut stream = input.execute(partition, context)?; + let mut batches = Vec::new(); + while let Some(batch) = stream.next().await { + batches.push(batch?); + } + + fs::create_dir_all(&shuffle_dir).map_err(|err| { + internal_datafusion_err!("dfx_engine: creating {shuffle_dir}: {err}") + })?; + let temp_path = final_path.with_extension("arrow.tmp"); + { + let file = fs::File::create(&temp_path).map_err(|err| { + internal_datafusion_err!("dfx_engine: creating {}: {err}", temp_path.display()) + })?; + let mut writer = StreamWriter::try_new(file, stream.schema().as_ref()) + .map_err(|err| internal_datafusion_err!("dfx_engine: ipc writer: {err}"))?; + for batch in &batches { + writer.write(batch).map_err(|err| { + internal_datafusion_err!("dfx_engine: writing batch: {err}") + })?; + } + writer + .finish() + .map_err(|err| internal_datafusion_err!("dfx_engine: ipc finish: {err}"))?; + } + fs::rename(&temp_path, &final_path).map_err(|err| { + internal_datafusion_err!("dfx_engine: publishing {}: {err}", final_path.display()) + })?; + + Ok::<_, DataFusionError>(batches) + }; + + // Written on first poll rather than here: `execute` must return + // promptly, so the work happens when the consumer drives the stream. + let stream = futures::stream::once(collected) + .map(|result| match result { + Ok(batches) => futures::stream::iter(batches.into_iter().map(Ok)).boxed(), + Err(err) => futures::stream::once(async move { Err(err) }).boxed(), + }) + .flatten(); + Ok(Box::pin(RecordBatchStreamAdapter::new(schema, stream))) + } + + fn read_partition(&self, partition: usize) -> Result { + let path = partition_path(&self.shuffle_dir, self.stage_id, partition); + let file = fs::File::open(&path).map_err(|err| { + internal_datafusion_err!("dfx_engine: opening {}: {err}", path.display()) + })?; + let reader = StreamReader::try_new(file, None).map_err(|err| { + internal_datafusion_err!("dfx_engine: reading {}: {err}", path.display()) + })?; + let schema = reader.schema(); + let batches = reader + .collect::>>() + .map_err(|err| { + internal_datafusion_err!("dfx_engine: reading {}: {err}", path.display()) + })?; + Ok(Box::pin(MemoryStream::try_new(batches, schema, None)?)) + } +} + +impl DisplayAs for ShuffleStageExec { + fn fmt_as(&self, _t: DisplayFormatType, f: &mut fmt::Formatter<'_>) -> fmt::Result { + write!(f, "ShuffleStageExec: stage={}", self.stage_id) + } +} + +impl ExecutionPlan for ShuffleStageExec { + fn name(&self) -> &str { + Self::static_name() + } + + fn properties(&self) -> &Arc { + &self.properties + } + + fn children(&self) -> Vec<&Arc> { + vec![&self.input] + } + + fn apply_expressions( + &self, + _f: &mut dyn FnMut(&Arc) -> Result, + ) -> Result { + // Owns no expressions of its own; the child holds them all. + Ok(TreeNodeRecursion::Continue) + } + + fn with_new_children( + self: Arc, + mut children: Vec>, + ) -> Result> { + if children.len() != 1 { + return internal_err!( + "ShuffleStageExec expects exactly one child, got {}", + children.len() + ); + } + Ok(Arc::new(Self::new( + self.stage_id, + self.shuffle_dir.clone(), + children.swap_remove(0), + ))) + } + + /// Read this partition's results if they exist, otherwise compute them + /// and leave them where the next reader will find them. + /// + /// The same code runs on a worker and on the driver, and which branch it + /// takes is decided by the filesystem rather than by a mode flag: + /// + /// - On a worker the file is absent, so the child runs and the output is + /// written on the way past. + /// - On the driver the workers have already been and gone, so the file is + /// there and the child is never touched. + /// + /// A query with no workers at all takes the second branch on every + /// partition and still gets the right answer, having done the work itself. + fn execute( + &self, + partition: usize, + context: Arc, + ) -> Result { + if partition_path(&self.shuffle_dir, self.stage_id, partition).exists() { + return self.read_partition(partition); + } + self.write_partition(partition, context) + } +} diff --git a/examples/distributed/engine-library/stage-1-part-0.arrow b/examples/distributed/engine-library/stage-1-part-0.arrow new file mode 100644 index 0000000000000000000000000000000000000000..9e29db43c155ec56fb03371a5713dd6c3dad22bd GIT binary patch literal 1032 zcmbtSO-sW-5PfN4jT*FvqD4FuPsv3PL_Ef$ApU_OMcP#YX_S~iz4fQ{NBCpB#W(vQ zHVRt0%gpA@ynQn}q-mPw$^&|kwTp>RZP|NrOYNO88Le=`hBh0)*osW zbfiwm&I9b}9#a$NtH92H$#Y#)@l8=(Uc~1~GSc`<_h;XczVB_BnaIm!gB5^$T&k)r zi^u%oNoX8j^TAWN5svxpBXXr1U`Si)F8E=VFOwZG8t+P~}D&1Y5p WZ?tKBQ{gQ&XJ$`>{^RBU72YRWz%^9> literal 0 HcmV?d00001 diff --git a/examples/distributed/engine-library/stage-1-part-1.arrow b/examples/distributed/engine-library/stage-1-part-1.arrow new file mode 100644 index 0000000000000000000000000000000000000000..84c8ac2cdea9d27b018f26a918d2af398ce39cee GIT binary patch literal 1032 zcmbtSO-sW-5PfN4jT*FvqD4FuPt8RTL_Ef$ApU_OMcP#YX_S~kz4fQ{NBCpBt#9^2 zY!$R*mzm9*dHZH|h}K$P1waR~c9GC_5Yf9>5FzFUW7Y#9(K~HU>(V+nr0wH?9FQei zJ=zHk2txS=+RU`MWwneyrGK4G2Sr}0yjIh*Vx;D^Dkp04QsrZn4w=0d{jSkm*Y6t@ zw4_eV&LixZ9y1%~tH92H@$;;#2RCJXc`-Oo)1kp%x7`?&O}}=8>|2nqe|7Y zs(dUSo`lBn=?|X5O>oS2ACW8F0Daohbioh1LYZuj*_bSu@hwg6+KD*Q8kXqy{BKxd z+v42vWiyBEn}CJGUT^8J=fi#Yw`UIf{9-My-TfwqNcw&@)w#ag^u{W`Y5%@&Kc7|g Xzfsfrrb7QMvuAEk1ODga{}tvZYJfGB literal 0 HcmV?d00001 diff --git a/examples/distributed/engine-library/stage-1-part-2.arrow b/examples/distributed/engine-library/stage-1-part-2.arrow new file mode 100644 index 0000000000000000000000000000000000000000..e5c813320726aad1437c7e224d9e6ed4119cee62 GIT binary patch literal 1032 zcmbtSO-sW-5PfN4jT*FvqD4FuPsv3PL_Ef$ApU_OMcP#YX_S~iz4fQ{NBCpB#W(vQ zHVRt0%gpA@ynQn}q-mPw$^&|kwTp>RZP|NrOYNO88Le=`hBh0)*osW zbfiwm&I9b}9#a$NtH92H$#Y#)@l8=(Uc~1~GSc`<_h;XczVB_BnaIm!gB5^$T&k)r zi^u%oNoX8j^TAWN5svxpBXXr1U`Si)F8E=VFOwZG8t+P~}D&1Y5p XZ?tKBT_Jr-&6$}~ga0}Ce}(x8dO$VR literal 0 HcmV?d00001 diff --git a/examples/distributed/storage-library/src/codec.rs b/examples/distributed/storage-library/src/codec.rs index 9868e0014..56c0bc748 100644 --- a/examples/distributed/storage-library/src/codec.rs +++ b/examples/distributed/storage-library/src/codec.rs @@ -42,20 +42,24 @@ //! type, including extension types and field metadata. use std::fmt; +use std::path::Path; use std::sync::Arc; use std::sync::atomic::{AtomicUsize, Ordering}; use arrow::datatypes::Schema; use arrow::ipc::reader::StreamReader; use arrow::ipc::writer::StreamWriter; -use datafusion::common::{Result, internal_datafusion_err, internal_err}; +use datafusion::catalog::TableProvider; +use datafusion::common::{Result, TableReference, internal_datafusion_err, internal_err}; use datafusion::execution::TaskContext; use datafusion::physical_plan::ExecutionPlan; +use datafusion_proto::logical_plan::{DefaultLogicalExtensionCodec, LogicalExtensionCodec}; use datafusion_proto::physical_plan::{ DefaultPhysicalExtensionCodec, PhysicalExtensionCodec, PhysicalProtoConverterExtension, }; use crate::exec::{FileSlice, PartitionedParquetExec}; +use crate::table_provider::PartitionedParquetTable; /// Framing magic. The trailing digit is the payload version. const MAGIC: &[u8; 8] = b"DFXSTOR1"; @@ -71,6 +75,8 @@ pub(crate) struct CodecCounters { pub(crate) encoded: AtomicUsize, pub(crate) decoded: AtomicUsize, pub(crate) declined: AtomicUsize, + pub(crate) provider_encoded: AtomicUsize, + pub(crate) provider_decoded: AtomicUsize, } pub(crate) struct DfxStoragePhysicalCodec { @@ -224,3 +230,105 @@ impl PhysicalExtensionCodec for DfxStoragePhysicalCodec { )?)) } } + +/// Framing magic for the logical payload; the digit is its version. +const LOGICAL_MAGIC: &[u8; 8] = b"DFXSTOL1"; + +/// Carries this library's *table provider*, which a query planner forces. +/// +/// A provider library might reasonably think a physical codec is enough -- +/// its scan node is a physical node, after all. It is not. Installing any FFI +/// query planner means the session hands that planner the **logical** plan as +/// protobuf, and a logical plan holds its tables as `Arc`. +/// Encoding one is `try_encode_table_provider`, and the default codec has no +/// implementation, so without this codec a session that has *both* this +/// provider and any engine installed fails at `execution_plan()` with +/// "Error serializing custom table". +/// +/// The payload is the directory, because everything else this provider holds +/// -- the file list, their sizes, the schema -- is read back from the +/// directory when it is rebuilt. Durable metadata again, for the same reason: +/// the process that decodes this has never seen the table registered. +pub(crate) struct DfxStorageLogicalCodec { + inner: DefaultLogicalExtensionCodec, + pub(crate) counters: Arc, +} + +impl DfxStorageLogicalCodec { + pub(crate) fn new(counters: Arc) -> Self { + Self { + inner: DefaultLogicalExtensionCodec {}, + counters, + } + } +} + +impl fmt::Debug for DfxStorageLogicalCodec { + fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result { + formatter + .debug_struct("DfxStorageLogicalCodec") + .finish_non_exhaustive() + } +} + +impl LogicalExtensionCodec for DfxStorageLogicalCodec { + fn try_decode( + &self, + buf: &[u8], + inputs: &[datafusion::logical_expr::LogicalPlan], + ctx: &TaskContext, + ) -> Result { + // This library defines no logical extension node. + self.inner.try_decode(buf, inputs, ctx) + } + + fn try_encode( + &self, + node: &datafusion::logical_expr::Extension, + buf: &mut Vec, + ) -> Result<()> { + self.inner.try_encode(node, buf) + } + + fn try_encode_table_provider( + &self, + table_ref: &TableReference, + node: Arc, + buf: &mut Vec, + ) -> Result<()> { + let Some(table) = node.downcast_ref::() else { + self.counters.declined.fetch_add(1, Ordering::SeqCst); + return self.inner.try_encode_table_provider(table_ref, node, buf); + }; + buf.extend_from_slice(LOGICAL_MAGIC); + buf.extend_from_slice(table.directory.as_bytes()); + self.counters + .provider_encoded + .fetch_add(1, Ordering::SeqCst); + Ok(()) + } + + fn try_decode_table_provider( + &self, + buf: &[u8], + table_ref: &TableReference, + schema: arrow::datatypes::SchemaRef, + ctx: &TaskContext, + ) -> Result> { + let Some(directory) = buf.strip_prefix(LOGICAL_MAGIC) else { + self.counters.declined.fetch_add(1, Ordering::SeqCst); + return self + .inner + .try_decode_table_provider(buf, table_ref, schema, ctx); + }; + let directory = std::str::from_utf8(directory).map_err(|err| { + internal_datafusion_err!("dfx_storage: bad directory in payload: {err}") + })?; + self.counters + .provider_decoded + .fetch_add(1, Ordering::SeqCst); + Ok(Arc::new(PartitionedParquetTable::try_new(Path::new( + directory, + ))?)) + } +} diff --git a/examples/distributed/storage-library/src/extension.rs b/examples/distributed/storage-library/src/extension.rs index b42944b3f..657d28088 100644 --- a/examples/distributed/storage-library/src/extension.rs +++ b/examples/distributed/storage-library/src/extension.rs @@ -26,15 +26,18 @@ use std::fmt; use std::sync::Arc; use std::sync::atomic::Ordering; +use datafusion_ffi::proto::logical_extension_codec::FFI_LogicalExtensionCodec; use datafusion_ffi::proto::physical_extension_codec::FFI_PhysicalExtensionCodec; +use datafusion_proto::logical_plan::LogicalExtensionCodec; use datafusion_proto::physical_plan::PhysicalExtensionCodec; use datafusion_python_util::{ - create_physical_extension_capsule, ffi_task_context_provider_from_pycapsule, get_tokio_runtime, + create_logical_extension_capsule, create_physical_extension_capsule, + ffi_task_context_provider_from_pycapsule, get_tokio_runtime, }; use pyo3::prelude::*; use pyo3::types::{PyCapsule, PyDict}; -use crate::codec::{CodecCounters, DfxStoragePhysicalCodec}; +use crate::codec::{CodecCounters, DfxStorageLogicalCodec, DfxStoragePhysicalCodec}; /// Wire id this codec's payloads carry. /// @@ -45,6 +48,33 @@ use crate::codec::{CodecCounters, DfxStoragePhysicalCodec}; /// every decode would fail. const PHYSICAL_CODEC_ID: &str = "dfx_storage.physical.v1"; +/// Logical companion to [`PHYSICAL_CODEC_ID`]. +const LOGICAL_CODEC_ID: &str = "dfx_storage.logical.v1"; + +/// Carries this library's logical codec. See [`BundledPhysicalCodec`]. +#[pyclass(name = "BundledLogicalCodec", module = "dfx_storage")] +pub(crate) struct BundledLogicalCodec { + codec: FFI_LogicalExtensionCodec, +} + +#[pymethods] +impl BundledLogicalCodec { + #[getter] + fn __datafusion_codec_id__(&self) -> &'static str { + LOGICAL_CODEC_ID + } + + #[pyo3(signature = (session=None))] + fn __datafusion_logical_extension_codec__<'py>( + &self, + py: Python<'py>, + session: Option>, + ) -> PyResult> { + let _ = session; + create_logical_extension_capsule(py, &self.codec) + } +} + /// Carries this library's physical codec as an object rather than a capsule. /// /// `with_extensions` requires an object: a codec's wire id is read off the @@ -141,15 +171,24 @@ impl DfxStorageExtension { let codec: Arc = Arc::new(DfxStoragePhysicalCodec::new(Arc::clone(&self.counters))); - let ffi = FFI_PhysicalExtensionCodec::new(codec, Some(runtime), provider); + let ffi = FFI_PhysicalExtensionCodec::new(codec, Some(runtime.clone()), provider.clone()); let physical = Py::new(py, BundledPhysicalCodec { codec: ffi })?; - // No logical codec: this library defines no logical extension node. - // Its table provider crosses FFI as a provider, not as a plan node. + // The logical codec is not optional, even though this library defines + // no logical extension *node*. Its table provider is held in the + // logical plan, and any installed query planner receives that plan as + // protobuf -- so without this the session fails to plan at all. See + // `DfxStorageLogicalCodec`. + let logical: Arc = + Arc::new(DfxStorageLogicalCodec::new(Arc::clone(&self.counters))); + let ffi_logical = FFI_LogicalExtensionCodec::new(logical, Some(runtime), provider); + let logical = Py::new(py, BundledLogicalCodec { codec: ffi_logical })?; + let components = py .import("datafusion")? .getattr("SessionExtensionComponents")?; let kwargs = PyDict::new(py); + kwargs.set_item("logical_extension_codecs", (logical,))?; kwargs.set_item("physical_extension_codecs", (physical,))?; components.call((), Some(&kwargs)) } diff --git a/examples/distributed/storage-library/src/table_provider.rs b/examples/distributed/storage-library/src/table_provider.rs index 4b980c03f..3c33c41df 100644 --- a/examples/distributed/storage-library/src/table_provider.rs +++ b/examples/distributed/storage-library/src/table_provider.rs @@ -45,6 +45,10 @@ use crate::exec::{FileSlice, PartitionedParquetExec}; /// Scans `*.parquet` under `directory`, one partition per file. #[derive(Debug)] pub(crate) struct PartitionedParquetTable { + /// Kept so the logical codec can write it down. Everything else here is + /// derived from the directory, so the path is the whole encoding -- see + /// [`crate::codec::DfxStorageLogicalCodec`]. + pub(crate) directory: String, files: Vec, schema: SchemaRef, } @@ -86,6 +90,7 @@ impl PartitionedParquetTable { let schema = Self::read_schema(&paths[0])?; Ok(Self { + directory: directory.to_string_lossy().into_owned(), files, schema: Arc::new(schema), }) From f8ca9b0d8acdaf046cdabad38598b6f2e4742acd Mon Sep 17 00:00:00 2001 From: Tim Saucer Date: Wed, 9 Sep 2026 13:22:06 -0400 Subject: [PATCH 04/24] Add the queries and the cross-library integration tests MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit 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) --- .../python/dfx_engine/driver.py | 22 +- .../python/tests/_test_three_libraries.py | 396 ++++++++++++++++++ .../engine-library/python/tests/conftest.py | 97 +++++ .../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(-) create mode 100644 examples/distributed/engine-library/python/tests/_test_three_libraries.py create mode 100644 examples/distributed/engine-library/python/tests/conftest.py delete mode 100644 examples/distributed/engine-library/stage-1-part-0.arrow delete mode 100644 examples/distributed/engine-library/stage-1-part-1.arrow delete mode 100644 examples/distributed/engine-library/stage-1-part-2.arrow create mode 100644 examples/distributed/run_tpch.py diff --git a/examples/distributed/engine-library/python/dfx_engine/driver.py b/examples/distributed/engine-library/python/dfx_engine/driver.py index baa145f15..804aadf7b 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 @@ 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 000000000..a72c897ac --- /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 + ``....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 000000000..6c0791ecd --- /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) + + +@pytest.fixture(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), +] + + +@pytest.fixture +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 + + +@pytest.fixture +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 4c308f57e..11de45039 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>RZP|NrOYNO88Le=`hBh0)*osW zbfiwm&I9b}9#a$NtH92H$#Y#)@l8=(Uc~1~GSc`<_h;XczVB_BnaIm!gB5^$T&k)r zi^u%oNoX8j^TAWN5svxpBXXr1U`Si)F8E=VFOwZG8t+P~}D&1Y5p WZ?tKBQ{gQ&XJ$`>{^RBU72YRWz%^9> 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 84c8ac2cdea9d27b018f26a918d2af398ce39cee..0000000000000000000000000000000000000000 GIT binary patch literal 0 HcmV?d00001 literal 1032 zcmbtSO-sW-5PfN4jT*FvqD4FuPt8RTL_Ef$ApU_OMcP#YX_S~kz4fQ{NBCpBt#9^2 zY!$R*mzm9*dHZH|h}K$P1waR~c9GC_5Yf9>5FzFUW7Y#9(K~HU>(V+nr0wH?9FQei zJ=zHk2txS=+RU`MWwneyrGK4G2Sr}0yjIh*Vx;D^Dkp04QsrZn4w=0d{jSkm*Y6t@ zw4_eV&LixZ9y1%~tH92H@$;;#2RCJXc`-Oo)1kp%x7`?&O}}=8>|2nqe|7Y zs(dUSo`lBn=?|X5O>oS2ACW8F0Daohbioh1LYZuj*_bSu@hwg6+KD*Q8kXqy{BKxd z+v42vWiyBEn}CJGUT^8J=fi#Yw`UIf{9-My-TfwqNcw&@)w#ag^u{W`Y5%@&Kc7|g Xzfsfrrb7QMvuAEk1ODga{}tvZYJfGB 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 e5c813320726aad1437c7e224d9e6ed4119cee62..0000000000000000000000000000000000000000 GIT binary patch literal 0 HcmV?d00001 literal 1032 zcmbtSO-sW-5PfN4jT*FvqD4FuPsv3PL_Ef$ApU_OMcP#YX_S~iz4fQ{NBCpB#W(vQ zHVRt0%gpA@ynQn}q-mPw$^&|kwTp>RZP|NrOYNO88Le=`hBh0)*osW zbfiwm&I9b}9#a$NtH92H$#Y#)@l8=(Uc~1~GSc`<_h;XczVB_BnaIm!gB5^$T&k)r zi^u%oNoX8j^TAWN5svxpBXXr1U`Si)F8E=VFOwZG8t+P~}D&1Y5p XZ?tKBT_Jr-&6$}~ga0}Ce}(x8dO$VR diff --git a/examples/distributed/run_tpch.py b/examples/distributed/run_tpch.py new file mode 100644 index 000000000..90a25deee --- /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()) From 9339ac88fb21bf79bf841c3211b4e48c62a6449b Mon Sep 17 00:00:00 2001 From: Tim Saucer Date: Wed, 9 Sep 2026 14:18:12 -0400 Subject: [PATCH 05/24] Wire the distributed example into CI, and say which example is which MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Three maturin builds and three test invocations for #1719, plus the documentation change that keeping three example trees requires. The plan had been to retire `datafusion-ffi-query-planner-example` and fold it into the new engine. That is now off, on evidence from building the engine: roughly fifteen of its forty-seven tests cover planner *layering*, and `dfx_engine`'s planner structurally cannot delegate to a `fallback`. Delegating hands physical planning back to the host, which returns opaque `ForeignExecutionPlan` nodes the engine can neither serialize nor split — a stage-splitting planner has to plan for itself. So the new example has nothing for those tests to nest, and deleting the crate would delete real coverage of the most subtle part of #1679's contract. Three trees then, with distinct jobs, which the guide now states up front rather than leaving a reader to infer: `examples/distributed` is the worked example and the place to start; `datafusion-ffi-example` is the capsule-protocol test bed, one of every hook exercised hard; and `datafusion-ffi-query-planner-example` is the planner-composition test bed. The guide's "three roles in a query" section described only the latter two. Two stale claims fixed while in there. `examples/README.md` linked three `sql-on-*.py` files that do not exist. The planner example's README said its planner "owns no serializable types of its own and deliberately uses only built-in physical nodes", which stopped being true when `DistributedExec` was added — and the sentence mattered, because owning a node is exactly why that library ships its codec and planner as one bundle. The `actionlint` pre-commit hook needs Docker and could not run here; the workflow files are otherwise lint-clean and parse as YAML. Co-Authored-By: Claude Opus 5 (1M context) --- .github/workflows/build.yml | 35 +++++++++++++++++++ .github/workflows/test.yml | 9 +++++ docs/source/extension-guide/index.md | 27 +++++++++++--- examples/README.md | 24 +++++++------ .../README.md | 4 +-- 5 files changed, 81 insertions(+), 18 deletions(-) diff --git a/.github/workflows/build.yml b/.github/workflows/build.yml index d7af9b663..4ed5782bd 100644 --- a/.github/workflows/build.yml +++ b/.github/workflows/build.yml @@ -206,6 +206,38 @@ jobs: args: --out dist rustup-components: rust-std + # The three libraries of the distributed example. Built in dependency + # order for readability only; they are independent cdylibs. + - name: Build distributed example storage library + if: matrix.python-tag == 'abi3' + uses: PyO3/maturin-action@v1 + with: + target: x86_64-unknown-linux-gnu + manylinux: "2_28" + working-directory: examples/distributed/storage-library + args: --out dist + rustup-components: rust-std + + - name: Build distributed example UDF library + if: matrix.python-tag == 'abi3' + uses: PyO3/maturin-action@v1 + with: + target: x86_64-unknown-linux-gnu + manylinux: "2_28" + working-directory: examples/distributed/udf-library + args: --out dist + rustup-components: rust-std + + - name: Build distributed example engine library + if: matrix.python-tag == 'abi3' + uses: PyO3/maturin-action@v1 + with: + target: x86_64-unknown-linux-gnu + manylinux: "2_28" + working-directory: examples/distributed/engine-library + args: --out dist + rustup-components: rust-std + - name: Archive wheels uses: actions/upload-artifact@v7 with: @@ -220,6 +252,9 @@ jobs: path: | examples/datafusion-ffi-example/dist/* examples/datafusion-ffi-query-planner-example/dist/* + examples/distributed/storage-library/dist/* + examples/distributed/udf-library/dist/* + examples/distributed/engine-library/dist/* # ============================================ # Build - Linux ARM64 diff --git a/.github/workflows/test.yml b/.github/workflows/test.yml index 047b35039..c1efe0bae 100644 --- a/.github/workflows/test.yml +++ b/.github/workflows/test.yml @@ -127,6 +127,15 @@ jobs: uv run --no-project pytest python/tests/_test*.py cd ../datafusion-ffi-query-planner-example uv run --no-project pytest python/tests/_test*.py + # The distributed example. Its tests spawn worker processes with + # `sys.executable`, so they need the same interpreter the wheels + # were installed into -- which `uv run` gives them. + cd ../distributed/storage-library + uv run --no-project pytest python/tests/_test*.py + cd ../udf-library + uv run --no-project pytest python/tests/_test*.py + cd ../engine-library + uv run --no-project pytest python/tests/_test*.py - name: Run tpchgen-cli to create 1 Gb dataset if: matrix.wheel-tag == 'abi3' diff --git a/docs/source/extension-guide/index.md b/docs/source/extension-guide/index.md index 91d42a3c5..dce944de4 100644 --- a/docs/source/extension-guide/index.md +++ b/docs/source/extension-guide/index.md @@ -59,11 +59,27 @@ this section only makes sense once they are distinct in your head: the codecs that serialize them. - **A planner library** — owns a query planner and the configuration it needs. -The worked examples in this repository use two separate crates, -[`datafusion-ffi-example`] and [`datafusion-ffi-query-planner-example`], so -each role has a distinct shared-library identity. A real library may play more -than one role; keeping them separate in the examples is what makes the -boundaries observable. +A real library may play more than one role. Keeping them in separate crates is +what makes the boundaries observable, because each one is then a distinct +shared library and the FFI conversions are real rather than same-library +downcasts. + +The examples in this repository are three trees, and it is worth knowing which +one answers your question: + +- [`examples/distributed`] is the **worked example**, and the place to start. + Three libraries — functions, tables, and an engine — cooperate on one query + whose leaf stage runs in separate worker processes. It is also the only + example whose plans genuinely leave the process, so it is where the codecs + encode durable metadata rather than tokens. +- [`datafusion-ffi-example`] is the **capsule-protocol test bed**: one of every + hook, exercised hard. Read it to see the shape of a getter, not to see a + library designed well. +- [`datafusion-ffi-query-planner-example`] is the **planner-composition test + bed**: what happens when more than one library contributes a query planner, + and how `fallback` nests them. The distributed example cannot cover this — + a planner that rewrites the plan into stages has to plan for itself, so it + has no use for a fallback. The session owns the codecs used for the exchange and supplies them to the foreign planner. That is what lets the planner decode provider-owned objects, @@ -134,3 +150,4 @@ checklist [`datafusion-ffi-example`]: https://github.com/apache/datafusion-python/tree/main/examples/datafusion-ffi-example [`datafusion-ffi-query-planner-example`]: https://github.com/apache/datafusion-python/tree/main/examples/datafusion-ffi-query-planner-example +[`examples/distributed`]: https://github.com/apache/datafusion-python/tree/main/examples/distributed diff --git a/examples/README.md b/examples/README.md index 7bbb45dcf..e38932aeb 100644 --- a/examples/README.md +++ b/examples/README.md @@ -51,23 +51,25 @@ Here is a direct link to the file used in the examples: ### Rust FFI Extensions -- [Table providers, functions, and codecs](./datafusion-ffi-example/) -- [Independent query planner and planner configuration](./datafusion-ffi-query-planner-example/) +Start with the worked example; the other two are focused test beds that +exercise one part of the protocol hard rather than reading as a tutorial. -These two crates form a three-library interoperability example with -`datafusion-python`. They are separate shared libraries so the tests exercise real FFI -type and codec boundaries rather than same-library Rust downcasts. +- [**Three libraries in one distributed query**](./distributed/) — a UDF + library, a table provider with its own scan node, and a toy engine that + splits the plan and runs each partition in a separate process. Read this one + first. +- [Capsule protocol conformance](./datafusion-ffi-example/) — table providers, + catalogs, functions, config, and codecs, one of each. +- [Query planner composition](./datafusion-ffi-query-planner-example/) — what + happens when more than one library contributes a planner, and how they nest. + +Each is a separate shared library, so the tests exercise real FFI type and +codec boundaries rather than same-library Rust downcasts. ### Substrait Support - [Serialize query plans using Substrait](./substrait.py) -### Executing SQL against DataFrame Libraries (Experimental) - -- [Executing SQL on Polars](./sql-on-polars.py) -- [Executing SQL on Pandas](./sql-on-pandas.py) -- [Executing SQL on cuDF](./sql-on-cudf.py) - ## TPC-H Examples Within the subdirectory `tpch` there are 22 examples that reproduce queries in diff --git a/examples/datafusion-ffi-query-planner-example/README.md b/examples/datafusion-ffi-query-planner-example/README.md index af128886f..e7c105e2e 100644 --- a/examples/datafusion-ffi-query-planner-example/README.md +++ b/examples/datafusion-ffi-query-planner-example/README.md @@ -76,8 +76,8 @@ ctx.register_udf(provider_udf) ctx.set_query_planner(MyQueryPlanner()) ``` -`MyPlannerConfig` is transferred through the foreign session. `MyQueryPlanner` reads `ffi_query_planner.max_rows`, creates the plan with `DefaultPhysicalPlanner`, and adds a built-in `GlobalLimitExec`. The test changes the setting with `SET` and verifies the new row limit. +`MyPlannerConfig` is transferred through the foreign session. `MyQueryPlanner` reads `ffi_query_planner.max_rows`, creates the plan with `DefaultPhysicalPlanner`, adds a built-in `GlobalLimitExec`, and wraps the result in a `DistributedExec` of its own. The test changes the setting with `SET` and verifies the new row limit. -The provider's codec chain is attached to the planner when it is installed and is also used to decode the returned physical plan in `datafusion-python`. Extension codecs compose: each `with_logical_extension_codec` / `with_physical_extension_codec` call appends to the session's codec chain, and each payload records which codec wrote it, so several libraries can install codecs on the same session and the order between them does not affect decoding. This planner owns no serializable types of its own and deliberately uses only built-in physical nodes. Install the codecs before the planner where possible; installing a codec afterwards rebuilds the planner against the new chain, but planner-last order is easier to audit. That rebuild is one level deep — a planner constructed with `fallback=` keeps the codecs its fallback was imported with — so codecs-first is a requirement rather than a preference once planners are layered. See [Rebinding a planner's codecs is one level deep](https://datafusion.apache.org/python/extension-guide/query-planners.html#install-codecs-before-a-layered-planner). +The provider's codec chain is attached to the planner when it is installed and is also used to decode the returned physical plan in `datafusion-python`. Extension codecs compose: each `with_logical_extension_codec` / `with_physical_extension_codec` call appends to the session's codec chain, and each payload records which codec wrote it, so several libraries can install codecs on the same session and the order between them does not affect decoding. This planner does own a node of its own — `DistributedExec`, which nothing else in the process can serialize — and that is why it ships its codec and its planner as one bundle. Install the codecs before the planner where possible; installing a codec afterwards rebuilds the planner against the new chain, but planner-last order is easier to audit. That rebuild is one level deep — a planner constructed with `fallback=` keeps the codecs its fallback was imported with — so codecs-first is a requirement rather than a preference once planners are layered. See [Rebinding a planner's codecs is one level deep](https://datafusion.apache.org/python/extension-guide/query-planners.html#install-codecs-before-a-layered-planner). For the limits behind that choice — how the codec chain dispatches, which node kinds survive the boundary, and what a derived context shares with the context it came from — see the [Extension Guide](https://datafusion.apache.org/python/extension-guide/index.html). From 467fef1bc6010f3442f9fc07d3cf449124bc9174 Mon Sep 17 00:00:00 2001 From: Tim Saucer Date: Wed, 9 Sep 2026 14:22:27 -0400 Subject: [PATCH 06/24] Document what a worker has to reproduce, and three findings that cost time MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The documentation half of #1719. Each section here exists because building the example ran into the thing it describes. **`distributing-work/query-engines.md` gains the worker-parity checklist.** This was the gap I most expected to find and did: there is no way to snapshot a SessionContext and restore it elsewhere — SessionConfig is write-only from Python, and `df_settings` is readable but lists `datafusion.runtime.*` keys with no namespace to set them back into — so parity has to be built the same way twice, and nothing said what "the same" covers. Nine items, each of which the example gets wrong somewhere on purpose to show the failure. **`extension-guide/query-planners.md` gains "plan against your own optimizer rules".** A planner returns protobuf rather than a plan handle, so every query serializes its output. Physical planning applies `session.physical_optimizers()`, which over FFI are the *host's* rules, so each one hands the library back a `ForeignExecutionPlan` — and a stock `CooperativeExec` wrapped that way has no reachable `try_to_proto`. A perfectly serializable node becomes unserializable by having crossed a boundary. It is also opaque to `downcast_ref`, so a planner that means to rewrite the plan cannot see what it was given. Wrapping the session with a locally-owned rule list fixes both, and the section says when to do that instead of delegating to a fallback: a planner does one or the other. **`extension-guide/codecs.md` gains "a table provider needs a logical codec".** A provider library reasonably concludes a physical codec is enough, since its scan is a physical node. It is enough until someone installs a query planner, which receives the logical plan — holding the provider as an `Arc` — and then the session fails while planning with "Error serializing custom table". Found by shipping the storage library without one. `extension_codec_durable_metadata` also now points at a codec that does encode durable metadata, which it previously could not: it described what to do, said the in-repo examples deliberately do not do it, and left the reader with no implementation to read. **`distributing-work/expressions.md` sharpens the UDF-portability rule.** The existing text said imports are captured by reference, which is true but not the useful distinction. What decides it is whether cloudpickle can resolve the name to an importable `module.qualname`: a module attribute like `pyarrow.compute` is stored as an import of that submodule and works, while a *function* in one of your modules becomes a pointer and requires your code installed on the worker. The same callable is around 1 kB from `__main__` and around 30 bytes from a package, so moving a helper into one silently changes what ships. Now a table, with the failure signature: a bare `ModuleNotFoundError` raised during plan decode, naming neither UDFs nor serialization. Also: a README for the example that says what it is not, and two checklist items — ship a logical codec with a provider, and decode in a different process in at least one test, since a token-registry codec passes every in-process round trip. Every `{ref}` added here resolves; checked by extracting defined labels and references across the docs tree. One dangling reference exists in `aggregations.md` (`spark-functions`) and predates this work. Co-Authored-By: Claude Opus 5 (1M context) --- docs/source/extension-guide/checklist.md | 16 +- docs/source/extension-guide/codecs.md | 50 ++++-- docs/source/extension-guide/query-planners.md | 47 ++++++ .../distributing-work/expressions.md | 32 +++- .../distributing-work/query-engines.md | 55 ++++++- examples/distributed/README.md | 147 ++++++++++++++++++ 6 files changed, 326 insertions(+), 21 deletions(-) create mode 100644 examples/distributed/README.md diff --git a/docs/source/extension-guide/checklist.md b/docs/source/extension-guide/checklist.md index 6e204b3f5..677ccbefb 100644 --- a/docs/source/extension-guide/checklist.md +++ b/docs/source/extension-guide/checklist.md @@ -56,6 +56,14 @@ publish. Each links to the page that explains it. - [ ] **You round-trip a plan in a test and assert *your* codec did the work.** Both being installed does not mean your node reached you. → {ref}`extension_codec_order` +- [ ] **You ship a logical codec too, if you contribute a table provider.** A + physical codec is not enough: an installed query planner receives the + logical plan, which holds your provider, and the session fails to plan + without one. → {ref}`extension_codec_provider_logical` +- [ ] **You decode in a *different process* in at least one test.** A codec + that parks the object in a process-global map passes every in-process + round trip and fails the first real one. + → {ref}`extension_codec_durable_metadata` ## Bundles and planners @@ -94,6 +102,8 @@ publish. Each links to the page that explains it. process-local token. The examples in this repository use tokens to make ownership observable; that is a demonstration, not a pattern. → {ref}`extension_codec_durable_metadata` -- [ ] **You have integration tests across a real FFI boundary.** The two - example crates in this repository are the pattern: build the cdylib, - install the wheel, then exercise it from Python. +- [ ] **You have integration tests across a real FFI boundary.** The example + trees in this repository are the pattern: build the cdylib, install the + wheel, then exercise it from Python. `examples/distributed` additionally + spawns worker processes, which is the only way to catch a codec that + only works in the process that wrote it. diff --git a/docs/source/extension-guide/codecs.md b/docs/source/extension-guide/codecs.md index 160439927..0bb7b9375 100644 --- a/docs/source/extension-guide/codecs.md +++ b/docs/source/extension-guide/codecs.md @@ -58,15 +58,47 @@ Your payload has to be enough to rebuild the object somewhere your process is not. Write the metadata a fresh instance can be constructed from — a path, a connection string, a schema, the options the object was created with. -The example codecs in this repository do not do this, and it is worth knowing -before copying them. They keep a process-local `HashMap` of live providers and -encode an integer token into it: encoding inserts, decoding removes. That makes -Rust type identity observable across three separately loaded libraries in one -test, which is what the examples exist to show. It also means a decode consumes -its token, so the same bytes cannot be decoded twice, one encoded plan cannot -fan out to several readers, and a plan that never reaches a decoder keeps its -provider alive for the life of the process. A real codec has none of those -properties because it does not park the object anywhere. +Two of the example codecs in this repository do not do this, and it is worth +knowing before copying them. `datafusion-ffi-example` and +`datafusion-ffi-query-planner-example` keep a process-local `HashMap` of live +providers and encode an integer token into it: encoding inserts, decoding +removes. That makes Rust type identity observable across three separately +loaded libraries in one test, which is what those examples exist to show. It +also means a decode consumes its token, so the same bytes cannot be decoded +twice, one encoded plan cannot fan out to several readers, and a plan that +never reaches a decoder keeps its provider alive for the life of the process. +A real codec has none of those properties because it does not park the object +anywhere. + +For one that does it properly, read +[`examples/distributed/storage-library`](https://github.com/apache/datafusion-python/tree/main/examples/distributed/storage-library). +Its payload is the file paths, the projection, the row limit, and the schema — +enough to rebuild the scan from nothing — and its tests decode a plan in a +separate interpreter that never registered the table. + +(extension_codec_provider_logical)= + +## A table provider needs a *logical* codec + +A provider library can reasonably conclude it needs only a physical codec: its +scan is a physical node, so that is where its own type appears. That holds +right up until someone installs a query planner. + +An FFI query planner is handed the **logical** plan, as protobuf. A logical +plan holds its tables as `Arc`, and the default codec's +`try_encode_table_provider` is unimplemented. So a session with your provider +and any engine installed fails while planning, before anything is executed, +with: + +```text +Error serializing custom table ... caused by +Execution error: No installed extension codec handled a table provider +``` + +Implement `try_encode_table_provider` and `try_decode_table_provider`, and +contribute the logical codec alongside the physical one. The payload can be +small — the storage library writes just the directory, because everything else +it holds is read back from there — but it has to exist. (extension_codec_ids)= diff --git a/docs/source/extension-guide/query-planners.md b/docs/source/extension-guide/query-planners.md index e373bb552..0a5fa94fa 100644 --- a/docs/source/extension-guide/query-planners.md +++ b/docs/source/extension-guide/query-planners.md @@ -38,6 +38,53 @@ against the codecs of the session that will run the query. `MyQueryPlanner` in [`datafusion-ffi-query-planner-example`] is the worked implementation. +(planner_host_optimizer_rules)= + +## Plan against your own optimizer rules + +Your planner returns its plan as **protobuf**, not as a handle. Every query +therefore serializes what you produce, and anything in it that cannot be +encoded is your problem rather than a distant one. + +That matters because of where physical optimization runs. Physical planning +applies `session.physical_optimizers()`, and when the session arrived over FFI +those rules are the *host's* — so each one runs back across the boundary and +hands you a `ForeignExecutionPlan` wrapping the result. `EnsureCooperative` is +on by default and will do exactly this. A stock `CooperativeExec` produced that +way has no reachable `try_to_proto`, so a node that is perfectly serializable +in the process that made it becomes unserializable in yours: + +```text +Internal error: Unsupported plan and extension codec failed with +[This feature is not implemented: PhysicalExtensionCodec is not provided]. +Plan: ForeignExecutionPlan { name: "CooperativeExec", ... } +``` + +A foreign node is also opaque to `downcast_ref`, so a planner that means to +*rewrite* the plan — inserting stages, say — cannot inspect what it was given. + +Both problems go away if the rules run on your side. Wrap the session you were +handed in one that delegates everything except `physical_optimizers()`, and +return the stock rule set from there: + +```rust +let local = LocalOptimizerSession::new(session); // owns PhysicalOptimizer::default().rules +DefaultPhysicalPlanner::default() + .create_physical_plan(logical_plan, &local) + .await? +``` + +`LocalOptimizerSession` in +[`examples/distributed/engine-library`](https://github.com/apache/datafusion-python/tree/main/examples/distributed/engine-library) +is about twenty delegating methods and one override. + +Delegating to a `fallback` avoids the problem differently, by not planning at +all: the plan comes back from whoever you delegated to, already concrete. That +is the right choice for a planner that only layers behaviour on another, and +the wrong one for a planner that needs to rewrite the result — you cannot +rewrite a subtree you hold an opaque handle to. A planner does one or the +other. + ## One planner per session A session holds exactly one query planner. Calling `set_query_planner` again diff --git a/docs/source/user-guide/distributing-work/expressions.md b/docs/source/user-guide/distributing-work/expressions.md index 252defc24..b817b3e49 100644 --- a/docs/source/user-guide/distributing-work/expressions.md +++ b/docs/source/user-guide/distributing-work/expressions.md @@ -119,14 +119,30 @@ requirements on the worker environment: stamps the sender's `(major, minor)`; mismatches raise a clear error naming both versions. Align the Python version on driver and workers. -- **Imported modules must be importable on the worker.** cloudpickle - captures the callable *by value* (bytecode and closure cells travel - whole), but names resolved through `import` are captured *by - reference* — module path only. A UDF doing - `from mylib import transform` requires `mylib` installed on the - worker. Same applies to bound methods of imported classes. - Self-contained UDFs (no imports beyond what the worker already has, - e.g. `pyarrow`) avoid this entirely. +- **Anything the callable names must be reachable on the worker.** + cloudpickle captures the function's own body *by value* — bytecode and + closure cells travel whole — but every global it refers to is captured *by + reference* if cloudpickle can resolve it to an importable + `module.qualname`. The worker then imports it by that path. + + So the rule is not "imports are bad", it is **whether the name has an + importable home**: + + | The callable refers to | Travels as | Worker needs | + | --- | --- | --- | + | a nested or `__main__`-level function | the function itself | nothing | + | a module, including a submodule like `pyarrow.compute` | an import of that module | the module installed | + | a function in an importable module of yours | a pointer to `yourmod.helper` | **your code installed** | + + The third row is the one that surprises people, and the size difference + makes it concrete: one small function is around 1 kB pickled from + `__main__` and around 30 bytes from an importable module, because the + second is only a pointer. Moving a helper out of a script and into a + package silently changes what gets shipped. + + It fails on the worker as a bare + `ModuleNotFoundError: No module named 'yourmod'`, raised while the plan is + being decoded, with nothing in the message about UDFs or serialization. ## Registering shared UDFs on workers diff --git a/docs/source/user-guide/distributing-work/query-engines.md b/docs/source/user-guide/distributing-work/query-engines.md index 0aa045c79..a67d99fae 100644 --- a/docs/source/user-guide/distributing-work/query-engines.md +++ b/docs/source/user-guide/distributing-work/query-engines.md @@ -71,11 +71,64 @@ If you install more than one library, pass them in one {ref}`user_guide_extensions` for the details of installing extension libraries, and {ref}`ffi` if you want to write an engine yourself. +(distributed_worker_parity)= + +## What a worker has to reproduce + +An engine ships your plan to a process that has never seen your session. That +process has to be able to rebuild everything the plan refers to, and there is +**no way to snapshot a {py:class}`~datafusion.SessionContext` and restore it +somewhere else**: {py:class}`~datafusion.SessionConfig` is write-only from +Python, and while `information_schema.df_settings` can be read back, it lists +`datafusion.runtime.*` keys that have no configuration namespace to set them +into again. + +So parity is not automatic. It is something you build the same way twice, and +these are the things that have to match. Most engines handle several of them +for you — check which. + +- **Codec ids.** A plan records which codec wrote each payload, and decoding + routes on that id. Pin ids with `__datafusion_codec_id__` rather than + letting them default to a class's import path, and compare + {py:meth}`~datafusion.SessionContext.physical_extension_codec_ids` on both + sides before shipping anything. See {ref}`extension_codec_ids`. +- **Functions the plan names.** A function resolves either from the receiving + session's registry or from a codec. Either is enough; neither is automatic. + A Python UDF is the exception — it travels inside the plan. +- **Object stores**, registered *before* the plan is decoded rather than + before it is executed. Decoding a Parquet scan resolves its store. +- **Config extensions**, installed before any namespaced key is set. Setting a + key in a namespace that has not been declared is an error, not a no-op. +- **The Python minor version**, if any inline Python UDF is involved. + Cloudpickle payloads are stamped with the sender's version and refuse to + load on another. Launching workers with `sys.executable` makes this true by + construction; a hardcoded `python` does not. +- **The `cloudpickle` version**, which is *not* stamped. Cross-version loading + usually works and is not guaranteed. Pin it. +- **`target_partitions`**, if a worker re-plans anything. Left to default it + follows the core count, so two differently-sized machines disagree. + +Two more that are about lifetime rather than configuration: + +- **One session per worker, alive for the whole process.** An FFI codec + resolves names against the session captured when its bundle was installed, + and that reference is weak — see {ref}`extension_sessions`. +- **Pass a context to `to_bytes`.** {py:meth}`~datafusion.ExecutionPlan.to_bytes` + takes an optional context, and without one it uses an empty codec chain that + cannot encode any library's nodes. The argument being optional makes this + easy to miss, because the failure only appears once an extension node is in + the plan. + ## Available engines Query-level distribution is being built upstream. Neither project below is usable from datafusion-python yet; both sections will fill -in as the integrations land. +in as the integrations land. In the meantime the repository contains a +worked example you can read and run: +[`examples/distributed`](https://github.com/apache/datafusion-python/tree/main/examples/distributed) +splits a query across worker processes using three separate extension +libraries, and is written to make each of the requirements above visible — +including the ways they fail. ### datafusion-distributed diff --git a/examples/distributed/README.md b/examples/distributed/README.md new file mode 100644 index 000000000..6d1343ed7 --- /dev/null +++ b/examples/distributed/README.md @@ -0,0 +1,147 @@ + + +# Three libraries, one distributed query + +A worked example of what `datafusion-python`'s extension protocol is *for*: +several independently compiled libraries, none of which knows about the +others, cooperating on a single query whose work runs in separate operating +system processes. + +Everything here is real. The workers are separate interpreters. The plan they +run was serialized by the driver and decoded by them. If you break the +serialization, the tests fail. + +## The three libraries + +| Crate | Owns | Installed with | +| --- | --- | --- | +| `udf-library` (`dfx_udfs`) | a scalar function, an aggregate, a window function | **by hand** — `register_udf` plus two `with_*_extension_codec` calls | +| `storage-library` (`dfx_storage`) | a Parquet table provider and its own scan node | `with_extensions` | +| `engine-library` (`dfx_engine`) | a query planner, a stage node, and the driver/worker machinery | `with_extensions` | + +One of them is deliberately old-fashioned. `dfx_udfs` exposes no +`__datafusion_session_components__`, so it cannot be installed as a bundle and +its caller has to do five things in the right order instead of one. That is +not a strawman: `SessionExtensionComponents` carries codec fields only, so a +library that contributes *functions* has nowhere to put them today. Mixed +setups are the normal case, and this example shows what one costs. + +## Running it + +```console +$ cd examples/distributed/engine-library +$ uv venv && uv pip install pytest pyarrow ../.. ../storage-library ../udf-library +$ uv run maturin develop +$ uv run pytest python/tests/_test*.py +``` + +Against the real TPC-H data — generate it as +[`examples/tpch`](../tpch/README.md) describes, then: + +```console +$ uv run python ../run_tpch.py --partitions 4 +``` + +## What actually happens + +The engine's planner splits the plan at the partial aggregate, which is where +DataFusion has already split it for its own reasons: a `GROUP BY` becomes a +partial pass per input partition and a final pass that merges them, and the +partial passes are independent by construction. + +``` +SortPreservingMergeExec + ProjectionExec + AggregateExec: mode=FinalPartitioned <- driver merges + RepartitionExec: Hash([l_returnflag], 2) + FFI_ExecutionPlan: ShuffleStageExec <- shipped to workers + AggregateExec: mode=Partial <- one worker per partition + FFI_ExecutionPlan: PartitionedParquetExec +``` + +The driver serializes the `ShuffleStageExec` subtree, starts one worker per +partition, and waits. Each worker rebuilds an equivalent session, decodes the +plan, runs *its* partition, and writes the result to an Arrow IPC file. The +driver then runs the whole query itself — and the stage node, finding the +files already there, streams them instead of recomputing. + +One node does both halves of that exchange, which is why nothing has to +rewrite the plan in between. It also means a query run with no workers at all +still gets the right answer; it just does the work itself. + +## The four things worth reading + +**`engine-library/python/dfx_engine/session.py`** is the point of the whole +example. There is no way to snapshot a `SessionContext` and restore it +elsewhere, so worker parity cannot be automated — it has to be *built the same +way twice*, from data small enough to put in a message. Both the driver and +every worker call one `build_session`. Anything a query depends on that is not +in the `SessionSpec` is a bug waiting for a worker to find it. + +**`storage-library/src/codec.rs`** is the repository's only codec that encodes +durable metadata. The others park the live object in a process-global map and +encode an integer token, which is fine for making Rust type identity +observable in a test and useless the moment the bytes leave the process. This +one writes the file paths, the projection, and the schema, so the same bytes +decode twice, decode on ten workers, and decode tomorrow. + +**`udf-library/python/tests/_test_udfs.py`** shows that installing a +library's codec is an *alternative* to registering its functions, not an +addition. Three workers, three configurations: + +| worker has | result | +| --- | --- | +| the codec, no registrations | works; the codec rebuilds each function from its name | +| the registrations, no codec | works; the registry answers first and the codec is never consulted | +| neither | fails, naming `dfx_net_revenue` | + +The middle row is the trap. On the driver, where the functions are always +registered, a broken or missing codec looks completely fine. + +**`engine-library/python/tests/_test_three_libraries.py`** runs the queries, +and pins the failure modes next to the successes — including a Python UDF that +works on the driver and fails on the worker. + +## Things this example is not + +It writes shuffle results to local files, so "distributed" means several +processes on one machine. Adding a network is a transport change and would not +alter anything above it. + +It holds one partition of results in memory before writing, because an Arrow +IPC stream needs its schema up front. A production engine would stream to the +file and track completion separately. + +It has one stage. A real engine chains them, and the interesting problems — +scheduling, retries, straggler handling, memory limits — all live in the part +this example replaces with `subprocess.Popen` and a `for` loop. + +It is slower than running the query in one process. Four processes on one +laptop cannot beat one process that skips a round trip through Arrow IPC +files. The comparison the tests make is *agreement*, not speed. + +## Further reading + +- [Distributed query engines](https://datafusion.apache.org/python/user-guide/distributing-work/query-engines.html) + — using an engine, and the checklist for what a worker has to reproduce. +- [Extension Guide](https://datafusion.apache.org/python/extension-guide/index.html) + — writing a library like these. +- [Encode metadata, not a handle to a live object](https://datafusion.apache.org/python/extension-guide/codecs.html#encode-metadata-not-a-handle-to-a-live-object) + — what a codec should put on the wire, and why. From 05b40fb8e2d900abda137e67fdb6861343f8e7e0 Mon Sep 17 00:00:00 2001 From: Tim Saucer Date: Thu, 10 Sep 2026 13:24:06 -0400 Subject: [PATCH 07/24] Build the example libraries on the host, and prime their deps once The five example libraries built in the Linux x86_64 job are test fixtures, not release artifacts. The `test-ffi-manylinux-x86_64` artifact they feed is consumed only by test.yml, which installs them on a runner of the same kind. Building them in the manylinux container therefore bought nothing, and cost a container start and an in-container rustup install apiece. They now build on the host with `container: off`, which also brings them under the Swatinem/rust-cache the job already sets up. Dropping the explicit `target` matters as much as dropping the container: naming a triple puts the artifacts under `target//debug`, where the priming build below cannot reach them. maturin runs one `cargo rustc` per library, and cargo resolves features per package, so each invocation could rebuild the shared datafusion crates under its own feature union -- measured at 78 seconds for a single transition between two of these libraries, on an otherwise warm target directory. A single `cargo build` spanning all five resolves that union once. Measured afterwards, the five per-package builds reuse every dependency and only compile their own leaf crate: about 6 seconds in total. Co-Authored-By: Claude Opus 5 (1M context) --- .github/workflows/build.yml | 49 +++++++++++++++++++++++++------------ 1 file changed, 33 insertions(+), 16 deletions(-) diff --git a/.github/workflows/build.yml b/.github/workflows/build.yml index 4ed5782bd..649a2c51f 100644 --- a/.github/workflows/build.yml +++ b/.github/workflows/build.yml @@ -185,58 +185,75 @@ jobs: features: "protoc,substrait" manylinux: "2_28" + # The example libraries below are test fixtures, not release artifacts: + # the `test-ffi-manylinux-x86_64` artifact they feed is consumed only by + # test.yml, which installs them on a runner like this one. They therefore + # build on the host (`container: off`) rather than in the manylinux + # container, which drops a container start and an in-container rustup + # install per library and lets the Swatinem/rust-cache above cover them. + # + # None of them pass `target`, so their artifacts land in `target/debug` + # alongside the priming build's. Naming a target would move them under + # `target//debug` and the priming build would be wasted. + # + # maturin runs one `cargo rustc` per library and cargo resolves features + # per package, so five separate invocations each rebuild the shared + # datafusion crates under their own feature union -- measured at over a + # minute for a single transition between two of these libraries. One + # build spanning all five resolves that union once; afterwards each + # maturin step only compiles its own leaf crate. + - name: Prime example crate dependencies + if: matrix.python-tag == 'abi3' + run: > + cargo build --lib + -p datafusion-ffi-example + -p datafusion-ffi-query-planner-example + -p dfx-storage + -p dfx-udfs + -p dfx-engine + # FFI test wheel only needs to be built once per platform; gate to abi3. - name: Build FFI provider test library if: matrix.python-tag == 'abi3' uses: PyO3/maturin-action@v1 with: - target: x86_64-unknown-linux-gnu - manylinux: "2_28" + container: "off" working-directory: examples/datafusion-ffi-example args: --out dist - rustup-components: rust-std - name: Build FFI query planner test library if: matrix.python-tag == 'abi3' uses: PyO3/maturin-action@v1 with: - target: x86_64-unknown-linux-gnu - manylinux: "2_28" + container: "off" working-directory: examples/datafusion-ffi-query-planner-example args: --out dist - rustup-components: rust-std - # The three libraries of the distributed example. Built in dependency + # The three libraries of the distributed example. Listed in dependency # order for readability only; they are independent cdylibs. - name: Build distributed example storage library if: matrix.python-tag == 'abi3' uses: PyO3/maturin-action@v1 with: - target: x86_64-unknown-linux-gnu - manylinux: "2_28" + container: "off" working-directory: examples/distributed/storage-library args: --out dist - rustup-components: rust-std - name: Build distributed example UDF library if: matrix.python-tag == 'abi3' uses: PyO3/maturin-action@v1 with: - target: x86_64-unknown-linux-gnu - manylinux: "2_28" + container: "off" working-directory: examples/distributed/udf-library args: --out dist - rustup-components: rust-std - name: Build distributed example engine library if: matrix.python-tag == 'abi3' uses: PyO3/maturin-action@v1 with: - target: x86_64-unknown-linux-gnu - manylinux: "2_28" + container: "off" working-directory: examples/distributed/engine-library args: --out dist - rustup-components: rust-std - name: Archive wheels uses: actions/upload-artifact@v7 From 016ad9d68924c55f2e95f51d50d686101fb5cfb5 Mon Sep 17 00:00:00 2001 From: Tim Saucer Date: Thu, 10 Sep 2026 14:21:31 -0400 Subject: [PATCH 08/24] Number the stages, and ship all of them `insert_stage` wrapped every topmost partial aggregate but gave each one the same `STAGE_ID`. Two stages sharing an id exchange results through the same `stage-1-part-N.arrow` paths, and a union drives its branches concurrently, so they race. A `UNION ALL` of two `GROUP BY`s fails outright: Internal error: dfx_engine: publishing .../stage-1-part-1.arrow: No such file or directory (os error 2) Three defects, one symptom: - Stage ids are now allocated from a counter threaded through `insert_stage`, not a constant. Threaded rather than global because the driver plans the same query twice -- once to ship the stages, once to read their results back -- and only per-call numbering has both calls agree on which stage is which. `stages_inserted` counts every stage instead of once per planning call. - The temporary file a writer builds a partition in is unique to that writer rather than derived from the final name. Two writers of one partition were interleaving their batches into a single file and then one rename found it already consumed, which is why the failure surfaced as a rename error. The rename that is supposed to make publishing atomic was the thing that broke. - The driver ships every stage. `find_stages` returns all of them in pre-order and `run_distributed` dispatches one worker per `(stage_id, partition)`; leaving one behind had the driver compute it locally while the plan it shipped claimed otherwise. `DistributedResult.partitions`/`worker_rows` become `tasks`/`task_rows` keyed by both, since partition 0 of two stages is two pieces of work. Python cannot read a stage id back off a plan -- the FFI wrapper's display replaces `ShuffleStageExec: stage=2` -- so the numbering is a convention shared with Rust, and `_internal.stage_id(index)` is the one place the arithmetic is written down. Co-Authored-By: Claude Opus 5 (1M context) --- examples/distributed/README.md | 4 +- .../python/dfx_engine/driver.py | 135 +++++++++++------- .../python/tests/_test_three_libraries.py | 72 +++++++++- .../distributed/engine-library/src/lib.rs | 16 ++- .../distributed/engine-library/src/planner.rs | 81 +++++++---- .../distributed/engine-library/src/stage.rs | 25 +++- examples/distributed/run_tpch.py | 4 +- 7 files changed, 245 insertions(+), 92 deletions(-) diff --git a/examples/distributed/README.md b/examples/distributed/README.md index 6d1343ed7..35fc1d5e8 100644 --- a/examples/distributed/README.md +++ b/examples/distributed/README.md @@ -129,7 +129,9 @@ It holds one partition of results in memory before writing, because an Arrow IPC stream needs its schema up front. A production engine would stream to the file and track completion separately. -It has one stage. A real engine chains them, and the interesting problems — +It has one *level* of stages. A query with an aggregate in each branch of a +union gets one stage per branch, and they run side by side — but no stage ever +feeds another. A real engine chains them, and the interesting problems — scheduling, retries, straggler handling, memory limits — all live in the part this example replaces with `subprocess.Popen` and a `for` loop. diff --git a/examples/distributed/engine-library/python/dfx_engine/driver.py b/examples/distributed/engine-library/python/dfx_engine/driver.py index 804aadf7b..8f1d8d6f9 100644 --- a/examples/distributed/engine-library/python/dfx_engine/driver.py +++ b/examples/distributed/engine-library/python/dfx_engine/driver.py @@ -18,10 +18,10 @@ """The driver: split a query into tasks, fan them out, collect the answer. The shape is deliberately boring, because the interesting part is not the -scheduling. What matters is the four things the driver has to get right, each +scheduling. What matters is the five things the driver has to get right, each of which is a way a real deployment goes wrong: -1. It serializes the stage **with** its session. ``to_bytes(None)`` uses an +1. It serializes each stage **with** its session. ``to_bytes(None)`` uses an empty codec chain and cannot encode any library's node. 2. It ships the codec ids it used, so a worker can refuse a plan it would misread rather than decode it with the wrong codec. @@ -30,6 +30,9 @@ cannot disagree. 4. It waits for every worker before reading, because the stage node decides whether to read or recompute by looking at the filesystem. +5. It ships *every* stage. A plan can hold more than one -- an aggregate in + each branch of a union, say -- and they are independent subtrees rather + than a chain. """ from __future__ import annotations @@ -49,7 +52,12 @@ from datafusion.plan import ExecutionPlan from datafusion.user_defined import ScalarUDF -__all__ = ["DistributedResult", "find_stage", "run_distributed"] +__all__ = [ + "DistributedResult", + "find_stage", + "find_stages", + "run_distributed", +] class DistributedResult: @@ -58,21 +66,26 @@ class DistributedResult: def __init__( self, batches: list[pa.RecordBatch], - partitions: list[int], - worker_rows: dict[int, int], + tasks: list[tuple[int, int]], + task_rows: dict[tuple[int, int], int], ) -> None: self.batches = batches - self.partitions = partitions - """Partition indices that were dispatched, one per worker.""" - self.worker_rows = worker_rows - """Rows each worker produced, keyed by partition index.""" + self.tasks = tasks + """``(stage_id, partition)`` pairs that were dispatched, one per worker. + + Keyed by both, not by partition alone: a query with an aggregate in + more than one branch has more than one stage, and partition 0 of each + is a different piece of work. + """ + self.task_rows = task_rows + """Rows each worker produced, keyed as :attr:`tasks` is.""" STAGE_NODE_NAME = "ShuffleStageExec" -def find_stage(plan: ExecutionPlan) -> ExecutionPlan | None: - """Locate the stage node the planner inserted. +def find_stages(plan: ExecutionPlan) -> list[ExecutionPlan]: + """Locate every stage node the planner inserted, in pre-order. Matched on the display string because a Python caller has no way to downcast a Rust plan node -- there is no ``isinstance`` across an FFI @@ -87,27 +100,39 @@ def find_stage(plan: ExecutionPlan) -> ExecutionPlan | None: A foreign node reports its own name nested inside the wrapper's, which makes anchored matches on plan text quietly wrong -- the kind of thing that works in a single-library test and fails the moment a real extension - is involved. + is involved. It is also why the stage *id* has to be recomputed here + rather than read: the wrapper dropped it. + + Pre-order is the contract with the planner, which numbers stages in the + same walk, so the nth node returned here has the id + ``_internal.stage_id(n)``. The recursion stops at a stage rather than + descending into it, because the planner never nests one inside another. """ if STAGE_NODE_NAME in plan.display(): - return plan - for child in plan.children(): - found = find_stage(child) - if found is not None: - return found - return None + return [plan] + return [stage for child in plan.children() for stage in find_stages(child)] + + +def find_stage(plan: ExecutionPlan) -> ExecutionPlan | None: + """The first stage in `plan`, or ``None``. + + For tests and callers that only care whether the planner split at all. + See :func:`find_stages`. + """ + stages = find_stages(plan) + return stages[0] if stages else None def _dispatch( - envelope: dict, envelope_dir: pathlib.Path, partition: int + envelope: dict, envelope_dir: pathlib.Path, stage_id: int, partition: int ) -> subprocess.Popen[str]: - """Start one worker for one partition. + """Start one worker for one ``(stage, partition)``. ``sys.executable``, not ``python``: a worker on a different Python minor version cannot load a cloudpickled inline UDF, and that failure is far from its cause. """ - path = envelope_dir / f"task-{partition}.json" + path = envelope_dir / f"task-{stage_id}-{partition}.json" path.write_text(json.dumps(envelope)) return subprocess.Popen( # noqa: S603 [sys.executable, "-m", "dfx_engine.worker", str(path)], @@ -120,7 +145,7 @@ def _dispatch( 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. + """Run `sql`, executing each stage partition in its own worker process. Requires ``spec.shuffle_dir``: without it the planner inserts no stage and there is nothing to distribute. @@ -135,13 +160,13 @@ def run_distributed( message = "run_distributed needs a shuffle_dir; build_session got none" raise ValueError(message) - ctx, engine, _storage = build_session(spec) + 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) - if stage is None: + stages = find_stages(plan) + if not stages: message = ( "no ShuffleStageExec in the plan; the engine's planner did not run, " "or its config extension was not registered" @@ -151,48 +176,56 @@ def run_distributed( shuffle_dir = pathlib.Path(spec.shuffle_dir) shuffle_dir.mkdir(parents=True, exist_ok=True) - # Encode the stage subtree, through the session that owns the codecs. - plan_path = shuffle_dir / "stage.plan" - plan_path.write_bytes(stage.to_bytes(ctx)) - - stage_id = _internal.stage_id() - partitions = list(range(stage.partition_count)) - envelopes = [ - { - "spec": spec.to_json(), - "plan": str(plan_path), - "stage_id": stage_id, - "partition": partition, - } - for partition in partitions - ] - - # One process per partition, all in flight together. This is the claim the + # One task per (stage, partition). Every stage is shipped, not just the + # first: a query with an aggregate in two branches has two independent + # subtrees, and leaving one behind would have the driver compute it + # locally while the plan it shipped claimed otherwise. + tasks: list[tuple[int, int]] = [] + envelopes = [] + for index, stage in enumerate(stages): + # The id the planner gave this stage, recomputed from its position + # because the FFI wrapper's display does not carry it. + stage_id = _internal.stage_id(index) + # Encode the stage subtree, through the session that owns the codecs. + plan_path = shuffle_dir / f"stage-{stage_id}.plan" + plan_path.write_bytes(stage.to_bytes(ctx)) + for partition in range(stage.partition_count): + tasks.append((stage_id, partition)) + envelopes.append( + { + "spec": spec.to_json(), + "plan": str(plan_path), + "stage_id": stage_id, + "partition": partition, + } + ) + + # One process per task, all in flight together. This is the claim the # example is making: each worker reads a different file and writes a # different result, so they need no coordination beyond the directory. workers = [ - _dispatch(envelope, shuffle_dir, partition) - for envelope, partition in zip(envelopes, partitions, strict=True) + _dispatch(envelope, shuffle_dir, stage_id, partition) + for envelope, (stage_id, partition) in zip(envelopes, tasks, strict=True) ] - worker_rows: dict[int, int] = {} + task_rows: dict[tuple[int, int], int] = {} failures = [] - for partition, worker in zip(partitions, workers, strict=True): + for task, worker in zip(tasks, workers, strict=True): stdout, stderr = worker.communicate() if worker.returncode != 0: - failures.append(f"partition {partition} failed:\n{stderr}") + stage_id, partition = task + failures.append(f"stage {stage_id} partition {partition} failed:\n{stderr}") continue - worker_rows[partition] = json.loads(stdout)["rows"] + task_rows[task] = json.loads(stdout)["rows"] if failures: raise RuntimeError("\n".join(failures)) # Now run the whole query here. Every stage partition has a file, so the - # stage node streams them instead of recomputing -- the driver does only + # stage nodes stream them instead of recomputing -- the driver does only # the final merge. batches = ctx.sql(sql).collect() - _ = engine - return DistributedResult(batches, partitions, worker_rows) + return DistributedResult(batches, tasks, task_rows) def run_local( diff --git a/examples/distributed/engine-library/python/tests/_test_three_libraries.py b/examples/distributed/engine-library/python/tests/_test_three_libraries.py index a72c897ac..36215b978 100644 --- a/examples/distributed/engine-library/python/tests/_test_three_libraries.py +++ b/examples/distributed/engine-library/python/tests/_test_three_libraries.py @@ -34,7 +34,7 @@ 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.driver import find_stage, find_stages, run_distributed, run_local from dfx_engine.session import SessionSpec, build_session, expected_codec_ids from dfx_engine.worker import run_task @@ -83,7 +83,8 @@ 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] + # One stage, so every task is a partition of stage 1. + assert result.tasks == [(1, 0), (1, 1), (1, 2), (1, 3)] assert _rows(result.batches) == _rows(run_local(Q1, spec)) # Checked by hand against the fixture. assert _rows(result.batches) == [ @@ -211,11 +212,11 @@ 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) + assert sorted(result.tasks) == [(1, 0), (1, 1), (1, 2), (1, 3)] + assert len(set(result.tasks)) == len(result.tasks) # 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) + assert sum(result.task_rows.values()) == 8 + assert set(result.task_rows) == set(result.tasks) def test_each_worker_published_its_own_file(spec: SessionSpec) -> None: @@ -231,6 +232,8 @@ def test_each_worker_published_its_own_file(spec: SessionSpec) -> None: for partition in range(4) ) assert produced == expected + # Nothing half-written: a temporary file is unique to its writer and is + # renamed away when the partition is complete. assert list(shuffle.glob("*.tmp")) == [] @@ -295,6 +298,63 @@ def test_the_plan_splits_at_the_partial_aggregate(spec: SessionSpec) -> None: assert stage.output_partitioning.scheme == "UnknownPartitioning" +TWO_BRANCHES = """ +select l_returnflag as g, sum(l_quantity) as qty +from lineitem group by l_returnflag +union all +select l_linestatus as g, sum(l_quantity) as qty +from other group by l_linestatus +""" + + +def test_two_branches_get_two_separately_numbered_stages( + lineitem_dir: pathlib.Path, tmp_path: pathlib.Path +) -> None: + """A partial aggregate in each branch is two stages, not one. + + Both branches are independent subtrees and both want shipping, so the + planner wraps each. They must not share a stage id: a stage exchanges + results through paths built from that id, so two stages numbered alike + would write to the same files -- and because a union drives its branches + concurrently, they would do it at the same time. That fails outright + rather than quietly, but only because the schemas happen to differ. + """ + spec = SessionSpec( + tables={"lineitem": str(lineitem_dir), "other": str(lineitem_dir)}, + shuffle_dir=str(tmp_path / "shuffle"), + target_partitions=2, + ) + + ctx, engine, _storage = build_session(spec) + plan = ctx.sql(TWO_BRANCHES).execution_plan() + + assert plan.display_indent().count("ShuffleStageExec") == 2 + assert len(find_stages(plan)) == 2 + assert engine.stages_inserted() == 2 + + result = run_distributed(TWO_BRANCHES, spec) + + # Four partitions apiece, under two distinct stage ids. + assert sorted(result.tasks) == [ + (stage, part) for stage in (1, 2) for part in range(4) + ] + # Sorted: a union has no ordering of its own, so the branches interleave + # differently from run to run. + assert sorted(_rows(result.batches)) == sorted(_rows(run_local(TWO_BRANCHES, spec))) + + # Each stage published its own files, so neither read the other's. + produced = sorted( + path.name for path in pathlib.Path(spec.shuffle_dir).glob("*.arrow") + ) + assert produced == sorted( + pathlib.Path( + _internal.partition_path(spec.shuffle_dir, _internal.stage_id(index), part) + ).name + for index in range(2) + for part in range(4) + ) + + # --- the ways it goes wrong ------------------------------------------------- diff --git a/examples/distributed/engine-library/src/lib.rs b/examples/distributed/engine-library/src/lib.rs index 2dd3b7e17..2a38f7eb4 100644 --- a/examples/distributed/engine-library/src/lib.rs +++ b/examples/distributed/engine-library/src/lib.rs @@ -28,6 +28,7 @@ //! //! One of three libraries in `examples/distributed`. This one owns execution. +use pyo3::exceptions::PyValueError; use pyo3::prelude::*; use crate::config::DfxEngineConfig; @@ -52,10 +53,19 @@ fn partition_path(shuffle_dir: &str, stage_id: u32, partition: usize) -> String .into_owned() } -/// The stage id this engine's planner produces. +/// Id of the `index`th stage in a plan, counting in pre-order from the root. +/// +/// Exported for the same reason as [`partition_path`]: a foreign node's +/// one-line display is replaced by the FFI wrapper's, so Python cannot read a +/// stage id back off a plan and has to agree with Rust about the numbering +/// instead. Keeping the arithmetic here means the two cannot drift. #[pyfunction] -fn stage_id() -> u32 { - planner::STAGE_ID +#[pyo3(signature = (index=0))] +fn stage_id(index: usize) -> PyResult { + u32::try_from(index) + .ok() + .and_then(|index| planner::FIRST_STAGE_ID.checked_add(index)) + .ok_or_else(|| PyValueError::new_err(format!("stage index {index} is out of range"))) } #[pymodule] diff --git a/examples/distributed/engine-library/src/planner.rs b/examples/distributed/engine-library/src/planner.rs index 11de45039..f81ed6e2a 100644 --- a/examples/distributed/engine-library/src/planner.rs +++ b/examples/distributed/engine-library/src/planner.rs @@ -25,8 +25,15 @@ //! [`ShuffleStageExec`] is the whole rewrite. //! //! A query with no aggregate gets its whole plan wrapped instead, so there is -//! always exactly one stage and the orchestration in Python has one shape to +//! always at least one stage and the orchestration in Python has one shape to //! deal with. +//! +//! There can be more than one. A `UNION ALL` of two `GROUP BY`s, or a join +//! between two of them, puts a partial aggregate in each branch, and the +//! branches are independent subtrees that both want shipping. Each stage +//! therefore gets its own id, assigned in the order a pre-order walk finds +//! them: stages exchange results through paths built from that id, so two +//! stages sharing one would write to the same files and race each other. use std::sync::Arc; use std::sync::atomic::{AtomicUsize, Ordering}; @@ -54,9 +61,15 @@ pub(crate) const SHUFFLE_DIR_KEY: &str = "dfx_engine.shuffle_dir"; /// every foreign config extension is namespaced under `datafusion_ffi`. const FFI_SHUFFLE_DIR_KEY: &str = "datafusion_ffi.dfx_engine.shuffle_dir"; -/// The one stage id this engine produces. A real engine would number a chain -/// of them; one is enough to show the mechanism. -pub(crate) const STAGE_ID: u32 = 1; +/// Id of the first stage in a plan. Later ones count up from here. +/// +/// Numbering is a *convention shared with Python*, which cannot read a stage +/// id back off a plan: a foreign node's one-line display is replaced by the +/// FFI wrapper's, so `ShuffleStageExec: stage=2` never reaches the driver. +/// Both sides instead agree that the nth stage found in a pre-order walk has +/// id `FIRST_STAGE_ID + n`, and [`crate::stage_id`] is the one place that +/// arithmetic is written down. +pub(crate) const FIRST_STAGE_ID: u32 = 1; /// What the planner did, so a test can assert it rather than infer it. #[derive(Default, Debug)] @@ -80,37 +93,46 @@ pub(crate) fn shuffle_dir_from_options(options: &ConfigOptions) -> Option, shuffle_dir: &str, -) -> Result<(Arc, bool)> { + next_id: &mut u32, +) -> Result<(Arc, usize)> { if let Some(aggregate) = plan.downcast_ref::() && matches!(aggregate.mode(), AggregateMode::Partial) { - let stage = ShuffleStageExec::new(STAGE_ID, shuffle_dir.to_string(), Arc::clone(&plan)); - return Ok((Arc::new(stage), true)); + let stage_id = *next_id; + *next_id += 1; + let stage = ShuffleStageExec::new(stage_id, shuffle_dir.to_string(), Arc::clone(&plan)); + return Ok((Arc::new(stage), 1)); } - let mut inserted = false; + let mut inserted = 0; let mut children = Vec::new(); for child in plan.children() { - let (child, child_inserted) = insert_stage(Arc::clone(child), shuffle_dir)?; - inserted |= child_inserted; + let (child, child_inserted) = insert_stage(Arc::clone(child), shuffle_dir, next_id)?; + inserted += child_inserted; children.push(child); } - if !inserted { - return Ok((plan, false)); + if inserted == 0 { + return Ok((plan, 0)); } // `Keep`: the replacement is a `ShuffleStageExec` wrapping the node it // replaced, and that node takes its properties from its child, so the // parent's view of its children is unchanged. let options = ReplaceChildrenOptions::new(ChildrenPropertiesMode::Keep); - Ok((plan.replace_children(children, options)?, true)) + Ok((plan.replace_children(children, options)?, inserted)) } #[derive(Debug)] @@ -160,18 +182,23 @@ impl QueryPlanner for DistributedQueryPlanner { return Ok(plan); }; - let (plan, inserted) = insert_stage(plan, &shuffle_dir)?; - if inserted { - self.observations - .stages_inserted - .fetch_add(1, Ordering::SeqCst); - return Ok(plan); - } + let mut next_id = FIRST_STAGE_ID; + let (plan, inserted) = insert_stage(plan, &shuffle_dir, &mut next_id)?; + + // Nothing to split at, so the whole plan is the one stage. Counted + // the same way as the rewritten case, so `stages_inserted` is the + // number of stages a test can expect to find in the plan. + let (plan, inserted) = match inserted { + 0 => ( + Arc::new(ShuffleStageExec::new(FIRST_STAGE_ID, shuffle_dir, plan)) as _, + 1, + ), + inserted => (plan, inserted), + }; - // Nothing to split at, so the whole plan is the stage. self.observations .stages_inserted - .fetch_add(1, Ordering::SeqCst); - Ok(Arc::new(ShuffleStageExec::new(STAGE_ID, shuffle_dir, plan))) + .fetch_add(inserted, Ordering::SeqCst); + Ok(plan) } } diff --git a/examples/distributed/engine-library/src/stage.rs b/examples/distributed/engine-library/src/stage.rs index 16bf43c91..e2cd51458 100644 --- a/examples/distributed/engine-library/src/stage.rs +++ b/examples/distributed/engine-library/src/stage.rs @@ -39,6 +39,7 @@ use std::path::{Path, PathBuf}; use std::sync::Arc; +use std::sync::atomic::{AtomicU64, Ordering}; use std::{fmt, fs}; use arrow::ipc::reader::StreamReader; @@ -61,6 +62,24 @@ pub(crate) fn partition_path(shuffle_dir: &str, stage_id: u32, partition: usize) Path::new(shuffle_dir).join(format!("stage-{stage_id}-part-{partition}.arrow")) } +/// Where a writer builds a partition before publishing it. +/// +/// Unique per writer, not merely per partition. Deriving the temporary name +/// from the final one alone would have two writers of the same partition +/// interleave their batches into one file, and then have one of the renames +/// fail because the other already consumed it -- so the rename that is +/// supposed to make publishing atomic would instead be the thing that broke. +/// The driver dispatches each partition once, but a node that is safe only +/// because of how its caller schedules work is not safe. +fn temp_partition_path(shuffle_dir: &str, stage_id: u32, partition: usize) -> PathBuf { + static NEXT: AtomicU64 = AtomicU64::new(0); + let unique = NEXT.fetch_add(1, Ordering::Relaxed); + Path::new(shuffle_dir).join(format!( + "stage-{stage_id}-part-{partition}.{}-{unique}.arrow.tmp", + std::process::id() + )) +} + /// Marks a subtree as one stage of a distributed query. #[derive(Debug)] pub(crate) struct ShuffleStageExec { @@ -93,7 +112,9 @@ impl ShuffleStageExec { /// /// Published by rename, so a reader can never observe a half-written /// file. The driver waits for workers to exit before reading, but relying - /// on that alone would break for anyone who overlapped the two. + /// on that alone would break for anyone who overlapped the two. The + /// writer side of the same argument is why [`temp_partition_path`] is + /// unique per writer rather than per partition. fn write_partition( &self, partition: usize, @@ -102,6 +123,7 @@ impl ShuffleStageExec { let input = Arc::clone(&self.input); let schema = input.schema(); let final_path = partition_path(&self.shuffle_dir, self.stage_id, partition); + let temp_path = temp_partition_path(&self.shuffle_dir, self.stage_id, partition); let shuffle_dir = self.shuffle_dir.clone(); let collected = async move { @@ -114,7 +136,6 @@ impl ShuffleStageExec { fs::create_dir_all(&shuffle_dir).map_err(|err| { internal_datafusion_err!("dfx_engine: creating {shuffle_dir}: {err}") })?; - let temp_path = final_path.with_extension("arrow.tmp"); { let file = fs::File::create(&temp_path).map_err(|err| { internal_datafusion_err!("dfx_engine: creating {}: {err}", temp_path.display()) diff --git a/examples/distributed/run_tpch.py b/examples/distributed/run_tpch.py index 90a25deee..b32381bc7 100644 --- a/examples/distributed/run_tpch.py +++ b/examples/distributed/run_tpch.py @@ -136,8 +136,8 @@ def main(argv: list[str] | None = None) -> int: 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}" + f"distributed: {distributed:.2f}s across {len(result.tasks)} " + f"worker process(es); rows per (stage, partition) {result.task_rows}" ) start = time.monotonic() From 3a61caf862082fd5ea4609efd7880e5f4d4f5b59 Mon Sep 17 00:00:00 2001 From: Tim Saucer Date: Thu, 10 Sep 2026 14:23:48 -0400 Subject: [PATCH 09/24] Read the third codec id instead of copying it `expected_codec_ids` said its ids were "read from the libraries rather than written out here", then wrote out `dfx_udfs.physical.v1`. Every other literal id in the example is in a test pinning that id, where writing it is the point; this was the only copy on a live path. A stale copy fails badly rather than obviously: on a version bump `build_session` would report "session codec ids [...] do not match the expected [...]" against a session that was in fact correct, which is the class of misdirected failure the check exists to prevent. `dfx_udfs` ships no bundle, so it has no bundle class to hang a `physical_codec_id()` on, and its id comes off a throwaway codec. `__datafusion_codec_id__` is a property of the codec object -- where the protocol puts it -- and the two static methods are a convenience their bundles need only because a bundle is not itself a codec. Adding one to `CodecObservations` for symmetry would have erased the asymmetry this library exists to show. Co-Authored-By: Claude Opus 5 (1M context) --- .../engine-library/python/dfx_engine/session.py | 14 +++++++++++--- 1 file changed, 11 insertions(+), 3 deletions(-) diff --git a/examples/distributed/engine-library/python/dfx_engine/session.py b/examples/distributed/engine-library/python/dfx_engine/session.py index 14390d01f..cdbb5399a 100644 --- a/examples/distributed/engine-library/python/dfx_engine/session.py +++ b/examples/distributed/engine-library/python/dfx_engine/session.py @@ -97,14 +97,22 @@ def from_json(payload: Mapping) -> SessionSpec: def expected_codec_ids() -> list[str]: """The physical codec ids a correctly-built session carries. - Read from the libraries rather than written out here, so adding a library - to :func:`build_session` and forgetting this list is not possible. + Read from the libraries rather than written out here. A copy of an id in + this file would go stale on a version bump, and the failure would be + :func:`build_session` accusing a session that was in fact correct. + + The third one looks different because ``dfx_udfs`` ships no bundle, so + there is no bundle class to hang a ``physical_codec_id()`` on. Its id + comes off a throwaway codec instead: ``__datafusion_codec_id__`` is a + property of the codec object, which is where the protocol puts it, and the + two static methods above are a convenience their bundles need only because + a bundle is not itself a codec. """ return sorted( [ dfx_storage.DfxStorageExtension.physical_codec_id(), _internal.DfxEngineExtension.physical_codec_id(), - "dfx_udfs.physical.v1", + dfx_udfs.CodecObservations().physical_codec().__datafusion_codec_id__, ] ) From e3ae58e015e0198c675c6ca9266b45e6bad3fd3c Mon Sep 17 00:00:00 2001 From: Tim Saucer Date: Thu, 10 Sep 2026 14:58:26 -0400 Subject: [PATCH 10/24] Make the example's run instructions work Three things stood between a reader and a passing test run. `uv pip install ... ../..` pointed at `examples/`, which has no `pyproject.toml`; the repository root is `../../..`. And `uv run maturin develop` fails with `Failed to spawn: maturin` unless maturin is in the venv, which nothing put there. Both were invisible to anyone who already had a working environment. The third was not in the README at all. None of the example projects declared `[tool.pytest.ini_options]`, so pytest's rootdir search walked up to the repository root and ran these suites under the root's pytest settings and its `conftest.py` -- which exists to inject a doctest namespace for `datafusion` and imports numpy. A reader following the README got `ModuleNotFoundError: No module named 'numpy'` at collection, from a package the example has no use for. It worked in CI only because the venv there is the root project's, synced with `--dev`. Each example now caps the search with its own config, which also drops the `Unknown config option: asyncio_mode` warnings every run was emitting. `python_files` picks up the underscore convention the test files already follow, so `uv run pytest` needs no path argument. No tracked example file contains a Python doctest, so nothing was relying on the root's `--doctest-modules`. Verified by running the instructions verbatim in a shell with neither the repository's venv on `PATH` nor `VIRTUAL_ENV` set: 18 passed, rootdir the example directory. The two older example crates have the same latent rootdir problem and the same one-block fix; left alone here. Co-Authored-By: Claude Opus 5 (1M context) --- .gitignore | 3 +++ examples/distributed/README.md | 18 ++++++++++++++++-- .../distributed/engine-library/pyproject.toml | 7 +++++++ .../distributed/storage-library/pyproject.toml | 14 ++++++++++++++ .../distributed/udf-library/pyproject.toml | 7 +++++++ 5 files changed, 47 insertions(+), 2 deletions(-) diff --git a/.gitignore b/.gitignore index ef00c0fd6..f5ade5698 100644 --- a/.gitignore +++ b/.gitignore @@ -37,3 +37,6 @@ docs/mdbook/book .pyo3_build_config examples/distributed/*/.venv/ +# Left behind by the `uv` commands in examples/distributed/README.md. The +# example projects are not locked -- only the root project is. +examples/distributed/*/uv.lock diff --git a/examples/distributed/README.md b/examples/distributed/README.md index 35fc1d5e8..28b354434 100644 --- a/examples/distributed/README.md +++ b/examples/distributed/README.md @@ -47,11 +47,25 @@ setups are the normal case, and this example shows what one costs. ```console $ cd examples/distributed/engine-library -$ uv venv && uv pip install pytest pyarrow ../.. ../storage-library ../udf-library +$ uv venv +$ uv pip install pytest maturin ../../.. ../storage-library ../udf-library $ uv run maturin develop -$ uv run pytest python/tests/_test*.py +$ uv run pytest ``` +`../../..` is the repository root. The example needs the `datafusion` built +from this checkout rather than a released wheel, because an extension library +and its host have to agree on the FFI ABI. + +The other two libraries are installed for a reason beyond running their own +tests: `dfx_engine.session` imports both, so `import dfx_engine` fails without +them. That first `uv pip install` builds three Rust libraries from source and +is not quick; `maturin develop` then builds the fourth, this one, in place. + +Each library's tests can be run from its own directory the same way, without +the sibling installs. `storage-library` and `udf-library` need only +`pytest maturin ../../..`. + Against the real TPC-H data — generate it as [`examples/tpch`](../tpch/README.md) describes, then: diff --git a/examples/distributed/engine-library/pyproject.toml b/examples/distributed/engine-library/pyproject.toml index ca9de9818..da786e50c 100644 --- a/examples/distributed/engine-library/pyproject.toml +++ b/examples/distributed/engine-library/pyproject.toml @@ -32,3 +32,10 @@ dynamic = ["version"] features = ["pyo3/extension-module"] python-source = "python" module-name = "dfx_engine._internal" + +# See the same block in `../storage-library/pyproject.toml`: this caps pytest's +# rootdir search so the suite does not inherit the repository root's pytest +# settings and its numpy-importing `conftest.py`. +[tool.pytest.ini_options] +testpaths = ["python/tests"] +python_files = ["_test_*.py"] diff --git a/examples/distributed/storage-library/pyproject.toml b/examples/distributed/storage-library/pyproject.toml index 24bc294fd..556b74c95 100644 --- a/examples/distributed/storage-library/pyproject.toml +++ b/examples/distributed/storage-library/pyproject.toml @@ -30,3 +30,17 @@ dynamic = ["version"] [tool.maturin] features = ["pyo3/extension-module"] + +# Here so pytest's rootdir search stops in this directory. With no config file +# of its own, the search walks up to the repository root, and this suite then +# runs under the root's pytest settings and its `conftest.py` -- which exists +# to inject a doctest namespace for `datafusion` and imports numpy. Running +# these tests would need a package the example has no use for, and would warn +# about pytest-asyncio options this environment has no reason to install. +# +# `python_files` is the underscore convention the test files already follow, so +# a bare `pytest` finds them; the leading underscore is what keeps a run from +# the repository root out of here. +[tool.pytest.ini_options] +testpaths = ["python/tests"] +python_files = ["_test_*.py"] diff --git a/examples/distributed/udf-library/pyproject.toml b/examples/distributed/udf-library/pyproject.toml index 8e87abf47..defe67942 100644 --- a/examples/distributed/udf-library/pyproject.toml +++ b/examples/distributed/udf-library/pyproject.toml @@ -30,3 +30,10 @@ dynamic = ["version"] [tool.maturin] features = ["pyo3/extension-module"] + +# See the same block in `../storage-library/pyproject.toml`: this caps pytest's +# rootdir search so the suite does not inherit the repository root's pytest +# settings and its numpy-importing `conftest.py`. +[tool.pytest.ini_options] +testpaths = ["python/tests"] +python_files = ["_test_*.py"] From 6e3dc962f63362c95147686071952e0b33b86ae0 Mon Sep 17 00:00:00 2001 From: Tim Saucer Date: Thu, 10 Sep 2026 15:08:21 -0400 Subject: [PATCH 11/24] Cover build_session's own consistency check `test_a_session_missing_a_library_is_rejected_at_build` was named for the check and described it in its docstring, but only built a bare `SessionContext` and asserted it carried no codec ids. `build_session` was never called, so the `RuntimeError` it exists to raise had no test. Neutering the guard leaves that test passing, which is the proof. The branch cannot be reached through the function's signature: which libraries get installed is written into `build_session`, not taken from the spec, so no argument can make the session come out wrong. The new test patches `expected_codec_ids` instead and says why -- the check guards this module against being edited inconsistently, a library added to one list and not the other, rather than anything a caller passes. The mismatched peer case is already covered by `test_a_worker_whose_codecs_disagree_refuses_the_plan`, which compares a worker's session against the driver's envelope. The old test keeps its (weaker, true) assertion under a name that says what it checks: a `SessionContext` carries none of the three libraries until `build_session` puts them there. Verified by mutation: with `if installed != expected` forced false the new test fails with "DID NOT RAISE RuntimeError" and the renamed one still passes. Co-Authored-By: Claude Opus 5 (1M context) --- .../python/tests/_test_three_libraries.py | 40 ++++++++++++++++--- 1 file changed, 35 insertions(+), 5 deletions(-) diff --git a/examples/distributed/engine-library/python/tests/_test_three_libraries.py b/examples/distributed/engine-library/python/tests/_test_three_libraries.py index 36215b978..ab5373ef2 100644 --- a/examples/distributed/engine-library/python/tests/_test_three_libraries.py +++ b/examples/distributed/engine-library/python/tests/_test_three_libraries.py @@ -389,16 +389,46 @@ def test_a_worker_whose_codecs_disagree_refuses_the_plan(spec: SessionSpec) -> N 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. +def test_a_bare_session_carries_none_of_the_three_libraries() -> None: + """Nothing about a `SessionContext` is installed by default. - 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. + Every codec in :func:`expected_codec_ids` is there because + :func:`build_session` put it there, which is why that function is the only + supported way to build a driver or a worker. """ ctx = SessionContext() installed = sorted(ctx.physical_extension_codec_ids()) - assert installed != expected_codec_ids() assert installed == [] + assert expected_codec_ids() != installed + + +def test_build_session_rejects_a_session_it_built_wrong( + spec: SessionSpec, monkeypatch: pytest.MonkeyPatch +) -> None: + """`build_session` checks its own work before handing the session over. + + 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. + + Induced by patching the expectation, because no argument can produce the + mismatch from the other side: *which* libraries get installed is written + into :func:`build_session`, not taken from the spec. That is what the + check guards -- this module being edited inconsistently, a library added + to one list and not the other -- rather than anything a caller passes. + Its counterpart for a genuinely mismatched peer is + `test_a_worker_whose_codecs_disagree_refuses_the_plan`, which compares a + worker's session against the driver's envelope. + """ + with_a_fourth = sorted([*expected_codec_ids(), "dfx_absent.physical.v1"]) + monkeypatch.setattr("dfx_engine.session.expected_codec_ids", lambda: with_a_fourth) + + with pytest.raises(RuntimeError, match="do not match the expected") as excinfo: + build_session(spec) + + # The message names both sides and says what breaks, so the reader does + # not have to guess which list is wrong. + assert "dfx_absent.physical.v1" in str(excinfo.value) + assert "will fail to decode" in str(excinfo.value) def test_a_plan_encoded_without_a_context_cannot_be_encoded( From be191ac5bf61289ccbec78e9f936cecd02bcdab6 Mon Sep 17 00:00:00 2001 From: Tim Saucer Date: Thu, 10 Sep 2026 15:18:35 -0400 Subject: [PATCH 12/24] Point local_session.rs at a section that exists The module doc said "See the 'Known gaps' section of the extension guide". No such section exists, and `git log -S` finds none on any branch under docs/ -- it was never a valid pointer rather than a rename casualty. The content it wanted is the section this PR adds to the query-planner guide, "Plan against your own optimizer rules", which names `LocalOptimizerSession` as its worked example. Linking to it makes the pair mutually discoverable instead of one-way. Rust cannot use a Sphinx `:ref:`, so this is the rendered URL, matching how the query-planner example's README links into the guide. Verified against a fresh `sphinx-build`: the page contains `section id="plan-against-your-own-optimizer-rules"`. The two other guide anchors this PR's READMEs link to resolve as well, and the build reports no undefined labels. Co-Authored-By: Claude Opus 5 (1M context) --- examples/distributed/engine-library/src/local_session.rs | 7 +++++-- 1 file changed, 5 insertions(+), 2 deletions(-) diff --git a/examples/distributed/engine-library/src/local_session.rs b/examples/distributed/engine-library/src/local_session.rs index 3660c5c84..dd5e8e894 100644 --- a/examples/distributed/engine-library/src/local_session.rs +++ b/examples/distributed/engine-library/src/local_session.rs @@ -23,8 +23,11 @@ //! `CooperativeExec` produced that way has no reachable `try_to_proto`, so a //! planner that must serialize its result -- and `FFI_QueryPlanner` always //! must, it returns proto bytes rather than a handle -- fails on a node that -//! is perfectly serializable in the process that made it. See the "Known gaps" -//! section of the extension guide. +//! is perfectly serializable in the process that made it. The guide spells +//! this out, with the error it produces, under [Plan against your own +//! optimizer rules](https://datafusion.apache.org/python/extension-guide/query-planners.html#plan-against-your-own-optimizer-rules) +//! -- which names this type as the worked example, so the two are meant to be +//! read together. //! //! Wrapping the session with a locally-owned copy of the same rule set keeps //! every rewrite inside this library, where the nodes stay concrete. That is From 4662e18338fbfe8a938056cf1d1d581795014a67 Mon Sep 17 00:00:00 2001 From: Tim Saucer Date: Thu, 10 Sep 2026 15:29:15 -0400 Subject: [PATCH 13/24] Say why table_options_mut panics on purpose The comment justified the panic empirically -- "Physical planning never calls this; verified in the spike" -- which reads as an observation that could stop holding, and invites the obvious "fix" of keeping a cloned `TableOptions` and lending out `&mut` to that. That fix would be worse than the panic. This type borrows the session it wraps and has no table options of its own, so a caller could set a Parquet option, watch it apply to nothing, and have no way to notice. Silently dropping a mutation is not an improvement on refusing one. The real reason the panic is unreachable is stronger than the comment claimed, and is type-level: the method takes `&mut self`, `create_physical_plan` takes `&dyn Session`, nothing in DataFusion asks for a `&mut dyn Session`, and `planner.rs` binds the wrapper immutably. Probed with the compiler rather than by reading -- adding `local.table_options_mut()` at the call site fails with E0596, "cannot borrow `local` as mutable". The trait declares the method with no default impl, so it has to be written either way. Behaviour unchanged. Co-Authored-By: Claude Opus 5 (1M context) --- .../engine-library/src/local_session.rs | 19 ++++++++++++++++--- 1 file changed, 16 insertions(+), 3 deletions(-) diff --git a/examples/distributed/engine-library/src/local_session.rs b/examples/distributed/engine-library/src/local_session.rs index dd5e8e894..ee971c2e5 100644 --- a/examples/distributed/engine-library/src/local_session.rs +++ b/examples/distributed/engine-library/src/local_session.rs @@ -153,9 +153,22 @@ impl Session for LocalOptimizerSession<'_> { } fn table_options_mut(&mut self) -> &mut TableOptions { - // The wrapper only borrows `inner`, so it cannot hand out a mutable - // reference. Physical planning never calls this; verified in the spike. - unimplemented!("LocalOptimizerSession does not support table_options_mut") + // Written because the trait requires it -- there is no default -- and + // unreachable by type rather than by luck: this takes `&mut self`, + // `create_physical_plan` takes `&dyn Session`, nothing in DataFusion + // asks for a `&mut dyn Session`, and `planner.rs` binds the wrapper + // immutably. A call here does not compile, so the guarantee is the + // borrow checker's rather than a comment's. + // + // Panicking is the answer on purpose, not an unfinished one. This type + // borrows the session it wraps and has no table options of its own to + // lend. Keeping a clone and returning `&mut` to that would compile and + // would be worse: a caller could set a Parquet option, watch it apply + // to nothing, and have no way to notice. Quietly dropping a mutation + // is not an improvement on refusing one. + unimplemented!( + "LocalOptimizerSession borrows its session and cannot lend out mutable table options" + ) } fn task_ctx(&self) -> Arc { From 60bb0f81acc29effbff1f4590b0d06ed05395a57 Mon Sep 17 00:00:00 2001 From: Tim Saucer Date: Thu, 10 Sep 2026 15:32:51 -0400 Subject: [PATCH 14/24] Link the two sections this PR added `planner_host_optimizer_rules` and `distributed_worker_parity` were defined and never referenced. Both rendered fine, so nothing was broken -- they were just unreachable from anywhere a reader would be standing. The planner one goes on the checklist, next to the `fallback` items, which needed it: "Your planner hook wraps `fallback` and delegates to it" was unqualified, and the new section says a planner that rewrites the plan structurally cannot delegate. The two items now name each other instead of contradicting each other. Worker parity gets two, because it has two audiences. Users arrive at it from `expressions.md`, whose portability section covers two of the nine things a worker has to reproduce and did not say there were seven more. Engine authors arrive from the checklist: the guide already says "most engines handle several of them for you -- check which", which is only checkable if engines say so. Verified against a fresh `sphinx-build`: all three render as links to `#planner-host-optimizer-rules` and `#distributed-worker-parity`, both ids present on their pages, link text taken from the section headings, and no undefined labels anywhere in the build. Co-Authored-By: Claude Opus 5 (1M context) --- docs/source/extension-guide/checklist.md | 14 +++++++++++++- .../user-guide/distributing-work/expressions.md | 6 ++++++ 2 files changed, 19 insertions(+), 1 deletion(-) diff --git a/docs/source/extension-guide/checklist.md b/docs/source/extension-guide/checklist.md index 677ccbefb..4b0088179 100644 --- a/docs/source/extension-guide/checklist.md +++ b/docs/source/extension-guide/checklist.md @@ -78,8 +78,15 @@ publish. Each links to the page that explains it. `with_extensions` refuses a capsule, because there would be nothing to name the codec by. → {ref}`extension_bundles_codecs_are_objects` - [ ] **Your planner hook wraps `fallback` and delegates to it.** Ignoring it - replaces every layer beneath you, which is legal but not composable. + replaces every layer beneath you, which is legal but not composable — + unless your planner rewrites the plan, as in the next item. → {ref}`extension_bundles` +- [ ] **Your planner plans for itself if it rewrites the plan**, leaving + `fallback` unused. The two are exclusive: delegating hands planning back + to the host and returns nodes you can neither downcast nor split. + Planning for yourself then means supplying your *own* optimizer rules, + because a session that arrived over FFI carries the host's. + → {ref}`planner_host_optimizer_rules` - [ ] **Your planner hook returns `None`, not `fallback`, when it has nothing to contribute.** Returning `fallback` installs the session's own planner as a foreign one and adds an FFI hop that was not there. @@ -107,3 +114,8 @@ publish. Each links to the page that explains it. wheel, then exercise it from Python. `examples/distributed` additionally spawns worker processes, which is the only way to catch a codec that only works in the process that wrote it. +- [ ] **If you ship an engine, say which worker-parity items you handle** and + which you leave to your users. A `SessionContext` cannot be snapshotted + and restored elsewhere, so every one of them is somebody's job, and your + users cannot tell whose from the outside. + → {ref}`distributed_worker_parity` diff --git a/docs/source/user-guide/distributing-work/expressions.md b/docs/source/user-guide/distributing-work/expressions.md index b817b3e49..4e30bc0f9 100644 --- a/docs/source/user-guide/distributing-work/expressions.md +++ b/docs/source/user-guide/distributing-work/expressions.md @@ -144,6 +144,12 @@ requirements on the worker environment: `ModuleNotFoundError: No module named 'yourmod'`, raised while the plan is being decoded, with nothing in the message about UDFs or serialization. +Both requirements above are entries on a longer list. A worker has to +reproduce more than the environment its UDFs need — codec ids, object stores, +config extensions, `target_partitions` — and none of it can be copied off a +running session, so each one is something you build the same way twice. See +{ref}`distributed_worker_parity`. + ## Registering shared UDFs on workers When an expression references an FFI capsule UDF (or any UDF the From 1774343c3f69b13531a0c6d1748e297c9b4fe8ce Mon Sep 17 00:00:00 2001 From: Tim Saucer Date: Thu, 10 Sep 2026 15:51:31 -0400 Subject: [PATCH 15/24] Decode a provider with the schema from the plan `try_decode_table_provider` ignored its `schema` argument and re-read the first Parquet footer instead. That argument is not advisory: a serialized `CustomScan` carries the table schema, and `from_proto` resolves the scan's projection from column *names* to indices against it -- let column_indices = columns.columns.iter() .map(|name| schema.index_of(name)) -- while `TableScanBuilder::build` then applies those indices to the schema the decoded provider reports: let schema = source.schema(); ... schema.fields()[*i] Unchecked. So a column added to the directory ahead of the others since the plan was written selects the wrong one, and a column removed indexes off the end and panics inside DataFusion. Taking the schema the plan carries makes both unreachable, and skips a footer read the caller had already done. `try_new` keeps reading the footer for the registration path, where nobody has supplied a schema yet. Also exposes `provider_encode_calls` / `provider_decode_calls`. Both counters were incremented in Rust and never reachable from Python, so the logical half of this codec had no assertion anywhere -- only the physical `encode_calls` / `decode_calls` were exposed. Two tests: one drives the logical round trip through `LogicalPlan.to_bytes` and asserts each half ran, and one drifts the directory between encode and decode and asserts the decoded scan still reports the plan's schema. Confirmed by mutation -- restoring the footer read makes the second fail with the drifted `[label, sensor_id, reading]`. Co-Authored-By: Claude Opus 5 (1M context) --- .../python/tests/_test_portable_codec.py | 65 ++++++++++++++++++- .../distributed/storage-library/src/codec.rs | 11 +++- .../storage-library/src/extension.rs | 16 +++++ .../storage-library/src/table_provider.rs | 37 +++++++++-- 4 files changed, 119 insertions(+), 10 deletions(-) diff --git a/examples/distributed/storage-library/python/tests/_test_portable_codec.py b/examples/distributed/storage-library/python/tests/_test_portable_codec.py index 716d92fd6..91acd90ab 100644 --- a/examples/distributed/storage-library/python/tests/_test_portable_codec.py +++ b/examples/distributed/storage-library/python/tests/_test_portable_codec.py @@ -32,9 +32,11 @@ import textwrap from typing import TYPE_CHECKING +import pyarrow as pa +import pyarrow.parquet as pq import pytest from datafusion import SessionContext -from datafusion.plan import ExecutionPlan +from datafusion.plan import ExecutionPlan, LogicalPlan from dfx_storage import DfxStorageExtension, PartitionedParquetTable if TYPE_CHECKING: @@ -157,6 +159,67 @@ def test_stock_nodes_never_reach_this_codec(readings_dir: pathlib.Path) -> None: assert bundle.declined_calls() == 0 +def test_the_logical_codec_carries_the_provider(readings_dir: pathlib.Path) -> None: + """The provider is held in the logical plan, so it needs its own codec. + + A physical codec is not enough. Nothing here installs a query planner, so + this is the only test that reaches the logical half directly -- but any + session with an engine installed takes this path on every query, which is + what makes a provider library shipping only a physical codec fail as soon + as it meets one. + """ + ctx, bundle = _configured(readings_dir) + plan = ctx.sql("select sensor_id from readings").logical_plan() + assert bundle.provider_encode_calls() == 0 + + blob = plan.to_bytes(ctx) + assert bundle.provider_encode_calls() == 1 + # The directory, under the logical payload's own magic. + assert b"DFXSTOL1" in blob + assert str(readings_dir).encode() in blob + + LogicalPlan.from_bytes(ctx, blob) + assert bundle.provider_decode_calls() == 1 + + +def test_a_decoded_provider_reports_the_schema_from_the_plan( + readings_dir: pathlib.Path, +) -> None: + """The decoder takes the plan's schema rather than re-reading a footer. + + A serialized scan carries its projection as column *names*, which the + decoder resolves to indices against the schema in the plan and then + applies to whatever this provider reports -- without a bounds check. So + the two have to be the same schema. + + Here the directory gains a column in front of the others after the plan + was written. Re-reading the footer on decode would make index 0 mean + `label` while the plan means `sensor_id`, and the query would quietly + return the wrong column. + """ + ctx, _ = _configured(readings_dir) + blob = ctx.sql("select sensor_id from readings").logical_plan().to_bytes(ctx) + + pq.write_table( + pa.table( + { + "label": ["a", "b", "c"], + "sensor_id": [0, 1, 2], + "reading": [1.5, 2.5, 3.5], + } + ), + readings_dir / "part-0.parquet", + ) + + restored = LogicalPlan.from_bytes(ctx, blob) + + # Still the schema the plan was built against: `label` is not in it, and + # the projected column is the one that was asked for. + schema_text = restored.display_indent_schema() + assert "sensor_id" in schema_text + assert "label" not in schema_text + + WORKER = textwrap.dedent( """ import sys diff --git a/examples/distributed/storage-library/src/codec.rs b/examples/distributed/storage-library/src/codec.rs index 56c0bc748..5a16e4006 100644 --- a/examples/distributed/storage-library/src/codec.rs +++ b/examples/distributed/storage-library/src/codec.rs @@ -327,8 +327,13 @@ impl LogicalExtensionCodec for DfxStorageLogicalCodec { self.counters .provider_decoded .fetch_add(1, Ordering::SeqCst); - Ok(Arc::new(PartitionedParquetTable::try_new(Path::new( - directory, - ))?)) + // `schema` is the table schema recorded in the plan, and it is the one + // the plan's projection indices were resolved against -- so it is the + // schema this provider has to report, not one re-read from a file that + // may have changed since. See `try_new_with_schema`. + Ok(Arc::new(PartitionedParquetTable::try_new_with_schema( + Path::new(directory), + schema, + )?)) } } diff --git a/examples/distributed/storage-library/src/extension.rs b/examples/distributed/storage-library/src/extension.rs index 657d28088..76d496177 100644 --- a/examples/distributed/storage-library/src/extension.rs +++ b/examples/distributed/storage-library/src/extension.rs @@ -151,6 +151,22 @@ impl DfxStorageExtension { self.counters.declined.load(Ordering::SeqCst) } + /// How often the *logical* codec wrote this library's table provider. + /// + /// Separate from [`Self::encode_calls`] because the two halves are reached + /// by different callers: the physical codec carries the scan node, and + /// this one carries the provider held in the logical plan. A library that + /// shipped only the physical half would report activity here of zero and + /// fail the moment a query planner was installed. + fn provider_encode_calls(&self) -> usize { + self.counters.provider_encoded.load(Ordering::SeqCst) + } + + /// How often it rebuilt the provider, which is the half a worker runs. + fn provider_decode_calls(&self) -> usize { + self.counters.provider_decoded.load(Ordering::SeqCst) + } + /// The wire id, so a driver can put it in a worker's task envelope and /// the worker can check it before decoding anything. #[staticmethod] diff --git a/examples/distributed/storage-library/src/table_provider.rs b/examples/distributed/storage-library/src/table_provider.rs index 3c33c41df..b18abded9 100644 --- a/examples/distributed/storage-library/src/table_provider.rs +++ b/examples/distributed/storage-library/src/table_provider.rs @@ -45,9 +45,10 @@ use crate::exec::{FileSlice, PartitionedParquetExec}; /// Scans `*.parquet` under `directory`, one partition per file. #[derive(Debug)] pub(crate) struct PartitionedParquetTable { - /// Kept so the logical codec can write it down. Everything else here is - /// derived from the directory, so the path is the whole encoding -- see - /// [`crate::codec::DfxStorageLogicalCodec`]. + /// Kept so the logical codec can write it down. The file list is derived + /// from the directory, so the path is the whole encoding -- see + /// [`crate::codec::DfxStorageLogicalCodec`]. The schema is not derived on + /// the decode path; see [`Self::try_new_with_schema`]. pub(crate) directory: String, files: Vec, schema: SchemaRef, @@ -56,11 +57,32 @@ pub(crate) struct PartitionedParquetTable { impl PartitionedParquetTable { /// Read the directory listing and the first file's schema, once. /// + /// For the registration path, where nobody has told us the schema yet. + pub(crate) fn try_new(directory: &Path) -> Result { + Self::open(directory, None) + } + + /// Open with the schema the plan was built against, rather than re-reading + /// it from a file. + /// + /// This is the decode path, and taking the caller's schema is not an + /// optimisation. A serialized `CustomScan` carries the table schema and + /// its projection as *column names*; the decoder turns those names into + /// indices against the encoded schema and then applies them to whatever + /// this provider reports. Re-reading a footer here would let the two + /// drift, and the indices are applied without a bounds check -- so a + /// column added to the directory since the plan was written silently + /// selects the wrong one, and a column removed panics inside DataFusion. + /// Neither is reachable if the schema in the plan is the schema used. + pub(crate) fn try_new_with_schema(directory: &Path, schema: SchemaRef) -> Result { + Self::open(directory, Some(schema)) + } + /// Sorted by path so that partition `i` means the same file in every /// process that opens the same directory. Directory iteration order is /// not specified, and a worker that disagreed with the driver about which /// file is partition 3 would silently produce wrong answers. - pub(crate) fn try_new(directory: &Path) -> Result { + fn open(directory: &Path, schema: Option) -> Result { let mut paths: Vec<_> = fs::read_dir(directory) .map_err(|err| DataFusionError::External(Box::new(err)))? .collect::>>() @@ -88,11 +110,14 @@ impl PartitionedParquetTable { }); } - let schema = Self::read_schema(&paths[0])?; + let schema = match schema { + Some(schema) => schema, + None => Arc::new(Self::read_schema(&paths[0])?), + }; Ok(Self { directory: directory.to_string_lossy().into_owned(), files, - schema: Arc::new(schema), + schema, }) } From 628c45175b7157dffcfdcdbd8662640cf26bc221 Mon Sep 17 00:00:00 2001 From: Tim Saucer Date: Thu, 10 Sep 2026 16:02:36 -0400 Subject: [PATCH 16/24] Refuse a malformed projection instead of dropping it `try_decode` parsed the projection with `filter_map`, so an element that would not convert was skipped and the decode succeeded with a shorter list. The indices are positional, so that is a different query rather than a degraded one: lose one and it reads the wrong columns, lose all of them and it reads none -- with a well-formed plan either way, and nothing downstream able to tell. Demonstrated on the real wire format before fixing it. Editing a payload's `"projection":[0,1]` to the same-width `"projection":[1e1]` -- valid JSON, but a float, so `as_u64` declines it -- decoded without complaint and produced batches whose schema was `[]` instead of `["sensor_id", "reading"]`. It now fails with "projection index 10.0 is not a column number". Both the projection and the limit are parsed fallibly, a non-array projection is rejected rather than treated as absent, and `usize::try_from` replaces `as usize` so a 64-bit index cannot truncate on a 32-bit worker. Absent and null still both mean "every column". The test edits the payload in place at constant length, because the JSON sits in a length-delimited protobuf field and a resized payload would corrupt the framing rather than exercise the codec. Co-Authored-By: Claude Opus 5 (1M context) --- .../python/tests/_test_portable_codec.py | 29 +++++++++++ .../distributed/storage-library/src/codec.rs | 49 ++++++++++++++++--- 2 files changed, 71 insertions(+), 7 deletions(-) diff --git a/examples/distributed/storage-library/python/tests/_test_portable_codec.py b/examples/distributed/storage-library/python/tests/_test_portable_codec.py index 91acd90ab..d59541235 100644 --- a/examples/distributed/storage-library/python/tests/_test_portable_codec.py +++ b/examples/distributed/storage-library/python/tests/_test_portable_codec.py @@ -159,6 +159,35 @@ def test_stock_nodes_never_reach_this_codec(readings_dir: pathlib.Path) -> None: assert bundle.declined_calls() == 0 +def test_a_malformed_projection_is_refused_not_dropped( + readings_dir: pathlib.Path, +) -> None: + """A projection index that will not parse has to be an error. + + Dropping it instead would hand back a shorter projection, and because the + indices are positional that is a different query rather than a degraded + one -- losing one reads the wrong columns and losing all of them reads + none, with a well-formed plan either way. A codec is the last place that + can tell a malformed payload from a valid one. + + The payload is edited in place, keeping its length: the JSON sits inside a + length-delimited protobuf field, so `[1,2]` is replaced by the same-width + `[1e1]`, which is a valid JSON *float* and so not a column number. Writing + a shorter or longer payload would corrupt the protobuf instead and test + the wrong thing. + """ + ctx, _ = _configured(readings_dir) + plan = ctx.sql("select sensor_id, reading from readings").execution_plan() + blob = plan.to_bytes(ctx) + assert b'"projection":[0,1]' in blob + + mangled = blob.replace(b'"projection":[0,1]', b'"projection":[1e1]') + assert len(mangled) == len(blob), "the edit has to preserve the protobuf framing" + + with pytest.raises(Exception, match="is not a column number"): + ExecutionPlan.from_bytes(ctx, mangled) + + def test_the_logical_codec_carries_the_provider(readings_dir: pathlib.Path) -> None: """The provider is held in the logical plan, so it needs its own codec. diff --git a/examples/distributed/storage-library/src/codec.rs b/examples/distributed/storage-library/src/codec.rs index 5a16e4006..e5e7cf6de 100644 --- a/examples/distributed/storage-library/src/codec.rs +++ b/examples/distributed/storage-library/src/codec.rs @@ -215,13 +215,48 @@ impl PhysicalExtensionCodec for DfxStoragePhysicalCodec { }) }) .collect::>>()?; - let projection = descriptor["projection"].as_array().map(|indices| { - indices - .iter() - .filter_map(|index| index.as_u64().map(|index| index as usize)) - .collect::>() - }); - let limit = descriptor["limit"].as_u64().map(|limit| limit as usize); + // Every element has to parse. `filter_map` here would drop the ones + // that did not and hand back a *shorter* projection, which is not a + // degraded answer but a different query: the indices are positional, + // so losing one silently reads the wrong columns, and losing all of + // them reads none. A codec is the last place that can tell a + // malformed payload from a valid one, because everything downstream + // sees a well-formed plan. + let projection = match &descriptor["projection"] { + // Absent and null both mean "every column". + serde_json::Value::Null => None, + serde_json::Value::Array(indices) => Some( + indices + .iter() + .map(|index| { + index + .as_u64() + .and_then(|index| usize::try_from(index).ok()) + .ok_or_else(|| { + internal_datafusion_err!( + "dfx_storage: projection index {index} is not a column number" + ) + }) + }) + .collect::>>()?, + ), + other => { + return internal_err!( + "dfx_storage: projection must be a list of column numbers or null, got {other}" + ); + } + }; + let limit = match &descriptor["limit"] { + serde_json::Value::Null => None, + value => Some( + value + .as_u64() + .and_then(|limit| usize::try_from(limit).ok()) + .ok_or_else(|| { + internal_datafusion_err!("dfx_storage: limit {value} is not a row count") + })?, + ), + }; let schema = Arc::new(schema_from_ipc_bytes(schema_bytes)?); self.counters.decoded.fetch_add(1, Ordering::SeqCst); From 22912acd258d4265d0554909fce1c6fb13ed2ebe Mon Sep 17 00:00:00 2001 From: Tim Saucer Date: Thu, 10 Sep 2026 16:22:47 -0400 Subject: [PATCH 17/24] Raise Execution, not Internal, for a bad payload Every payload error in both codecs was `internal_datafusion_err!`, so a worker handed a truncated or version-skewed payload got: Internal error: dfx_storage: projection index 10.0 is not a column number. This issue was likely caused by a bug in DataFusion's code. Please help us to resolve this by filing a bug report in our issue tracker The advice is wrong and the class is wrong. Upstream documents `Internal` as "due to bugs in DataFusion", says "a user should not be able to trigger internal errors under normal circumstances by feeding in malformed queries, bad data, etc.", and adds that I/O errors "do NOT fall under this category". A codec reads bytes written by another process, so nearly everything it can hit is bad data. Reclassified to `Execution`: payload framing and parsing in both codecs, and the shuffle file I/O in `stage.rs`. `Internal` survives in the three places it is the correct class, each now saying why -- `with_new_children` and `execute(partition)`, whose arguments come from an optimizer rule rather than a payload and so cannot be reached by any input, and encoding an object this process already holds in memory. The class does survive FFI, which was worth checking: `df_result!` wraps the error as `DataFusionError::Ffi` carrying the inner error's `Display`, so the message a Python caller sees goes from "FFI error: Internal error: ... file a bug report" to "FFI error: Execution error: ...". A test pins it, because nothing else would notice a regression: the malformed-projection case now asserts the message names no bug report. The storage codec's module doc records the rule, since two classes now coexist in one file. Co-Authored-By: Claude Opus 5 (1M context) --- .../distributed/engine-library/src/codec.rs | 14 ++--- .../distributed/engine-library/src/stage.rs | 40 ++++++------- .../python/tests/_test_portable_codec.py | 9 ++- .../distributed/storage-library/src/codec.rs | 60 +++++++++++++------ .../distributed/storage-library/src/exec.rs | 3 + 5 files changed, 79 insertions(+), 47 deletions(-) diff --git a/examples/distributed/engine-library/src/codec.rs b/examples/distributed/engine-library/src/codec.rs index f4188ab78..fa34f9a66 100644 --- a/examples/distributed/engine-library/src/codec.rs +++ b/examples/distributed/engine-library/src/codec.rs @@ -35,7 +35,7 @@ use std::fmt; use std::sync::Arc; use std::sync::atomic::{AtomicUsize, Ordering}; -use datafusion::common::{Result, internal_datafusion_err, internal_err}; +use datafusion::common::{Result, exec_datafusion_err, exec_err}; use datafusion::execution::TaskContext; use datafusion::physical_plan::ExecutionPlan; use datafusion_proto::physical_plan::{ @@ -107,20 +107,20 @@ impl PhysicalExtensionCodec for DfxEnginePhysicalCodec { return self.inner.try_decode(buf, inputs, ctx, proto_converter); }; - let (stage_id, shuffle_dir) = rest.split_at_checked(4).ok_or_else(|| { - internal_datafusion_err!("dfx_engine: payload truncated before stage id") - })?; + let (stage_id, shuffle_dir) = rest + .split_at_checked(4) + .ok_or_else(|| exec_datafusion_err!("dfx_engine: payload truncated before stage id"))?; let stage_id = u32::from_le_bytes( stage_id .try_into() - .map_err(|_| internal_datafusion_err!("dfx_engine: bad stage id"))?, + .map_err(|_| exec_datafusion_err!("dfx_engine: bad stage id"))?, ); let shuffle_dir = std::str::from_utf8(shuffle_dir) - .map_err(|err| internal_datafusion_err!("dfx_engine: bad shuffle dir: {err}"))?; + .map_err(|err| exec_datafusion_err!("dfx_engine: bad shuffle dir: {err}"))?; // The child arrives already decoded, by the host's chain. let [input] = inputs else { - return internal_err!( + return exec_err!( "ShuffleStageExec expects exactly one input, got {}", inputs.len() ); diff --git a/examples/distributed/engine-library/src/stage.rs b/examples/distributed/engine-library/src/stage.rs index e2cd51458..80ca72944 100644 --- a/examples/distributed/engine-library/src/stage.rs +++ b/examples/distributed/engine-library/src/stage.rs @@ -45,7 +45,7 @@ use std::{fmt, fs}; use arrow::ipc::reader::StreamReader; use arrow::ipc::writer::StreamWriter; use datafusion::common::tree_node::TreeNodeRecursion; -use datafusion::common::{DataFusionError, Result, internal_datafusion_err, internal_err}; +use datafusion::common::{DataFusionError, Result, exec_datafusion_err, internal_err}; use datafusion::execution::{SendableRecordBatchStream, TaskContext}; use datafusion::physical_expr::PhysicalExpr; use datafusion::physical_plan::memory::MemoryStream; @@ -133,26 +133,25 @@ impl ShuffleStageExec { batches.push(batch?); } - fs::create_dir_all(&shuffle_dir).map_err(|err| { - internal_datafusion_err!("dfx_engine: creating {shuffle_dir}: {err}") - })?; + fs::create_dir_all(&shuffle_dir) + .map_err(|err| exec_datafusion_err!("dfx_engine: creating {shuffle_dir}: {err}"))?; { let file = fs::File::create(&temp_path).map_err(|err| { - internal_datafusion_err!("dfx_engine: creating {}: {err}", temp_path.display()) + exec_datafusion_err!("dfx_engine: creating {}: {err}", temp_path.display()) })?; let mut writer = StreamWriter::try_new(file, stream.schema().as_ref()) - .map_err(|err| internal_datafusion_err!("dfx_engine: ipc writer: {err}"))?; + .map_err(|err| exec_datafusion_err!("dfx_engine: ipc writer: {err}"))?; for batch in &batches { - writer.write(batch).map_err(|err| { - internal_datafusion_err!("dfx_engine: writing batch: {err}") - })?; + writer + .write(batch) + .map_err(|err| exec_datafusion_err!("dfx_engine: writing batch: {err}"))?; } writer .finish() - .map_err(|err| internal_datafusion_err!("dfx_engine: ipc finish: {err}"))?; + .map_err(|err| exec_datafusion_err!("dfx_engine: ipc finish: {err}"))?; } fs::rename(&temp_path, &final_path).map_err(|err| { - internal_datafusion_err!("dfx_engine: publishing {}: {err}", final_path.display()) + exec_datafusion_err!("dfx_engine: publishing {}: {err}", final_path.display()) })?; Ok::<_, DataFusionError>(batches) @@ -171,18 +170,14 @@ impl ShuffleStageExec { fn read_partition(&self, partition: usize) -> Result { let path = partition_path(&self.shuffle_dir, self.stage_id, partition); - let file = fs::File::open(&path).map_err(|err| { - internal_datafusion_err!("dfx_engine: opening {}: {err}", path.display()) - })?; - let reader = StreamReader::try_new(file, None).map_err(|err| { - internal_datafusion_err!("dfx_engine: reading {}: {err}", path.display()) - })?; + let file = fs::File::open(&path) + .map_err(|err| exec_datafusion_err!("dfx_engine: opening {}: {err}", path.display()))?; + let reader = StreamReader::try_new(file, None) + .map_err(|err| exec_datafusion_err!("dfx_engine: reading {}: {err}", path.display()))?; let schema = reader.schema(); let batches = reader .collect::>>() - .map_err(|err| { - internal_datafusion_err!("dfx_engine: reading {}: {err}", path.display()) - })?; + .map_err(|err| exec_datafusion_err!("dfx_engine: reading {}: {err}", path.display()))?; Ok(Box::pin(MemoryStream::try_new(batches, schema, None)?)) } } @@ -218,6 +213,11 @@ impl ExecutionPlan for ShuffleStageExec { self: Arc, mut children: Vec>, ) -> Result> { + // The one `Internal` in this file, and the reason the rest are not: + // children here come from an optimizer rule, not from a payload. No + // input a user supplies can reach this, so if it fires the caller has + // a bug and a bug report is the right advice. Everything driven by + // bytes or by the filesystem is `Execution`. if children.len() != 1 { return internal_err!( "ShuffleStageExec expects exactly one child, got {}", diff --git a/examples/distributed/storage-library/python/tests/_test_portable_codec.py b/examples/distributed/storage-library/python/tests/_test_portable_codec.py index d59541235..4f4efb902 100644 --- a/examples/distributed/storage-library/python/tests/_test_portable_codec.py +++ b/examples/distributed/storage-library/python/tests/_test_portable_codec.py @@ -184,9 +184,16 @@ def test_a_malformed_projection_is_refused_not_dropped( mangled = blob.replace(b'"projection":[0,1]', b'"projection":[1e1]') assert len(mangled) == len(blob), "the edit has to preserve the protobuf framing" - with pytest.raises(Exception, match="is not a column number"): + with pytest.raises(Exception, match="is not a column number") as excinfo: ExecutionPlan.from_bytes(ctx, mangled) + # A bad payload is `Execution`, not `Internal`. DataFusion appends a + # "please file a bug report" line to every internal error, and sending + # someone to DataFusion's issue tracker over a corrupt payload of ours + # wastes their time and the maintainers'. + assert "bug report" not in str(excinfo.value) + assert "Internal error" not in str(excinfo.value) + def test_the_logical_codec_carries_the_provider(readings_dir: pathlib.Path) -> None: """The provider is held in the logical plan, so it needs its own codec. diff --git a/examples/distributed/storage-library/src/codec.rs b/examples/distributed/storage-library/src/codec.rs index e5e7cf6de..70edfc5ee 100644 --- a/examples/distributed/storage-library/src/codec.rs +++ b/examples/distributed/storage-library/src/codec.rs @@ -40,6 +40,23 @@ //! fields because a human debugging a worker can read it; the schema is Arrow //! IPC because that is the only encoding guaranteed to round-trip every Arrow //! type, including extension types and field metadata. +//! +//! # Which error to raise +//! +//! Two classes appear below, and the split is DataFusion's own rule rather +//! than a preference. `DataFusionError::Internal` is documented as "due to +//! bugs in DataFusion", it appends *"please help us to resolve this by filing +//! a bug report"* to every message, and "a user should not be able to trigger +//! internal errors under normal circumstances by feeding in malformed +//! queries, bad data, etc." +//! +//! A codec reads bytes from somewhere else, so almost everything that can go +//! wrong here is bad data: a truncated payload, a version skew, a projection +//! index that is not a number. Those are `Execution`, because the person +//! reading the message needs to look at the payload, not at DataFusion's +//! issue tracker. `Internal` is left for the two things that really would be +//! this library's fault -- failing to encode an object it holds in memory, +//! and being handed a plan node that violates its own contract. use std::fmt; use std::path::Path; @@ -50,7 +67,9 @@ use arrow::datatypes::Schema; use arrow::ipc::reader::StreamReader; use arrow::ipc::writer::StreamWriter; use datafusion::catalog::TableProvider; -use datafusion::common::{Result, TableReference, internal_datafusion_err, internal_err}; +use datafusion::common::{ + Result, TableReference, exec_datafusion_err, exec_err, internal_datafusion_err, +}; use datafusion::execution::TaskContext; use datafusion::physical_plan::ExecutionPlan; use datafusion_proto::logical_plan::{DefaultLogicalExtensionCodec, LogicalExtensionCodec}; @@ -103,6 +122,8 @@ impl fmt::Debug for DfxStoragePhysicalCodec { } } +/// `Internal` on purpose: the schema being written is one this process is +/// already holding, so a failure here is this library's bug and not bad input. fn schema_to_ipc_bytes(schema: &Schema) -> Result> { let mut buf: Vec = Vec::new(); { @@ -117,7 +138,7 @@ fn schema_to_ipc_bytes(schema: &Schema) -> Result> { fn schema_from_ipc_bytes(bytes: &[u8]) -> Result { let reader = StreamReader::try_new(std::io::Cursor::new(bytes), None) - .map_err(|err| internal_datafusion_err!("dfx_storage: reading schema: {err}"))?; + .map_err(|err| exec_datafusion_err!("dfx_storage: reading schema: {err}"))?; Ok(reader.schema().as_ref().clone()) } @@ -144,13 +165,15 @@ impl PhysicalExtensionCodec for DfxStoragePhysicalCodec { "projection": exec.projection, "limit": exec.limit, }); + // `Internal`, like `schema_to_ipc_bytes`: serializing a value built + // two lines up cannot fail on anything but a bug here. let json = serde_json::to_vec(&descriptor) .map_err(|err| internal_datafusion_err!("dfx_storage: encoding descriptor: {err}"))?; let schema = schema_to_ipc_bytes(&exec.table_schema)?; buf.extend_from_slice(MAGIC); let json_len = u32::try_from(json.len()) - .map_err(|_| internal_datafusion_err!("dfx_storage: descriptor too large to encode"))?; + .map_err(|_| exec_datafusion_err!("dfx_storage: descriptor too large to encode"))?; buf.extend_from_slice(&json_len.to_le_bytes()); buf.extend_from_slice(&json); buf.extend_from_slice(&schema); @@ -175,39 +198,39 @@ impl PhysicalExtensionCodec for DfxStoragePhysicalCodec { return self.inner.try_decode(buf, inputs, ctx, proto_converter); }; if !inputs.is_empty() { - return internal_err!( + return exec_err!( "PartitionedParquetExec is a leaf, got {} input(s)", inputs.len() ); } let (len_bytes, rest) = rest.split_at_checked(4).ok_or_else(|| { - internal_datafusion_err!("dfx_storage: payload truncated before descriptor length") + exec_datafusion_err!("dfx_storage: payload truncated before descriptor length") })?; let json_len = u32::from_le_bytes( len_bytes .try_into() - .map_err(|_| internal_datafusion_err!("dfx_storage: bad descriptor length"))?, + .map_err(|_| exec_datafusion_err!("dfx_storage: bad descriptor length"))?, ) as usize; let (json, schema_bytes) = rest.split_at_checked(json_len).ok_or_else(|| { - internal_datafusion_err!( + exec_datafusion_err!( "dfx_storage: descriptor claims {json_len} bytes, {} remain", rest.len() ) })?; let descriptor: serde_json::Value = serde_json::from_slice(json) - .map_err(|err| internal_datafusion_err!("dfx_storage: bad descriptor: {err}"))?; + .map_err(|err| exec_datafusion_err!("dfx_storage: bad descriptor: {err}"))?; let files = descriptor["files"] .as_array() - .ok_or_else(|| internal_datafusion_err!("dfx_storage: descriptor has no file list"))? + .ok_or_else(|| exec_datafusion_err!("dfx_storage: descriptor has no file list"))? .iter() .map(|file| { - let path = file["path"].as_str().ok_or_else(|| { - internal_datafusion_err!("dfx_storage: file entry has no path") - })?; + let path = file["path"] + .as_str() + .ok_or_else(|| exec_datafusion_err!("dfx_storage: file entry has no path"))?; let size = file["size"].as_u64().ok_or_else(|| { - internal_datafusion_err!("dfx_storage: file entry {path} has no size") + exec_datafusion_err!("dfx_storage: file entry {path} has no size") })?; Ok(FileSlice { path: path.to_string(), @@ -233,7 +256,7 @@ impl PhysicalExtensionCodec for DfxStoragePhysicalCodec { .as_u64() .and_then(|index| usize::try_from(index).ok()) .ok_or_else(|| { - internal_datafusion_err!( + exec_datafusion_err!( "dfx_storage: projection index {index} is not a column number" ) }) @@ -241,7 +264,7 @@ impl PhysicalExtensionCodec for DfxStoragePhysicalCodec { .collect::>>()?, ), other => { - return internal_err!( + return exec_err!( "dfx_storage: projection must be a list of column numbers or null, got {other}" ); } @@ -253,7 +276,7 @@ impl PhysicalExtensionCodec for DfxStoragePhysicalCodec { .as_u64() .and_then(|limit| usize::try_from(limit).ok()) .ok_or_else(|| { - internal_datafusion_err!("dfx_storage: limit {value} is not a row count") + exec_datafusion_err!("dfx_storage: limit {value} is not a row count") })?, ), }; @@ -356,9 +379,8 @@ impl LogicalExtensionCodec for DfxStorageLogicalCodec { .inner .try_decode_table_provider(buf, table_ref, schema, ctx); }; - let directory = std::str::from_utf8(directory).map_err(|err| { - internal_datafusion_err!("dfx_storage: bad directory in payload: {err}") - })?; + let directory = std::str::from_utf8(directory) + .map_err(|err| exec_datafusion_err!("dfx_storage: bad directory in payload: {err}"))?; self.counters .provider_decoded .fetch_add(1, Ordering::SeqCst); diff --git a/examples/distributed/storage-library/src/exec.rs b/examples/distributed/storage-library/src/exec.rs index 3f09f0391..c33036a67 100644 --- a/examples/distributed/storage-library/src/exec.rs +++ b/examples/distributed/storage-library/src/exec.rs @@ -105,6 +105,9 @@ impl PartitionedParquetExec { /// exists to own the *description* of the scan across a process boundary, /// not to reimplement Parquet. fn scan_for(&self, partition: usize) -> Result> { + // `Internal`: `execute` is only ever called with a partition this node + // said it had, so an out-of-range one is a caller's bug rather than + // bad data. Contrast the payload errors in [`crate::codec`]. let slice = self.files.get(partition).ok_or_else(|| { datafusion::common::internal_datafusion_err!( "PartitionedParquetExec has {} partition(s), asked for {partition}", From 9f9f688a6c3b351058374a607fd1d063ce15e72b Mon Sep 17 00:00:00 2001 From: Tim Saucer Date: Thu, 10 Sep 2026 16:29:24 -0400 Subject: [PATCH 18/24] Stop converting a fallback this planner never calls `__datafusion_session_planner__` ran `ffi_query_planner_from_pycapsule` over its `fallback` and then threw the result away with `let _ = fallback;`. The hook reads as though the fallback matters and then says it does not, which is the wrong thing to copy from an example. Nothing is lost by dropping it. I had assumed the conversion was at least serving as validation, and it is not: `_export_query_planner` already runs the same function over whatever the previous hook returned -- its doc says that is deliberate, "so a malformed planner surfaces at the hook that produced it rather than at the final install". So the second conversion repeated a getter call, a capsule check and an ABI check to produce a value with no reader. `DistributedQueryPlanner::fallback` goes too. It was only ever constructed `None`, which made the `Some(fallback)` arm of `create_physical_plan` unreachable, and unreachable code in an example is a liability -- it reads as a supported path and cannot rot loudly. What that arm documented is now prose on the struct: a planner delegates or rewrites, never both, and if you write the layering kind then hold the fallback and call it directly rather than going through `Session::create_physical_plan`, which dispatches through the installed planner and recurses until the stack overflows. `datafusion-ffi-query-planner-example` remains the crate that demonstrates layering for real, so the redundant illustration here is no loss. No behaviour change: the fallback was already unused. Co-Authored-By: Claude Opus 5 (1M context) --- .../engine-library/src/extension.rs | 32 ++++++------ .../distributed/engine-library/src/planner.rs | 49 +++++++++---------- 2 files changed, 42 insertions(+), 39 deletions(-) diff --git a/examples/distributed/engine-library/src/extension.rs b/examples/distributed/engine-library/src/extension.rs index a53541d00..d0f6d30ea 100644 --- a/examples/distributed/engine-library/src/extension.rs +++ b/examples/distributed/engine-library/src/extension.rs @@ -35,7 +35,7 @@ use datafusion_proto::physical_plan::PhysicalExtensionCodec; use datafusion_python_util::{ create_physical_extension_capsule, create_query_planner_capsule, ffi_logical_codec_from_pycapsule, ffi_physical_codec_from_pycapsule, - ffi_query_planner_from_pycapsule, ffi_task_context_provider_from_pycapsule, get_tokio_runtime, + ffi_task_context_provider_from_pycapsule, get_tokio_runtime, }; use datafusion_session::QueryPlanner; use pyo3::prelude::*; @@ -163,30 +163,34 @@ impl DfxEngineExtension { components.call((), Some(&kwargs)) } - /// Contribute this engine's planner, nesting it on whatever came before. + /// Contribute this engine's planner, replacing whatever came before. /// /// Runs after every bundle's codecs are installed, so `ctx` carries the /// final chains and the planner is not left encoding through a partial - /// set. `fallback` is the planner assembled so far; delegating to it is - /// what makes several planner-shipping libraries composable, and - /// returning a planner that ignored it would discard every layer beneath. + /// set. + /// + /// `_fallback` is the planner assembled so far, and this engine does not + /// use it: delegating would hand physical planning to the host and bring + /// the plan back as opaque foreign nodes, which cannot be split, and + /// splitting is the whole point. A planner that only rearranged stock + /// nodes would keep it -- see [`DistributedQueryPlanner`]. + /// + /// Not converted, either, which is worth saying because converting it and + /// dropping the result is an easy line to write. There is nothing for + /// this hook to validate: `SessionContext._export_query_planner` already + /// ran `ffi_query_planner_from_pycapsule` over whatever the previous hook + /// returned, precisely so a malformed planner surfaces at the hook that + /// produced it. Converting again would repeat a getter call, a capsule + /// check and an ABI check to reach a value this planner never calls. fn __datafusion_session_planner__<'py>( &self, py: Python<'py>, ctx: Bound<'py, PyAny>, - fallback: Bound<'py, PyAny>, + _fallback: Bound<'py, PyAny>, ) -> PyResult> { - let fallback = ffi_query_planner_from_pycapsule(&fallback, Some(&ctx))?; let planner: Arc = Arc::new(DistributedQueryPlanner { observations: Arc::clone(&self.observations), - // Deliberately not layered. Delegating would hand physical - // planning to the host and bring the plan back as opaque foreign - // nodes, which this engine cannot split -- so it plans for itself - // and the fallback goes unused. A planner that only rearranged - // stock nodes would keep it. - fallback: None, }); - let _ = fallback; // The planner takes the *host's* codecs, not ones built here. By now // those are the final chains, and this library has no business diff --git a/examples/distributed/engine-library/src/planner.rs b/examples/distributed/engine-library/src/planner.rs index f81ed6e2a..da0a95147 100644 --- a/examples/distributed/engine-library/src/planner.rs +++ b/examples/distributed/engine-library/src/planner.rs @@ -135,16 +135,24 @@ fn insert_stage( Ok((plan.replace_children(children, options)?, inserted)) } +/// Holds no `fallback`, and that is the design rather than an omission. +/// +/// A planner either delegates or rewrites. Delegating hands physical planning +/// to whoever is underneath, including the host, and brings the plan back as +/// opaque foreign nodes -- which cannot be split, and splitting is the only +/// thing this library exists to do. So the hook is handed a fallback and +/// leaves it alone; see `DfxEngineExtension::__datafusion_session_planner__`. +/// +/// Two things worth knowing if you write the layering kind instead. Hold the +/// fallback as an `Option>` and call it +/// directly: `Session::create_physical_plan` looks like the way to delegate +/// and is not, because it dispatches through the session's *installed* +/// planner, so calling it from inside that planner recurses until the stack +/// overflows. And `datafusion-ffi-query-planner-example` is the crate that +/// demonstrates layering for real, including how `fallback` nests. #[derive(Debug)] pub(crate) struct DistributedQueryPlanner { pub(crate) observations: Arc, - /// Planner to layer on top of, if the session already had one. - /// - /// Held so several planner-shipping libraries compose. Note that - /// `Session::create_physical_plan` cannot be used for this: it dispatches - /// through the session's installed planner, so calling it from inside that - /// planner recurses until the stack overflows. - pub(crate) fallback: Option>, } #[async_trait] @@ -156,24 +164,15 @@ impl QueryPlanner for DistributedQueryPlanner { ) -> Result> { self.observations.plan_calls.fetch_add(1, Ordering::SeqCst); - let plan = match self.fallback.as_ref() { - // Delegating hands physical planning to whoever is underneath, - // including the host. That is correct for composition, but it - // means the plan comes back as opaque foreign nodes this engine - // cannot split -- so a fallback and a split are exclusive, and - // the split is what this library is for. - Some(fallback) => return fallback.create_physical_plan(logical_plan, session).await, - None => { - // Plan against a session that owns the stock rule set locally - // instead of reaching back over FFI for the host's. Without - // this the plan contains `ForeignExecutionPlan` wrappers that - // cannot be serialized and cannot be rewritten. - let local = LocalOptimizerSession::new(session); - DefaultPhysicalPlanner::default() - .create_physical_plan(logical_plan, &local) - .await? - } - }; + // Plan against a session that owns the stock rule set locally instead + // of reaching back over FFI for the host's. Without this the plan + // contains `ForeignExecutionPlan` wrappers that cannot be serialized + // and cannot be rewritten -- and rewriting is the next thing that + // happens here. + let local = LocalOptimizerSession::new(session); + let plan = DefaultPhysicalPlanner::default() + .create_physical_plan(logical_plan, &local) + .await?; let Some(shuffle_dir) = shuffle_dir_from_options(session.config_options()) else { // No shuffle directory configured: leave the plan alone and let it From 4a8cb2ce28aaa50b3d0f17e81dd2374e67de0f16 Mon Sep 17 00:00:00 2001 From: Tim Saucer Date: Thu, 10 Sep 2026 16:41:15 -0400 Subject: [PATCH 19/24] Clear the small stuff off the distributed example Six unrelated tidy-ups, no behaviour change to any query. `driver.py` reads a worker's row count off the last line of its stdout rather than parsing the whole buffer. A worker shares stdout with everything loaded into it, so one stray print turns `json.loads` into a failure a long way from its cause; an unreadable or absent report now says which stage and partition produced it. `dfx_storage` exports `BundledLogicalCodec`. It was constructed with `module = "dfx_storage"` but never added to the module, so that claim was untrue and the physical half was exported while the logical half was not. `write_partition` writes the file with the schema the stream declares instead of the one the child's stream reports. They agree for any well-behaved child, and pinning it to the declaration makes a disagreement loud: `StreamWriter::write` rejects a batch whose schema differs, so a child contradicting its own `schema()` fails there rather than publishing a file readers were told to expect something else from. `run_tpch.py` raises instead of asserting. That comparison is the only thing making the script a check rather than a demo, and `python -O` drops an `assert`, which would leave it printing a table it never verified. The tolerance rationale moves to the new function's docstring, where it was otherwise duplicated. The CI artifact is `test-example-wheels-x86_64`. It stopped being manylinux when these builds moved to the host, and it carries five projects rather than the two the old name and step titles implied. `dfx_engine` documents why it declares no dependency on the sibling libraries. I tried declaring them -- `dfx_engine.session` imports both, so `import dfx_engine` fails without them -- and it breaks the build outright, because neither is published: Because dfx-storage was not found in the package registry and your project depends on dfx-storage, we can conclude that your project's requirements are unsatisfiable. So the note records the dead end next to the field someone will otherwise fill in again. The requirement itself stays in the README, beside the install command that satisfies it. Verified: all three wheels build, install together the way test.yml does, and their suites pass 13 / 11 / 19. Both workflow files parse as YAML; actionlint could not run locally, as it needs Docker. Co-Authored-By: Claude Opus 5 (1M context) --- .github/workflows/build.yml | 6 +- .github/workflows/test.yml | 10 ++-- .../distributed/engine-library/pyproject.toml | 12 ++++ .../python/dfx_engine/driver.py | 26 ++++++++- .../distributed/engine-library/src/stage.rs | 10 +++- examples/distributed/run_tpch.py | 56 +++++++++++++------ .../distributed/storage-library/src/lib.rs | 5 +- 7 files changed, 96 insertions(+), 29 deletions(-) diff --git a/.github/workflows/build.yml b/.github/workflows/build.yml index 649a2c51f..7b3303b51 100644 --- a/.github/workflows/build.yml +++ b/.github/workflows/build.yml @@ -186,7 +186,7 @@ jobs: manylinux: "2_28" # The example libraries below are test fixtures, not release artifacts: - # the `test-ffi-manylinux-x86_64` artifact they feed is consumed only by + # the `test-example-wheels-x86_64` artifact they feed is consumed only by # test.yml, which installs them on a runner like this one. They therefore # build on the host (`container: off`) rather than in the manylinux # container, which drops a container start and an in-container rustup @@ -261,11 +261,11 @@ jobs: name: dist-manylinux-x86_64-${{ matrix.python-tag }} path: dist/* - - name: Archive FFI test wheel + - name: Archive example test wheels if: matrix.python-tag == 'abi3' uses: actions/upload-artifact@v7 with: - name: test-ffi-manylinux-x86_64 + name: test-example-wheels-x86_64 path: | examples/datafusion-ffi-example/dist/* examples/datafusion-ffi-query-planner-example/dist/* diff --git a/.github/workflows/test.yml b/.github/workflows/test.yml index c1efe0bae..550f2c386 100644 --- a/.github/workflows/test.yml +++ b/.github/workflows/test.yml @@ -73,11 +73,11 @@ jobs: path: wheels/ # FFI test wheel only built once (under the abi3 matrix entry in build.yml). - - name: Download pre-built FFI test wheel + - name: Download pre-built example test wheels if: matrix.wheel-tag == 'abi3' uses: actions/download-artifact@v8 with: - name: test-ffi-manylinux-x86_64 + name: test-example-wheels-x86_64 path: wheels/ - name: Install from pre-built wheels @@ -93,9 +93,9 @@ jobs: uv venv --python "${{ steps.setup-python.outputs.python-path }}" VENV_PY="$PWD/.venv/bin/python" uv sync --python "$VENV_PY" --dev --no-install-package datafusion - # Search recursively: the FFI artifact bundles more than one - # project, so upload-artifact keeps a `/dist/` prefix - # and the wheels are not all at the top of wheels/. + # Search recursively: the example artifact bundles five projects, + # so upload-artifact keeps a `/dist/` prefix and the + # wheels are not all at the top of wheels/. WHEELS=$(find wheels/ -name "*.whl") if [ -n "$WHEELS" ]; then echo "Installing wheels:" diff --git a/examples/distributed/engine-library/pyproject.toml b/examples/distributed/engine-library/pyproject.toml index da786e50c..038a6d710 100644 --- a/examples/distributed/engine-library/pyproject.toml +++ b/examples/distributed/engine-library/pyproject.toml @@ -27,6 +27,18 @@ classifiers = [ "Programming Language :: Python :: Implementation :: CPython", ] dynamic = ["version"] +# No `dependencies`, deliberately, though `dfx_engine.session` imports +# `dfx_storage` and `dfx_udfs` at module scope and `import dfx_engine` fails +# without them. Declaring them does not work: neither is published, so +# resolution fails before a wheel is built at all -- +# +# Because dfx-storage was not found in the package registry and your +# project depends on dfx-storage, we can conclude that your project's +# requirements are unsatisfiable. +# +# which would break the build rather than document the requirement. The +# requirement is in the README instead, next to the install command that +# satisfies it. [tool.maturin] features = ["pyo3/extension-module"] diff --git a/examples/distributed/engine-library/python/dfx_engine/driver.py b/examples/distributed/engine-library/python/dfx_engine/driver.py index 8f1d8d6f9..65d13c97d 100644 --- a/examples/distributed/engine-library/python/dfx_engine/driver.py +++ b/examples/distributed/engine-library/python/dfx_engine/driver.py @@ -142,6 +142,30 @@ def _dispatch( ) +def _report(task: tuple[int, int], stdout: str) -> int: + """Read a worker's row count off its last line of output. + + The last line, not the whole stream: a worker's stdout is shared with + everything loaded into it, and one stray `print` from a library -- or a + warning some future dependency decides to write there -- would turn + `json.loads` on the whole buffer into a confusing failure a long way from + its cause. + """ + stage_id, partition = task + lines = [line for line in stdout.splitlines() if line.strip()] + if not lines: + message = f"stage {stage_id} partition {partition} printed no report" + raise RuntimeError(message) + try: + return json.loads(lines[-1])["rows"] + except (ValueError, KeyError) as err: + message = ( + f"stage {stage_id} partition {partition} printed an unreadable " + f"report {lines[-1]!r}" + ) + raise RuntimeError(message) from err + + def run_distributed( sql: str, spec: SessionSpec, extra_udfs: list[ScalarUDF] | None = None ) -> DistributedResult: @@ -216,7 +240,7 @@ def run_distributed( stage_id, partition = task failures.append(f"stage {stage_id} partition {partition} failed:\n{stderr}") continue - task_rows[task] = json.loads(stdout)["rows"] + task_rows[task] = _report(task, stdout) if failures: raise RuntimeError("\n".join(failures)) diff --git a/examples/distributed/engine-library/src/stage.rs b/examples/distributed/engine-library/src/stage.rs index 80ca72944..37454b71c 100644 --- a/examples/distributed/engine-library/src/stage.rs +++ b/examples/distributed/engine-library/src/stage.rs @@ -125,6 +125,14 @@ impl ShuffleStageExec { let final_path = partition_path(&self.shuffle_dir, self.stage_id, partition); let temp_path = temp_partition_path(&self.shuffle_dir, self.stage_id, partition); let shuffle_dir = self.shuffle_dir.clone(); + // The schema the file is written with is the one this stream declares, + // not the one the child's stream happens to report. They agree for any + // well-behaved child, and pinning it to the declaration is what makes + // a disagreement loud: `StreamWriter::write` rejects a batch whose + // schema differs, so a child that contradicts its own `schema()` fails + // here instead of publishing a file readers were told to expect + // something else from. + let written_schema = Arc::clone(&schema); let collected = async move { let mut stream = input.execute(partition, context)?; @@ -139,7 +147,7 @@ impl ShuffleStageExec { let file = fs::File::create(&temp_path).map_err(|err| { exec_datafusion_err!("dfx_engine: creating {}: {err}", temp_path.display()) })?; - let mut writer = StreamWriter::try_new(file, stream.schema().as_ref()) + let mut writer = StreamWriter::try_new(file, written_schema.as_ref()) .map_err(|err| exec_datafusion_err!("dfx_engine: ipc writer: {err}"))?; for batch in &batches { writer diff --git a/examples/distributed/run_tpch.py b/examples/distributed/run_tpch.py index b32381bc7..08b00014d 100644 --- a/examples/distributed/run_tpch.py +++ b/examples/distributed/run_tpch.py @@ -92,6 +92,39 @@ def reshard( return written +def compare(table: pa.Table, reference: pa.Table) -> None: + """Raise unless `table` matches `reference`, floats to 1e-6 relative. + + Floats get a tolerance rather than 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. Every distributed engine has this property; + it is worth knowing before someone diffs two runs and concludes the split + is broken. + """ + if table.column_names != reference.column_names: + message = ( + f"column names differ: {table.column_names} vs {reference.column_names}" + ) + raise ValueError(message) + + for name in table.column_names: + got = table.column(name).to_pylist() + want = reference.column(name).to_pylist() + if len(got) != len(want): + message = f"{name}: {len(got)} rows distributed, {len(want)} local" + raise ValueError(message) + for lhs, rhs in zip(got, want, strict=True): + close = ( + abs(lhs - rhs) <= 1e-6 * max(1.0, abs(rhs)) + if isinstance(lhs, float) + else lhs == rhs + ) + if not close: + message = f"{name}: {lhs!r} distributed, {rhs!r} local" + raise ValueError(message) + + def main(argv: list[str] | None = None) -> int: parser = argparse.ArgumentParser(description=__doc__) parser.add_argument( @@ -150,24 +183,11 @@ def main(argv: list[str] | None = None) -> int: 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) + # Raises rather than asserts. This comparison is the only thing that + # makes the script a check rather than a demo, and `python -O` removes + # an `assert` -- which would leave it printing a table it never + # verified. + compare(table, reference) print("\nsame answer both ways (floats to within 1e-6 relative):\n") names = table.column_names diff --git a/examples/distributed/storage-library/src/lib.rs b/examples/distributed/storage-library/src/lib.rs index bb01cf9c7..b94b9f1ca 100644 --- a/examples/distributed/storage-library/src/lib.rs +++ b/examples/distributed/storage-library/src/lib.rs @@ -22,7 +22,7 @@ use pyo3::prelude::*; -use crate::extension::{BundledPhysicalCodec, DfxStorageExtension}; +use crate::extension::{BundledLogicalCodec, BundledPhysicalCodec, DfxStorageExtension}; use crate::table_provider::PyPartitionedParquetTable; mod codec; @@ -33,6 +33,9 @@ mod table_provider; #[pymodule] fn dfx_storage(m: &Bound<'_, PyModule>) -> PyResult<()> { pyo3_log::init(); + // Both bundled codecs, so that `module = "dfx_storage"` on each is true + // and a caller inspecting a session's codecs sees a type it can look up. + m.add_class::()?; m.add_class::()?; m.add_class::()?; m.add_class::()?; From b0030c79759e0b4f5680e001e1b31d1271cee6fc Mon Sep 17 00:00:00 2001 From: Tim Saucer Date: Fri, 11 Sep 2026 08:19:13 -0400 Subject: [PATCH 20/24] Refuse a shuffle directory that already holds stage output A stage node reads partition `i` if its file exists and computes it otherwise, which is what lets one node be both halves of the exchange. The cost is that the filesystem is the state, and stage ids restart at `FIRST_STAGE_ID` for every plan -- so a second query pointed at the same shuffle directory finds the first one's files and reads them, having computed nothing. Nothing downstream can catch it. When the two plans' stage schemas differ it surfaces as an unrelated-looking schema error; when they agree -- the same query over data that has since changed -- the previous answer arrives silently. `run_distributed` now checks before building a session. The glob lives next to `partition_path` in `stage.rs` and is exported like it, so the question "does this directory hold stage output?" is asked with the same naming convention the node answers it with. The check is scoped to stage output rather than to an empty directory, because the driver writes its own encoded plans and task envelopes there too. Co-Authored-By: Claude Opus 5 (1M context) --- .../distributing-work/query-engines.md | 10 ++++ examples/distributed/README.md | 6 +++ .../python/dfx_engine/driver.py | 44 ++++++++++++++++- .../python/dfx_engine/session.py | 9 +++- .../python/tests/_test_three_libraries.py | 47 +++++++++++++++++++ .../distributed/engine-library/src/lib.rs | 11 +++++ .../distributed/engine-library/src/stage.rs | 9 ++++ 7 files changed, 133 insertions(+), 3 deletions(-) diff --git a/docs/source/user-guide/distributing-work/query-engines.md b/docs/source/user-guide/distributing-work/query-engines.md index a67d99fae..c2f4881a2 100644 --- a/docs/source/user-guide/distributing-work/query-engines.md +++ b/docs/source/user-guide/distributing-work/query-engines.md @@ -119,6 +119,16 @@ Two more that are about lifetime rather than configuration: easy to miss, because the failure only appears once an extension node is in the plan. +And one that is about the engine rather than the session: + +- **Where stages exchange results, scoped to one query.** An engine that + publishes stage output to a location you name, and numbers its stages from + scratch per plan, will read a previous query's results back if you point two + queries at the same location. Ask your engine whether it isolates that for + you; if it does not, give each query its own. This one is worth confirming + rather than assuming, because when the two plans happen to agree on schema + the wrong answer arrives without an error. + ## Available engines Query-level distribution is being built upstream. Neither project diff --git a/examples/distributed/README.md b/examples/distributed/README.md index 28b354434..c9bb23cc3 100644 --- a/examples/distributed/README.md +++ b/examples/distributed/README.md @@ -100,6 +100,12 @@ One node does both halves of that exchange, which is why nothing has to rewrite the plan in between. It also means a query run with no workers at all still gets the right answer; it just does the work itself. +The cost of that trick is that the filesystem is the state, and stage ids +restart at 1 for every plan — so a second query pointed at the same shuffle +directory would find the first one's files and read them. The driver refuses +a directory that already holds stage output rather than letting that through: +one shuffle directory per query. + ## The four things worth reading **`engine-library/python/dfx_engine/session.py`** is the point of the whole diff --git a/examples/distributed/engine-library/python/dfx_engine/driver.py b/examples/distributed/engine-library/python/dfx_engine/driver.py index 65d13c97d..1dc78ccde 100644 --- a/examples/distributed/engine-library/python/dfx_engine/driver.py +++ b/examples/distributed/engine-library/python/dfx_engine/driver.py @@ -18,7 +18,7 @@ """The driver: split a query into tasks, fan them out, collect the answer. The shape is deliberately boring, because the interesting part is not the -scheduling. What matters is the five things the driver has to get right, each +scheduling. What matters is the six things the driver has to get right, each of which is a way a real deployment goes wrong: 1. It serializes each stage **with** its session. ``to_bytes(None)`` uses an @@ -33,6 +33,9 @@ 5. It ships *every* stage. A plan can hold more than one -- an aggregate in each branch of a union, say -- and they are independent subtrees rather than a chain. +6. It refuses a shuffle directory that already holds stage output. That same + "read it if it is there" rule is what makes a *second* query in the same + directory read the first one's results -- see :func:`require_empty_shuffle`. """ from __future__ import annotations @@ -54,9 +57,12 @@ __all__ = [ "DistributedResult", + "dataframe_for", "find_stage", "find_stages", + "require_empty_shuffle", "run_distributed", + "run_local", ] @@ -166,13 +172,43 @@ def _report(task: tuple[int, int], stdout: str) -> int: raise RuntimeError(message) from err +def require_empty_shuffle(shuffle_dir: pathlib.Path) -> None: + """Refuse a directory that already holds stage output. + + A stage node reads partition `i` if the file for it exists and computes it + otherwise, which is what lets one node be both halves of the exchange. The + cost is that the filesystem *is* the state, and stage ids restart at + ``stage_id(0)`` for every plan -- so a second query in the same directory + finds files left by the first and reads them, having computed nothing. + + Nothing downstream can catch that. If the two plans' stage schemas differ + the failure is an unrelated-looking schema error a long way from here, and + if they agree -- the same query over data that has since changed, say -- + the answer is simply the old one, silently. + + So the rule this enforces is **one shuffle directory per query**, and the + check belongs here rather than in the node: the node cannot tell a file + this run's worker wrote from one last run's worker wrote, but the driver + knows it has not dispatched anything yet. + """ + stale = sorted(path.name for path in shuffle_dir.glob(_internal.partition_glob())) + if stale: + message = ( + f"{shuffle_dir} already holds stage output {stale}; a stage reads " + f"a partition file if it finds one, so this query would return the " + f"previous query's results. Use one shuffle directory per query." + ) + raise RuntimeError(message) + + def run_distributed( sql: str, spec: SessionSpec, extra_udfs: list[ScalarUDF] | None = None ) -> DistributedResult: """Run `sql`, executing each stage partition in its own worker process. Requires ``spec.shuffle_dir``: without it the planner inserts no stage and - there is nothing to distribute. + there is nothing to distribute. The directory must not already hold stage + output -- see :func:`require_empty_shuffle`. ``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 @@ -184,6 +220,10 @@ def run_distributed( message = "run_distributed needs a shuffle_dir; build_session got none" raise ValueError(message) + # Before building a session, so a reused directory costs nothing to + # diagnose. `glob` on a directory that does not exist yet yields nothing. + require_empty_shuffle(pathlib.Path(spec.shuffle_dir)) + ctx, _engine, _storage = build_session(spec) for function in extra_udfs or []: ctx.register_udf(function) diff --git a/examples/distributed/engine-library/python/dfx_engine/session.py b/examples/distributed/engine-library/python/dfx_engine/session.py index cdbb5399a..ce96299bf 100644 --- a/examples/distributed/engine-library/python/dfx_engine/session.py +++ b/examples/distributed/engine-library/python/dfx_engine/session.py @@ -62,7 +62,14 @@ class SessionSpec: """Table name to the directory ``dfx_storage`` should scan for it.""" shuffle_dir: str - """Where stages exchange results. Empty means "run in this process".""" + """Where stages exchange results. Empty means "run in this process". + + One directory per query. Stage ids restart at 1 for every plan, and a + stage reads a partition file if it finds one, so a directory reused across + two queries hands the second one the first one's results. + :func:`~dfx_engine.driver.run_distributed` refuses that rather than + letting it through. + """ target_partitions: int = 2 """Pinned rather than defaulted to the core count. diff --git a/examples/distributed/engine-library/python/tests/_test_three_libraries.py b/examples/distributed/engine-library/python/tests/_test_three_libraries.py index ab5373ef2..b8bd512eb 100644 --- a/examples/distributed/engine-library/python/tests/_test_three_libraries.py +++ b/examples/distributed/engine-library/python/tests/_test_three_libraries.py @@ -377,6 +377,53 @@ def test_without_a_shuffle_dir_nothing_is_distributed( assert _rows(ctx.sql(Q1).collect())[0] == ("A", "F", 3, 10.0, 1000.0) +def test_a_reused_shuffle_directory_is_refused( + spec: SessionSpec, tmp_path: pathlib.Path +) -> None: + """A second query in one directory would read the first one's results. + + Stage ids restart at 1 for every plan and a stage reads a partition file + if it finds one, so the files `Q1` leaves behind are exactly the files + `REVENUE`'s stage looks for. Nothing downstream can catch that: here the + two schemas differ so it would surface as an unrelated-looking error, but + re-running *the same* query over changed data would simply return the old + answer. + """ + run_distributed(Q1, spec) + + with pytest.raises(RuntimeError, match="already holds stage output") as excinfo: + run_distributed(REVENUE, spec) + + # Names the files and the rule, so the reader does not have to work out + # why a directory that "looks fine" was rejected. + assert "stage-1-part-0.arrow" in str(excinfo.value) + assert "one shuffle directory per query" in str(excinfo.value) + + # And a fresh directory is all it takes. + elsewhere = dataclasses.replace(spec, shuffle_dir=str(tmp_path / "second")) + assert _rows(run_distributed(REVENUE, elsewhere).batches) == _rows( + run_local(REVENUE, spec) + ) + + +def test_the_guard_is_scoped_to_stage_output(spec: SessionSpec) -> None: + """The driver's own scratch in the same directory is not stage output. + + `run_distributed` writes each stage's encoded plan and each worker's task + envelope beside the results, so a check that rejected any non-empty + directory would reject every second call for the wrong reason -- and the + obvious fix, deleting what it found, would delete those too. + """ + shuffle = pathlib.Path(spec.shuffle_dir) + shuffle.mkdir(parents=True) + (shuffle / "stage-1.plan").write_bytes(b"leftover") + (shuffle / "task-1-0.json").write_text("{}") + + result = run_distributed(Q1, spec) + + assert _rows(result.batches) == _rows(run_local(Q1, spec)) + + 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 = { diff --git a/examples/distributed/engine-library/src/lib.rs b/examples/distributed/engine-library/src/lib.rs index 2a38f7eb4..37e73802e 100644 --- a/examples/distributed/engine-library/src/lib.rs +++ b/examples/distributed/engine-library/src/lib.rs @@ -53,6 +53,16 @@ fn partition_path(shuffle_dir: &str, stage_id: u32, partition: usize) -> String .into_owned() } +/// Glob matching every file [`partition_path`] can produce, for any stage. +/// +/// Exported for the same reason: the driver refuses a shuffle directory that +/// already holds stage output, and it has to ask that question with the same +/// naming convention the node answers it with. +#[pyfunction] +fn partition_glob() -> &'static str { + stage::PARTITION_GLOB +} + /// Id of the `index`th stage in a plan, counting in pre-order from the root. /// /// Exported for the same reason as [`partition_path`]: a foreign node's @@ -74,6 +84,7 @@ fn _internal(m: &Bound<'_, PyModule>) -> PyResult<()> { m.add_class::()?; m.add_class::()?; m.add_class::()?; + m.add_function(wrap_pyfunction!(partition_glob, m)?)?; m.add_function(wrap_pyfunction!(partition_path, m)?)?; m.add_function(wrap_pyfunction!(stage_id, m)?)?; Ok(()) diff --git a/examples/distributed/engine-library/src/stage.rs b/examples/distributed/engine-library/src/stage.rs index 37454b71c..ee7ed750d 100644 --- a/examples/distributed/engine-library/src/stage.rs +++ b/examples/distributed/engine-library/src/stage.rs @@ -62,6 +62,15 @@ pub(crate) fn partition_path(shuffle_dir: &str, stage_id: u32, partition: usize) Path::new(shuffle_dir).join(format!("stage-{stage_id}-part-{partition}.arrow")) } +/// Matches every file [`partition_path`] can produce, for any stage. +/// +/// Here rather than in Python because a caller asking "does this directory +/// already hold stage output?" must not re-spell the naming convention: the +/// answer decides whether [`ShuffleStageExec::execute`] reads or recomputes, +/// so a pattern that drifted from the path would report an empty directory +/// that is not one. The two are deliberately adjacent for that reason. +pub(crate) const PARTITION_GLOB: &str = "stage-*-part-*.arrow"; + /// Where a writer builds a partition before publishing it. /// /// Unique per writer, not merely per partition. Deriving the temporary name From 089a8563236852dfc3aef88b5997e8f404040631 Mon Sep 17 00:00:00 2001 From: Tim Saucer Date: Fri, 11 Sep 2026 08:21:01 -0400 Subject: [PATCH 21/24] Check a shuffle file's schema instead of adopting it `write_partition` pins what it writes to the node's declared schema so a child that contradicts its own `schema()` fails at the writer. The read path did the opposite: it took `reader.schema()` off the file and handed it to the stream unexamined. Reaching that code means a file was found, and finding one says nothing about who wrote it -- an older build, or a query whose stage happened to be numbered the same, both leave something readable behind. Adopting its schema pushes the disagreement up to whichever operator first uses the batches, where the message no longer names the file; where the field lists differ only in type, there may be no error at all. Compare fields and refuse a mismatch, naming the file and both field lists. Fields rather than the whole schema: a name, a type or a nullability that disagrees changes how a batch is read, and schema-level metadata does not. Co-Authored-By: Claude Opus 5 (1M context) --- .../python/tests/_test_three_libraries.py | 32 +++++++++++ .../distributed/engine-library/src/stage.rs | 55 ++++++++++++++++++- 2 files changed, 84 insertions(+), 3 deletions(-) diff --git a/examples/distributed/engine-library/python/tests/_test_three_libraries.py b/examples/distributed/engine-library/python/tests/_test_three_libraries.py index b8bd512eb..79c59eb29 100644 --- a/examples/distributed/engine-library/python/tests/_test_three_libraries.py +++ b/examples/distributed/engine-library/python/tests/_test_three_libraries.py @@ -255,6 +255,38 @@ def test_the_driver_reads_the_workers_output(spec: SessionSpec) -> None: ctx.sql(Q1).collect() +def test_a_shuffle_file_with_the_wrong_schema_is_refused(spec: SessionSpec) -> None: + """A *readable* file is not the same as one this stage wrote. + + The previous test corrupts the bytes, which the IPC reader rejects on its + own. This one leaves a perfectly valid Arrow stream carrying the wrong + columns -- what an older build, or a query whose stage was numbered the + same, would leave behind. Adopting its schema would push the disagreement + up to whichever operator first used the batches, with nothing in the + message about the file it came from. + """ + run_distributed(Q1, spec) + + victim = pathlib.Path( + _internal.partition_path(spec.shuffle_dir, _internal.stage_id(), 1) + ) + table = pa.table({"unrelated": [1, 2, 3]}) + sink = pa.BufferOutputStream() + with pa.ipc.new_stream(sink, table.schema) as writer: + writer.write_table(table) + victim.write_bytes(sink.getvalue().to_pybytes()) + + ctx, _engine, _storage = build_session(spec) + with pytest.raises(Exception, match="this stage produces") as excinfo: + ctx.sql(Q1).collect() + + # Names the file and both field lists, so the mismatch is readable + # without opening either side. + assert victim.name in str(excinfo.value) + assert "unrelated: Int64" in str(excinfo.value) + assert "l_returnflag: Utf8" in str(excinfo.value) + + 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) diff --git a/examples/distributed/engine-library/src/stage.rs b/examples/distributed/engine-library/src/stage.rs index ee7ed750d..9bcd81a84 100644 --- a/examples/distributed/engine-library/src/stage.rs +++ b/examples/distributed/engine-library/src/stage.rs @@ -42,10 +42,11 @@ use std::sync::Arc; use std::sync::atomic::{AtomicU64, Ordering}; use std::{fmt, fs}; +use arrow::datatypes::Schema; use arrow::ipc::reader::StreamReader; use arrow::ipc::writer::StreamWriter; use datafusion::common::tree_node::TreeNodeRecursion; -use datafusion::common::{DataFusionError, Result, exec_datafusion_err, internal_err}; +use datafusion::common::{DataFusionError, Result, exec_datafusion_err, exec_err, internal_err}; use datafusion::execution::{SendableRecordBatchStream, TaskContext}; use datafusion::physical_expr::PhysicalExpr; use datafusion::physical_plan::memory::MemoryStream; @@ -89,6 +90,20 @@ fn temp_partition_path(shuffle_dir: &str, stage_id: u32, partition: usize) -> Pa )) } +/// `name: type` per field, for an error a reader can act on. +/// +/// The `Debug` of a `Fields` runs to a screenful for even a small schema and +/// buries the names among nullability flags and empty metadata maps, which is +/// the opposite of what someone comparing two of them needs. +fn describe(schema: &Schema) -> String { + schema + .fields() + .iter() + .map(|field| format!("{}: {}", field.name(), field.data_type())) + .collect::>() + .join(", ") +} + /// Marks a subtree as one stage of a distributed query. #[derive(Debug)] pub(crate) struct ShuffleStageExec { @@ -185,17 +200,51 @@ impl ShuffleStageExec { Ok(Box::pin(RecordBatchStreamAdapter::new(schema, stream))) } + /// Stream a partition somebody already published. + /// + /// The file's schema is checked rather than adopted, which is the other + /// half of the argument in [`Self::write_partition`]. That one pins what + /// it writes to this node's declared schema so a child contradicting its + /// own `schema()` fails at the writer; this one refuses a file that + /// disagrees, so a file *this process did not write* fails at the reader. + /// + /// Worth doing because reaching here at all means a file was found, and + /// nothing about finding it says who wrote it -- an older build of this + /// library, or a query whose stage happened to be numbered the same, both + /// leave something readable behind. Adopting its schema instead pushes + /// the disagreement up to whichever operator first tries to use the + /// batches, where the error no longer names the file it came from; and + /// where the two field lists differ only in type, there may be no error + /// at all. + /// + /// Fields only, not the whole schema: a name, a type or a nullability + /// that disagrees changes how a batch is read, and schema-level metadata + /// does not. fn read_partition(&self, partition: usize) -> Result { let path = partition_path(&self.shuffle_dir, self.stage_id, partition); let file = fs::File::open(&path) .map_err(|err| exec_datafusion_err!("dfx_engine: opening {}: {err}", path.display()))?; let reader = StreamReader::try_new(file, None) .map_err(|err| exec_datafusion_err!("dfx_engine: reading {}: {err}", path.display()))?; - let schema = reader.schema(); + + let expected = self.schema(); + let found = reader.schema(); + if found.fields() != expected.fields() { + return exec_err!( + "dfx_engine: {} holds [{}] but this stage produces [{}]", + path.display(), + describe(&found), + describe(&expected) + ); + } + let batches = reader .collect::>>() .map_err(|err| exec_datafusion_err!("dfx_engine: reading {}: {err}", path.display()))?; - Ok(Box::pin(MemoryStream::try_new(batches, schema, None)?)) + // The declared schema, not the file's, for the same reason the writer + // used it: the two agree on every field by the check above, and the + // one this node advertises is the one its parent was planned against. + Ok(Box::pin(MemoryStream::try_new(batches, expected, None)?)) } } From d76822f93bf93e7d5f74502965e99c79669db27e Mon Sep 17 00:00:00 2001 From: Tim Saucer Date: Fri, 11 Sep 2026 08:29:36 -0400 Subject: [PATCH 22/24] Size a scalar function's output from the batch, not its arguments `values_to_arrays` infers its length from the arguments it is given, and when they are all scalars there is nothing to infer from, so it returns one row. `dfx_net_revenue` then produced a single value for a batch of a hundred, into a column the rest of the plan sizes at a hundred. Use `to_array(args.number_rows)` per argument instead, which is the pattern `crates/core/src/udf.rs` already follows. Iterating to `number_rows` rather than to whichever argument was examined first also makes the output length a stated contract. No test: an all-literal call is constant, so `SimplifyExpressions` folds it before execution and the function only ever sees a one-row batch that agrees with `number_rows`. There is no reachable path from Python, and this repository does not run `cargo test`. The invisibility is the reason to get the shape right in a crate written to be copied. Co-Authored-By: Claude Opus 5 (1M context) --- .../distributed/udf-library/src/functions.rs | 28 ++++++++++++++++--- 1 file changed, 24 insertions(+), 4 deletions(-) diff --git a/examples/distributed/udf-library/src/functions.rs b/examples/distributed/udf-library/src/functions.rs index ec7e95eb5..9ef5c258d 100644 --- a/examples/distributed/udf-library/src/functions.rs +++ b/examples/distributed/udf-library/src/functions.rs @@ -72,16 +72,36 @@ impl ScalarUDFImpl for NetRevenue { Ok(DataType::Float64) } + /// # Why `to_array(number_rows)` and not `values_to_arrays` + /// + /// `values_to_arrays` infers its length from the arguments, and when they + /// are *all* scalars there is nothing to infer from, so it produces one + /// row. This function would then hand back a single value for a batch of + /// a hundred, into a column the rest of the plan sizes at a hundred. + /// + /// `number_rows` is the batch's own count and is the only argument- + /// independent answer, which is why it is passed. Taking it also makes + /// the output length a stated contract rather than something read off + /// whichever argument happened to be examined first. + /// + /// Not reachable from SQL today: an all-literal call is constant, so + /// `SimplifyExpressions` folds it before execution and this code sees a + /// one-row batch that agrees with `number_rows`. That is precisely what + /// makes the shape worth copying correctly -- the bug is invisible until + /// something suppresses the fold. fn invoke_with_args(&self, args: ScalarFunctionArgs) -> Result { - let arrays = ColumnarValue::values_to_arrays(&args.args)?; - let [price, discount, tax] = arrays.as_slice() else { - return exec_err!("{NET_REVENUE} takes 3 arguments, got {}", arrays.len()); + let rows = args.number_rows; + let [price, discount, tax] = args.args.as_slice() else { + return exec_err!("{NET_REVENUE} takes 3 arguments, got {}", args.args.len()); }; + let price = price.to_array(rows)?; + let discount = discount.to_array(rows)?; + let tax = tax.to_array(rows)?; let price = price.as_primitive::(); let discount = discount.as_primitive::(); let tax = tax.as_primitive::(); - let values: Float64Array = (0..price.len()) + let values: Float64Array = (0..rows) .map(|row| { if price.is_null(row) || discount.is_null(row) || tax.is_null(row) { return None; From 936cb1bd918adff375595e150bd4f62821f37f1a Mon Sep 17 00:00:00 2001 From: Tim Saucer Date: Fri, 11 Sep 2026 08:34:47 -0400 Subject: [PATCH 23/24] Say what the example-test gate actually does The comment claimed these steps "only need to run once", but `wheel-tag == 'abi3'` matches five of the six matrix entries, not one. Anyone reading it to decide whether adding a suite here is expensive got the wrong answer by a factor of five. Running five times is correct and worth stating as intent rather than leaving as an accident: the example wheels are abi3, so the interpreter underneath them is the only thing that varies between those entries, and these suites are what exercise the capsule protocol from Python. The cost is four seconds for the two FFI examples, measured, against a job that takes four minutes. Renamed to match, since the step now runs the distributed example too. actionlint not run: it needs a Docker daemon, which is unavailable here. The file parses as YAML and the step list is unchanged apart from the rename. Co-Authored-By: Claude Opus 5 (1M context) --- .github/workflows/test.yml | 10 ++++++++-- 1 file changed, 8 insertions(+), 2 deletions(-) diff --git a/.github/workflows/test.yml b/.github/workflows/test.yml index 550f2c386..186f78098 100644 --- a/.github/workflows/test.yml +++ b/.github/workflows/test.yml @@ -119,8 +119,14 @@ jobs: # free-threaded build and re-pick the system 3.12 (see install step). uv run --python "$PWD/.venv/bin/python" --no-project pytest -v --import-mode=importlib - # FFI + TPC-H examples only need to run once; gate to abi3 entries. - - name: FFI unit tests + # Gated to the abi3 entries because the free-threaded build has no + # example wheels to test against -- not to run once. This is five of + # the six matrix entries, and running the suites against 3.10 through + # 3.14 is the point: the example wheels are abi3, so the interpreter + # underneath them is the only thing that varies, and these suites are + # what exercise the capsule protocol from Python. Four seconds for the + # two FFI examples as of this writing. + - name: Example library tests if: matrix.wheel-tag == 'abi3' run: | cd examples/datafusion-ffi-example From 857a9ff782ade58999557aca7bed0f720d31df36 Mon Sep 17 00:00:00 2001 From: Tim Saucer Date: Fri, 11 Sep 2026 08:38:06 -0400 Subject: [PATCH 24/24] Keep the driver's scratch out of the results directory Three small things the shuffle-directory work left behind. `run_distributed` wrote each stage's encoded plan and each worker's task envelope beside the partition files. Harmless, because the stage node builds exact paths rather than globbing, but it meant "does this directory hold stage output?" -- the question `require_empty_shuffle` now asks -- had to be answered by pattern rather than by looking. They move to a `tasks/` subdirectory, so the invariant is that everything directly under the shuffle directory is stage output. The guard stays scoped to the stage node's own naming regardless: a caller pointing it at a directory of their own should get an answer about the files that would actually be read. `.gitignore` gains the shuffle output pattern. Every supported path writes these under a temporary directory, so one in the tree means a relative shuffle_dir reached a stage node -- which is how three of them were committed on this branch once already. `compare` in run_tpch.py did `abs(lhs - rhs)` before testing for None, so a null in a float column raised TypeError and reported a genuine null-handling difference between the two runs as a crash in the checker. Co-Authored-By: Claude Opus 5 (1M context) --- .gitignore | 5 +++ .../python/dfx_engine/driver.py | 18 +++++++-- .../python/tests/_test_three_libraries.py | 40 +++++++++++++++---- examples/distributed/run_tpch.py | 14 ++++--- 4 files changed, 62 insertions(+), 15 deletions(-) diff --git a/.gitignore b/.gitignore index f5ade5698..7afc67b7b 100644 --- a/.gitignore +++ b/.gitignore @@ -40,3 +40,8 @@ examples/distributed/*/.venv/ # Left behind by the `uv` commands in examples/distributed/README.md. The # example projects are not locked -- only the root project is. examples/distributed/*/uv.lock +# Shuffle output. Every supported path writes these under a temporary +# directory, so one appearing in the tree means a relative shuffle_dir +# reached a stage node and the files landed in the working directory +# instead -- which is how three of them were committed once already. +examples/distributed/**/stage-*-part-*.arrow diff --git a/examples/distributed/engine-library/python/dfx_engine/driver.py b/examples/distributed/engine-library/python/dfx_engine/driver.py index 1dc78ccde..f1b5ef089 100644 --- a/examples/distributed/engine-library/python/dfx_engine/driver.py +++ b/examples/distributed/engine-library/python/dfx_engine/driver.py @@ -190,6 +190,12 @@ def require_empty_shuffle(shuffle_dir: pathlib.Path) -> None: check belongs here rather than in the node: the node cannot tell a file this run's worker wrote from one last run's worker wrote, but the driver knows it has not dispatched anything yet. + + Scoped to stage output, not to "is the directory empty". The two are the + same thing only because :func:`run_distributed` keeps its own scratch in + a ``tasks/`` subdirectory; a caller who points this at a directory of + their own should get an answer about the files that would actually be + read, not about the ones they put there. """ stale = sorted(path.name for path in shuffle_dir.glob(_internal.partition_glob())) if stale: @@ -238,7 +244,13 @@ def run_distributed( raise RuntimeError(message) shuffle_dir = pathlib.Path(spec.shuffle_dir) - shuffle_dir.mkdir(parents=True, exist_ok=True) + # The driver's scratch lives in a subdirectory rather than beside the + # results. The stage node owns the `stage-*-part-*.arrow` namespace in + # `shuffle_dir` and nothing else should write there, so that + # :func:`require_empty_shuffle` is asking about stage output rather than + # about whatever else the driver happens to have left lying around. + task_dir = shuffle_dir / "tasks" + task_dir.mkdir(parents=True, exist_ok=True) # One task per (stage, partition). Every stage is shipped, not just the # first: a query with an aggregate in two branches has two independent @@ -251,7 +263,7 @@ def run_distributed( # because the FFI wrapper's display does not carry it. stage_id = _internal.stage_id(index) # Encode the stage subtree, through the session that owns the codecs. - plan_path = shuffle_dir / f"stage-{stage_id}.plan" + plan_path = task_dir / f"stage-{stage_id}.plan" plan_path.write_bytes(stage.to_bytes(ctx)) for partition in range(stage.partition_count): tasks.append((stage_id, partition)) @@ -268,7 +280,7 @@ def run_distributed( # example is making: each worker reads a different file and writes a # different result, so they need no coordination beyond the directory. workers = [ - _dispatch(envelope, shuffle_dir, stage_id, partition) + _dispatch(envelope, task_dir, stage_id, partition) for envelope, (stage_id, partition) in zip(envelopes, tasks, strict=True) ] diff --git a/examples/distributed/engine-library/python/tests/_test_three_libraries.py b/examples/distributed/engine-library/python/tests/_test_three_libraries.py index 79c59eb29..31dfb36f9 100644 --- a/examples/distributed/engine-library/python/tests/_test_three_libraries.py +++ b/examples/distributed/engine-library/python/tests/_test_three_libraries.py @@ -439,23 +439,49 @@ def test_a_reused_shuffle_directory_is_refused( def test_the_guard_is_scoped_to_stage_output(spec: SessionSpec) -> None: - """The driver's own scratch in the same directory is not stage output. + """Only the files a stage would read count, not everything in the way. - `run_distributed` writes each stage's encoded plan and each worker's task - envelope beside the results, so a check that rejected any non-empty - directory would reject every second call for the wrong reason -- and the - obvious fix, deleting what it found, would delete those too. + The question the guard asks is "would this query read somebody else's + results", and the only files that can do that are the ones matching the + stage node's own naming. Rejecting any non-empty directory would answer a + different question, and the obvious follow-on -- deleting what it found -- + would then delete things it never looked at properly. """ shuffle = pathlib.Path(spec.shuffle_dir) shuffle.mkdir(parents=True) - (shuffle / "stage-1.plan").write_bytes(b"leftover") - (shuffle / "task-1-0.json").write_text("{}") + (shuffle / "notes.txt").write_text("mine, not the engine's") + (shuffle / "stage-1.plan").write_bytes(b"looks close, reads nothing") result = run_distributed(Q1, spec) assert _rows(result.batches) == _rows(run_local(Q1, spec)) +def test_the_driver_keeps_its_scratch_out_of_the_results(spec: SessionSpec) -> None: + """Encoded plans and task envelopes go in a subdirectory. + + Not tidiness: it is what makes "does this directory hold stage output?" + answerable by looking, which is what + :func:`~dfx_engine.driver.require_empty_shuffle` does. + """ + run_distributed(Q1, spec) + + shuffle = pathlib.Path(spec.shuffle_dir) + beside_the_results = sorted( + path.name for path in shuffle.iterdir() if path.is_file() + ) + assert beside_the_results == sorted( + pathlib.Path( + _internal.partition_path(spec.shuffle_dir, _internal.stage_id(), partition) + ).name + for partition in range(4) + ) + + scratch = sorted(path.name for path in (shuffle / "tasks").iterdir()) + assert "stage-1.plan" in scratch + assert "task-1-0.json" in scratch + + 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 = { diff --git a/examples/distributed/run_tpch.py b/examples/distributed/run_tpch.py index 08b00014d..dae88db18 100644 --- a/examples/distributed/run_tpch.py +++ b/examples/distributed/run_tpch.py @@ -115,11 +115,15 @@ def compare(table: pa.Table, reference: pa.Table) -> None: message = f"{name}: {len(got)} rows distributed, {len(want)} local" raise ValueError(message) for lhs, rhs in zip(got, want, strict=True): - close = ( - abs(lhs - rhs) <= 1e-6 * max(1.0, abs(rhs)) - if isinstance(lhs, float) - else lhs == rhs - ) + if lhs is None or rhs is None: + # Checked before the float branch, which would raise + # `TypeError` on `None - None` and report a null-handling + # difference between the two runs as a crash in the checker. + close = lhs is None and rhs is None + elif isinstance(lhs, float): + close = abs(lhs - rhs) <= 1e-6 * max(1.0, abs(rhs)) + else: + close = lhs == rhs if not close: message = f"{name}: {lhs!r} distributed, {rhs!r} local" raise ValueError(message)