You are viewing a plain text version of this content. The canonical link for it is here.
Posted to commits@datafusion.apache.org by ti...@apache.org on 2026/09/09 18:23:59 UTC
(datafusion-python) 02/06: Add dfx_udfs: the library that cannot be installed as a bundle
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 <ti...@gmail.com>
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) <no...@anthropic.com>
---
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,
+ )
+
+
+@pytest.fixture
+def revenue_plan(lineitem: pathlib.Path) -> bytes:
+ ctx, _ = _session(lineitem)
+ return (
+ ctx.sql(
+ "select dfx_net_revenue(l_extendedprice, l_discount, l_tax) from lineitem"
+ )
+ .execution_plan()
+ .to_bytes(ctx)
+ )
+
+
+def test_a_worker_with_only_the_codec_can_run_the_plan(
+ revenue_plan: bytes, lineitem: pathlib.Path, tmp_path: pathlib.Path
+) -> None:
+ """Installing the codec is an alternative to registering the functions.
+
+ A separate process that has never registered `dfx_net_revenue` rebuilds it
+ from the name in the plan.
+ """
+ result = _run_worker(tmp_path, revenue_plan, "codec", lineitem)
+
+ assert result.returncode == 0, result.stderr
+ assert "total=530.0" in result.stdout
+ # The codec, not a registry hit, is what answered.
+ assert "decoded=1" in result.stdout
+
+
+def test_a_worker_with_only_the_registrations_can_run_the_plan(
+ revenue_plan: bytes, lineitem: pathlib.Path, tmp_path: pathlib.Path
+) -> None:
+ """And registering the functions is an alternative to the codec."""
+ result = _run_worker(tmp_path, revenue_plan, "registry", lineitem)
+
+ assert result.returncode == 0, result.stderr
+ assert "total=530.0" in result.stdout
+ assert "decoded=0" in result.stdout
+
+
+def test_a_worker_with_neither_names_the_function_it_cannot_find(
+ revenue_plan: bytes, lineitem: pathlib.Path, tmp_path: pathlib.Path
+) -> None:
+ """The failure is legible, and it arrives at decode rather than at execute.
+
+ This is the whole cost of a function library that ships without a codec
+ and without documenting what a worker must register.
+ """
+ result = _run_worker(tmp_path, revenue_plan, "neither", lineitem)
+
+ assert result.returncode != 0
+ assert "dfx_net_revenue" in result.stderr
+
+
+def test_the_codec_declines_names_it_does_not_own(lineitem: pathlib.Path) -> None:
+ """A name-only payload reaches every codec, so declining matters.
+
+ With no bytes there is no codec id to route on. A codec that answered for
+ any name it was handed would hijack another library's functions.
+ """
+ ctx, observations = _session(lineitem, with_codecs=True)
+ # `abs` is a built-in, so the plan references a name this library does not
+ # own; decoding offers it around.
+ blob = (
+ ctx.sql("select abs(l_discount) from lineitem").execution_plan().to_bytes(ctx)
+ )
+ ExecutionPlan.from_bytes(ctx, blob)
+
+ assert observations.decode_calls() == 0
+
+
+def test_the_codec_ids_are_pinned() -> None:
+ """Renaming the exporting class must not invalidate written plans."""
+ observations = CodecObservations()
+
+ assert observations.logical_codec().__datafusion_codec_id__ == "dfx_udfs.logical.v1"
+ assert (
+ observations.physical_codec().__datafusion_codec_id__ == "dfx_udfs.physical.v1"
+ )
+
+
+def test_this_library_cannot_be_installed_as_a_bundle() -> None:
+ """The mixed-workflow case, asserted rather than described.
+
+ `SessionExtensionComponents` carries codec fields only, so a function
+ library has nowhere to put its functions and this one does not pretend
+ otherwise. `with_extensions` rejects it by name.
+ """
+ observations = CodecObservations()
+
+ assert not hasattr(observations, "__datafusion_session_components__")
+ assert not hasattr(observations, "__datafusion_session_planner__")
+
+ with pytest.raises(TypeError, match=re.escape("__datafusion_session_components__")):
+ SessionContext().with_extensions(observations)
diff --git a/examples/distributed/udf-library/python/tests/conftest.py b/examples/distributed/udf-library/python/tests/conftest.py
new file mode 100644
index 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)
+
+
+@pytest.fixture(autouse=True)
+def fail_on_log_warnings() -> Generator[None, Any, None]:
+ handler = _FailOnWarning()
+ logging.root.addHandler(handler)
+ yield
+ logging.root.removeHandler(handler)
+
+
+@pytest.fixture
+def lineitem(tmp_path: pathlib.Path) -> pathlib.Path:
+ """A TPC-H-shaped slice, small enough to check the arithmetic by hand.
+
+ Two files, because the aggregate has to be correct when DataFusion splits
+ it into a partial pass per partition and merges the results.
+ """
+ directory = tmp_path / "lineitem"
+ directory.mkdir()
+ pq.write_table(
+ pa.table(
+ {
+ "l_extendedprice": [100.0, 200.0],
+ "l_discount": [0.0, 0.5],
+ "l_tax": [0.0, 0.0],
+ "l_quantity": [1.0, 3.0],
+ }
+ ),
+ directory / "part-0.parquet",
+ )
+ pq.write_table(
+ pa.table(
+ {
+ "l_extendedprice": [400.0],
+ "l_discount": [0.25],
+ "l_tax": [0.1],
+ "l_quantity": [4.0],
+ }
+ ),
+ directory / "part-1.parquet",
+ )
+ return directory
diff --git a/examples/distributed/udf-library/src/codec.rs b/examples/distributed/udf-library/src/codec.rs
new file mode 100644
index 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: commits-unsubscribe@datafusion.apache.org
For additional commands, e-mail: commits-help@datafusion.apache.org