diff --git a/Cargo.lock b/Cargo.lock index c7632732a..6a7f68438 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -1209,6 +1209,7 @@ dependencies = [ "datafusion-catalog", "datafusion-common", "datafusion-ffi", + "datafusion-proto", "datafusion-python-util", "datafusion-session", "pyo3", diff --git a/crates/core/src/context.rs b/crates/core/src/context.rs index 84182ff19..04427f49c 100644 --- a/crates/core/src/context.rs +++ b/crates/core/src/context.rs @@ -1608,6 +1608,81 @@ impl PySessionContext { derived.set_session_query_planner(None); derived } + + /// Create the destination context for a `with_extensions` transaction. + /// + /// Private support method for `SessionContext.with_extensions`. The + /// returned context is the single `Arc` that every FFI + /// task-context provider created during the transaction must target; + /// `_install_extensions` later mutates its state in place rather than + /// deriving a new context. + pub fn _derive_for_extensions(&self) -> Self { + Self { + ctx: Arc::new(SessionContext::new_with_state(self.ctx.state())), + logical_codec: Arc::clone(&self.logical_codec), + physical_codec: Arc::clone(&self.physical_codec), + } + } + + /// Commit a `with_extensions` transaction onto this context. + /// + /// Private support method for `SessionContext.with_extensions`; `self` + /// must be a context produced by `_derive_for_extensions`. Codec capsules + /// are imported and validated before any state change, so a failure + /// leaves the context untouched. The final state is written through this + /// context's own `state_ref()`, never a derived context, so FFI + /// task-context providers bound to it stay valid. + #[pyo3(signature = (logical_codecs, physical_codecs, planner=None))] + pub fn _install_extensions<'py>( + slf: &Bound<'py, Self>, + logical_codecs: Vec>, + physical_codecs: Vec>, + planner: Option>, + ) -> PyDataFusionResult { + // Chains are built as local values, so a codec that fails to import -- + // or that collides with an id already installed -- leaves the session + // untouched. Nothing is borrowed across a call back into Python. + let (mut logical_codec, mut physical_codec) = { + let this = slf.borrow(); + ( + this.logical_codec.as_ref().clone(), + this.physical_codec.as_ref().clone(), + ) + }; + + for codec in logical_codecs { + let id = resolve_codec_id(&codec, None, &logical_codec.codec_ids())?; + let inner_ffi = ffi_logical_codec_from_pycapsule(codec, Some(slf.as_any()))?; + let inner: Arc = (&inner_ffi).into(); + logical_codec = logical_codec.with_additional_codec(id, inner); + } + let logical_codec = Arc::new(logical_codec); + + for codec in physical_codecs { + let id = resolve_codec_id(&codec, None, &physical_codec.codec_ids())?; + let inner_ffi = ffi_physical_codec_from_pycapsule(codec, Some(slf.as_any()))?; + let inner: Arc = (&inner_ffi).into(); + physical_codec = physical_codec.with_additional_codec(id, inner); + } + let physical_codec = Arc::new(physical_codec); + + let planner = planner + .map(|planner| ffi_query_planner_from_pycapsule(&planner, Some(slf.as_any()))) + .transpose()?; + + let derived = Self { + ctx: Arc::clone(&slf.borrow().ctx), + logical_codec, + physical_codec, + }; + // Bind the planner only once the codec chains are final, and through + // the derived handle so it carries them. Passing `None` still rebuilds + // whichever planner the session already holds against the new chains, + // exactly as `with_logical_extension_codec` does. + derived.set_session_query_planner(planner); + + Ok(derived) + } } impl PySessionContext { diff --git a/docs/source/contributor-guide/ffi.md b/docs/source/contributor-guide/ffi.md index d86858a83..c3d781126 100644 --- a/docs/source/contributor-guide/ffi.md +++ b/docs/source/contributor-guide/ffi.md @@ -343,6 +343,61 @@ The current FFI logical codec supports providers and UDFs but not arbitrary cust `LogicalPlan::Extension` nodes. See both example READMEs for the supported flow and local build commands. +### Extension bundles: `with_extensions` + +The chaining above works, but it makes the caller responsible for two things that are +easy to get wrong: keeping every intermediate context alive, and installing the codecs +before the planner. Every codec and planner capsule carries an +`FFI_TaskContextProvider` holding a *weak* reference to the context it was built +against, so a component bound to a `with_*` result that is then discarded fails at +query time with `TaskContextProvider went out of scope over FFI boundary`. + +`SessionContext.with_extensions` removes both hazards. An extension library exposes a +bundle object implementing `__datafusion_session_extension__`: + +```python +class MyEngineExtension: + def __datafusion_session_extension__(self, ctx: SessionContext) -> SessionExtensionComponents: + # Create fresh components bound to `ctx` on every call. `ctx` is the + # exact context the host will return from with_extensions. + return SessionExtensionComponents( + logical_extension_codecs=(self._make_logical_codec(ctx),), + physical_extension_codecs=(self._make_physical_codec(ctx),), + query_planner=self._make_planner(ctx), + ) +``` + +The host creates one destination context, passes it to every factory, installs all the +codecs, binds the planner against the final codec chains, and returns that context in +a single step: + +```python +ctx = SessionContext(config).with_extensions(lib_a.Extension(), lib_b.Extension()) +ctx.register_table("t", lib_a.TableProvider()) +ctx.register_udf(udf(lib_b.SomeUDF())) +``` + +Extensions are processed left to right and their codecs are appended to the chain in +that order. As above, order affects only encoding — decoding routes by id. At most one +extension per call may supply a query planner. If any factory raises, the source +context is left exactly as it was. + +Bundle objects must be configuration-only: create fresh components on each call, never +cache bound components, and do not retain the context passed in. Catalogs are shared +with the source context, so registrations made during binding are not rolled back on +failure. + +The returned context is the strong owner of every installed component's task-context +provider, and dependent objects do not extend its lifetime. A `DataFrame`, logical +plan, or capsule can outlive the context, but any operation that reaches an FFI codec +after the context is collected fails with `TaskContextProvider went out of scope over +FFI boundary`. Keep the context alive for as long as objects derived from it are in +use. + +`MyPlannerExtension` in [`datafusion-ffi-query-planner-example`] is a complete Rust +implementation of the protocol, including taking the task-context provider off the +supplied context and constructing a Python `SessionExtensionComponents`. + ### Capsule getters receive the session they are installed on `__datafusion_query_planner__`, `__datafusion_logical_extension_codec__`, and @@ -516,6 +571,9 @@ the original handle rebinds the session's planner back to the original handle's instead, which is the trap `test_reinstalling_a_planner_rebinds_the_session_to_that_handles_codecs` pins. +`with_extensions` sidesteps the ordering question entirely: it installs every codec +before it binds the planner, so there is no "afterwards" for a bundle's own planner. + ## Alternative Approach Suppose you needed to expose some other features of DataFusion and you could not wait diff --git a/examples/datafusion-ffi-query-planner-example/Cargo.toml b/examples/datafusion-ffi-query-planner-example/Cargo.toml index 4d02c69f1..263f034b8 100644 --- a/examples/datafusion-ffi-query-planner-example/Cargo.toml +++ b/examples/datafusion-ffi-query-planner-example/Cargo.toml @@ -31,6 +31,7 @@ datafusion = { workspace = true } datafusion-catalog = { workspace = true, default-features = false } datafusion-common = { workspace = true, default-features = false } datafusion-ffi = { workspace = true } +datafusion-proto = { workspace = true } datafusion-session = { workspace = true } async-trait = { workspace = true } datafusion-python-util.workspace = true diff --git a/examples/datafusion-ffi-query-planner-example/README.md b/examples/datafusion-ffi-query-planner-example/README.md index 66bc45196..597a25142 100644 --- a/examples/datafusion-ffi-query-planner-example/README.md +++ b/examples/datafusion-ffi-query-planner-example/README.md @@ -41,7 +41,22 @@ uv run pytest \ examples/datafusion-ffi-query-planner-example/python/tests/_test*.py ``` -The integration test follows this setup: +The preferred setup uses `SessionContext.with_extensions` with extension bundles: + +```python +config = SessionConfig().with_extension(MyPlannerConfig(max_rows=3)) +ctx = SessionContext(config).with_extensions(provider_bundle, MyPlannerExtension()) +ctx.register_table("numbers", provider) +ctx.register_udf(provider_udf) +``` + +`MyPlannerExtension` implements the `__datafusion_session_extension__` protocol: it +receives the destination context, binds fresh codec and planner components to that +context's task-context provider, and returns them as `SessionExtensionComponents`. +The host installs everything in one step, so no component can end up bound to an +intermediate context that is later collected. + +The integration tests also cover the low-level chaining setup: ```python config = SessionConfig().with_extension(MyPlannerConfig(max_rows=3)) diff --git a/examples/datafusion-ffi-query-planner-example/python/tests/_test_three_library_query_planner.py b/examples/datafusion-ffi-query-planner-example/python/tests/_test_three_library_query_planner.py index 97453b279..adfbf4fb9 100644 --- a/examples/datafusion-ffi-query-planner-example/python/tests/_test_three_library_query_planner.py +++ b/examples/datafusion-ffi-query-planner-example/python/tests/_test_three_library_query_planner.py @@ -21,7 +21,14 @@ import pyarrow as pa import pytest -from datafusion import Expr, SessionConfig, SessionContext, col, udf +from datafusion import ( + Expr, + SessionConfig, + SessionContext, + SessionExtensionComponents, + col, + udf, +) from datafusion_ffi_example import ( IsNullUDF, MyCatalogProvider, @@ -30,7 +37,11 @@ MyPhysicalOptimizerRule, MyTableProvider, ) -from datafusion_ffi_query_planner_example import MyPlannerConfig, MyQueryPlanner +from datafusion_ffi_query_planner_example import ( + MyPlannerConfig, + MyPlannerExtension, + MyQueryPlanner, +) def configured_context(max_rows: int): @@ -688,6 +699,275 @@ def test_query_planner_rejects_invalid_config(max_rows: str): ctx.sql(f"SET ffi_query_planner.max_rows = '{max_rows}'").collect() +class ProviderCodecsExtension: + """Bundles the provider library's codecs for ``with_extensions``. + + These codecs keep their own private task-context provider, so they only + need to be created once; the bundle can hand out the same exporters on + every call. + """ + + def __init__(self) -> None: + self.logical_codec = MyLogicalExtensionCodec() + self.physical_codec = MyPhysicalExtensionCodec() + + def __datafusion_session_extension__( + self, ctx: SessionContext + ) -> SessionExtensionComponents: + return SessionExtensionComponents( + logical_extension_codecs=(self.logical_codec,), + physical_extension_codecs=(self.physical_codec,), + ) + + +class _NamedCodec: + """Forwards a codec's capsule getters under a declared id. + + ``with_extensions`` takes no ``codec_id=``, so an extension that ships a + codec class another extension also ships declares + ``__datafusion_codec_id__`` on the object it hands over. Both getters are + forwarded because one wrapper stands in for whichever kind it wraps. + """ + + def __init__(self, codec: object, codec_id: str) -> None: + self._codec = codec + self.__datafusion_codec_id__ = codec_id + + def __datafusion_logical_extension_codec__(self, session: object = None) -> object: + return self._codec.__datafusion_logical_extension_codec__(session) + + def __datafusion_physical_extension_codec__(self, session: object = None) -> object: + return self._codec.__datafusion_physical_extension_codec__(session) + + +class IdentifiedProviderCodecsExtension(ProviderCodecsExtension): + """``ProviderCodecsExtension`` whose codecs carry ids of their own.""" + + def __init__(self, prefix: str) -> None: + super().__init__() + self.logical = _NamedCodec(self.logical_codec, f"{prefix}.logical") + self.physical = _NamedCodec(self.physical_codec, f"{prefix}.physical") + + def __datafusion_session_extension__( + self, ctx: SessionContext + ) -> SessionExtensionComponents: + return SessionExtensionComponents( + logical_extension_codecs=(self.logical,), + physical_extension_codecs=(self.physical,), + ) + + +def test_with_extensions_three_library_query(): + """One with_extensions call installs provider codecs and a planner bundle, + and a real non-empty plan flows across the three libraries.""" + config = SessionConfig().with_extension(MyPlannerConfig(max_rows=3)) + provider_ext = ProviderCodecsExtension() + planner_ext = MyPlannerExtension() + ctx = SessionContext(config).with_extensions(provider_ext, planner_ext) + ctx.register_table("numbers", MyTableProvider(1, 6, 1)) + ctx.register_udf(udf(IsNullUDF())) + + batches = ctx.sql( + 'SELECT "A", my_custom_is_null("A") AS is_null FROM numbers ORDER BY "A"' + ).collect() + assert batches[0].column(0).to_pylist() == [0, 1, 2] + assert batches[0].column(1).to_pylist() == [False, False, False] + assert planner_ext.plan_calls() >= 1 + assert planner_ext.last_max_rows() == 3 + assert planner_ext.foreign_session_observed() + assert planner_ext.foreign_provider_observed() + assert planner_ext.foreign_plan_observed() + assert provider_ext.logical_codec.table_provider_encode_calls() > 0 + assert provider_ext.logical_codec.table_provider_decode_calls() > 0 + assert provider_ext.physical_codec.execution_plan_encode_calls() > 0 + assert provider_ext.physical_codec.execution_plan_decode_calls() > 0 + + +def test_with_extensions_provider_targets_returned_context(): + """The bundle's task-context provider reads current state from the + returned context, not the source it was derived from.""" + config = SessionConfig().with_extension(MyPlannerConfig(max_rows=3)) + source = SessionContext(config) + source.register_table("numbers", MyTableProvider(1, 6, 1)) + planner_ext = MyPlannerExtension() + result = source.with_extensions(ProviderCodecsExtension(), planner_ext) + + # Diverge the two live contexts. Config state is copied at derivation, + # so after these statements source and result disagree. + source.sql("SET ffi_query_planner.max_rows = 5").collect() + result.sql("SET ffi_query_planner.max_rows = 2").collect() + + batches = result.sql('SELECT "A" FROM numbers ORDER BY "A"').collect() + assert batches[0].column(0).to_pylist() == [0, 1] + assert planner_ext.last_max_rows() == 2 + + # Resolving the provider bound during with_extensions is what a codec's + # decode callback does. Seeing 2 (never 5) proves the provider targets the + # returned context rather than the source. + assert planner_ext.max_rows_through_provider() == 2 + + +def test_with_extensions_survives_dropping_source_and_bundles(): + """Neither the source context nor the bundle objects are needed to keep + the installed components' task-context provider alive.""" + config = SessionConfig().with_extension(MyPlannerConfig(max_rows=2)) + ctx = SessionContext(config).with_extensions( + ProviderCodecsExtension(), MyPlannerExtension() + ) + ctx.register_table("numbers", MyTableProvider(1, 6, 1)) + gc.collect() + + batches = ctx.sql('SELECT "A" FROM numbers ORDER BY "A"').collect() + assert batches[0].column(0).to_pylist() == [0, 1] + + +def test_with_extensions_sees_state_changes_after_install(): + """Tables, UDFs, and config changes made after installation are visible + to the planner and to provider callbacks.""" + config = SessionConfig().with_extension(MyPlannerConfig(max_rows=4)) + planner_ext = MyPlannerExtension() + ctx = SessionContext(config).with_extensions(ProviderCodecsExtension(), planner_ext) + + ctx.register_table("numbers", MyTableProvider(1, 6, 1)) + ctx.register_udf(udf(IsNullUDF())) + ctx.sql("SET ffi_query_planner.max_rows = 2").collect() + + batches = ctx.sql( + 'SELECT "A", my_custom_is_null("A") AS is_null FROM numbers ORDER BY "A"' + ).collect() + assert batches[0].column(0).to_pylist() == [0, 1] + assert planner_ext.last_max_rows() == 2 + assert planner_ext.max_rows_through_provider() == 2 + + +def test_with_extensions_bundle_is_reusable(): + """Installing the same bundle into two contexts binds fresh components to + each destination.""" + planner_ext = MyPlannerExtension() + + config_a = SessionConfig().with_extension(MyPlannerConfig(max_rows=2)) + ctx_a = SessionContext(config_a).with_extensions( + ProviderCodecsExtension(), planner_ext + ) + ctx_a.register_table("numbers", MyTableProvider(1, 6, 1)) + + config_b = SessionConfig().with_extension(MyPlannerConfig(max_rows=3)) + ctx_b = SessionContext(config_b).with_extensions( + ProviderCodecsExtension(), planner_ext + ) + ctx_b.register_table("numbers", MyTableProvider(1, 6, 1)) + + batches = ctx_a.sql('SELECT "A" FROM numbers ORDER BY "A"').collect() + assert batches[0].column(0).to_pylist() == [0, 1] + assert planner_ext.last_max_rows() == 2 + + batches = ctx_b.sql('SELECT "A" FROM numbers ORDER BY "A"').collect() + assert batches[0].column(0).to_pylist() == [0, 1, 2] + assert planner_ext.last_max_rows() == 3 + + +def test_with_extensions_failure_leaves_source_usable(): + """A failing factory after a successful one leaves the source context + fully functional.""" + + class BoomExtension: + def __datafusion_session_extension__( + self, ctx: SessionContext + ) -> SessionExtensionComponents: + msg = "boom" + raise RuntimeError(msg) + + config = SessionConfig().with_extension(MyPlannerConfig(max_rows=2)) + source = SessionContext(config) + source.register_table("numbers", MyTableProvider(1, 6, 1)) + + with pytest.raises(RuntimeError, match="boom"): + source.with_extensions(MyPlannerExtension(), BoomExtension()) + + # No planner was installed, so the default planner runs unrestricted. + batches = source.sql('SELECT "A" FROM numbers ORDER BY "A"').collect() + assert batches[0].column(0).to_pylist() == [0, 1, 2, 3, 4, 5] + + +def test_with_extensions_rebinds_existing_planner(): + """Codec-only bundles installed on a context that already has an FFI + planner rebind that planner to the new codec chains.""" + config = SessionConfig().with_extension(MyPlannerConfig(max_rows=2)) + planner = MyQueryPlanner() + ctx = SessionContext(config) + ctx.set_query_planner(planner) + provider_ext = ProviderCodecsExtension() + ctx = ctx.with_extensions(provider_ext) + ctx.register_table("numbers", MyTableProvider(1, 6, 1)) + + batches = ctx.sql('SELECT "A" FROM numbers ORDER BY "A"').collect() + assert batches[0].column(0).to_pylist() == [0, 1] + assert planner.last_max_rows() == 2 + # The planner only sees these codecs if it was rebound to the chains + # built during with_extensions. + assert provider_ext.logical_codec.table_provider_decode_calls() > 0 + assert provider_ext.physical_codec.execution_plan_decode_calls() > 0 + + +def test_with_extensions_rejects_two_bundles_of_the_same_codec_class(): + """Two bundles contributing the same codec class collide on id. + + Ids are derived from the exporting class, so two instances of one class + claim the same id. A payload names its codec by id when it is decoded, so + the ambiguity is refused at install time rather than resolved by position. + """ + config = SessionConfig().with_extension(MyPlannerConfig(max_rows=2)) + with pytest.raises(ValueError, match="is already installed on this session"): + SessionContext(config).with_extensions( + ProviderCodecsExtension(), ProviderCodecsExtension(), MyPlannerExtension() + ) + + +def test_with_extensions_accepts_distinct_codec_ids(): + """Declaring ``__datafusion_codec_id__`` resolves the collision above. + + Both codec pairs then install, and the query still runs end to end: only + the codec that wrote a payload is asked to decode it, so the second pair + is simply never consulted. + """ + config = SessionConfig().with_extension(MyPlannerConfig(max_rows=2)) + ext_a = ProviderCodecsExtension() + ext_b = IdentifiedProviderCodecsExtension("second") + ctx = SessionContext(config).with_extensions(ext_a, ext_b, MyPlannerExtension()) + ctx.register_table("numbers", MyTableProvider(1, 6, 1)) + + batches = ctx.sql('SELECT "A" FROM numbers ORDER BY "A"').collect() + assert batches[0].column(0).to_pylist() == [0, 1] + + ids = ctx.logical_extension_codec_ids() + assert "datafusion_ffi_example.MyLogicalExtensionCodec" in ids + assert "second.logical" in ids + + # The first pair wrote the payloads, so decoding routes back to it alone. + assert ext_a.logical_codec.table_provider_encode_calls() > 0 + assert ext_a.logical_codec.table_provider_decode_calls() > 0 + assert ext_b.logical_codec.table_provider_decode_calls() == 0 + + +def test_dataframe_outliving_context_fails_cleanly(): + """A DataFrame does not keep its SessionContext alive. FFI components + resolve the task context through a weak reference, so using the + DataFrame after dropping the context raises a clean error instead of + crashing. This locks in the documented ownership contract: the context + must outlive DataFrames that depend on FFI codecs.""" + config = SessionConfig().with_extension(MyPlannerConfig(max_rows=2)) + ctx = SessionContext(config).with_extensions( + ProviderCodecsExtension(), MyPlannerExtension() + ) + ctx.register_table("numbers", MyTableProvider(1, 6, 1)) + df = ctx.sql('SELECT "A" FROM numbers ORDER BY "A"') + del ctx + gc.collect() + + with pytest.raises(Exception, match="went out of scope"): + df.collect() + + def test_composed_codecs_with_query_planner(): """A second pair of codecs installed on top of the provider codecs composes with them instead of replacing them. The extra codecs diff --git a/examples/datafusion-ffi-query-planner-example/src/extension.rs b/examples/datafusion-ffi-query-planner-example/src/extension.rs new file mode 100644 index 000000000..3f60cc819 --- /dev/null +++ b/examples/datafusion-ffi-query-planner-example/src/extension.rs @@ -0,0 +1,299 @@ +// 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. + +use std::fmt; +use std::sync::atomic::Ordering; +use std::sync::{Arc, Mutex}; + +use datafusion::arrow::datatypes::SchemaRef; +use datafusion::common::{Result, TableReference}; +use datafusion::datasource::TableProvider; +use datafusion::execution::TaskContext; +use datafusion::logical_expr::{Extension, LogicalPlan}; +use datafusion::physical_plan::ExecutionPlan; +use datafusion_ffi::execution::FFI_TaskContextProvider; +use datafusion_ffi::proto::logical_extension_codec::FFI_LogicalExtensionCodec; +use datafusion_ffi::proto::physical_extension_codec::FFI_PhysicalExtensionCodec; +use datafusion_ffi::query_planner::FFI_QueryPlanner; +use datafusion_proto::logical_plan::{DefaultLogicalExtensionCodec, LogicalExtensionCodec}; +use datafusion_proto::physical_plan::{ + DefaultPhysicalExtensionCodec, PhysicalExtensionCodec, PhysicalProtoConverterExtension, +}; +use datafusion_python_util::{ + create_logical_extension_capsule, create_physical_extension_capsule, + create_query_planner_capsule, ffi_logical_codec_from_pycapsule, + ffi_physical_codec_from_pycapsule, ffi_task_context_provider_from_pycapsule, get_tokio_runtime, +}; +use datafusion_session::QueryPlanner; +use pyo3::prelude::*; +use pyo3::types::PyDict; + +use crate::planner::{DistributedQueryPlanner, PlannerObservations, planner_config_from_options}; + +/// Values of `ffi_query_planner.max_rows` observed through the task-context +/// provider bound at installation time. +/// +/// Only populated when a codec in this bundle is actually consulted. The host +/// dispatches a framed payload straight to the codec whose id it names, so a +/// decline-all codec like the ones here is normally never asked to decode. The +/// binding itself is proved by [`MyPlannerExtension::max_rows_through_provider`], +/// which reads the provider directly rather than waiting for a callback. +type ObservedMaxRows = Arc>>; + +/// The task-context provider handed to this bundle's components, if it has been +/// installed. `FFI_TaskContextProvider` holds its session weakly, so keeping one +/// here does not keep the destination context alive. +type BoundProvider = Arc>>; + +fn record_task_ctx(observed: &ObservedMaxRows, ctx: &TaskContext) { + if let Ok(config) = planner_config_from_options(ctx.session_config().options()) + && let Ok(mut observed) = observed.lock() + { + observed.push(config.max_rows); + } +} + +/// Records the task context resolved by the FFI wrapper, then declines by +/// delegating to the default codec so the host's codec chain falls through to +/// the codec that owns the payload. +struct ObservingLogicalExtensionCodec { + inner: DefaultLogicalExtensionCodec, + observed: ObservedMaxRows, +} + +impl fmt::Debug for ObservingLogicalExtensionCodec { + fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result { + formatter + .debug_struct("ObservingLogicalExtensionCodec") + .finish_non_exhaustive() + } +} + +impl LogicalExtensionCodec for ObservingLogicalExtensionCodec { + fn try_decode( + &self, + buf: &[u8], + inputs: &[LogicalPlan], + ctx: &TaskContext, + ) -> Result { + record_task_ctx(&self.observed, ctx); + self.inner.try_decode(buf, inputs, ctx) + } + + fn try_encode(&self, node: &Extension, buf: &mut Vec) -> Result<()> { + self.inner.try_encode(node, buf) + } + + fn try_decode_table_provider( + &self, + buf: &[u8], + table_ref: &TableReference, + schema: SchemaRef, + ctx: &TaskContext, + ) -> Result> { + record_task_ctx(&self.observed, ctx); + self.inner + .try_decode_table_provider(buf, table_ref, schema, ctx) + } + + fn try_encode_table_provider( + &self, + table_ref: &TableReference, + node: Arc, + buf: &mut Vec, + ) -> Result<()> { + self.inner.try_encode_table_provider(table_ref, node, buf) + } +} + +/// Physical companion to [`ObservingLogicalExtensionCodec`]. +struct ObservingPhysicalExtensionCodec { + inner: DefaultPhysicalExtensionCodec, + observed: ObservedMaxRows, +} + +impl fmt::Debug for ObservingPhysicalExtensionCodec { + fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result { + formatter + .debug_struct("ObservingPhysicalExtensionCodec") + .finish_non_exhaustive() + } +} + +impl PhysicalExtensionCodec for ObservingPhysicalExtensionCodec { + fn try_decode( + &self, + buf: &[u8], + inputs: &[Arc], + ctx: &TaskContext, + proto_converter: &dyn PhysicalProtoConverterExtension, + ) -> Result> { + record_task_ctx(&self.observed, ctx); + self.inner.try_decode(buf, inputs, ctx, proto_converter) + } + + fn try_encode( + &self, + node: Arc, + buf: &mut Vec, + proto_converter: &dyn PhysicalProtoConverterExtension, + ) -> Result<()> { + self.inner.try_encode(node, buf, proto_converter) + } +} + +/// Extension bundle for `SessionContext.with_extensions`. +/// +/// Mirrors how a distributed engine such as Ballista packages its session +/// extensions: the object itself is reusable configuration, and every +/// `__datafusion_session_extension__` call creates fresh codec and planner +/// components bound to the task-context provider of the context it receives. +#[pyclass( + from_py_object, + name = "MyPlannerExtension", + module = "datafusion_ffi_query_planner_example", + subclass +)] +#[derive(Default, Clone)] +pub(crate) struct MyPlannerExtension { + observations: Arc, + observed_max_rows: ObservedMaxRows, + bound_provider: BoundProvider, +} + +impl fmt::Debug for MyPlannerExtension { + fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result { + formatter + .debug_struct("MyPlannerExtension") + .field("observations", &self.observations) + .finish_non_exhaustive() + } +} + +#[pymethods] +impl MyPlannerExtension { + #[new] + fn new() -> Self { + Self::default() + } + + fn plan_calls(&self) -> usize { + self.observations.plan_calls.load(Ordering::SeqCst) + } + + fn last_max_rows(&self) -> usize { + self.observations.last_max_rows.load(Ordering::SeqCst) + } + + fn foreign_session_observed(&self) -> bool { + self.observations.foreign_session.load(Ordering::SeqCst) + } + + fn foreign_provider_observed(&self) -> bool { + self.observations.foreign_provider.load(Ordering::SeqCst) + } + + fn foreign_plan_observed(&self) -> bool { + self.observations.foreign_plan.load(Ordering::SeqCst) + } + + /// `ffi_query_planner.max_rows` values seen through the bound + /// task-context provider during codec decode calls. + /// + /// Usually empty: the host routes a framed payload to the codec named in + /// it, so codecs that own nothing are not consulted. + fn decode_max_rows_seen(&self) -> Vec { + self.observed_max_rows + .lock() + .map(|observed| observed.clone()) + .unwrap_or_default() + } + + /// `ffi_query_planner.max_rows` read through the task-context provider + /// this bundle was last bound to. + /// + /// Resolving the provider is what a codec's decode callback does, so this + /// answers which session those callbacks would resolve against -- the + /// context `with_extensions` returned, not the one it was called on. + /// Returns ``None`` if the bundle was never installed, or if the context it + /// was bound to has been dropped: the provider holds it weakly. + fn max_rows_through_provider(&self) -> Option { + let provider = self.bound_provider.lock().ok()?.clone()?; + let task_ctx = Arc::::try_from(&provider).ok()?; + planner_config_from_options(task_ctx.session_config().options()) + .ok() + .map(|config| config.max_rows) + } + + fn __datafusion_session_extension__<'py>( + &self, + py: Python<'py>, + ctx: Bound<'py, PyAny>, + ) -> PyResult> { + // Bind every component to the destination context supplied by the + // host. Components must not be cached across calls: each installation + // targets a different context. + // + // The task-context provider comes off that context rather than from a + // `SessionContext` built here, so the codecs' decode callbacks resolve + // names against the session that will actually run the query. + let provider = ffi_task_context_provider_from_pycapsule(&ctx)?; + if let Ok(mut bound) = self.bound_provider.lock() { + *bound = Some(provider.clone()); + } + let runtime = get_tokio_runtime().handle().clone(); + + let logical: Arc = Arc::new(ObservingLogicalExtensionCodec { + inner: DefaultLogicalExtensionCodec {}, + observed: Arc::clone(&self.observed_max_rows), + }); + let ffi_logical = + FFI_LogicalExtensionCodec::new(logical, Some(runtime.clone()), provider.clone()); + let logical_capsule = create_logical_extension_capsule(py, &ffi_logical)?; + + let physical: Arc = + Arc::new(ObservingPhysicalExtensionCodec { + inner: DefaultPhysicalExtensionCodec {}, + observed: Arc::clone(&self.observed_max_rows), + }); + let ffi_physical = + FFI_PhysicalExtensionCodec::new(physical, Some(runtime), provider.clone()); + let physical_capsule = create_physical_extension_capsule(py, &ffi_physical)?; + + let planner: Arc = Arc::new(DistributedQueryPlanner { + observations: Arc::clone(&self.observations), + fallback: None, + }); + // The planner takes the host's codecs, not ones built here. Installing + // the codecs above rebuilds the planner against them anyway, and this + // library has no business minting a provider of its own. + let host_logical = ffi_logical_codec_from_pycapsule(ctx.clone(), None)?; + let host_physical = ffi_physical_codec_from_pycapsule(ctx, None)?; + let ffi_planner = + FFI_QueryPlanner::new_with_ffi_codecs(planner, host_logical, host_physical); + let planner_capsule = create_query_planner_capsule(py, &ffi_planner)?; + + let components = py + .import("datafusion")? + .getattr("SessionExtensionComponents")?; + let kwargs = PyDict::new(py); + kwargs.set_item("logical_extension_codecs", (logical_capsule,))?; + kwargs.set_item("physical_extension_codecs", (physical_capsule,))?; + kwargs.set_item("query_planner", planner_capsule)?; + components.call((), Some(&kwargs)) + } +} diff --git a/examples/datafusion-ffi-query-planner-example/src/lib.rs b/examples/datafusion-ffi-query-planner-example/src/lib.rs index c505c1ce7..70d4a42c5 100644 --- a/examples/datafusion-ffi-query-planner-example/src/lib.rs +++ b/examples/datafusion-ffi-query-planner-example/src/lib.rs @@ -18,15 +18,18 @@ use pyo3::prelude::*; use crate::config::MyPlannerConfig; +use crate::extension::MyPlannerExtension; use crate::planner::MyQueryPlanner; mod config; +mod extension; mod planner; #[pymodule] fn datafusion_ffi_query_planner_example(m: &Bound<'_, PyModule>) -> PyResult<()> { pyo3_log::init(); m.add_class::()?; + m.add_class::()?; m.add_class::()?; Ok(()) } diff --git a/examples/datafusion-ffi-query-planner-example/src/planner.rs b/examples/datafusion-ffi-query-planner-example/src/planner.rs index 67262e39c..733536d21 100644 --- a/examples/datafusion-ffi-query-planner-example/src/planner.rs +++ b/examples/datafusion-ffi-query-planner-example/src/planner.rs @@ -52,14 +52,14 @@ use crate::config::MyPlannerConfig; /// most recent plan would be answering a different question than the one its /// accessor name asks. #[derive(Default)] -struct PlannerObservations { - plan_calls: AtomicUsize, - last_max_rows: AtomicUsize, - foreign_session: AtomicBool, - foreign_provider: AtomicBool, - foreign_plan: AtomicBool, +pub(crate) struct PlannerObservations { + pub(crate) plan_calls: AtomicUsize, + pub(crate) last_max_rows: AtomicUsize, + pub(crate) foreign_session: AtomicBool, + pub(crate) foreign_provider: AtomicBool, + pub(crate) foreign_plan: AtomicBool, /// Only ever set to `true`, so it is already cumulative. - used_fallback: AtomicBool, + pub(crate) used_fallback: AtomicBool, } impl fmt::Debug for PlannerObservations { @@ -104,8 +104,12 @@ const MAX_ROWS_KEY: &str = "ffi_query_planner.max_rows"; const FFI_MAX_ROWS_KEY: &str = "datafusion_ffi.ffi_query_planner.max_rows"; fn planner_config(session: &dyn Session) -> datafusion::common::Result { - let options = session.config_options(); + planner_config_from_options(session.config_options()) +} +pub(crate) fn planner_config_from_options( + options: &datafusion::common::config::ConfigOptions, +) -> datafusion::common::Result { // Prefer the raw entry. `local_or_ffi_extension` discards a value it cannot // parse and hands back `MyPlannerConfig::default()`, which would quietly turn // a typo into a different row limit instead of reporting it. @@ -143,8 +147,8 @@ fn planner_config(session: &dyn Session) -> datafusion::common::Result, +pub(crate) struct DistributedQueryPlanner { + pub(crate) observations: Arc, /// Planner to hand the work to instead of planning here. /// /// This is how a real planner layers on top of an existing one. The capsule @@ -156,7 +160,7 @@ struct DistributedQueryPlanner { /// Note that `Session::create_physical_plan` cannot be used for this. It /// dispatches through the session's installed query planner, so calling it /// from inside that planner recurses until the stack overflows. - fallback: Option>, + pub(crate) fallback: Option>, } #[async_trait] diff --git a/python/datafusion/__init__.py b/python/datafusion/__init__.py index 9c55f446c..4b02a383e 100644 --- a/python/datafusion/__init__.py +++ b/python/datafusion/__init__.py @@ -92,6 +92,10 @@ ) from .dataframe_formatter import configure_formatter from .expr import Expr, WindowFrame +from .extensions import ( + SessionExtensionComponents, + SessionExtensionExportable, +) from .io import read_avro, read_csv, read_json, read_parquet from .options import CsvReadOptions from .plan import ExecutionPlan, LogicalPlan, Metric, MetricsSet @@ -134,6 +138,8 @@ "ScalarUDF", "SessionConfig", "SessionContext", + "SessionExtensionComponents", + "SessionExtensionExportable", "Table", "TableFunction", "TableProviderFactory", @@ -146,6 +152,7 @@ "common", "configure_formatter", "expr", + "extensions", "functions", "ipc", "lit", diff --git a/python/datafusion/context.py b/python/datafusion/context.py index 644c7b445..202ff14ba 100644 --- a/python/datafusion/context.py +++ b/python/datafusion/context.py @@ -69,6 +69,11 @@ ) from datafusion.dataframe import DataFrame from datafusion.expr import sort_list_to_raw_sort_list +from datafusion.extensions import ( + QueryPlannerExportable, + SessionExtensionComponents, + SessionExtensionExportable, +) from datafusion.options import ( DEFAULT_MAX_INFER_SCHEMA, CsvReadOptions, @@ -145,18 +150,6 @@ class PhysicalOptimizerRuleExportable(Protocol): def __datafusion_physical_optimizer_rule__(self) -> object: ... # noqa: D105 -class QueryPlannerExportable(Protocol): - """Type hint for object that has a __datafusion_query_planner__ PyCapsule. - - The method returns a PyCapsule wrapping an ``FFI_QueryPlanner``, typically - produced by a separate compiled extension. ``session`` is the - :py:class:`SessionContext` the planner is being installed on; take the - extension codecs from it rather than building your own. - """ - - def __datafusion_query_planner__(self, session: Any) -> object: ... # noqa: D105 - - class SessionConfig: """Session configuration options.""" @@ -1817,6 +1810,113 @@ def set_query_planner(self, planner: QueryPlannerExportable | _PyCapsule) -> Non """ self.ctx.set_query_planner(planner) + def with_extensions( + self, *extensions: SessionExtensionExportable + ) -> SessionContext: + """Create a new session context with the given extension bundles. + + This is the preferred way to install FFI extensions that need a + task-context provider (extension codecs and query planners). Each + extension's ``__datafusion_session_extension__`` method is called with + the destination context so it can bind its components to that exact + context, then all components are installed in one step. This avoids + the pitfalls of chaining :py:meth:`with_logical_extension_codec`, + :py:meth:`with_physical_extension_codec`, and + :py:meth:`set_query_planner` by hand, where components can end up + bound to an intermediate context that is later garbage collected. + + Codecs compose with the existing chain and with each other: extensions + are processed left to right and their codecs are appended to the chain + in that order. Decoding routes by codec id, so the order matters only + for encoding. At most one extension may supply a query planner. If none + does, an existing FFI planner on the source context is rebound to the + final codec chains. + + If any extension raises or returns invalid components, the source + context's state is left unchanged and the partially built destination + is discarded. Extension factories must treat the context they receive + as configuration-only: catalogs are shared with the source context, so + registering tables or otherwise mutating the context during binding is + not rolled back on failure. + + The returned context is the strong owner of the installed components' + task-context providers. Keep it alive for as long as DataFrames or + plans derived from it are in use; FFI operations after the context is + collected raise an error. + + Args: + extensions: One or more objects implementing + ``__datafusion_session_extension__`` (see + :py:class:`SessionExtensionExportable`). + + Returns: + A new context with all extension components installed. + + Raises: + TypeError: If an argument does not implement the protocol or + returns something other than a + :py:class:`SessionExtensionComponents`. + ValueError: If no extensions are given, more than one extension + supplies a query planner, or two codecs claim the same id. Ids + are derived the same way :py:meth:`with_logical_extension_codec` + derives them, so an extension that contributes two instances of + one codec class must declare ``__datafusion_codec_id__`` on at + least one of them. + + Examples: + >>> from my_extension import DistributedEngineExtension # doctest: +SKIP + >>> ctx = SessionContext().with_extensions( + ... DistributedEngineExtension("scheduler:50050") + ... ) # doctest: +SKIP + >>> ctx.sql("SELECT 1").collect() # doctest: +SKIP + """ + if not extensions: + msg = "with_extensions requires at least one extension" + raise ValueError(msg) + for extension in extensions: + if not hasattr(extension, "__datafusion_session_extension__"): + msg = ( + "Extension does not implement __datafusion_session_extension__: " + f"{extension!r}" + ) + raise TypeError(msg) + + # Single destination context. Every component the extensions create + # must bind to this context; _install_extensions later mutates its + # state in place so those bindings stay valid. + destination = SessionContext.__new__(SessionContext) + destination.ctx = self.ctx._derive_for_extensions() + + logical_codecs: list[LogicalExtensionCodecExportable | _PyCapsule] = [] + physical_codecs: list[PhysicalExtensionCodecExportable | _PyCapsule] = [] + planner: QueryPlannerExportable | _PyCapsule | None = None + for extension in extensions: + components = extension.__datafusion_session_extension__(destination) + if not isinstance(components, SessionExtensionComponents): + msg = ( + "__datafusion_session_extension__ must return " + "SessionExtensionComponents, got " + f"{type(components).__name__} from {extension!r}" + ) + raise TypeError(msg) + logical_codecs.extend(components.logical_extension_codecs) + physical_codecs.extend(components.physical_extension_codecs) + if components.query_planner is not None: + if planner is not None: + msg = ( + "Multiple extensions supplied a query planner; a " + "session context has exactly one. Layer planners " + "explicitly instead." + ) + raise ValueError(msg) + planner = components.query_planner + + new = SessionContext.__new__(SessionContext) + new.ctx = destination.ctx._install_extensions( + logical_codecs, physical_codecs, planner + ) + return new + def table_provider(self, name: str) -> Table: """Return the :py:class:`~datafusion.catalog.Table` for the given table name. diff --git a/python/datafusion/extensions.py b/python/datafusion/extensions.py new file mode 100644 index 000000000..e228a96de --- /dev/null +++ b/python/datafusion/extensions.py @@ -0,0 +1,114 @@ +# 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. + +"""Protocols and value types for installing extensions on a session context. + +An *extension* is a reusable configuration object — typically shipped by a +separate compiled library — that contributes components to a +:py:class:`~datafusion.context.SessionContext`. It implements +:py:class:`SessionExtensionExportable` by returning a +:py:class:`SessionExtensionComponents` describing what it contributes, and is +installed with :py:meth:`~datafusion.context.SessionContext.with_extensions`:: + + ctx = SessionContext().with_extensions(MyLibraryExtension()) + +Installing through ``with_extensions`` rather than by chaining the individual +``with_*`` methods matters for components that hold a task-context provider: +the extension is handed the destination context so every component binds to +the session that is actually returned. See the FFI extensions guide in the +contributor documentation for the full rationale. +""" + +from __future__ import annotations + +from dataclasses import dataclass +from typing import TYPE_CHECKING, Any, Protocol + +if TYPE_CHECKING: + from _typeshed import CapsuleType as _PyCapsule + + from datafusion.context import SessionContext + from datafusion.user_defined import ( + LogicalExtensionCodecExportable, + PhysicalExtensionCodecExportable, + ) + +__all__ = [ + "QueryPlannerExportable", + "SessionExtensionComponents", + "SessionExtensionExportable", +] + + +class QueryPlannerExportable(Protocol): + """Type hint for object that has a __datafusion_query_planner__ PyCapsule. + + The method returns a PyCapsule wrapping an ``FFI_QueryPlanner``, typically + produced by a separate compiled extension. ``session`` is the + :py:class:`~datafusion.context.SessionContext` the planner is being + installed on; take the extension codecs from it rather than building your + own. + """ + + def __datafusion_query_planner__(self, session: Any) -> object: ... # noqa: D105 + + +@dataclass(frozen=True) +class SessionExtensionComponents: + """Components an extension contributes to a session context. + + Returned by :py:meth:`SessionExtensionExportable.__datafusion_session_extension__` + and consumed by + :py:meth:`~datafusion.context.SessionContext.with_extensions`. Every + component must be created against the context passed to that method; + components bound to any other context hold a task-context provider for the + wrong session and cannot be rebound. + """ + + logical_extension_codecs: tuple[ + LogicalExtensionCodecExportable | _PyCapsule, ... + ] = () + """Logical codecs to add to the session's codec chain, in declaration order.""" + + physical_extension_codecs: tuple[ + PhysicalExtensionCodecExportable | _PyCapsule, ... + ] = () + """Physical codecs to add to the session's codec chain, in declaration order.""" + + query_planner: QueryPlannerExportable | _PyCapsule | None = None + """Optional query planner. + + At most one extension per + :py:meth:`~datafusion.context.SessionContext.with_extensions` call may + supply one. + """ + + +class SessionExtensionExportable(Protocol): + """Type hint for extension bundles installable via ``with_extensions``. + + Implementations are reusable configuration objects: they must not retain a + :py:class:`~datafusion.context.SessionContext` and must create fresh + components on every call using the context supplied by + :py:meth:`~datafusion.context.SessionContext.with_extensions`. They should + also avoid mutating global state during binding, since a failed + installation discards the destination context. + """ + + def __datafusion_session_extension__( # noqa: D105 + self, ctx: SessionContext + ) -> SessionExtensionComponents: ... diff --git a/python/tests/test_context.py b/python/tests/test_context.py index 3c95835af..f509d8afa 100644 --- a/python/tests/test_context.py +++ b/python/tests/test_context.py @@ -16,6 +16,7 @@ # under the License. import ctypes import datetime as dt +import gc import gzip import pathlib import shutil @@ -29,6 +30,7 @@ RuntimeEnvBuilder, SessionConfig, SessionContext, + SessionExtensionComponents, SQLOptions, Table, column, @@ -879,6 +881,121 @@ def test_contexts_sharing_a_session_share_the_planner(ctx): assert sibling.session_id() == ctx.session_id() +class _CodecOnlyExtension: + """Contributes decline-all codecs exported from an unrelated session.""" + + def __init__(self): + self.exporter = SessionContext() + self.bound_ctx = None + + def __datafusion_session_extension__(self, ctx): + self.bound_ctx = ctx + return SessionExtensionComponents( + logical_extension_codecs=( + self.exporter.__datafusion_logical_extension_codec__(), + ), + physical_extension_codecs=( + self.exporter.__datafusion_physical_extension_codec__(), + ), + ) + + +class _PlannerExtension: + """Contributes the destination context's own exported planner.""" + + def __datafusion_session_extension__(self, ctx): + return SessionExtensionComponents( + query_planner=ctx.__datafusion_query_planner__() + ) + + +def test_with_extensions_requires_an_extension(ctx): + with pytest.raises(ValueError, match="at least one extension"): + ctx.with_extensions() + + +def test_with_extensions_rejects_non_extension(ctx): + with pytest.raises(TypeError, match="__datafusion_session_extension__"): + ctx.with_extensions(object()) + + +def test_with_extensions_rejects_bad_components(ctx): + class BadExtension: + def __datafusion_session_extension__(self, ctx): + return 42 + + with pytest.raises(TypeError, match="SessionExtensionComponents"): + ctx.with_extensions(BadExtension()) + + +def test_with_extensions_rejects_multiple_planners(ctx): + with pytest.raises(ValueError, match="query planner"): + ctx.with_extensions(_PlannerExtension(), _PlannerExtension()) + + +def test_with_extensions_rejects_bad_codec_capsule(ctx): + class BadCodecExtension: + def __datafusion_session_extension__(self, ctx): + return SessionExtensionComponents( + logical_extension_codecs=(ctx.__datafusion_task_context_provider__(),), + ) + + with pytest.raises( + ValueError, match="Expected name 'datafusion_logical_extension_codec'" + ): + ctx.with_extensions(BadCodecExtension()) + + +def test_with_extensions_installs_codecs_and_planner(ctx): + ctx.register_record_batches( + "extensions_test", + [[pa.RecordBatch.from_pydict({"value": [1, 2, 3]})]], + ) + extension = _CodecOnlyExtension() + result = ctx.with_extensions(extension, _PlannerExtension()) + + assert result.table_exist("extensions_test") + # In-memory tables need a real extension codec to round-trip through the + # FFI planner, so query plans that don't serialize a table provider. + batches = result.sql("SELECT 1 AS value").collect() + assert batches[0].column(0) == pa.array([1]) + + +def test_with_extensions_binds_to_returned_context(ctx): + extension = _CodecOnlyExtension() + result = ctx.with_extensions(extension) + + # The context passed to the factory shares the same underlying session + # as the returned context: registrations made through it are visible. + extension.bound_ctx.register_record_batches( + "bound_test", + [[pa.RecordBatch.from_pydict({"value": [1]})]], + ) + assert result.table_exist("bound_test") + + +def test_with_extensions_survives_source_collection(): + extension = _CodecOnlyExtension() + result = SessionContext().with_extensions(extension, _PlannerExtension()) + gc.collect() + + batches = result.sql("SELECT 1 AS value").collect() + assert batches[0].column(0) == pa.array([1]) + + +def test_with_extensions_failure_leaves_source_usable(ctx): + class BoomExtension: + def __datafusion_session_extension__(self, ctx): + msg = "boom" + raise RuntimeError(msg) + + with pytest.raises(RuntimeError, match="boom"): + ctx.with_extensions(_CodecOnlyExtension(), BoomExtension()) + + batches = ctx.sql("SELECT 1 AS value").collect() + assert batches[0].column(0) == pa.array([1]) + + def test_table_provider(ctx): batch = pa.RecordBatch.from_pydict({"x": [10, 20, 30]}) ctx.register_record_batches("provider_test", [[batch]]) diff --git a/python/tests/test_wrapper_coverage.py b/python/tests/test_wrapper_coverage.py index cf6719ecf..b1afd6832 100644 --- a/python/tests/test_wrapper_coverage.py +++ b/python/tests/test_wrapper_coverage.py @@ -67,6 +67,14 @@ def missing_exports(internal_obj, wrapped_obj) -> None: pytest.fail(f"Missing __repr__: {internal_obj.__name__}") for internal_attr_name in dir(internal_obj): + # Single-underscore names are private support methods for the + # wrappers (e.g. SessionContext._install_extensions) and are not + # part of the public surface that requires a wrapper. + if internal_attr_name.startswith("_") and not internal_attr_name.startswith( + "__" + ): + continue + wrapped_attr_name = internal_attr_name.removeprefix("Raw") assert wrapped_attr_name in dir(wrapped_obj)