This is an automated email from the ASF dual-hosted git repository. timsaucer pushed a commit to branch feat/ffi-with-extensions in repository https://gitbox.apache.org/repos/asf/datafusion-python.git
commit b644f97428ab57bf52a157103395f22bc99d4f88 Author: Tim Saucer <[email protected]> AuthorDate: Fri Aug 7 14:17:45 2026 -0400 Add extension-bundle example and with_extensions FFI tests MyPlannerExtension in the query-planner example crate implements the __datafusion_session_extension__ protocol from Rust: it extracts the destination context's task-context provider, binds fresh observing codecs and a planner to it, and returns SessionExtensionComponents. Its codecs record the max_rows config value resolved through the weak provider, letting tests prove the provider targets the returned context rather than the source. Documents with_extensions as the preferred API in the FFI guide. Co-Authored-By: Claude Fable 5 <[email protected]> --- docs/source/contributor-guide/ffi.md | 62 ++++- .../datafusion-ffi-query-planner-example/README.md | 17 +- .../tests/_test_three_library_query_planner.py | 170 ++++++++++++- .../src/extension.rs | 263 +++++++++++++++++++++ .../src/lib.rs | 3 + .../src/planner.rs | 22 +- 6 files changed, 523 insertions(+), 14 deletions(-) diff --git a/docs/source/contributor-guide/ffi.md b/docs/source/contributor-guide/ffi.md index 3b5b8b91..17c58826 100644 --- a/docs/source/contributor-guide/ffi.md +++ b/docs/source/contributor-guide/ffi.md @@ -283,8 +283,60 @@ planner are rebound automatically, but a planner wrapped inside another planner fallback is opaque and keeps the codecs it was exported with. **Install all extension codecs before exporting or chaining planners.** -Putting it together for a session using two extension libraries that each provide -tables, functions, and a query planner: +### Extension bundles: `with_extensions` + +Codec and planner capsules carry an `FFI_TaskContextProvider` holding a *weak* +reference to the `SessionContext` they were created against. A capsule does not keep +that context alive, and a component bound to one context cannot be rebound to another. +Chaining the low-level `with_*` methods by hand therefore risks binding components to +an intermediate context that is later garbage collected, which fails at query time +with `TaskContextProvider went out of scope over FFI boundary` — or worse, silently +reads stale session state. + +`SessionContext.with_extensions` avoids this by construction. An extension library +exposes a bundle object implementing the `__datafusion_session_extension__` protocol: + +```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 +codecs, binds the planner against the final codec chains, and returns the 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 prepend to the codec chain, so codecs from +later extensions are consulted first. At most one extension may supply a query +planner. If any factory fails, the source context's state is unchanged. + +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. + +`MyPlannerExtension` in [`datafusion-ffi-query-planner-example`] is a complete Rust +implementation of this protocol, including extracting the task-context provider from +the supplied context and constructing a Python `SessionExtensionComponents`. + +### Advanced: chaining the low-level methods + +The `with_logical_extension_codec`, `with_physical_extension_codec`, and +`with_query_planner` methods remain available for advanced use. Putting them together +for a session using two extension libraries that each provide tables, functions, and +a query planner: ```python ctx = SessionContext(config) @@ -307,6 +359,12 @@ ctx.register_table("t", lib_a.TableProvider()) ctx.register_udf(udf(lib_b.SomeUDF())) ``` +When chaining by hand, keep the final context assigned to `ctx` as the single owner: +components created against earlier intermediate contexts (for example a codec +constructed with a context that is later discarded) hold weak references that break +once that intermediate context is collected. Prefer `with_extensions` whenever the +extension library provides a bundle. + ## 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/README.md b/examples/datafusion-ffi-query-planner-example/README.md index 7e1ec899..55e89869 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(PlannerConfig(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(PlannerConfig(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 0991c5c9..5846661f 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 @@ -20,14 +20,23 @@ from __future__ import annotations import gc import pytest -from datafusion import SessionConfig, SessionContext, udf +from datafusion import ( + SessionConfig, + SessionContext, + SessionExtensionComponents, + udf, +) from datafusion_ffi_example import ( IsNullUDF, MyLogicalExtensionCodec, MyPhysicalExtensionCodec, MyTableProvider, ) -from datafusion_ffi_query_planner_example import MyQueryPlanner, PlannerConfig +from datafusion_ffi_query_planner_example import ( + MyPlannerExtension, + MyQueryPlanner, + PlannerConfig, +) def configured_context(max_rows: int): @@ -120,6 +129,163 @@ 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,), + ) + + +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(PlannerConfig(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(PlannerConfig(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 + + # Codec decode calls resolve the task context through the weak provider + # bound during with_extensions. Seeing 2 (never 5) proves the provider + # targets the returned context rather than the source. + seen = planner_ext.decode_max_rows_seen() + assert seen, "expected the bundle codecs to observe at least one decode" + assert set(seen) == {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(PlannerConfig(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(PlannerConfig(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 + seen = planner_ext.decode_max_rows_seen() + assert seen + assert set(seen) == {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(PlannerConfig(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(PlannerConfig(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(PlannerConfig(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_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 00000000..714d3b32 --- /dev/null +++ b/examples/datafusion-ffi-query-planner-example/src/extension.rs @@ -0,0 +1,263 @@ +// 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::ptr::NonNull; +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::get_tokio_runtime; +use datafusion_session::QueryPlanner; +use pyo3::prelude::*; +use pyo3::types::{PyCapsule, 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. Decoding reads the provider's current +/// session state, so these prove which context the provider targets. +type ObservedMaxRows = Arc<Mutex<Vec<usize>>>; + +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<Extension> { + record_task_ctx(&self.observed, ctx); + self.inner.try_decode(buf, inputs, ctx) + } + + fn try_encode(&self, node: &Extension, buf: &mut Vec<u8>) -> Result<()> { + self.inner.try_encode(node, buf) + } + + fn try_decode_table_provider( + &self, + buf: &[u8], + table_ref: &TableReference, + schema: SchemaRef, + ctx: &TaskContext, + ) -> Result<Arc<dyn TableProvider>> { + 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<dyn TableProvider>, + buf: &mut Vec<u8>, + ) -> 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<dyn ExecutionPlan>], + ctx: &TaskContext, + proto_converter: &dyn PhysicalProtoConverterExtension, + ) -> Result<Arc<dyn ExecutionPlan>> { + record_task_ctx(&self.observed, ctx); + self.inner.try_decode(buf, inputs, ctx, proto_converter) + } + + fn try_encode( + &self, + node: Arc<dyn ExecutionPlan>, + buf: &mut Vec<u8>, + proto_converter: &dyn PhysicalProtoConverterExtension, + ) -> Result<()> { + self.inner.try_encode(node, buf, proto_converter) + } +} + +fn task_ctx_provider_from_session(ctx: &Bound<'_, PyAny>) -> PyResult<FFI_TaskContextProvider> { + let capsule = ctx.call_method0("__datafusion_task_context_provider__")?; + let capsule = capsule.cast::<PyCapsule>()?; + let provider: NonNull<FFI_TaskContextProvider> = capsule + .pointer_checked(Some(c"datafusion_task_context_provider"))? + .cast(); + // The FFI provider holds a weak reference; cloning it does not keep the + // session context alive. The host's returned context is the strong owner. + Ok(unsafe { provider.as_ref() }.clone()) +} + +/// 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(Debug, Default, Clone)] +pub(crate) struct MyPlannerExtension { + observations: Arc<PlannerObservations>, + observed_max_rows: ObservedMaxRows, +} + +#[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. + fn decode_max_rows_seen(&self) -> Vec<usize> { + self.observed_max_rows + .lock() + .map(|observed| observed.clone()) + .unwrap_or_default() + } + + fn __datafusion_session_extension__<'py>( + &self, + py: Python<'py>, + ctx: Bound<'py, PyAny>, + ) -> PyResult<Bound<'py, PyAny>> { + // Bind every component to the destination context supplied by the + // host. Components must not be cached across calls: each installation + // targets a different context. + let provider = task_ctx_provider_from_session(&ctx)?; + let runtime = get_tokio_runtime().handle().clone(); + + let logical: Arc<dyn LogicalExtensionCodec> = 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 = + PyCapsule::new_with_value(py, ffi_logical, cr"datafusion_logical_extension_codec")?; + + let physical: Arc<dyn PhysicalExtensionCodec + Send> = + Arc::new(ObservingPhysicalExtensionCodec { + inner: DefaultPhysicalExtensionCodec {}, + observed: Arc::clone(&self.observed_max_rows), + }); + let ffi_physical = + FFI_PhysicalExtensionCodec::new(physical, Some(runtime.clone()), provider.clone()); + let physical_capsule = + PyCapsule::new_with_value(py, ffi_physical, cr"datafusion_physical_extension_codec")?; + + let planner: Arc<dyn QueryPlanner + Send + Sync> = Arc::new(DistributedQueryPlanner { + observations: Arc::clone(&self.observations), + }); + let ffi_planner = FFI_QueryPlanner::new( + planner, + Some(runtime), + provider, + Arc::new(DefaultLogicalExtensionCodec {}), + Arc::new(DefaultPhysicalExtensionCodec {}), + ); + let planner_capsule = + PyCapsule::new_with_value(py, ffi_planner, cr"datafusion_query_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 7635c299..ff8ce784 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::PlannerConfig; +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::<PlannerConfig>()?; + m.add_class::<MyPlannerExtension>()?; m.add_class::<MyQueryPlanner>()?; Ok(()) } diff --git a/examples/datafusion-ffi-query-planner-example/src/planner.rs b/examples/datafusion-ffi-query-planner-example/src/planner.rs index cb767ffa..3477f8c0 100644 --- a/examples/datafusion-ffi-query-planner-example/src/planner.rs +++ b/examples/datafusion-ffi-query-planner-example/src/planner.rs @@ -41,12 +41,12 @@ use pyo3::types::PyCapsule; use crate::config::PlannerConfig; #[derive(Debug, 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, } fn logical_plan_has_foreign_provider(plan: &LogicalPlan) -> bool { @@ -70,8 +70,12 @@ fn physical_plan_has_foreign_plan(plan: &Arc<dyn ExecutionPlan>) -> bool { } fn planner_config(session: &dyn Session) -> datafusion::common::Result<PlannerConfig> { - 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<PlannerConfig> { // Read the flattened entry first. Some DataFusion revisions add an extra // `datafusion_ffi` namespace while reconstructing a ForeignSession. Parsing // it directly also ensures malformed values are reported instead of being @@ -107,8 +111,8 @@ fn planner_config(session: &dyn Session) -> datafusion::common::Result<PlannerCo } #[derive(Debug)] -struct DistributedQueryPlanner { - observations: Arc<PlannerObservations>, +pub(crate) struct DistributedQueryPlanner { + pub(crate) observations: Arc<PlannerObservations>, } #[async_trait] --------------------------------------------------------------------- To unsubscribe, e-mail: [email protected] For additional commands, e-mail: [email protected]
