This is an automated email from the ASF dual-hosted git repository. timsaucer pushed a commit to branch feat/distributed-extensions-example in repository https://gitbox.apache.org/repos/asf/datafusion-python.git
commit 6571d893e179fd8b32c740a2f4b08a3e80296bf9 Author: Tim Saucer <[email protected]> AuthorDate: Wed Sep 9 13:00:51 2026 -0400 Add dfx_engine: a toy distributed engine in both halves 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<dyn TableProvider>`. 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) <[email protected]> --- Cargo.lock | 18 ++ Cargo.toml | 1 + examples/distributed/engine-library/Cargo.toml | 51 +++++ examples/distributed/engine-library/build.rs | 20 ++ examples/distributed/engine-library/pyproject.toml | 34 +++ .../engine-library/python/dfx_engine/__init__.py | 40 ++++ .../engine-library/python/dfx_engine/driver.py | 206 ++++++++++++++++++ .../engine-library/python/dfx_engine/session.py | 170 +++++++++++++++ .../engine-library/python/dfx_engine/worker.py | 113 ++++++++++ examples/distributed/engine-library/src/codec.rs | 136 ++++++++++++ examples/distributed/engine-library/src/config.rs | 113 ++++++++++ .../distributed/engine-library/src/extension.rs | 200 +++++++++++++++++ examples/distributed/engine-library/src/lib.rs | 70 ++++++ .../engine-library/src/local_session.rs | 161 ++++++++++++++ examples/distributed/engine-library/src/planner.rs | 170 +++++++++++++++ examples/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 examples/distributed/storage-library/src/codec.rs | 110 +++++++++- .../distributed/storage-library/src/extension.rs | 49 ++++- .../storage-library/src/table_provider.rs | 5 + 22 files changed, 1897 insertions(+), 6 deletions(-) diff --git a/Cargo.lock b/Cargo.lock index 6ae37a13..d946352c 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 fb7db274..14f34e59 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 00000000..c26d326d --- /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 00000000..4878d8b0 --- /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 00000000..ca9de981 --- /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 00000000..d98cbf1d --- /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 00000000..baa145f1 --- /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 00000000..14390d01 --- /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"] + + [email protected](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 00000000..0250bf10 --- /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 <envelope.json>``. 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 <envelope.json>\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 00000000..f4188ab7 --- /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<CodecCounters>, +} + +impl DfxEnginePhysicalCodec { + pub(crate) fn new(counters: Arc<CodecCounters>) -> 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<dyn ExecutionPlan>, + buf: &mut Vec<u8>, + 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::<ShuffleStageExec>() 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<dyn ExecutionPlan>], + ctx: &TaskContext, + proto_converter: &dyn PhysicalProtoConverterExtension, + ) -> Result<Arc<dyn ExecutionPlan>> { + 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 00000000..05199905 --- /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<Bound<'py, PyCapsule>> { + 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<dyn ExtensionOptions> { + Box::new(self.clone()) + } + + fn set(&mut self, key: &str, value: &str) -> datafusion_common::Result<()> { + ConfigField::set(self, key, value) + } + + fn entries(&self) -> Vec<ConfigEntry> { + 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<V: 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 00000000..a53541d0 --- /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<Bound<'py, PyAny>>, + ) -> PyResult<Bound<'py, PyCapsule>> { + 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<PlannerObservations>, + counters: Arc<CodecCounters>, +} + +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<Bound<'py, PyAny>> { + // 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<dyn PhysicalExtensionCodec + Send> = + 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<Bound<'py, PyCapsule>> { + let fallback = ffi_query_planner_from_pycapsule(&fallback, Some(&ctx))?; + let planner: Arc<dyn QueryPlanner + Send + Sync> = 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 00000000..2dd3b7e1 --- /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::<BundledPhysicalCodec>()?; + m.add_class::<DfxEngineConfig>()?; + m.add_class::<DfxEngineExtension>()?; + 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 00000000..3660c5c8 --- /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<Arc<dyn PhysicalOptimizerRule + Send + Sync>>, +} + +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<dyn PhysicalOptimizerRule + Send + Sync>] { + &self.rules + } + + fn session_id(&self) -> &str { + self.inner.session_id() + } + + fn config(&self) -> &SessionConfig { + self.inner.config() + } + + fn catalog_list(&self) -> Arc<dyn CatalogProviderList> { + self.inner.catalog_list() + } + + fn config_options(&self) -> &ConfigOptions { + self.inner.config_options() + } + + fn optimize(&self, plan: &LogicalPlan) -> Result<LogicalPlan> { + self.inner.optimize(plan) + } + + async fn create_physical_plan( + &self, + logical_plan: &LogicalPlan, + ) -> Result<Arc<dyn ExecutionPlan>> { + self.inner.create_physical_plan(logical_plan).await + } + + fn create_physical_expr( + &self, + expr: Expr, + df_schema: &DFSchema, + ) -> Result<Arc<dyn PhysicalExpr>> { + self.inner.create_physical_expr(expr, df_schema) + } + + fn scalar_functions(&self) -> &HashMap<String, Arc<ScalarUDF>> { + self.inner.scalar_functions() + } + + fn higher_order_functions(&self) -> &HashMap<String, Arc<HigherOrderUDF>> { + self.inner.higher_order_functions() + } + + fn aggregate_functions(&self) -> &HashMap<String, Arc<AggregateUDF>> { + self.inner.aggregate_functions() + } + + fn window_functions(&self) -> &HashMap<String, Arc<WindowUDF>> { + self.inner.window_functions() + } + + fn extension_type_registry(&self) -> &ExtensionTypeRegistryRef { + self.inner.extension_type_registry() + } + + fn runtime_env(&self) -> &Arc<RuntimeEnv> { + 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::<ForeignSession>()` 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<TaskContext> { + 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 00000000..4c308f57 --- /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<String> { + 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<dyn ExecutionPlan>, + shuffle_dir: &str, +) -> Result<(Arc<dyn ExecutionPlan>, bool)> { + if let Some(aggregate) = plan.downcast_ref::<AggregateExec>() + && 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<PlannerObservations>, + /// 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<Arc<dyn QueryPlanner + Send + Sync>>, +} + +#[async_trait] +impl QueryPlanner for DistributedQueryPlanner { + async fn create_physical_plan( + &self, + logical_plan: &LogicalPlan, + session: &dyn Session, + ) -> Result<Arc<dyn ExecutionPlan>> { + 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 00000000..16bf43c9 --- /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<dyn ExecutionPlan>, + properties: Arc<PlanProperties>, +} + +impl ShuffleStageExec { + pub(crate) fn new(stage_id: u32, shuffle_dir: String, input: Arc<dyn ExecutionPlan>) -> 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<TaskContext>, + ) -> Result<SendableRecordBatchStream> { + 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<SendableRecordBatchStream> { + 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::<arrow::error::Result<Vec<_>>>() + .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<PlanProperties> { + &self.properties + } + + fn children(&self) -> Vec<&Arc<dyn ExecutionPlan>> { + vec![&self.input] + } + + fn apply_expressions( + &self, + _f: &mut dyn FnMut(&Arc<dyn PhysicalExpr>) -> Result<TreeNodeRecursion>, + ) -> Result<TreeNodeRecursion> { + // Owns no expressions of its own; the child holds them all. + Ok(TreeNodeRecursion::Continue) + } + + fn with_new_children( + self: Arc<Self>, + mut children: Vec<Arc<dyn ExecutionPlan>>, + ) -> Result<Arc<dyn ExecutionPlan>> { + 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<TaskContext>, + ) -> Result<SendableRecordBatchStream> { + 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 00000000..9e29db43 Binary files /dev/null and b/examples/distributed/engine-library/stage-1-part-0.arrow differ 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 00000000..84c8ac2c Binary files /dev/null and b/examples/distributed/engine-library/stage-1-part-1.arrow differ 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 00000000..e5c81332 Binary files /dev/null and b/examples/distributed/engine-library/stage-1-part-2.arrow differ diff --git a/examples/distributed/storage-library/src/codec.rs b/examples/distributed/storage-library/src/codec.rs index 9868e001..56c0bc74 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<dyn TableProvider>`. +/// 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<CodecCounters>, +} + +impl DfxStorageLogicalCodec { + pub(crate) fn new(counters: Arc<CodecCounters>) -> 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<datafusion::logical_expr::Extension> { + // 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<u8>, + ) -> Result<()> { + self.inner.try_encode(node, buf) + } + + fn try_encode_table_provider( + &self, + table_ref: &TableReference, + node: Arc<dyn TableProvider>, + buf: &mut Vec<u8>, + ) -> Result<()> { + let Some(table) = node.downcast_ref::<PartitionedParquetTable>() 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<Arc<dyn TableProvider>> { + 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 b42944b3..657d2808 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<Bound<'py, PyAny>>, + ) -> PyResult<Bound<'py, PyCapsule>> { + 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<dyn PhysicalExtensionCodec + Send> = 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<dyn LogicalExtensionCodec> = + 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 4b980c03..3c33c41d 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<FileSlice>, 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), }) --------------------------------------------------------------------- To unsubscribe, e-mail: [email protected] For additional commands, e-mail: [email protected]
