From 3a2476861570a25b7cee41ae79a40066b88c18e8 Mon Sep 17 00:00:00 2001 From: Tim Saucer Date: Tue, 15 Sep 2026 10:01:31 -0400 Subject: [PATCH 1/3] feat: let extension bundles declare physical optimizer rules MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit `SessionExtensionComponents.physical_optimizer_rules` completes the group of components whose capsule getter takes no argument, so a library shipping a rule alongside anything else no longer asks the caller for a separate `add_physical_optimizer_rule` call. Rules are the one kind with no collision rule: they accumulate rather than replace, so two bundles contributing one each is the normal case and there is nothing to refuse. Installing them splits across the two new private primitives `_resolve_extension_physical_optimizer_rules` and `_install_extension_physical_optimizer_rules`, keeping the commit step infallible: the capsules are imported during resolution, so a rule that fails to import cannot leave the session with a planner already bound. All the rules in one call go on in a single `SessionState` rebuild. `add_physical_optimizer_rule` rebuilds per call, which for a bundle with several would clone the whole state that many times and leave the earlier ones installed if a later one failed. The session id is carried across the rebuild for the same reason that method carries it. `PhysicalOptimizerRuleExportable` moves from `datafusion.context` to `datafusion.extensions`, alongside the rest of the protocol family, and is now exported from the package root. It stays importable from `datafusion.context`. `MyRuleExtension` in `datafusion-ffi-example` declares two rules, which is what makes accumulation observable — each carries its own counter and both fire. Co-Authored-By: Claude Opus 5 (1M context) --- crates/core/src/context.rs | 65 ++++++++++++++++ docs/source/extension-guide/bundles.md | 4 + .../extension-guide/other-components.md | 16 ++++ .../python/tests/_test_session_extension.py | 74 ++++++++++++++++++- .../datafusion-ffi-example/src/extension.rs | 59 +++++++++++++++ examples/datafusion-ffi-example/src/lib.rs | 3 +- .../src/physical_optimizer.rs | 4 +- python/datafusion/__init__.py | 2 + python/datafusion/context.py | 34 ++++----- python/datafusion/extensions.py | 49 ++++++++++++ python/tests/test_context.py | 53 ++++++++++++- python/tests/test_wrapper_coverage.py | 4 + 12 files changed, 343 insertions(+), 24 deletions(-) diff --git a/crates/core/src/context.rs b/crates/core/src/context.rs index fe75668fb..cc8d59079 100644 --- a/crates/core/src/context.rs +++ b/crates/core/src/context.rs @@ -44,6 +44,7 @@ use datafusion::execution::options::{ArrowReadOptions, ReadOptions}; use datafusion::execution::runtime_env::RuntimeEnvBuilder; use datafusion::execution::session_state::SessionStateBuilder; use datafusion::execution::{FunctionRegistry, TaskContextProvider}; +use datafusion::physical_optimizer::PhysicalOptimizerRule; use datafusion::physical_plan::ExecutionPlanProperties; use datafusion::prelude::{ AvroReadOptions, CsvReadOptions, DataFrame, JsonReadOptions, ParquetReadOptions, @@ -1762,6 +1763,70 @@ impl PySessionContext { slf.borrow().set_session_query_planner(planner); Ok(()) } + + /// Import the physical optimizer rules a `with_extensions` call declared. + /// + /// The fallible half of installing them, run while the call can still fail + /// harmlessly. Every capsule is imported here so that + /// [`Self::_install_extension_physical_optimizer_rules`] has nothing left + /// that can raise — a rule that failed to import after the planner was + /// bound would leave the session half-installed, and there is no derived + /// context to roll back to. + /// + /// **Writes nothing.** + pub fn _resolve_extension_physical_optimizer_rules( + &self, + rules: Vec>, + ) -> PyDataFusionResult { + let rules = rules + .iter() + .map(physical_optimizer_rule_from_pycapsule) + .collect::, _>>()?; + Ok(PyPhysicalOptimizerRules { rules }) + } + + /// Commit the physical optimizer rules for a `with_extensions` call. + /// + /// Rules accumulate rather than replace, so unlike a planner there is no + /// composition order to get right and no collision to refuse. + /// + /// All of them go on in **one** `SessionState` rebuild. + /// [`Self::add_physical_optimizer_rule`] rebuilds per call, which for a + /// bundle contributing several would clone the whole state that many times + /// and, worse, leave the earlier rules installed if a later one failed. + /// Nothing here can fail: the capsules were imported by + /// [`Self::_resolve_extension_physical_optimizer_rules`]. + pub fn _install_extension_physical_optimizer_rules( + &self, + resolved: PyRef<'_, PyPhysicalOptimizerRules>, + ) { + if resolved.rules.is_empty() { + return; + } + let state_ref = self.ctx.state_ref(); + let mut guard = state_ref.write(); + // The session id has to be carried over for the same reason + // `add_physical_optimizer_rule` carries it: the builder mints a fresh + // one, and losing it leaves `session_id()` disagreeing with every + // `TaskContext` the session has already handed out. + let mut builder = SessionStateBuilder::new_from_existing(guard.clone()) + .with_session_id(guard.session_id().to_string()); + for rule in resolved.rules.iter().cloned() { + builder = builder.with_physical_optimizer_rule(rule); + } + *guard = builder.build(); + } +} + +/// Physical optimizer rules imported for a `with_extensions` call. +/// +/// Opaque to Python, and deliberately not added to the module: it exists only +/// to carry imported rules from the resolve step to the commit step, so the +/// import can fail before anything is written. `with_extensions` is its only +/// producer and its only consumer. +#[pyclass(name = "PhysicalOptimizerRules", module = "datafusion._internal")] +pub struct PyPhysicalOptimizerRules { + rules: Vec>, } impl PySessionContext { diff --git a/docs/source/extension-guide/bundles.md b/docs/source/extension-guide/bundles.md index c3f04d1e1..57ee1da0b 100644 --- a/docs/source/extension-guide/bundles.md +++ b/docs/source/extension-guide/bundles.md @@ -324,6 +324,10 @@ Two cases this does *not* catch: DataFusion function, and replacing one by name is a supported thing to do — `enable_spark_functions` works that way. +Physical optimizer rules are exempt from all of this: they accumulate rather +than replace, so two bundles contributing one each is the normal case and there +is nothing to refuse. See {doc}`other-components`. + Your caller cannot rename your function, so stay out of the way: prefix the names with something tied to your library. diff --git a/docs/source/extension-guide/other-components.md b/docs/source/extension-guide/other-components.md index 919b8718e..83b770d56 100644 --- a/docs/source/extension-guide/other-components.md +++ b/docs/source/extension-guide/other-components.md @@ -47,6 +47,22 @@ fn __datafusion_physical_optimizer_rule__<'py>( } ``` +If your library ships a rule alongside anything else, declare it on your bundle +as `physical_optimizer_rules` rather than asking the caller for a separate +`add_physical_optimizer_rule` call: + +```python +return SessionExtensionComponents(physical_optimizer_rules=(MyRule(),)) +``` + +Rules are the one kind of component with **no collision rule at all**: they +accumulate, so two libraries may each contribute one and neither has to know +about the other. Every rule in a call installs in a single `SessionState` +rebuild, where `add_physical_optimizer_rule` rebuilds once per call — which for +a bundle contributing several would clone the whole state that many times, and +would leave the earlier ones installed if a later one failed. See +{ref}`extension_bundles_transaction`. + ## Typed configuration **`__datafusion_extension_options__`** contributes typed configuration entries diff --git a/examples/datafusion-ffi-example/python/tests/_test_session_extension.py b/examples/datafusion-ffi-example/python/tests/_test_session_extension.py index fadc1fd70..58a205497 100644 --- a/examples/datafusion-ffi-example/python/tests/_test_session_extension.py +++ b/examples/datafusion-ffi-example/python/tests/_test_session_extension.py @@ -22,7 +22,7 @@ import pyarrow as pa import pytest from datafusion import SessionContext, SessionExtensionComponents -from datafusion_ffi_example import MyFunctionExtension +from datafusion_ffi_example import MyFunctionExtension, MyRuleExtension def _session(): @@ -112,6 +112,78 @@ def __datafusion_session_planner__(self, ctx, fallback) -> None: ctx.udf("my_custom_is_null") +def _query(ctx): + batch = pa.RecordBatch.from_arrays([pa.array([1, 2, 3])], names=["a"]) + ctx.register_record_batches("t", [[batch]]) + return ctx.sql("SELECT a FROM t").collect() + + +def test_declared_rules_all_fire(): + """Rules accumulate, so both of a bundle's two rules run. + + Nothing about installing the second displaces the first, which is what + makes rules different from a planner and why there is no collision to + refuse. + """ + extension = MyRuleExtension() + ctx = SessionContext().with_extensions(extension) + + assert _query(ctx)[0].column(0).to_pylist() == [1, 2, 3] + assert extension.first_calls() > 0 + assert extension.second_calls() > 0 + + +def test_rules_install_without_changing_the_session_id(): + """Installing rules rebuilds ``SessionState``; the id has to survive it. + + A fresh id would leave ``session_id()`` disagreeing with every + ``TaskContext`` the session already handed out, which is exactly what a + codec's decode callbacks resolve against. + """ + ctx = SessionContext() + before = ctx.session_id() + result = ctx.with_extensions(MyRuleExtension()) + + assert result.session_id() == before + assert ctx.session_id() == before + + +def test_rules_and_functions_install_together(): + """Two bundles, one contributing functions and one rules, in one call.""" + rules = MyRuleExtension() + ctx = SessionContext().with_extensions(MyFunctionExtension(), rules) + batch = pa.RecordBatch.from_arrays([pa.array([1, 2, None])], names=["a"]) + ctx.register_record_batches("t", [[batch]]) + + result = ctx.sql("SELECT my_custom_is_null(a) FROM t").collect() + + assert result[0].column(0).to_pylist() == [False, False, True] + assert rules.first_calls() > 0 + + +def test_a_failure_leaves_no_rule_installed(): + """The transaction covers rules, which write through a state rebuild. + + A rule reaching the session before the failing hook would be invisible to + ``session_id()`` and to the function registry, so this asserts on the + counter instead: an installed rule fires on the next query. + """ + rules = MyRuleExtension() + ctx = SessionContext() + + class BoomPlanner: + def __datafusion_session_planner__(self, ctx, fallback) -> None: + msg = "boom" + raise RuntimeError(msg) + + with pytest.raises(RuntimeError, match="boom"): + ctx.with_extensions(rules, BoomPlanner()) + + _query(ctx) + assert rules.first_calls() == 0 + assert rules.second_calls() == 0 + + def test_the_hook_returns_the_components_type(): """The bundle builds a real dataclass, not a duck-typed stand-in. diff --git a/examples/datafusion-ffi-example/src/extension.rs b/examples/datafusion-ffi-example/src/extension.rs index fe94a667f..4f6d114cc 100644 --- a/examples/datafusion-ffi-example/src/extension.rs +++ b/examples/datafusion-ffi-example/src/extension.rs @@ -19,6 +19,7 @@ use pyo3::types::{PyAnyMethods, PyDict, PyDictMethods}; use pyo3::{Bound, Py, PyAny, PyResult, Python, pyclass, pymethods}; use crate::aggregate_udf::MySumUDF; +use crate::physical_optimizer::MyPhysicalOptimizerRule; use crate::scalar_udf::IsNullUDF; use crate::window_udf::MyRankUDF; @@ -62,3 +63,61 @@ impl MyFunctionExtension { components.call((), Some(&kwargs)) } } + +/// A bundle contributing two physical optimizer rules. +/// +/// Two, because that is what makes accumulation observable: rules never +/// collide the way function names do, so both of these install and both fire. +/// Each carries its own counter, which is how a test tells them apart. +#[pyclass( + from_py_object, + name = "MyRuleExtension", + module = "datafusion_ffi_example", + subclass +)] +#[derive(Debug, Clone, Default)] +pub(crate) struct MyRuleExtension { + first: MyPhysicalOptimizerRule, + second: MyPhysicalOptimizerRule, +} + +#[pymethods] +impl MyRuleExtension { + #[new] + fn new() -> Self { + Self::default() + } + + /// How many times the first declared rule has run. + fn first_calls(&self) -> usize { + self.first.optimize_calls() + } + + /// How many times the second declared rule has run. + fn second_calls(&self) -> usize { + self.second.optimize_calls() + } + + /// `ctx` is unused: a rule getter takes no argument, so there is nothing + /// session-scoped to bind. + fn __datafusion_session_components__<'py>( + &self, + py: Python<'py>, + ctx: Bound<'py, PyAny>, + ) -> PyResult> { + let _ = ctx; + + let components = py + .import("datafusion")? + .getattr("SessionExtensionComponents")?; + let kwargs = PyDict::new(py); + kwargs.set_item( + "physical_optimizer_rules", + ( + Py::new(py, self.first.clone())?, + Py::new(py, self.second.clone())?, + ), + )?; + components.call((), Some(&kwargs)) + } +} diff --git a/examples/datafusion-ffi-example/src/lib.rs b/examples/datafusion-ffi-example/src/lib.rs index b680c84de..097c8ef87 100644 --- a/examples/datafusion-ffi-example/src/lib.rs +++ b/examples/datafusion-ffi-example/src/lib.rs @@ -20,7 +20,7 @@ use pyo3::prelude::*; use crate::aggregate_udf::MySumUDF; use crate::catalog_provider::{FixedSchemaProvider, MyCatalogProvider, MyCatalogProviderList}; use crate::config::MyConfig; -use crate::extension::MyFunctionExtension; +use crate::extension::{MyFunctionExtension, MyRuleExtension}; use crate::logical_extension_codec::MyLogicalExtensionCodec; use crate::name_only_codec::{NameOnlyFunction, NameOnlyUdfCodec}; use crate::physical_extension_codec::MyPhysicalExtensionCodec; @@ -66,5 +66,6 @@ fn datafusion_ffi_example(m: &Bound<'_, PyModule>) -> PyResult<()> { m.add_class::()?; m.add_class::()?; m.add_class::()?; + m.add_class::()?; Ok(()) } diff --git a/examples/datafusion-ffi-example/src/physical_optimizer.rs b/examples/datafusion-ffi-example/src/physical_optimizer.rs index e17510495..26caf8d65 100644 --- a/examples/datafusion-ffi-example/src/physical_optimizer.rs +++ b/examples/datafusion-ffi-example/src/physical_optimizer.rs @@ -72,11 +72,11 @@ pub(crate) struct MyPhysicalOptimizerRule { #[pymethods] impl MyPhysicalOptimizerRule { #[new] - fn new() -> Self { + pub(crate) fn new() -> Self { Self::default() } - fn optimize_calls(&self) -> usize { + pub(crate) fn optimize_calls(&self) -> usize { self.optimize_calls.load(Ordering::SeqCst) } diff --git a/python/datafusion/__init__.py b/python/datafusion/__init__.py index 1b44f8a73..f5a2b3849 100644 --- a/python/datafusion/__init__.py +++ b/python/datafusion/__init__.py @@ -93,6 +93,7 @@ from .dataframe_formatter import configure_formatter from .expr import Expr, WindowFrame from .extensions import ( + PhysicalOptimizerRuleExportable, QueryPlannerExportable, SessionComponentsExportable, SessionExtensionComponents, @@ -139,6 +140,7 @@ "MetricsSet", "ParquetColumnOptions", "ParquetWriterOptions", + "PhysicalOptimizerRuleExportable", "PhysicalPartitioning", "QueryPlannerExportable", "RecordBatch", diff --git a/python/datafusion/context.py b/python/datafusion/context.py index 75d8e66ab..a29dd6fa2 100644 --- a/python/datafusion/context.py +++ b/python/datafusion/context.py @@ -70,6 +70,7 @@ from datafusion.dataframe import DataFrame from datafusion.expr import sort_list_to_raw_sort_list from datafusion.extensions import ( + PhysicalOptimizerRuleExportable, QueryPlannerExportable, SessionComponentsExportable, SessionExtensionComponents, @@ -151,16 +152,6 @@ class TableProviderExportable(Protocol): def __datafusion_table_provider__(self, session: Any) -> object: ... # noqa: D105 -class PhysicalOptimizerRuleExportable(Protocol): - """Type hint for object that has __datafusion_physical_optimizer_rule__ PyCapsule. - - The method returns a PyCapsule wrapping an ``FFI_PhysicalOptimizerRule``, - typically produced by a separate compiled extension. - """ - - def __datafusion_physical_optimizer_rule__(self) -> object: ... # noqa: D105 - - class _FunctionKind(NamedTuple): """How one kind of declared function is resolved and registered. @@ -245,9 +236,9 @@ def _collect_contributions( context derived from it. Returns: - The logical codecs, the physical codecs, and the declared functions as - ``(position, extension, function)`` triples, keyed by the - :py:data:`_FUNCTION_KINDS` field they arrived in. + The logical codecs, the physical codecs, and the declared functions and + optimizer rules as ``(position, extension, declaration)`` triples, keyed + by the ``SessionExtensionComponents`` field they arrived in. Raises: TypeError: If an argument implements neither hook, or a hook returns @@ -269,6 +260,9 @@ def _collect_contributions( declared: dict[str, list[tuple[int, object, Any]]] = { kind.field: [] for kind in _FUNCTION_KINDS } + # Rules are not a function kind: they accumulate rather than replace, so + # they carry no collision rule and install through their own primitive. + declared["physical_optimizer_rules"] = [] for position, extension in enumerate(extensions): if not isinstance(extension, SessionComponentsExportable): continue @@ -2117,10 +2111,10 @@ def with_extensions( Nothing is written to the session until every hook has returned and every component has been validated, so a hook that raises leaves the - session as it was. Declared functions register after the planner is - bound, and are visible on every handle sharing this session. A hook - that *mutates* the context it is handed — registering a table, say — is - not rolled back, which is why bundle objects must be + session as it was. Declared functions and optimizer rules install after + the planner is bound, and are visible on every handle sharing this + session. A hook that *mutates* the context it is handed — registering a + table, say — is not rolled back, which is why bundle objects must be configuration-only. Shares its session with this context — see :py:class:`SessionContext`. @@ -2214,6 +2208,11 @@ def with_extensions( ) for kind in _FUNCTION_KINDS ] + # Rules accumulate, so there is no name to check and nothing to refuse + # -- only the capsules to import while failing is still free. + resolved_rules = new.ctx._resolve_extension_physical_optimizer_rules( + [rule for _, _, rule in declared["physical_optimizer_rules"]] + ) # Phase two: nest the planners, outermost last. Each hook runs against # `new`, which carries the final chains, so a planner captured here @@ -2254,6 +2253,7 @@ def with_extensions( for register, functions in resolved: for function in functions: register(function) + new.ctx._install_extension_physical_optimizer_rules(resolved_rules) return new def table_provider(self, name: str) -> Table: diff --git a/python/datafusion/extensions.py b/python/datafusion/extensions.py index 1bb413237..403526cc8 100644 --- a/python/datafusion/extensions.py +++ b/python/datafusion/extensions.py @@ -63,6 +63,7 @@ ) __all__ = [ + "PhysicalOptimizerRuleExportable", "QueryPlannerExportable", "SessionComponentsExportable", "SessionExtensionComponents", @@ -70,6 +71,43 @@ ] +class PhysicalOptimizerRuleExportable(Protocol): + """Type hint for object that has a __datafusion_physical_optimizer_rule__ capsule. + + The method returns a PyCapsule wrapping an ``FFI_PhysicalOptimizerRule``, + typically produced by a separate compiled extension. It takes **no + argument**: a rule needs neither a codec nor a task-context provider, so + there is nothing session-scoped to hand it. + + Rules accumulate rather than replace, so several libraries may each + contribute one and none of them has to know about the others. Install one + with :py:meth:`~datafusion.context.SessionContext.add_physical_optimizer_rule`, + or declare it on a bundle as + :py:attr:`SessionExtensionComponents.physical_optimizer_rules`. + + Examples: + The getter is the whole protocol, and a capsule is what it must return + — anything else is refused where it is installed rather than at plan + time: + + >>> from datafusion import SessionContext + >>> ctx = SessionContext() + >>> ctx.add_physical_optimizer_rule(object()) + Traceback (most recent call last): + ... + RuntimeError: "Invalid datafusion_physical_optimizer_rule... + + Real usage. Skipped here (needs a built extension library); run for + real by ``test_ffi_physical_optimizer_rule`` in + ``datafusion-ffi-example``. + + >>> from datafusion_ffi_example import MyPhysicalOptimizerRule # doctest: +SKIP + >>> ctx.add_physical_optimizer_rule(MyPhysicalOptimizerRule()) # doctest: +SKIP + """ + + def __datafusion_physical_optimizer_rule__(self) -> object: ... # noqa: D105 + + class QueryPlannerExportable(Protocol): """Type hint for object that has a __datafusion_query_planner__ PyCapsule. @@ -255,6 +293,17 @@ class SessionExtensionComponents: :py:func:`~datafusion.udwf`. """ + physical_optimizer_rules: tuple[PhysicalOptimizerRuleExportable, ...] = _components( + "optimizer rule" + ) + """Physical optimizer rules to install on the session. + + Objects exposing ``__datafusion_physical_optimizer_rule__``. Unlike + functions these never collide: rules accumulate, so two extensions may each + contribute one without either having to know about the other. All the rules + in one call install together, in declaration order. + """ + def __post_init__(self) -> None: """Normalize each component field, rejecting what cannot become a tuple.""" # A bundle that writes `logical_extension_codecs=codec` instead of diff --git a/python/tests/test_context.py b/python/tests/test_context.py index a0276063a..5740df5c4 100644 --- a/python/tests/test_context.py +++ b/python/tests/test_context.py @@ -1670,6 +1670,52 @@ def test_session_extension_components_rejects_a_single_function(field): SessionExtensionComponents(**{field: _doubler()}) +def test_session_extension_components_rejects_a_single_optimizer_rule(): + """The same for rules, naming what that field holds.""" + with pytest.raises( + TypeError, match=r"must be an iterable of optimizer rule objects" + ): + SessionExtensionComponents(physical_optimizer_rules=object()) + + +def test_with_extensions_rejects_a_rule_that_is_not_a_capsule(ctx): + """A rule that will not import is refused, and nothing is installed. + + Importing the capsules is the only part of installing a rule that can + fail, so it happens during resolution. A failure here has to leave the + session alone even though the extension ahead of it declared a function + that was perfectly good. + """ + + class RuleExtension: + def __datafusion_session_components__(self, ctx): + return SessionExtensionComponents(physical_optimizer_rules=(object(),)) + + with pytest.raises(RuntimeError, match="datafusion_physical_optimizer_rule"): + ctx.with_extensions( + _FunctionExtension(udfs=(_doubler(),)), + RuleExtension(), + ) + + with pytest.raises(KeyError): + ctx.udf("double") + + +def test_with_extensions_declaring_no_rules_leaves_the_session_id(ctx): + """Installing rules rebuilds ``SessionState``; declaring none must not. + + The rebuild mints a fresh session id unless it is carried over, and a + changed id would break every ``TaskContext`` the session has handed out. + Asserted for the empty case too, because that is the one where the rebuild + would be pure cost. + """ + before = ctx.session_id() + result = ctx.with_extensions(_FunctionExtension(udfs=(_doubler(),))) + + assert result.session_id() == before + assert ctx.session_id() == before + + def test_every_component_field_has_an_installer(): """A field added to the components dataclass must be wired into the install. @@ -1680,9 +1726,9 @@ def test_every_component_field_has_an_installer(): Reaching into private names on purpose: the two sides answer different questions. The metadata says which fields are collections to normalize; - ``_FUNCTION_KINDS`` and the codec pair say which of them ``with_extensions`` - knows how to install. Nothing observable from outside can tell you they - have drifted, because the symptom is silence. + ``_FUNCTION_KINDS``, the codec pair, and the rules say which of them + ``with_extensions`` knows how to install. Nothing observable from outside + can tell you they have drifted, because the symptom is silence. """ from datafusion.context import _FUNCTION_KINDS @@ -1695,6 +1741,7 @@ def test_every_component_field_has_an_installer(): assert by_noun == { "codec": {"logical_extension_codecs", "physical_extension_codecs"}, "function": {kind.field for kind in _FUNCTION_KINDS}, + "optimizer rule": {"physical_optimizer_rules"}, } diff --git a/python/tests/test_wrapper_coverage.py b/python/tests/test_wrapper_coverage.py index 6927632b9..903fd4a83 100644 --- a/python/tests/test_wrapper_coverage.py +++ b/python/tests/test_wrapper_coverage.py @@ -39,6 +39,10 @@ "_install_extension_codecs", "_export_query_planner", "_install_extension_planner", + # Physical optimizer rules, split so the capsule import happens while + # the call can still fail without leaving the session half-installed. + "_resolve_extension_physical_optimizer_rules", + "_install_extension_physical_optimizer_rules", } ) From ed68913ef580e139a4157e8e0c374514b23eaab1 Mon Sep 17 00:00:00 2001 From: Tim Saucer Date: Tue, 15 Sep 2026 16:54:08 -0400 Subject: [PATCH 2/3] fix: name the bundle that declared something that is not a rule A rule the importer refuses used to raise from Rust with nothing but the capsule name. A caller who passed four bundles could not tell which one was at fault, and the resolve step is the last place that is known. `_resolve_declared_rules` looks for the getter in Python first, mirroring what `_resolve_declared_functions` already does for a declared function: TypeError A declared optimizer rule must expose __datafusion_physical_optimizer_rule__, got from Scoped to match the sibling rather than to go past it. A getter that is present but returns a non-capsule still falls through to the importer's `RuntimeError`, exactly as a declared function does, and `with_extensions` now documents that case instead of listing only the two errors it raises itself. `MyRuleExtension`'s two rules append to a run log they share, so the order they installed in is observable. The counters cannot show it: each rule has its own, so they say how often a rule ran but not when. `ffi-internals.md` describes the four-step commit order as a rule for the next field added to `SessionExtensionComponents`. `physical_optimizer_rules` is that field, so steps three and four name it rather than leaving the enumeration stale on the commit that invoked it. Why rules never collide is now argued once, in the extension guide. The protocol docstring and the field docstring state it and link there, and the docstring naming the test that runs its skipped example names the whole test. Co-Authored-By: Claude Opus 5 (1M context) --- .../source/contributor-guide/ffi-internals.md | 7 +-- .../python/tests/_test_session_extension.py | 15 ++++++ .../datafusion-ffi-example/src/extension.rs | 22 +++++++-- .../src/physical_optimizer.rs | 28 ++++++++++- python/datafusion/context.py | 47 +++++++++++++++++-- python/datafusion/extensions.py | 15 +++--- python/tests/test_context.py | 47 ++++++++++++++----- 7 files changed, 150 insertions(+), 31 deletions(-) diff --git a/docs/source/contributor-guide/ffi-internals.md b/docs/source/contributor-guide/ffi-internals.md index f25e14091..dfe6d6b17 100644 --- a/docs/source/contributor-guide/ffi-internals.md +++ b/docs/source/contributor-guide/ffi-internals.md @@ -128,9 +128,10 @@ A call therefore splits into a part that may fail and a part that may not: live on that handle rather than on the session, so this step writes nothing even though it can fail on a bad capsule or a duplicate id. 3. **Resolve.** Every declared function is wrapped and every name is checked, - and every `__datafusion_session_planner__` runs against the completed - chains. -4. **Commit.** The planner is bound and the functions are registered. + every declared physical optimizer rule has its capsule imported, and every + `__datafusion_session_planner__` runs against the completed chains. +4. **Commit.** The planner is bound, the functions are registered, and the + optimizer rules are installed in a single `SessionState` rebuild. Only step 4 touches the session, and every step that can fail happens before it. This is a rule for the next field added to diff --git a/examples/datafusion-ffi-example/python/tests/_test_session_extension.py b/examples/datafusion-ffi-example/python/tests/_test_session_extension.py index 58a205497..57ef722da 100644 --- a/examples/datafusion-ffi-example/python/tests/_test_session_extension.py +++ b/examples/datafusion-ffi-example/python/tests/_test_session_extension.py @@ -133,6 +133,21 @@ def test_declared_rules_all_fire(): assert extension.second_calls() > 0 +def test_declared_rules_run_in_declaration_order(): + """The order a bundle lists its rules in is the order they install in. + + Rules rewrite the plan one after another, so the order is part of what a + bundle declares. The counters cannot show it — each rule has its own — so + the two here append to a log they share. + """ + extension = MyRuleExtension() + ctx = SessionContext().with_extensions(extension) + + _query(ctx) + + assert extension.run_order() == [0, 1] + + def test_rules_install_without_changing_the_session_id(): """Installing rules rebuilds ``SessionState``; the id has to survive it. diff --git a/examples/datafusion-ffi-example/src/extension.rs b/examples/datafusion-ffi-example/src/extension.rs index 4f6d114cc..c0033b48c 100644 --- a/examples/datafusion-ffi-example/src/extension.rs +++ b/examples/datafusion-ffi-example/src/extension.rs @@ -15,6 +15,8 @@ // specific language governing permissions and limitations // under the License. +use std::sync::{Arc, Mutex}; + use pyo3::types::{PyAnyMethods, PyDict, PyDictMethods}; use pyo3::{Bound, Py, PyAny, PyResult, Python, pyclass, pymethods}; @@ -68,24 +70,32 @@ impl MyFunctionExtension { /// /// Two, because that is what makes accumulation observable: rules never /// collide the way function names do, so both of these install and both fire. -/// Each carries its own counter, which is how a test tells them apart. +/// Each carries its own counter, which is how a test tells them apart, and +/// both append to one run log, which is how a test sees the order they +/// installed in. #[pyclass( from_py_object, name = "MyRuleExtension", module = "datafusion_ffi_example", subclass )] -#[derive(Debug, Clone, Default)] +#[derive(Debug, Clone)] pub(crate) struct MyRuleExtension { first: MyPhysicalOptimizerRule, second: MyPhysicalOptimizerRule, + run_log: Arc>>, } #[pymethods] impl MyRuleExtension { #[new] fn new() -> Self { - Self::default() + let run_log = Arc::new(Mutex::new(Vec::new())); + Self { + first: MyPhysicalOptimizerRule::with_run_log(0, Arc::clone(&run_log)), + second: MyPhysicalOptimizerRule::with_run_log(1, Arc::clone(&run_log)), + run_log, + } } /// How many times the first declared rule has run. @@ -98,6 +108,12 @@ impl MyRuleExtension { self.second.optimize_calls() } + /// The labels of the two declared rules, in the order they ran: `0` is the + /// first declared and `1` the second. + fn run_order(&self) -> Vec { + self.run_log.lock().expect("run log poisoned").clone() + } + /// `ctx` is unused: a rule getter takes no argument, so there is nothing /// session-scoped to bind. fn __datafusion_session_components__<'py>( diff --git a/examples/datafusion-ffi-example/src/physical_optimizer.rs b/examples/datafusion-ffi-example/src/physical_optimizer.rs index 26caf8d65..38cb67376 100644 --- a/examples/datafusion-ffi-example/src/physical_optimizer.rs +++ b/examples/datafusion-ffi-example/src/physical_optimizer.rs @@ -15,8 +15,8 @@ // specific language governing permissions and limitations // under the License. -use std::sync::Arc; use std::sync::atomic::{AtomicUsize, Ordering}; +use std::sync::{Arc, Mutex}; use datafusion::common::Result; use datafusion::common::config::ConfigOptions; @@ -31,9 +31,15 @@ use pyo3::types::PyCapsule; /// shared counter each time it runs. Tests use the counter to prove that a /// session built with this rule actually routed physical planning through a /// user-supplied [`PhysicalOptimizerRule`] over FFI. +/// +/// A rule declared alongside siblings also appends its label to a log they +/// all share, which is what lets a test see the order they ran in. Counters +/// alone cannot: each rule has its own, so they say how often but not when. #[derive(Debug)] struct CountingPhysicalOptimizerRule { optimize_calls: Arc, + label: usize, + run_log: Option>>>, } impl PhysicalOptimizerRule for CountingPhysicalOptimizerRule { @@ -43,6 +49,9 @@ impl PhysicalOptimizerRule for CountingPhysicalOptimizerRule { _config: &ConfigOptions, ) -> Result> { self.optimize_calls.fetch_add(1, Ordering::SeqCst); + if let Some(run_log) = &self.run_log { + run_log.lock().expect("run log poisoned").push(self.label); + } Ok(plan) } @@ -67,6 +76,21 @@ impl PhysicalOptimizerRule for CountingPhysicalOptimizerRule { #[derive(Debug, Default, Clone)] pub(crate) struct MyPhysicalOptimizerRule { optimize_calls: Arc, + label: usize, + run_log: Option>>>, +} + +impl MyPhysicalOptimizerRule { + /// A rule that records where it ran relative to the siblings sharing + /// `run_log`. `label` is what it appends. Not exposed to Python: only a + /// bundle declaring several rules at once has siblings to order against. + pub(crate) fn with_run_log(label: usize, run_log: Arc>>) -> Self { + Self { + optimize_calls: Arc::new(AtomicUsize::new(0)), + label, + run_log: Some(run_log), + } + } } #[pymethods] @@ -87,6 +111,8 @@ impl MyPhysicalOptimizerRule { let rule: Arc = Arc::new(CountingPhysicalOptimizerRule { optimize_calls: Arc::clone(&self.optimize_calls), + label: self.label, + run_log: self.run_log.clone(), }); let runtime = get_tokio_runtime().handle().clone(); diff --git a/python/datafusion/context.py b/python/datafusion/context.py index a29dd6fa2..7f39d3d36 100644 --- a/python/datafusion/context.py +++ b/python/datafusion/context.py @@ -356,6 +356,40 @@ def _resolve_declared_functions( return resolved +def _resolve_declared_rules( + declared: list[tuple[int, object, Any]], resolve: Any +) -> Any: + """Import the capsule of every physical optimizer rule an extension declared. + + There is no name to check — rules accumulate — so unlike + :py:func:`_resolve_declared_functions` this only refuses a declaration that + cannot be a rule. The getter is looked for here rather than left to the + importer so that the error names the bundle that declared it; a caller who + passed four bundles cannot otherwise tell which one is at fault. + + Args: + declared: ``(position, extension, rule)`` triples in declaration order. + resolve: The primitive that imports a list of rules at once, returning + an opaque object for the commit step. + + Returns: + The imported rules, opaque, in declaration order. + + Raises: + TypeError: If a declaration does not expose + ``__datafusion_physical_optimizer_rule__``. + """ + for _, extension, rule in declared: + if not hasattr(rule, "__datafusion_physical_optimizer_rule__"): + msg = ( + "A declared optimizer rule must expose " + f"__datafusion_physical_optimizer_rule__, got {rule!r} " + f"from {extension!r}" + ) + raise TypeError(msg) + return resolve([rule for _, _, rule in declared]) + + class SessionConfig: """Session configuration options.""" @@ -2137,13 +2171,17 @@ def with_extensions( TypeError: If an argument implements neither hook, if a hook returns the wrong type, if a codec is contributed as a bare ``PyCapsule`` rather than an object exposing the getter, or if - a declared function is neither a wrapper nor exposes its - capsule getter. + a declared function or optimizer rule does not expose its + capsule getter and is not already a wrapper. ValueError: If two codecs claim the same id, if two extensions declare a function of one kind under the same name, or if a getter returns a capsule of the wrong kind. See :py:meth:`with_logical_extension_codec` for how ids are assigned. + RuntimeError: If a getter is present but returns something that is + not a ``PyCapsule`` at all. The message comes from the importer + and does not name the bundle, because by then the declaration + has already been accepted as the right shape. Examples: The returned handle is a different object sharing one session, and @@ -2210,8 +2248,9 @@ def with_extensions( ] # Rules accumulate, so there is no name to check and nothing to refuse # -- only the capsules to import while failing is still free. - resolved_rules = new.ctx._resolve_extension_physical_optimizer_rules( - [rule for _, _, rule in declared["physical_optimizer_rules"]] + resolved_rules = _resolve_declared_rules( + declared["physical_optimizer_rules"], + new.ctx._resolve_extension_physical_optimizer_rules, ) # Phase two: nest the planners, outermost last. Each hook runs against diff --git a/python/datafusion/extensions.py b/python/datafusion/extensions.py index 403526cc8..5964f4c12 100644 --- a/python/datafusion/extensions.py +++ b/python/datafusion/extensions.py @@ -79,11 +79,11 @@ class PhysicalOptimizerRuleExportable(Protocol): argument**: a rule needs neither a codec nor a task-context provider, so there is nothing session-scoped to hand it. - Rules accumulate rather than replace, so several libraries may each - contribute one and none of them has to know about the others. Install one - with :py:meth:`~datafusion.context.SessionContext.add_physical_optimizer_rule`, + Rules accumulate rather than replace. Install one with + :py:meth:`~datafusion.context.SessionContext.add_physical_optimizer_rule`, or declare it on a bundle as - :py:attr:`SessionExtensionComponents.physical_optimizer_rules`. + :py:attr:`SessionExtensionComponents.physical_optimizer_rules` — see + :ref:`extension_other_hooks`. Examples: The getter is the whole protocol, and a capsule is what it must return @@ -98,7 +98,7 @@ class PhysicalOptimizerRuleExportable(Protocol): RuntimeError: "Invalid datafusion_physical_optimizer_rule... Real usage. Skipped here (needs a built extension library); run for - real by ``test_ffi_physical_optimizer_rule`` in + real by ``test_ffi_physical_optimizer_rule_runs_during_planning`` in ``datafusion-ffi-example``. >>> from datafusion_ffi_example import MyPhysicalOptimizerRule # doctest: +SKIP @@ -299,9 +299,8 @@ class SessionExtensionComponents: """Physical optimizer rules to install on the session. Objects exposing ``__datafusion_physical_optimizer_rule__``. Unlike - functions these never collide: rules accumulate, so two extensions may each - contribute one without either having to know about the other. All the rules - in one call install together, in declaration order. + functions these never collide — they accumulate. All the rules in one call + install together, in declaration order. See :ref:`extension_other_hooks`. """ def __post_init__(self) -> None: diff --git a/python/tests/test_context.py b/python/tests/test_context.py index 5740df5c4..fadb61937 100644 --- a/python/tests/test_context.py +++ b/python/tests/test_context.py @@ -1678,20 +1678,20 @@ def test_session_extension_components_rejects_a_single_optimizer_rule(): SessionExtensionComponents(physical_optimizer_rules=object()) -def test_with_extensions_rejects_a_rule_that_is_not_a_capsule(ctx): - """A rule that will not import is refused, and nothing is installed. +def test_with_extensions_rejects_a_rule_that_is_not_a_rule(ctx): + """A declaration that is not a rule at all names the bundle that made it. - Importing the capsules is the only part of installing a rule that can - fail, so it happens during resolution. A failure here has to leave the - session alone even though the extension ahead of it declared a function - that was perfectly good. + Which of several bundles is at fault is the whole content of the message, + and the only place it is still known is here. A failure also has to leave + the session alone even though the extension ahead of it declared a + function that was perfectly good. """ class RuleExtension: def __datafusion_session_components__(self, ctx): return SessionExtensionComponents(physical_optimizer_rules=(object(),)) - with pytest.raises(RuntimeError, match="datafusion_physical_optimizer_rule"): + with pytest.raises(TypeError, match=r"got .* from .*RuleExtension"): ctx.with_extensions( _FunctionExtension(udfs=(_doubler(),)), RuleExtension(), @@ -1701,13 +1701,36 @@ def __datafusion_session_components__(self, ctx): ctx.udf("double") +def test_with_extensions_rejects_a_rule_whose_getter_returns_a_non_capsule(ctx): + """A rule shaped right but returning junk is refused by the importer. + + The bundle is past the point where it can be named — it declared the right + shape — so this is the one rule failure that surfaces as a ``RuntimeError`` + from the import rather than a ``TypeError`` from the resolve. + """ + + class NotACapsule: + def __datafusion_physical_optimizer_rule__(self): + return object() + + class RuleExtension: + def __datafusion_session_components__(self, ctx): + return SessionExtensionComponents(physical_optimizer_rules=(NotACapsule(),)) + + with pytest.raises(RuntimeError, match="datafusion_physical_optimizer_rule"): + ctx.with_extensions(_FunctionExtension(udfs=(_doubler(),)), RuleExtension()) + + with pytest.raises(KeyError): + ctx.udf("double") + + def test_with_extensions_declaring_no_rules_leaves_the_session_id(ctx): - """Installing rules rebuilds ``SessionState``; declaring none must not. + """Installing rules rebuilds ``SessionState``; the id has to survive it. - The rebuild mints a fresh session id unless it is carried over, and a - changed id would break every ``TaskContext`` the session has handed out. - Asserted for the empty case too, because that is the one where the rebuild - would be pure cost. + The rebuild mints a fresh id unless it is carried over, and a changed id + would break every ``TaskContext`` the session has handed out. Asserted for + a call declaring no rules as well, so the guarantee does not depend on + whether the rebuild was skipped. """ before = ctx.session_id() result = ctx.with_extensions(_FunctionExtension(udfs=(_doubler(),))) From c7c4bbff021934cb4d50f86a2e3cea169745dbdc Mon Sep 17 00:00:00 2001 From: Tim Saucer Date: Wed, 16 Sep 2026 08:57:25 -0400 Subject: [PATCH 3/3] docs: note that installing a rule drops prepared statements Installing a physical optimizer rule rebuilds the session state, and `SessionStateBuilder::build` starts the new state with an empty prepared-plan map, so a session that has run PREPARE reports the statement missing afterwards. It hits every handle sharing the session, not just the one a call returned. There is no fix in this repo. `SessionState` exposes `physical_optimizers` read-only and only the builder can append to it, so the rebuild is the only public path, and `prepared_plans` has no builder setter. The canonical home for the claim is a new `extension_rule_rebuild` section in the extension guide. `add_physical_optimizer_rule` and `with_extensions` each state it in one sentence and link there. The enumeration already on `add_physical_optimizer_rule` -- "tables, UDFs, and catalogs are preserved" -- was incomplete and now names the exception. Three smaller corrections alongside it: `test_with_extensions_declaring_no_rules_leaves_the_session_id` claimed to pin the id surviving the rebuild, but it declares no rules, so the rebuild is skipped and the assertion cannot reach that guarantee. Its docstring now says it is the no-op control and names the FFI test that does cover the rebuild. `test_declared_rules_run_in_declaration_order` asserted `run_order() == [0, 1]`, which ties the ordering claim to the optimizer running exactly once per query -- a count its sibling FFI test deliberately avoids asserting. It now checks the first pass only. `PyPhysicalOptimizerRules` is `frozen`. Nothing mutates it between the resolve and commit steps, so it has no reason to carry a borrow flag. Co-Authored-By: Claude Opus 5 (1M context) --- crates/core/src/context.rs | 10 ++++++- .../extension-guide/other-components.md | 26 +++++++++++++++++++ .../python/tests/_test_session_extension.py | 6 ++++- python/datafusion/context.py | 7 ++++- python/tests/test_context.py | 15 ++++++----- 5 files changed, 55 insertions(+), 9 deletions(-) diff --git a/crates/core/src/context.rs b/crates/core/src/context.rs index cc8d59079..5f238c1d5 100644 --- a/crates/core/src/context.rs +++ b/crates/core/src/context.rs @@ -1824,7 +1824,15 @@ impl PySessionContext { /// to carry imported rules from the resolve step to the commit step, so the /// import can fail before anything is written. `with_extensions` is its only /// producer and its only consumer. -#[pyclass(name = "PhysicalOptimizerRules", module = "datafusion._internal")] +/// +/// `frozen` because nothing mutates it between those two steps: the commit +/// only reads the rules back out, so there is no reason to pay for the runtime +/// borrow flag a mutable pyclass carries. +#[pyclass( + frozen, + name = "PhysicalOptimizerRules", + module = "datafusion._internal" +)] pub struct PyPhysicalOptimizerRules { rules: Vec>, } diff --git a/docs/source/extension-guide/other-components.md b/docs/source/extension-guide/other-components.md index 83b770d56..fe48682a6 100644 --- a/docs/source/extension-guide/other-components.md +++ b/docs/source/extension-guide/other-components.md @@ -63,6 +63,32 @@ a bundle contributing several would clone the whole state that many times, and would leave the earlier ones installed if a later one failed. See {ref}`extension_bundles_transaction`. +(extension_rule_rebuild)= + +### Installing a rule rebuilds the session state + +There is no way to append to a live `SessionState`: DataFusion exposes +`physical_optimizers` on one read-only, and only `SessionStateBuilder` can add +to the list. Installing a rule therefore rebuilds the state in place, and the +rebuild carries over the tables, functions, catalogs, and session id the old +one held. + +**Prepared statements are the exception.** `SessionStateBuilder::build` starts +the new state with an empty prepared-plan map, so a session that has run +`PREPARE` reports the statement missing once a rule is installed: + +```python +ctx.sql("PREPARE p AS SELECT a FROM t").collect() +ctx.with_extensions(MyRuleBundle()) +ctx.sql("EXECUTE p") # ValueError: Prepared statement 'p' does not exist +``` + +This is not specific to bundles — `add_physical_optimizer_rule` drops them the +same way, and both hit every handle sharing the session rather than only the +one the call returned. Install your rules before preparing anything. Batching a +bundle's rules into one rebuild is what keeps the cost to once per call instead +of once per rule. + ## Typed configuration **`__datafusion_extension_options__`** contributes typed configuration entries diff --git a/examples/datafusion-ffi-example/python/tests/_test_session_extension.py b/examples/datafusion-ffi-example/python/tests/_test_session_extension.py index 57ef722da..6d5a03bf0 100644 --- a/examples/datafusion-ffi-example/python/tests/_test_session_extension.py +++ b/examples/datafusion-ffi-example/python/tests/_test_session_extension.py @@ -139,13 +139,17 @@ def test_declared_rules_run_in_declaration_order(): Rules rewrite the plan one after another, so the order is part of what a bundle declares. The counters cannot show it — each rule has its own — so the two here append to a log they share. + + Only the first pass is asserted on. How many times a query optimizes is a + separate claim from what order the rules run in, and pinning both here + would report a changed pass count as an ordering bug. """ extension = MyRuleExtension() ctx = SessionContext().with_extensions(extension) _query(ctx) - assert extension.run_order() == [0, 1] + assert extension.run_order()[:2] == [0, 1] def test_rules_install_without_changing_the_session_id(): diff --git a/python/datafusion/context.py b/python/datafusion/context.py index 7f39d3d36..1ebfd804f 100644 --- a/python/datafusion/context.py +++ b/python/datafusion/context.py @@ -2055,7 +2055,8 @@ def add_physical_optimizer_rule( PyCapsule, typically produced by a separate compiled extension. The underlying :class:`SessionState` is rebuilt from its current state with the new rule appended, so previously registered tables, UDFs, - and catalogs are preserved. + and catalogs are preserved. Prepared statements are not — see + :ref:`extension_rule_rebuild`. Args: rule: Object exposing ``__datafusion_physical_optimizer_rule__``, @@ -2151,6 +2152,10 @@ def with_extensions( table, say — is not rolled back, which is why bundle objects must be configuration-only. + A call that installs optimizer rules rebuilds the session state, which + drops the session's prepared statements — see + :ref:`extension_rule_rebuild`. + Shares its session with this context — see :py:class:`SessionContext`. See :ref:`extension_bundles` in the online documentation for why the diff --git a/python/tests/test_context.py b/python/tests/test_context.py index fadb61937..84a143a64 100644 --- a/python/tests/test_context.py +++ b/python/tests/test_context.py @@ -1725,12 +1725,15 @@ def __datafusion_session_components__(self, ctx): def test_with_extensions_declaring_no_rules_leaves_the_session_id(ctx): - """Installing rules rebuilds ``SessionState``; the id has to survive it. - - The rebuild mints a fresh id unless it is carried over, and a changed id - would break every ``TaskContext`` the session has handed out. Asserted for - a call declaring no rules as well, so the guarantee does not depend on - whether the rebuild was skipped. + """A call declaring no rules leaves the session id alone. + + This is the control for the no-op path: with nothing to install the state + rebuild is skipped, so the id is untouched rather than carried over. The + carry-over itself is not reachable from here — the rebuild needs a real + rule capsule, which only a compiled extension can hand over. That half is + pinned by ``test_rules_install_without_changing_the_session_id`` in + ``datafusion-ffi-example``, where a fresh id would leave ``session_id()`` + disagreeing with every ``TaskContext`` the session has handed out. """ before = ctx.session_id() result = ctx.with_extensions(_FunctionExtension(udfs=(_doubler(),)))