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 a63583ce31f41f9fbe8b7c8a09a0577ed6f8dbb9 Author: Tim Saucer <[email protected]> AuthorDate: Wed Sep 9 12:41:13 2026 -0400 Add dfx_udfs: the library that cannot be installed as a bundle 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) <[email protected]> --- Cargo.lock | 18 ++ Cargo.toml | 1 + examples/distributed/udf-library/Cargo.toml | 51 ++++ examples/distributed/udf-library/build.rs | 20 ++ examples/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 +++++++++++++++++ examples/distributed/udf-library/src/functions.rs | 288 ++++++++++++++++++++ examples/distributed/udf-library/src/lib.rs | 46 ++++ examples/distributed/udf-library/src/python.rs | 231 ++++++++++++++++ 11 files changed, 1304 insertions(+) diff --git a/Cargo.lock b/Cargo.lock index 13ac8b64..6ae37a13 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 929c3bba..fb7db274 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 00000000..dac3eeff --- /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 00000000..4878d8b0 --- /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 00000000..8e87abf4 --- /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 00000000..54044578 --- /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, + ) + + [email protected] +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 00000000..2d30c435 --- /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) + + [email protected](autouse=True) +def fail_on_log_warnings() -> Generator[None, Any, None]: + handler = _FailOnWarning() + logging.root.addHandler(handler) + yield + logging.root.removeHandler(handler) + + [email protected] +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 00000000..c19df4db --- /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<CodecCounters>, +} + +impl DfxUdfsLogicalCodec { + pub(crate) fn new(counters: Arc<CodecCounters>) -> 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<datafusion::logical_expr::Extension> { + 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_decode_table_provider( + &self, + buf: &[u8], + table_ref: &datafusion::common::TableReference, + schema: arrow_schema::SchemaRef, + ctx: &datafusion::execution::TaskContext, + ) -> Result<Arc<dyn datafusion::datasource::TableProvider>> { + self.inner + .try_decode_table_provider(buf, table_ref, schema, ctx) + } + + fn try_encode_table_provider( + &self, + table_ref: &datafusion::common::TableReference, + node: Arc<dyn datafusion::datasource::TableProvider>, + buf: &mut Vec<u8>, + ) -> 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<u8>) -> Result<()> { + Ok(()) + } + + fn try_decode_udf(&self, name: &str, buf: &[u8]) -> Result<Arc<ScalarUDF>> { + decode_by_name!(self, name, buf, scalar_by_name, "scalar function") + } + + fn try_encode_udaf(&self, _node: &AggregateUDF, _buf: &mut Vec<u8>) -> Result<()> { + Ok(()) + } + + fn try_decode_udaf(&self, name: &str, buf: &[u8]) -> Result<Arc<AggregateUDF>> { + decode_by_name!(self, name, buf, aggregate_by_name, "aggregate function") + } + + fn try_encode_udwf(&self, _node: &WindowUDF, _buf: &mut Vec<u8>) -> Result<()> { + Ok(()) + } + + fn try_decode_udwf(&self, name: &str, buf: &[u8]) -> Result<Arc<WindowUDF>> { + 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<CodecCounters>, +} + +impl DfxUdfsPhysicalCodec { + pub(crate) fn new(counters: Arc<CodecCounters>) -> 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<dyn datafusion::physical_plan::ExecutionPlan>], + ctx: &datafusion::execution::TaskContext, + proto_converter: &dyn datafusion_proto::physical_plan::PhysicalProtoConverterExtension, + ) -> Result<Arc<dyn datafusion::physical_plan::ExecutionPlan>> { + // This library owns no execution plan nodes, only functions. + self.inner.try_decode(buf, inputs, ctx, proto_converter) + } + + fn try_encode( + &self, + node: Arc<dyn datafusion::physical_plan::ExecutionPlan>, + buf: &mut Vec<u8>, + 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<u8>) -> Result<()> { + Ok(()) + } + + fn try_decode_udf(&self, name: &str, buf: &[u8]) -> Result<Arc<ScalarUDF>> { + decode_by_name!(self, name, buf, scalar_by_name, "scalar function") + } + + fn try_encode_udaf(&self, _node: &AggregateUDF, _buf: &mut Vec<u8>) -> Result<()> { + Ok(()) + } + + fn try_decode_udaf(&self, name: &str, buf: &[u8]) -> Result<Arc<AggregateUDF>> { + decode_by_name!(self, name, buf, aggregate_by_name, "aggregate function") + } + + fn try_encode_udwf(&self, _node: &WindowUDF, _buf: &mut Vec<u8>) -> Result<()> { + Ok(()) + } + + fn try_decode_udwf(&self, name: &str, buf: &[u8]) -> Result<Arc<WindowUDF>> { + 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 00000000..ec7e95eb --- /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<DataType> { + Ok(DataType::Float64) + } + + fn invoke_with_args(&self, args: ScalarFunctionArgs) -> Result<ColumnarValue> { + 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::<Float64Type>(); + let discount = discount.as_primitive::<Float64Type>(); + let tax = tax.as_primitive::<Float64Type>(); + + 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<DataType> { + Ok(DataType::Float64) + } + + fn accumulator(&self, _args: AccumulatorArgs) -> Result<Box<dyn Accumulator>> { + Ok(Box::new(WeightedAvgAccumulator::default())) + } + + /// The two partial sums, in the order [`WeightedAvgAccumulator::state`] + /// returns them. + fn state_fields(&self, args: StateFieldsArgs) -> Result<Vec<FieldRef>> { + 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::<Float64Type>(); + let weight = weight.as_primitive::<Float64Type>(); + 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::<Float64Type>(); + let weight_sum = weight_sum.as_primitive::<Float64Type>(); + 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<Vec<ScalarValue>> { + Ok(vec![ + ScalarValue::Float64(Some(self.weighted_sum)), + ScalarValue::Float64(Some(self.weight_sum)), + ]) + } + + fn evaluate(&mut self) -> Result<ScalarValue> { + // 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<WindowUDF>, +} + +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<H: std::hash::Hasher>(&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<Box<dyn PartitionEvaluator>> { + self.inner.inner().partition_evaluator(args) + } + + fn field(&self, field_args: WindowUDFFieldArgs) -> Result<FieldRef> { + self.inner.inner().field(field_args) + } + + fn coerce_types(&self, arg_types: &[DataType]) -> Result<Vec<DataType>> { + 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<Arc<ScalarUDF>> { + (name == NET_REVENUE).then(|| Arc::new(ScalarUDF::from(NetRevenue::default()))) +} + +pub(crate) fn aggregate_by_name(name: &str) -> Option<Arc<AggregateUDF>> { + (name == WEIGHTED_AVG).then(|| Arc::new(AggregateUDF::from(WeightedAvg::default()))) +} + +pub(crate) fn window_by_name(name: &str) -> Option<Arc<WindowUDF>> { + (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 00000000..95fa87a3 --- /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::<PyCodecObservations>()?; + m.add_class::<PyLogicalCodec>()?; + m.add_class::<PyNetRevenue>()?; + m.add_class::<PyPhysicalCodec>()?; + m.add_class::<PyRevenueRank>()?; + m.add_class::<PyWeightedAvg>()?; + Ok(()) +} diff --git a/examples/distributed/udf-library/src/python.rs b/examples/distributed/udf-library/src/python.rs new file mode 100644 index 00000000..397aa672 --- /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<Bound<'py, PyCapsule>> { + 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<Bound<'py, PyCapsule>> { + 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<Bound<'py, PyCapsule>> { + 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<CodecCounters>, +} + +#[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<CodecCounters>, +} + +#[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<Bound<'py, PyCapsule>> { + let provider = ffi_task_context_provider_from_pycapsule(&session)?; + let runtime = get_tokio_runtime().handle().clone(); + let codec: Arc<dyn LogicalExtensionCodec> = + 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<CodecCounters>, +} + +#[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<Bound<'py, PyCapsule>> { + let provider = ffi_task_context_provider_from_pycapsule(&session)?; + let runtime = get_tokio_runtime().handle().clone(); + let codec: Arc<dyn PhysicalExtensionCodec + Send> = + Arc::new(DfxUdfsPhysicalCodec::new(Arc::clone(&self.counters))); + let ffi = FFI_PhysicalExtensionCodec::new(codec, Some(runtime), provider); + create_physical_extension_capsule(py, &ffi) + } +} --------------------------------------------------------------------- To unsubscribe, e-mail: [email protected] For additional commands, e-mail: [email protected]
