diff --git a/crates/core/src/context.rs b/crates/core/src/context.rs index fe75668fb..5f238c1d5 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,78 @@ 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. +/// +/// `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>, } impl PySessionContext { 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/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..fe48682a6 100644 --- a/docs/source/extension-guide/other-components.md +++ b/docs/source/extension-guide/other-components.md @@ -47,6 +47,48 @@ 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`. + +(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 fadc1fd70..6d5a03bf0 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,97 @@ 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_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. + + 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()[:2] == [0, 1] + + +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..c0033b48c 100644 --- a/examples/datafusion-ffi-example/src/extension.rs +++ b/examples/datafusion-ffi-example/src/extension.rs @@ -15,10 +15,13 @@ // 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}; use crate::aggregate_udf::MySumUDF; +use crate::physical_optimizer::MyPhysicalOptimizerRule; use crate::scalar_udf::IsNullUDF; use crate::window_udf::MyRankUDF; @@ -62,3 +65,75 @@ 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, 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)] +pub(crate) struct MyRuleExtension { + first: MyPhysicalOptimizerRule, + second: MyPhysicalOptimizerRule, + run_log: Arc>>, +} + +#[pymethods] +impl MyRuleExtension { + #[new] + fn new() -> Self { + 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. + 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() + } + + /// 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>( + &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..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,16 +76,31 @@ 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] 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) } @@ -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/__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..1ebfd804f 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 @@ -362,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.""" @@ -2027,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__``, @@ -2117,12 +2146,16 @@ 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. + 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 @@ -2143,13 +2176,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 @@ -2214,6 +2251,12 @@ 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 = _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 # `new`, which carries the final chains, so a planner captured here @@ -2254,6 +2297,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..5964f4c12 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. 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` — see + :ref:`extension_other_hooks`. + + 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_runs_during_planning`` 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,16 @@ 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 — they accumulate. All the rules in one call + install together, in declaration order. See :ref:`extension_other_hooks`. + """ + 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..84a143a64 100644 --- a/python/tests/test_context.py +++ b/python/tests/test_context.py @@ -1670,6 +1670,78 @@ 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_rule(ctx): + """A declaration that is not a rule at all names the bundle that made it. + + 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(TypeError, match=r"got .* from .*RuleExtension"): + ctx.with_extensions( + _FunctionExtension(udfs=(_doubler(),)), + RuleExtension(), + ) + + with pytest.raises(KeyError): + 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): + """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(),))) + + 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 +1752,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 +1767,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", } )