You are viewing a plain text version of this content. The canonical link for it is here.
Posted to github@datafusion.apache.org by "timsaucer (via GitHub)" <gi...@apache.org> on 2026/08/07 19:21:18 UTC
[PR] Add FFI query planner support [datafusion-python]
timsaucer opened a new pull request, #1677:
URL: https://github.com/apache/datafusion-python/pull/1677
# Which issue does this PR close?
Related to #1612. This PR does not close it, but provides the FFI query planner plumbing that a `datafusion-distributed` integration can build on.
This is part 1 of 3 in the split of #1672 (part 2: composable extension codecs; part 3: atomic `SessionContext.with_extensions`).
# Rationale for this change
Extension libraries (for example distributed execution engines) need to supply their own `QueryPlanner` to a `SessionContext` without compiling against the `datafusion-python` crate. This PR exposes the query planner over the FFI boundary, following the same PyCapsule pattern used for table providers and catalogs.
# What changes are included in this PR?
- `SessionContext.with_query_planner(planner)` installs a planner exported via a `__datafusion_query_planner__` PyCapsule, preserving existing session state and codec settings.
- `SessionContext.__datafusion_query_planner__()` exports the current planner so another planner can wrap it as an explicit fallback (a session holds exactly one planner; layering is explicit delegation).
- A `RuntimeAwareQueryPlanner` adapter binds foreign planners to the Tokio runtime owned by datafusion-python.
- New example crate `datafusion-ffi-query-planner-example` demonstrating a real three-library plan exchange (host, provider library, planner library as separate cdylibs), including session config transfer via `SessionConfig.with_extension`.
- New `docs/source/contributor-guide/ffi.md` sections covering the query planner capsule protocol and the three-library setup.
# Are there any user-facing changes?
New public APIs: `SessionContext.with_query_planner` and `SessionContext.__datafusion_query_planner__`. A new example crate ships under `examples/`. No breaking changes to existing APIs.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org
Re: [PR] Add FFI query planner support [datafusion-python]
Posted by "milenkovicm (via GitHub)" <gi...@apache.org>.
milenkovicm commented on code in PR #1677:
URL: https://github.com/apache/datafusion-python/pull/1677#discussion_r3875695314
##########
examples/datafusion-ffi-query-planner-example/src/planner.rs:
##########
@@ -0,0 +1,206 @@
+// 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::sync::Arc;
+use std::sync::atomic::{AtomicBool, AtomicUsize, Ordering};
+
+use async_trait::async_trait;
+use datafusion::execution::TaskContextProvider;
+use datafusion::execution::context::SessionContext;
+use datafusion::logical_expr::LogicalPlan;
+use datafusion::physical_plan::ExecutionPlan;
+use datafusion::physical_plan::limit::GlobalLimitExec;
+use datafusion::physical_planner::{DefaultPhysicalPlanner, PhysicalPlanner};
+use datafusion_catalog::default_table_source::source_as_provider;
+use datafusion_ffi::config::ExtensionOptionsFFIProvider;
+use datafusion_ffi::execution_plan::ForeignExecutionPlan;
+use datafusion_ffi::query_planner::FFI_QueryPlanner;
+use datafusion_ffi::session::ForeignSession;
+use datafusion_ffi::table_provider::ForeignTableProvider;
+use datafusion_proto::logical_plan::DefaultLogicalExtensionCodec;
+use datafusion_proto::physical_plan::DefaultPhysicalExtensionCodec;
+use datafusion_python_util::get_tokio_runtime;
+use datafusion_session::{QueryPlanner, Session};
+use pyo3::prelude::*;
+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,
+}
+
+fn logical_plan_has_foreign_provider(plan: &LogicalPlan) -> bool {
+ if let LogicalPlan::TableScan(scan) = plan
+ && let Ok(provider) = source_as_provider(&scan.source)
+ && provider.downcast_ref::<ForeignTableProvider>().is_some()
+ {
+ return true;
+ }
+ plan.inputs()
+ .iter()
+ .any(|input| logical_plan_has_foreign_provider(input))
+}
+
+fn physical_plan_has_foreign_plan(plan: &Arc<dyn ExecutionPlan>) -> bool {
+ plan.is::<ForeignExecutionPlan>()
+ || plan
+ .children()
+ .iter()
+ .any(|child| physical_plan_has_foreign_plan(child))
+}
+
+fn planner_config(session: &dyn Session) -> datafusion::common::Result<PlannerConfig> {
+ let options = session.config_options();
+
+ // 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
+ // replaced silently by PlannerConfig::default().
+ if let Some(entry) = options
+ .entries()
+ .into_iter()
+ .find(|entry| entry.key.ends_with("ffi_query_planner.max_rows"))
+ {
+ let value = entry.value.ok_or_else(|| {
+ datafusion::common::DataFusionError::Configuration(format!(
+ "{} must have a value",
+ entry.key
+ ))
+ })?;
+ let max_rows = value.parse::<usize>().map_err(|err| {
+ datafusion::common::DataFusionError::Configuration(format!(
+ "Invalid value '{value}' for {}: {err}",
+ entry.key
+ ))
+ })?;
+ if max_rows == 0 {
+ return Err(datafusion::common::DataFusionError::Configuration(
+ "ffi_query_planner.max_rows must be greater than zero".to_owned(),
+ ));
+ }
+ return Ok(PlannerConfig { max_rows });
+ }
+
+ Ok(options
+ .local_or_ffi_extension::<PlannerConfig>()
+ .unwrap_or_default())
+}
+
+#[derive(Debug)]
+struct DistributedQueryPlanner {
+ observations: Arc<PlannerObservations>,
+}
+
+#[async_trait]
+impl QueryPlanner for DistributedQueryPlanner {
+ async fn create_physical_plan(
+ &self,
+ logical_plan: &LogicalPlan,
+ session: &dyn Session,
+ ) -> datafusion::common::Result<Arc<dyn ExecutionPlan>> {
+ self.observations.plan_calls.fetch_add(1, Ordering::SeqCst);
+ self.observations
+ .foreign_session
+ .store(session.as_any().is::<ForeignSession>(), Ordering::SeqCst);
+ self.observations.foreign_provider.store(
+ logical_plan_has_foreign_provider(logical_plan),
+ Ordering::SeqCst,
+ );
+
+ let config = planner_config(session)?;
+ self.observations
+ .last_max_rows
+ .store(config.max_rows, Ordering::SeqCst);
+
+ let plan = DefaultPhysicalPlanner::default()
+ .create_physical_plan(logical_plan, session)
+ .await?;
+ self.observations
+ .foreign_plan
+ .store(physical_plan_has_foreign_plan(&plan), Ordering::SeqCst);
+
+ Ok(Arc::new(GlobalLimitExec::new(
+ plan,
+ 0,
+ Some(config.max_rows),
+ )))
+ }
+}
+
+#[pyclass(
+ from_py_object,
+ name = "MyQueryPlanner",
+ module = "datafusion_ffi_query_planner_example",
+ subclass
+)]
+#[derive(Debug, Default, Clone)]
+pub(crate) struct MyQueryPlanner {
+ observations: Arc<PlannerObservations>,
+}
+
+#[pymethods]
+impl MyQueryPlanner {
+ #[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)
+ }
+
+ fn __datafusion_query_planner__<'py>(
+ &self,
+ py: Python<'py>,
+ ) -> PyResult<Bound<'py, PyCapsule>> {
+ let planner: Arc<dyn QueryPlanner + Send + Sync> = Arc::new(DistributedQueryPlanner {
+ observations: Arc::clone(&self.observations),
+ });
+ let runtime = get_tokio_runtime().handle().clone();
+ let ctx_provider = Arc::new(SessionContext::new()) as Arc<dyn TaskContextProvider>;
Review Comment:
Thank you @timsaucer
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org
Re: [PR] Add FFI query planner support [datafusion-python]
Posted by "timsaucer (via GitHub)" <gi...@apache.org>.
timsaucer commented on code in PR #1677:
URL: https://github.com/apache/datafusion-python/pull/1677#discussion_r3789783139
##########
python/datafusion/context.py:
##########
@@ -1759,6 +1769,36 @@ def add_physical_optimizer_rule(
"""
self.ctx.add_physical_optimizer_rule(rule)
+ def with_query_planner(
Review Comment:
Good point! Also worth updating the skill to match this pattern
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org
Re: [PR] Add FFI query planner support [datafusion-python]
Posted by "milenkovicm (via GitHub)" <gi...@apache.org>.
milenkovicm commented on code in PR #1677:
URL: https://github.com/apache/datafusion-python/pull/1677#discussion_r3795851454
##########
examples/datafusion-ffi-query-planner-example/src/config.rs:
##########
@@ -0,0 +1,112 @@
+// 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::any::Any;
+
+use datafusion_common::config::{
+ ConfigEntry, ConfigExtension, ConfigField, ExtensionOptions, Visit,
+};
+use datafusion_common::{DataFusionError, config_err};
+use datafusion_ffi::config::extension_options::FFI_ExtensionOptions;
+use pyo3::exceptions::PyRuntimeError;
+use pyo3::prelude::*;
+use pyo3::types::PyCapsule;
+
+#[pyclass(
+ from_py_object,
+ name = "PlannerConfig",
+ module = "datafusion_ffi_query_planner_example",
+ subclass
+)]
+#[derive(Clone, Debug)]
+pub(crate) struct PlannerConfig {
+ pub max_rows: usize,
+}
+
+#[pymethods]
+impl PlannerConfig {
Review Comment:
Nit, `MyPlannerConfig` to have names aligned,
##########
examples/datafusion-ffi-query-planner-example/src/planner.rs:
##########
@@ -0,0 +1,206 @@
+// 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::sync::Arc;
+use std::sync::atomic::{AtomicBool, AtomicUsize, Ordering};
+
+use async_trait::async_trait;
+use datafusion::execution::TaskContextProvider;
+use datafusion::execution::context::SessionContext;
+use datafusion::logical_expr::LogicalPlan;
+use datafusion::physical_plan::ExecutionPlan;
+use datafusion::physical_plan::limit::GlobalLimitExec;
+use datafusion::physical_planner::{DefaultPhysicalPlanner, PhysicalPlanner};
+use datafusion_catalog::default_table_source::source_as_provider;
+use datafusion_ffi::config::ExtensionOptionsFFIProvider;
+use datafusion_ffi::execution_plan::ForeignExecutionPlan;
+use datafusion_ffi::query_planner::FFI_QueryPlanner;
+use datafusion_ffi::session::ForeignSession;
+use datafusion_ffi::table_provider::ForeignTableProvider;
+use datafusion_proto::logical_plan::DefaultLogicalExtensionCodec;
+use datafusion_proto::physical_plan::DefaultPhysicalExtensionCodec;
+use datafusion_python_util::get_tokio_runtime;
+use datafusion_session::{QueryPlanner, Session};
+use pyo3::prelude::*;
+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,
+}
+
+fn logical_plan_has_foreign_provider(plan: &LogicalPlan) -> bool {
+ if let LogicalPlan::TableScan(scan) = plan
+ && let Ok(provider) = source_as_provider(&scan.source)
+ && provider.downcast_ref::<ForeignTableProvider>().is_some()
+ {
+ return true;
+ }
+ plan.inputs()
+ .iter()
+ .any(|input| logical_plan_has_foreign_provider(input))
+}
+
+fn physical_plan_has_foreign_plan(plan: &Arc<dyn ExecutionPlan>) -> bool {
+ plan.is::<ForeignExecutionPlan>()
+ || plan
+ .children()
+ .iter()
+ .any(|child| physical_plan_has_foreign_plan(child))
+}
+
+fn planner_config(session: &dyn Session) -> datafusion::common::Result<PlannerConfig> {
+ let options = session.config_options();
+
+ // 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
+ // replaced silently by PlannerConfig::default().
+ if let Some(entry) = options
+ .entries()
+ .into_iter()
+ .find(|entry| entry.key.ends_with("ffi_query_planner.max_rows"))
+ {
+ let value = entry.value.ok_or_else(|| {
+ datafusion::common::DataFusionError::Configuration(format!(
+ "{} must have a value",
+ entry.key
+ ))
+ })?;
+ let max_rows = value.parse::<usize>().map_err(|err| {
+ datafusion::common::DataFusionError::Configuration(format!(
+ "Invalid value '{value}' for {}: {err}",
+ entry.key
+ ))
+ })?;
+ if max_rows == 0 {
+ return Err(datafusion::common::DataFusionError::Configuration(
+ "ffi_query_planner.max_rows must be greater than zero".to_owned(),
+ ));
+ }
+ return Ok(PlannerConfig { max_rows });
+ }
+
+ Ok(options
+ .local_or_ffi_extension::<PlannerConfig>()
+ .unwrap_or_default())
+}
+
+#[derive(Debug)]
+struct DistributedQueryPlanner {
+ observations: Arc<PlannerObservations>,
+}
+
+#[async_trait]
+impl QueryPlanner for DistributedQueryPlanner {
+ async fn create_physical_plan(
+ &self,
+ logical_plan: &LogicalPlan,
+ session: &dyn Session,
+ ) -> datafusion::common::Result<Arc<dyn ExecutionPlan>> {
+ self.observations.plan_calls.fetch_add(1, Ordering::SeqCst);
+ self.observations
+ .foreign_session
+ .store(session.as_any().is::<ForeignSession>(), Ordering::SeqCst);
+ self.observations.foreign_provider.store(
+ logical_plan_has_foreign_provider(logical_plan),
+ Ordering::SeqCst,
+ );
+
+ let config = planner_config(session)?;
+ self.observations
+ .last_max_rows
+ .store(config.max_rows, Ordering::SeqCst);
+
+ let plan = DefaultPhysicalPlanner::default()
+ .create_physical_plan(logical_plan, session)
+ .await?;
+ self.observations
+ .foreign_plan
+ .store(physical_plan_has_foreign_plan(&plan), Ordering::SeqCst);
+
+ Ok(Arc::new(GlobalLimitExec::new(
+ plan,
+ 0,
+ Some(config.max_rows),
+ )))
+ }
+}
+
+#[pyclass(
+ from_py_object,
+ name = "MyQueryPlanner",
+ module = "datafusion_ffi_query_planner_example",
+ subclass
+)]
+#[derive(Debug, Default, Clone)]
+pub(crate) struct MyQueryPlanner {
+ observations: Arc<PlannerObservations>,
+}
+
+#[pymethods]
+impl MyQueryPlanner {
+ #[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)
+ }
+
+ fn __datafusion_query_planner__<'py>(
+ &self,
+ py: Python<'py>,
+ ) -> PyResult<Bound<'py, PyCapsule>> {
+ let planner: Arc<dyn QueryPlanner + Send + Sync> = Arc::new(DistributedQueryPlanner {
+ observations: Arc::clone(&self.observations),
+ });
+ let runtime = get_tokio_runtime().handle().clone();
+ let ctx_provider = Arc::new(SessionContext::new()) as Arc<dyn TaskContextProvider>;
Review Comment:
is this session context be parameter of method call on the line `119` ? are those two different sessions ?
##########
examples/datafusion-ffi-query-planner-example/src/config.rs:
##########
@@ -0,0 +1,112 @@
+// 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::any::Any;
+
+use datafusion_common::config::{
+ ConfigEntry, ConfigExtension, ConfigField, ExtensionOptions, Visit,
+};
+use datafusion_common::{DataFusionError, config_err};
+use datafusion_ffi::config::extension_options::FFI_ExtensionOptions;
+use pyo3::exceptions::PyRuntimeError;
+use pyo3::prelude::*;
+use pyo3::types::PyCapsule;
+
+#[pyclass(
+ from_py_object,
+ name = "PlannerConfig",
+ module = "datafusion_ffi_query_planner_example",
+ subclass
+)]
+#[derive(Clone, Debug)]
+pub(crate) struct PlannerConfig {
+ pub max_rows: usize,
+}
Review Comment:
this is needed for ballista, thanks Tim for example
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org
Re: [PR] Add FFI query planner support [datafusion-python]
Posted by "ntjohnson1 (via GitHub)" <gi...@apache.org>.
ntjohnson1 commented on code in PR #1677:
URL: https://github.com/apache/datafusion-python/pull/1677#discussion_r3786169274
##########
python/datafusion/context.py:
##########
@@ -1759,6 +1769,36 @@ def add_physical_optimizer_rule(
"""
self.ctx.add_physical_optimizer_rule(rule)
+ def with_query_planner(
Review Comment:
Generally wonder if this builder pattern feels pythonic. Consistent with what's already here so no action requested. Didn't look at how many withs there are but
```python
ctx = SessionContext(config, planner)
```
feels a little more intuitive than
```python
ctx = SessionContext().with_query_planner(planner)
```
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org
Re: [PR] Add FFI query planner support [datafusion-python]
Posted by "timsaucer (via GitHub)" <gi...@apache.org>.
timsaucer commented on code in PR #1677:
URL: https://github.com/apache/datafusion-python/pull/1677#discussion_r3874565111
##########
examples/datafusion-ffi-query-planner-example/src/planner.rs:
##########
@@ -0,0 +1,206 @@
+// 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::sync::Arc;
+use std::sync::atomic::{AtomicBool, AtomicUsize, Ordering};
+
+use async_trait::async_trait;
+use datafusion::execution::TaskContextProvider;
+use datafusion::execution::context::SessionContext;
+use datafusion::logical_expr::LogicalPlan;
+use datafusion::physical_plan::ExecutionPlan;
+use datafusion::physical_plan::limit::GlobalLimitExec;
+use datafusion::physical_planner::{DefaultPhysicalPlanner, PhysicalPlanner};
+use datafusion_catalog::default_table_source::source_as_provider;
+use datafusion_ffi::config::ExtensionOptionsFFIProvider;
+use datafusion_ffi::execution_plan::ForeignExecutionPlan;
+use datafusion_ffi::query_planner::FFI_QueryPlanner;
+use datafusion_ffi::session::ForeignSession;
+use datafusion_ffi::table_provider::ForeignTableProvider;
+use datafusion_proto::logical_plan::DefaultLogicalExtensionCodec;
+use datafusion_proto::physical_plan::DefaultPhysicalExtensionCodec;
+use datafusion_python_util::get_tokio_runtime;
+use datafusion_session::{QueryPlanner, Session};
+use pyo3::prelude::*;
+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,
+}
+
+fn logical_plan_has_foreign_provider(plan: &LogicalPlan) -> bool {
+ if let LogicalPlan::TableScan(scan) = plan
+ && let Ok(provider) = source_as_provider(&scan.source)
+ && provider.downcast_ref::<ForeignTableProvider>().is_some()
+ {
+ return true;
+ }
+ plan.inputs()
+ .iter()
+ .any(|input| logical_plan_has_foreign_provider(input))
+}
+
+fn physical_plan_has_foreign_plan(plan: &Arc<dyn ExecutionPlan>) -> bool {
+ plan.is::<ForeignExecutionPlan>()
+ || plan
+ .children()
+ .iter()
+ .any(|child| physical_plan_has_foreign_plan(child))
+}
+
+fn planner_config(session: &dyn Session) -> datafusion::common::Result<PlannerConfig> {
+ let options = session.config_options();
+
+ // 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
+ // replaced silently by PlannerConfig::default().
+ if let Some(entry) = options
+ .entries()
+ .into_iter()
+ .find(|entry| entry.key.ends_with("ffi_query_planner.max_rows"))
+ {
+ let value = entry.value.ok_or_else(|| {
+ datafusion::common::DataFusionError::Configuration(format!(
+ "{} must have a value",
+ entry.key
+ ))
+ })?;
+ let max_rows = value.parse::<usize>().map_err(|err| {
+ datafusion::common::DataFusionError::Configuration(format!(
+ "Invalid value '{value}' for {}: {err}",
+ entry.key
+ ))
+ })?;
+ if max_rows == 0 {
+ return Err(datafusion::common::DataFusionError::Configuration(
+ "ffi_query_planner.max_rows must be greater than zero".to_owned(),
+ ));
+ }
+ return Ok(PlannerConfig { max_rows });
+ }
+
+ Ok(options
+ .local_or_ffi_extension::<PlannerConfig>()
+ .unwrap_or_default())
+}
+
+#[derive(Debug)]
+struct DistributedQueryPlanner {
+ observations: Arc<PlannerObservations>,
+}
+
+#[async_trait]
+impl QueryPlanner for DistributedQueryPlanner {
+ async fn create_physical_plan(
+ &self,
+ logical_plan: &LogicalPlan,
+ session: &dyn Session,
+ ) -> datafusion::common::Result<Arc<dyn ExecutionPlan>> {
+ self.observations.plan_calls.fetch_add(1, Ordering::SeqCst);
+ self.observations
+ .foreign_session
+ .store(session.as_any().is::<ForeignSession>(), Ordering::SeqCst);
+ self.observations.foreign_provider.store(
+ logical_plan_has_foreign_provider(logical_plan),
+ Ordering::SeqCst,
+ );
+
+ let config = planner_config(session)?;
+ self.observations
+ .last_max_rows
+ .store(config.max_rows, Ordering::SeqCst);
+
+ let plan = DefaultPhysicalPlanner::default()
+ .create_physical_plan(logical_plan, session)
+ .await?;
+ self.observations
+ .foreign_plan
+ .store(physical_plan_has_foreign_plan(&plan), Ordering::SeqCst);
+
+ Ok(Arc::new(GlobalLimitExec::new(
+ plan,
+ 0,
+ Some(config.max_rows),
+ )))
+ }
+}
+
+#[pyclass(
+ from_py_object,
+ name = "MyQueryPlanner",
+ module = "datafusion_ffi_query_planner_example",
+ subclass
+)]
+#[derive(Debug, Default, Clone)]
+pub(crate) struct MyQueryPlanner {
+ observations: Arc<PlannerObservations>,
+}
+
+#[pymethods]
+impl MyQueryPlanner {
+ #[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)
+ }
+
+ fn __datafusion_query_planner__<'py>(
+ &self,
+ py: Python<'py>,
+ ) -> PyResult<Bound<'py, PyCapsule>> {
+ let planner: Arc<dyn QueryPlanner + Send + Sync> = Arc::new(DistributedQueryPlanner {
+ observations: Arc::clone(&self.observations),
+ });
+ let runtime = get_tokio_runtime().handle().clone();
+ let ctx_provider = Arc::new(SessionContext::new()) as Arc<dyn TaskContextProvider>;
Review Comment:
Really good catch! This led me down a rabbit hole and I ended up needing two upstream fixes:
- https://github.com/apache/datafusion/issues/24688
- https://github.com/apache/datafusion/issues/24722
In the latest push we no longer create this session context just for the codecs.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org
Re: [PR] Add FFI query planner support [datafusion-python]
Posted by "timsaucer (via GitHub)" <gi...@apache.org>.
timsaucer commented on code in PR #1677:
URL: https://github.com/apache/datafusion-python/pull/1677#discussion_r3738310240
##########
.github/workflows/build.yml:
##########
@@ -196,6 +196,16 @@ jobs:
args: --out dist
rustup-components: rust-std
+ - name: Build FFI query planner test library
+ if: matrix.python-tag == 'abi3'
+ uses: PyO3/maturin-action@v1
+ with:
+ target: x86_64-unknown-linux-gnu
+ manylinux: "2_28"
+ working-directory: examples/datafusion-ffi-query-planner-example
+ args: --out dist
+ rustup-components: rust-std
Review Comment:
In order to prove that the 3 library approach works where we have different codecs and different execution plans provided, we are adding a second test library. This way we can make sure there is no accidental ability to reach into a foreign code block.
##########
crates/core/src/context.rs:
##########
@@ -1385,6 +1453,19 @@ impl PySessionContext {
create_logical_extension_capsule(py, ffi.as_ref())
}
+ pub fn __datafusion_query_planner__<'py>(
+ &self,
+ py: Python<'py>,
+ ) -> PyResult<Bound<'py, PyCapsule>> {
Review Comment:
We need our session context to export it's own query planner because we have a use case where one query planner can depend on another. This is already supported by datafusion-distributed, so we want to be certain we support it here.
##########
crates/core/src/context.rs:
##########
@@ -221,6 +227,42 @@ impl PySessionConfig {
}
}
+/// Adapts an FFI planner to the Tokio runtime owned by datafusion-python.
+///
+/// Upstream's `ForeignQueryPlanner` cannot recover the runtime handle from the
+/// `QueryPlanner` trait, so embedders that own the runtime must call
+/// `create_physical_plan_with_session_runtime` directly.
+#[derive(Debug, Clone)]
+struct RuntimeAwareQueryPlanner {
+ planner: FFI_QueryPlanner,
+}
Review Comment:
As the docstring says, the purpose of this is to make sure we attach the runtime handle when needed.
##########
crates/core/src/context.rs:
##########
@@ -1211,6 +1253,32 @@ impl PySessionContext {
Ok(())
}
+ pub fn with_query_planner(&self, planner: Bound<'_, PyAny>) -> PyDataFusionResult<Self> {
Review Comment:
This API is the main reason for this PR. Here we allow changing out the default query planner with a user provided query planner.
##########
examples/datafusion-ffi-query-planner-example/src/config.rs:
##########
@@ -0,0 +1,112 @@
+// 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::any::Any;
+
+use datafusion_common::config::{
+ ConfigEntry, ConfigExtension, ConfigField, ExtensionOptions, Visit,
+};
+use datafusion_common::{DataFusionError, config_err};
+use datafusion_ffi::config::extension_options::FFI_ExtensionOptions;
+use pyo3::exceptions::PyRuntimeError;
+use pyo3::prelude::*;
+use pyo3::types::PyCapsule;
+
+#[pyclass(
+ from_py_object,
+ name = "PlannerConfig",
+ module = "datafusion_ffi_query_planner_example",
+ subclass
+)]
+#[derive(Clone, Debug)]
+pub(crate) struct PlannerConfig {
+ pub max_rows: usize,
+}
Review Comment:
I'm adding this to the query planner example because it's a very common pattern that we will need custom configs for the query planner, so it is reasonable to need insurance that configs pass over the FFI boundary properly and to use as a demonstration to anyone who is providing such a library.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org
Re: [PR] Add FFI query planner support [datafusion-python]
Posted by "timsaucer (via GitHub)" <gi...@apache.org>.
timsaucer commented on PR #1677:
URL: https://github.com/apache/datafusion-python/pull/1677#issuecomment-5457153455
Definitely a bit of scope creep, but I think we've got a huge addition here! I'll ship once CI is green.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org
Re: [PR] Add FFI query planner support [datafusion-python]
Posted by "timsaucer (via GitHub)" <gi...@apache.org>.
timsaucer merged PR #1677:
URL: https://github.com/apache/datafusion-python/pull/1677
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org