gabotechs commented on code in PR #24028:
URL: https://github.com/apache/datafusion/pull/24028#discussion_r3711145744
##########
datafusion/ffi/src/session/mod.rs:
##########
@@ -355,6 +421,21 @@ impl FFI_SessionRef {
session: &(dyn Session + Send + Sync),
runtime: Option<Handle>,
logical_codec: FFI_LogicalExtensionCodec,
+ ) -> Self {
+ let physical_codec = FFI_PhysicalExtensionCodec::new(
+ Arc::new(DefaultPhysicalExtensionCodec {}),
+ runtime.clone(),
+ logical_codec.task_ctx_provider.clone(),
+ );
+ Self::new_with_ffi_codecs(session, runtime, logical_codec,
physical_codec)
+ }
+
+ /// Creates a new [`FFI_SessionRef`] using existing FFI codecs.
+ pub fn new_with_ffi_codecs(
+ session: &(dyn Session + Send + Sync),
Review Comment:
Session is already `Send` and `Sync`, do we need these extra constraints
here?
##########
datafusion/ffi/src/query_planner.rs:
##########
@@ -0,0 +1,416 @@
+// 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.
+
+//! FFI support for [`QueryPlanner`].
+//!
+//! A typical deployment has three libraries. Library A (for example,
+//! `datafusion-python`) owns the [`Session`] and codec registry. Library B
owns
+//! a custom table provider and its extension nodes. Library C (for example,
+//! Ballista or `datafusion-distributed`) owns the query planner. A serializes
a
+//! logical plan and invokes C, while `FFI_SessionRef` lets C call session
+//! services in A. C deserializes the logical plan, creates a physical plan,
+//! serializes that result, and returns it for A to deserialize. The logical
and
+//! physical extension codecs preserve nodes supplied by B.
+//!
+//! The physical result is serialized instead of returned as an
+//! [`crate::execution_plan::FFI_ExecutionPlan`]. An FFI execution-plan handle
is
+//! a foreign trait-object proxy, so even a built-in plan created in C cannot
be
+//! downcast to its concrete
+//! type in A. Serialization reconstructs known plan nodes with A's local Rust
+//! type identities, allowing A's optimizers and other consumers to downcast
+//! them. Extension codecs control how custom nodes are reconstructed.
+//!
+//! A node returned by B while C is planning is still foreign to C unless a
+//! codec boundary reconstructs it in C. The query-planner boundary guarantees
+//! that C-local serializable nodes, and extension nodes understood by the
+//! configured codecs, are reconstructed for A when the completed plan returns.
+//!
+//! # Delegating back to library A
+//!
+//! C commonly wants A's built-in planning as a starting point, then rewrites
the
+//! result. A must export its planner *before* installing C's planner on the
+//! session, and C must retain that handle: after the swap,
+//! [`Session::query_planner`] reports C's own planner, and
+//! [`Session::create_physical_plan`] dispatches to it, so either one is a
+//! self-call. Delegating to the retained handle is safe, because DataFusion's
+//! built-in physical planner never re-dispatches through [`Session`].
+//!
+//! Retain the planner rather than the session. [`FFI_QueryPlanner`] owns a
+//! reference-counted planner, so it outlives A's original session, whereas
+//! `FFI_SessionRef` borrows its session with the lifetime erased.
+
+use std::ffi::c_void;
+use std::sync::Arc;
+
+use async_ffi::{FfiFuture, FutureExt};
+use async_trait::async_trait;
+use datafusion_common::error::{DataFusionError, Result};
+use datafusion_expr::LogicalPlan;
+use datafusion_physical_plan::ExecutionPlan;
+use datafusion_proto::bytes::{
+ logical_plan_from_bytes_with_extension_codec,
+ logical_plan_to_bytes_with_extension_codec,
+ physical_plan_from_bytes_with_extension_codec,
+ physical_plan_to_bytes_with_extension_codec,
+};
+use datafusion_proto::logical_plan::{
+ DefaultLogicalExtensionCodec, LogicalExtensionCodec,
+};
+use datafusion_proto::physical_plan::{
+ DefaultPhysicalExtensionCodec, PhysicalExtensionCodec,
+};
+use datafusion_session::{QueryPlanner, Session};
+use stabby::vec::Vec as SVec;
+use tokio::runtime::Handle;
+
+use crate::execution::FFI_TaskContextProvider;
+use crate::proto::logical_extension_codec::FFI_LogicalExtensionCodec;
+use crate::proto::physical_extension_codec::FFI_PhysicalExtensionCodec;
+use crate::session::{FFI_SessionRef, ForeignSession};
+use crate::util::FFI_Result;
+use crate::{df_result, sresult_return};
+
+/// An ABI-stable handle to a [`QueryPlanner`] owned by another library.
+///
+/// The Rust-facing adapters serialize the input [`LogicalPlan`] and resulting
+/// [`ExecutionPlan`]; callers do not invoke the byte-oriented function pointer
+/// directly.
+#[repr(C)]
+#[derive(Debug)]
+pub struct FFI_QueryPlanner {
+ create_physical_plan: unsafe extern "C" fn(
+ &Self,
+ logical_plan_serialized: SVec<u8>,
+ session: FFI_SessionRef,
+ ) -> FfiFuture<FFI_Result<SVec<u8>>>,
+
+ /// Codec used to encode and decode logical plans and extension nodes.
+ pub logical_codec: FFI_LogicalExtensionCodec,
+
+ /// Codec used to encode and decode physical plans and extension nodes.
+ pub physical_codec: FFI_PhysicalExtensionCodec,
+
+ /// Used to create a clone of the query planner.
+ clone: unsafe extern "C" fn(planner: &Self) -> Self,
+
+ /// Release the memory of the private data when it is no longer being used.
+ release: unsafe extern "C" fn(arg: &mut Self),
+
+ /// Return the major DataFusion version number of this planner.
+ pub version: unsafe extern "C" fn() -> u64,
+
+ /// Internal data. This is only to be accessed by the provider of the
planner.
+ /// A [`ForeignQueryPlanner`] should never attempt to access this data.
+ private_data: *mut c_void,
+
+ /// Utility to identify when FFI objects are accessed locally through
+ /// the foreign interface. See [`crate::get_library_marker_id`].
+ pub library_marker_id: extern "C" fn() -> usize,
+}
+
+unsafe impl Send for FFI_QueryPlanner {}
+unsafe impl Sync for FFI_QueryPlanner {}
+
+struct QueryPlannerPrivateData {
+ planner: Arc<dyn QueryPlanner + Send + Sync>,
+}
+
+impl FFI_QueryPlanner {
+ fn inner(&self) -> &Arc<dyn QueryPlanner + Send + Sync> {
+ let private_data = self.private_data as *const QueryPlannerPrivateData;
+ unsafe { &(*private_data).planner }
+ }
+}
+
+unsafe extern "C" fn create_physical_plan_fn_wrapper(
+ planner: &FFI_QueryPlanner,
+ logical_plan_serialized: SVec<u8>,
+ session: FFI_SessionRef,
+) -> FfiFuture<FFI_Result<SVec<u8>>> {
+ let internal_planner = Arc::clone(planner.inner());
+ let logical_codec: Arc<dyn LogicalExtensionCodec> =
(&planner.logical_codec).into();
+ let physical_codec: Arc<dyn PhysicalExtensionCodec> =
+ (&planner.physical_codec).into();
+
+ async move {
+ let mut foreign_session = None;
+ let session = sresult_return!(
+ session
+ .as_local()
+ .map(Ok::<&(dyn Session + Send + Sync), DataFusionError>)
+ .unwrap_or_else(|| {
+ foreign_session =
Some(ForeignSession::try_from(&session)?);
+ Ok(foreign_session.as_ref().unwrap())
+ })
+ );
+
+ let logical_plan =
sresult_return!(logical_plan_from_bytes_with_extension_codec(
+ logical_plan_serialized.as_slice(),
+ session.task_ctx().as_ref(),
+ logical_codec.as_ref(),
+ ));
+
+ let physical_plan = sresult_return!(
+ internal_planner
+ .create_physical_plan(&logical_plan, session)
+ .await
+ );
+ let physical_plan =
sresult_return!(physical_plan_to_bytes_with_extension_codec(
+ physical_plan,
+ physical_codec.as_ref(),
+ ));
+
+ FFI_Result::Ok(physical_plan.iter().copied().collect())
+ }
+ .into_ffi()
+}
+
+unsafe extern "C" fn release_fn_wrapper(planner: &mut FFI_QueryPlanner) {
+ unsafe {
+ debug_assert!(!planner.private_data.is_null());
+ let private_data =
+ Box::from_raw(planner.private_data as *mut
QueryPlannerPrivateData);
+ drop(private_data);
+ planner.private_data = std::ptr::null_mut();
+ }
+}
+
+unsafe extern "C" fn clone_fn_wrapper(planner: &FFI_QueryPlanner) ->
FFI_QueryPlanner {
+ let old_planner = Arc::clone(planner.inner());
+
+ let private_data = Box::into_raw(Box::new(QueryPlannerPrivateData {
+ planner: old_planner,
+ })) as *mut c_void;
+
+ FFI_QueryPlanner {
+ create_physical_plan: create_physical_plan_fn_wrapper,
+ logical_codec: planner.logical_codec.clone(),
+ physical_codec: planner.physical_codec.clone(),
+ clone: clone_fn_wrapper,
+ release: release_fn_wrapper,
+ version: super::version,
+ private_data,
+ library_marker_id: crate::get_library_marker_id,
+ }
+}
+
+impl Drop for FFI_QueryPlanner {
+ fn drop(&mut self) {
+ unsafe { (self.release)(self) }
+ }
+}
+
+impl Clone for FFI_QueryPlanner {
+ fn clone(&self) -> Self {
+ unsafe { (self.clone)(self) }
+ }
+}
+
+impl FFI_QueryPlanner {
+ /// Creates an [`FFI_QueryPlanner`] with native extension codecs.
+ ///
+ /// Missing codecs use DataFusion's defaults. `runtime` and
+ /// `task_ctx_provider` support codec callbacks across the FFI boundary.
+ pub fn new(
+ planner: Arc<dyn QueryPlanner + Send + Sync>,
+ runtime: Option<Handle>,
+ task_ctx_provider: impl Into<FFI_TaskContextProvider>,
+ logical_codec: Option<Arc<dyn LogicalExtensionCodec>>,
+ physical_codec: Option<Arc<dyn PhysicalExtensionCodec + Send>>,
+ ) -> Self {
+ let task_ctx_provider = task_ctx_provider.into();
+ let logical_codec =
+ logical_codec.unwrap_or_else(||
Arc::new(DefaultLogicalExtensionCodec {}));
+ let logical_codec = FFI_LogicalExtensionCodec::new(
+ logical_codec,
+ runtime.clone(),
+ task_ctx_provider.clone(),
+ );
+ let physical_codec =
+ physical_codec.unwrap_or_else(||
Arc::new(DefaultPhysicalExtensionCodec {}));
+ let physical_codec =
+ FFI_PhysicalExtensionCodec::new(physical_codec, runtime,
task_ctx_provider);
+ Self::new_with_ffi_codecs(planner, logical_codec, physical_codec)
+ }
+
+ /// Creates an [`FFI_QueryPlanner`] using prebuilt FFI extension codecs.
+ ///
+ /// If `planner` is already foreign, this returns its original FFI handle
+ /// rather than adding another wrapper layer.
+ pub fn new_with_ffi_codecs(
+ planner: Arc<dyn QueryPlanner + Send + Sync>,
+ logical_codec: FFI_LogicalExtensionCodec,
+ physical_codec: FFI_PhysicalExtensionCodec,
+ ) -> Self {
+ let any_ref: &dyn std::any::Any = planner.as_ref();
+ if let Some(planner) = any_ref.downcast_ref::<ForeignQueryPlanner>() {
+ return planner.0.clone();
Review Comment:
Isn't this dropping the provided `logical_codec` and `physical_codec`? Isn't
this actually a problem?
##########
datafusion/ffi/src/query_planner.rs:
##########
@@ -0,0 +1,416 @@
+// 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.
+
+//! FFI support for [`QueryPlanner`].
+//!
+//! A typical deployment has three libraries. Library A (for example,
+//! `datafusion-python`) owns the [`Session`] and codec registry. Library B
owns
+//! a custom table provider and its extension nodes. Library C (for example,
+//! Ballista or `datafusion-distributed`) owns the query planner. A serializes
a
+//! logical plan and invokes C, while `FFI_SessionRef` lets C call session
+//! services in A. C deserializes the logical plan, creates a physical plan,
+//! serializes that result, and returns it for A to deserialize. The logical
and
+//! physical extension codecs preserve nodes supplied by B.
+//!
+//! The physical result is serialized instead of returned as an
+//! [`crate::execution_plan::FFI_ExecutionPlan`]. An FFI execution-plan handle
is
+//! a foreign trait-object proxy, so even a built-in plan created in C cannot
be
+//! downcast to its concrete
+//! type in A. Serialization reconstructs known plan nodes with A's local Rust
+//! type identities, allowing A's optimizers and other consumers to downcast
+//! them. Extension codecs control how custom nodes are reconstructed.
+//!
+//! A node returned by B while C is planning is still foreign to C unless a
+//! codec boundary reconstructs it in C. The query-planner boundary guarantees
+//! that C-local serializable nodes, and extension nodes understood by the
+//! configured codecs, are reconstructed for A when the completed plan returns.
+//!
+//! # Delegating back to library A
+//!
+//! C commonly wants A's built-in planning as a starting point, then rewrites
the
+//! result. A must export its planner *before* installing C's planner on the
+//! session, and C must retain that handle: after the swap,
+//! [`Session::query_planner`] reports C's own planner, and
+//! [`Session::create_physical_plan`] dispatches to it, so either one is a
+//! self-call. Delegating to the retained handle is safe, because DataFusion's
+//! built-in physical planner never re-dispatches through [`Session`].
+//!
+//! Retain the planner rather than the session. [`FFI_QueryPlanner`] owns a
+//! reference-counted planner, so it outlives A's original session, whereas
+//! `FFI_SessionRef` borrows its session with the lifetime erased.
+
+use std::ffi::c_void;
+use std::sync::Arc;
+
+use async_ffi::{FfiFuture, FutureExt};
+use async_trait::async_trait;
+use datafusion_common::error::{DataFusionError, Result};
+use datafusion_expr::LogicalPlan;
+use datafusion_physical_plan::ExecutionPlan;
+use datafusion_proto::bytes::{
+ logical_plan_from_bytes_with_extension_codec,
+ logical_plan_to_bytes_with_extension_codec,
+ physical_plan_from_bytes_with_extension_codec,
+ physical_plan_to_bytes_with_extension_codec,
+};
+use datafusion_proto::logical_plan::{
+ DefaultLogicalExtensionCodec, LogicalExtensionCodec,
+};
+use datafusion_proto::physical_plan::{
+ DefaultPhysicalExtensionCodec, PhysicalExtensionCodec,
+};
+use datafusion_session::{QueryPlanner, Session};
+use stabby::vec::Vec as SVec;
+use tokio::runtime::Handle;
+
+use crate::execution::FFI_TaskContextProvider;
+use crate::proto::logical_extension_codec::FFI_LogicalExtensionCodec;
+use crate::proto::physical_extension_codec::FFI_PhysicalExtensionCodec;
+use crate::session::{FFI_SessionRef, ForeignSession};
+use crate::util::FFI_Result;
+use crate::{df_result, sresult_return};
+
+/// An ABI-stable handle to a [`QueryPlanner`] owned by another library.
+///
+/// The Rust-facing adapters serialize the input [`LogicalPlan`] and resulting
+/// [`ExecutionPlan`]; callers do not invoke the byte-oriented function pointer
+/// directly.
+#[repr(C)]
+#[derive(Debug)]
+pub struct FFI_QueryPlanner {
+ create_physical_plan: unsafe extern "C" fn(
+ &Self,
+ logical_plan_serialized: SVec<u8>,
+ session: FFI_SessionRef,
+ ) -> FfiFuture<FFI_Result<SVec<u8>>>,
+
+ /// Codec used to encode and decode logical plans and extension nodes.
+ pub logical_codec: FFI_LogicalExtensionCodec,
+
+ /// Codec used to encode and decode physical plans and extension nodes.
+ pub physical_codec: FFI_PhysicalExtensionCodec,
+
+ /// Used to create a clone of the query planner.
+ clone: unsafe extern "C" fn(planner: &Self) -> Self,
+
+ /// Release the memory of the private data when it is no longer being used.
+ release: unsafe extern "C" fn(arg: &mut Self),
+
+ /// Return the major DataFusion version number of this planner.
+ pub version: unsafe extern "C" fn() -> u64,
+
+ /// Internal data. This is only to be accessed by the provider of the
planner.
+ /// A [`ForeignQueryPlanner`] should never attempt to access this data.
+ private_data: *mut c_void,
+
+ /// Utility to identify when FFI objects are accessed locally through
+ /// the foreign interface. See [`crate::get_library_marker_id`].
+ pub library_marker_id: extern "C" fn() -> usize,
+}
+
+unsafe impl Send for FFI_QueryPlanner {}
+unsafe impl Sync for FFI_QueryPlanner {}
+
+struct QueryPlannerPrivateData {
+ planner: Arc<dyn QueryPlanner + Send + Sync>,
+}
+
+impl FFI_QueryPlanner {
+ fn inner(&self) -> &Arc<dyn QueryPlanner + Send + Sync> {
+ let private_data = self.private_data as *const QueryPlannerPrivateData;
+ unsafe { &(*private_data).planner }
+ }
+}
+
+unsafe extern "C" fn create_physical_plan_fn_wrapper(
+ planner: &FFI_QueryPlanner,
+ logical_plan_serialized: SVec<u8>,
+ session: FFI_SessionRef,
+) -> FfiFuture<FFI_Result<SVec<u8>>> {
+ let internal_planner = Arc::clone(planner.inner());
+ let logical_codec: Arc<dyn LogicalExtensionCodec> =
(&planner.logical_codec).into();
+ let physical_codec: Arc<dyn PhysicalExtensionCodec> =
+ (&planner.physical_codec).into();
+
+ async move {
+ let mut foreign_session = None;
+ let session = sresult_return!(
+ session
+ .as_local()
+ .map(Ok::<&(dyn Session + Send + Sync), DataFusionError>)
+ .unwrap_or_else(|| {
+ foreign_session =
Some(ForeignSession::try_from(&session)?);
+ Ok(foreign_session.as_ref().unwrap())
+ })
+ );
+
+ let logical_plan =
sresult_return!(logical_plan_from_bytes_with_extension_codec(
+ logical_plan_serialized.as_slice(),
+ session.task_ctx().as_ref(),
+ logical_codec.as_ref(),
+ ));
+
+ let physical_plan = sresult_return!(
+ internal_planner
+ .create_physical_plan(&logical_plan, session)
+ .await
+ );
+ let physical_plan =
sresult_return!(physical_plan_to_bytes_with_extension_codec(
+ physical_plan,
+ physical_codec.as_ref(),
+ ));
+
+ FFI_Result::Ok(physical_plan.iter().copied().collect())
+ }
+ .into_ffi()
+}
+
+unsafe extern "C" fn release_fn_wrapper(planner: &mut FFI_QueryPlanner) {
+ unsafe {
+ debug_assert!(!planner.private_data.is_null());
+ let private_data =
+ Box::from_raw(planner.private_data as *mut
QueryPlannerPrivateData);
+ drop(private_data);
+ planner.private_data = std::ptr::null_mut();
+ }
+}
+
+unsafe extern "C" fn clone_fn_wrapper(planner: &FFI_QueryPlanner) ->
FFI_QueryPlanner {
+ let old_planner = Arc::clone(planner.inner());
+
+ let private_data = Box::into_raw(Box::new(QueryPlannerPrivateData {
+ planner: old_planner,
+ })) as *mut c_void;
+
+ FFI_QueryPlanner {
+ create_physical_plan: create_physical_plan_fn_wrapper,
+ logical_codec: planner.logical_codec.clone(),
+ physical_codec: planner.physical_codec.clone(),
+ clone: clone_fn_wrapper,
+ release: release_fn_wrapper,
+ version: super::version,
+ private_data,
+ library_marker_id: crate::get_library_marker_id,
+ }
+}
+
+impl Drop for FFI_QueryPlanner {
+ fn drop(&mut self) {
+ unsafe { (self.release)(self) }
+ }
+}
+
+impl Clone for FFI_QueryPlanner {
+ fn clone(&self) -> Self {
+ unsafe { (self.clone)(self) }
+ }
+}
+
+impl FFI_QueryPlanner {
+ /// Creates an [`FFI_QueryPlanner`] with native extension codecs.
+ ///
+ /// Missing codecs use DataFusion's defaults. `runtime` and
+ /// `task_ctx_provider` support codec callbacks across the FFI boundary.
+ pub fn new(
+ planner: Arc<dyn QueryPlanner + Send + Sync>,
+ runtime: Option<Handle>,
+ task_ctx_provider: impl Into<FFI_TaskContextProvider>,
+ logical_codec: Option<Arc<dyn LogicalExtensionCodec>>,
+ physical_codec: Option<Arc<dyn PhysicalExtensionCodec + Send>>,
Review Comment:
This `Send` requirement should not be necessary, `PhysicalExtensionCodec` is
already constrained to `Send`. It might actually give some problems to people
that already have a type erased `Arc<dyn PhysicalExtensionCodec>`.
##########
datafusion/ffi/src/query_planner.rs:
##########
@@ -0,0 +1,416 @@
+// 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.
+
+//! FFI support for [`QueryPlanner`].
+//!
+//! A typical deployment has three libraries. Library A (for example,
+//! `datafusion-python`) owns the [`Session`] and codec registry. Library B
owns
+//! a custom table provider and its extension nodes. Library C (for example,
+//! Ballista or `datafusion-distributed`) owns the query planner. A serializes
a
+//! logical plan and invokes C, while `FFI_SessionRef` lets C call session
+//! services in A. C deserializes the logical plan, creates a physical plan,
+//! serializes that result, and returns it for A to deserialize. The logical
and
+//! physical extension codecs preserve nodes supplied by B.
+//!
+//! The physical result is serialized instead of returned as an
+//! [`crate::execution_plan::FFI_ExecutionPlan`]. An FFI execution-plan handle
is
+//! a foreign trait-object proxy, so even a built-in plan created in C cannot
be
+//! downcast to its concrete
+//! type in A. Serialization reconstructs known plan nodes with A's local Rust
+//! type identities, allowing A's optimizers and other consumers to downcast
+//! them. Extension codecs control how custom nodes are reconstructed.
+//!
+//! A node returned by B while C is planning is still foreign to C unless a
+//! codec boundary reconstructs it in C. The query-planner boundary guarantees
+//! that C-local serializable nodes, and extension nodes understood by the
+//! configured codecs, are reconstructed for A when the completed plan returns.
+//!
+//! # Delegating back to library A
+//!
+//! C commonly wants A's built-in planning as a starting point, then rewrites
the
+//! result. A must export its planner *before* installing C's planner on the
+//! session, and C must retain that handle: after the swap,
+//! [`Session::query_planner`] reports C's own planner, and
+//! [`Session::create_physical_plan`] dispatches to it, so either one is a
+//! self-call. Delegating to the retained handle is safe, because DataFusion's
+//! built-in physical planner never re-dispatches through [`Session`].
+//!
+//! Retain the planner rather than the session. [`FFI_QueryPlanner`] owns a
+//! reference-counted planner, so it outlives A's original session, whereas
+//! `FFI_SessionRef` borrows its session with the lifetime erased.
+
+use std::ffi::c_void;
+use std::sync::Arc;
+
+use async_ffi::{FfiFuture, FutureExt};
+use async_trait::async_trait;
+use datafusion_common::error::{DataFusionError, Result};
+use datafusion_expr::LogicalPlan;
+use datafusion_physical_plan::ExecutionPlan;
+use datafusion_proto::bytes::{
+ logical_plan_from_bytes_with_extension_codec,
+ logical_plan_to_bytes_with_extension_codec,
+ physical_plan_from_bytes_with_extension_codec,
+ physical_plan_to_bytes_with_extension_codec,
+};
+use datafusion_proto::logical_plan::{
+ DefaultLogicalExtensionCodec, LogicalExtensionCodec,
+};
+use datafusion_proto::physical_plan::{
+ DefaultPhysicalExtensionCodec, PhysicalExtensionCodec,
+};
+use datafusion_session::{QueryPlanner, Session};
+use stabby::vec::Vec as SVec;
+use tokio::runtime::Handle;
+
+use crate::execution::FFI_TaskContextProvider;
+use crate::proto::logical_extension_codec::FFI_LogicalExtensionCodec;
+use crate::proto::physical_extension_codec::FFI_PhysicalExtensionCodec;
+use crate::session::{FFI_SessionRef, ForeignSession};
+use crate::util::FFI_Result;
+use crate::{df_result, sresult_return};
+
+/// An ABI-stable handle to a [`QueryPlanner`] owned by another library.
+///
+/// The Rust-facing adapters serialize the input [`LogicalPlan`] and resulting
+/// [`ExecutionPlan`]; callers do not invoke the byte-oriented function pointer
+/// directly.
+#[repr(C)]
+#[derive(Debug)]
+pub struct FFI_QueryPlanner {
+ create_physical_plan: unsafe extern "C" fn(
+ &Self,
+ logical_plan_serialized: SVec<u8>,
+ session: FFI_SessionRef,
+ ) -> FfiFuture<FFI_Result<SVec<u8>>>,
+
+ /// Codec used to encode and decode logical plans and extension nodes.
+ pub logical_codec: FFI_LogicalExtensionCodec,
+
+ /// Codec used to encode and decode physical plans and extension nodes.
+ pub physical_codec: FFI_PhysicalExtensionCodec,
+
+ /// Used to create a clone of the query planner.
+ clone: unsafe extern "C" fn(planner: &Self) -> Self,
+
+ /// Release the memory of the private data when it is no longer being used.
+ release: unsafe extern "C" fn(arg: &mut Self),
+
+ /// Return the major DataFusion version number of this planner.
+ pub version: unsafe extern "C" fn() -> u64,
+
+ /// Internal data. This is only to be accessed by the provider of the
planner.
+ /// A [`ForeignQueryPlanner`] should never attempt to access this data.
+ private_data: *mut c_void,
+
+ /// Utility to identify when FFI objects are accessed locally through
+ /// the foreign interface. See [`crate::get_library_marker_id`].
+ pub library_marker_id: extern "C" fn() -> usize,
+}
+
+unsafe impl Send for FFI_QueryPlanner {}
+unsafe impl Sync for FFI_QueryPlanner {}
+
+struct QueryPlannerPrivateData {
+ planner: Arc<dyn QueryPlanner + Send + Sync>,
+}
+
+impl FFI_QueryPlanner {
+ fn inner(&self) -> &Arc<dyn QueryPlanner + Send + Sync> {
+ let private_data = self.private_data as *const QueryPlannerPrivateData;
+ unsafe { &(*private_data).planner }
+ }
+}
+
+unsafe extern "C" fn create_physical_plan_fn_wrapper(
+ planner: &FFI_QueryPlanner,
+ logical_plan_serialized: SVec<u8>,
+ session: FFI_SessionRef,
+) -> FfiFuture<FFI_Result<SVec<u8>>> {
+ let internal_planner = Arc::clone(planner.inner());
+ let logical_codec: Arc<dyn LogicalExtensionCodec> =
(&planner.logical_codec).into();
+ let physical_codec: Arc<dyn PhysicalExtensionCodec> =
+ (&planner.physical_codec).into();
+
+ async move {
+ let mut foreign_session = None;
+ let session = sresult_return!(
+ session
+ .as_local()
+ .map(Ok::<&(dyn Session + Send + Sync), DataFusionError>)
+ .unwrap_or_else(|| {
+ foreign_session =
Some(ForeignSession::try_from(&session)?);
+ Ok(foreign_session.as_ref().unwrap())
+ })
+ );
+
+ let logical_plan =
sresult_return!(logical_plan_from_bytes_with_extension_codec(
+ logical_plan_serialized.as_slice(),
+ session.task_ctx().as_ref(),
+ logical_codec.as_ref(),
+ ));
+
+ let physical_plan = sresult_return!(
+ internal_planner
+ .create_physical_plan(&logical_plan, session)
+ .await
+ );
+ let physical_plan =
sresult_return!(physical_plan_to_bytes_with_extension_codec(
+ physical_plan,
+ physical_codec.as_ref(),
+ ));
+
+ FFI_Result::Ok(physical_plan.iter().copied().collect())
+ }
+ .into_ffi()
+}
+
+unsafe extern "C" fn release_fn_wrapper(planner: &mut FFI_QueryPlanner) {
+ unsafe {
+ debug_assert!(!planner.private_data.is_null());
+ let private_data =
+ Box::from_raw(planner.private_data as *mut
QueryPlannerPrivateData);
+ drop(private_data);
+ planner.private_data = std::ptr::null_mut();
+ }
+}
+
+unsafe extern "C" fn clone_fn_wrapper(planner: &FFI_QueryPlanner) ->
FFI_QueryPlanner {
+ let old_planner = Arc::clone(planner.inner());
+
+ let private_data = Box::into_raw(Box::new(QueryPlannerPrivateData {
+ planner: old_planner,
+ })) as *mut c_void;
+
+ FFI_QueryPlanner {
+ create_physical_plan: create_physical_plan_fn_wrapper,
+ logical_codec: planner.logical_codec.clone(),
+ physical_codec: planner.physical_codec.clone(),
+ clone: clone_fn_wrapper,
+ release: release_fn_wrapper,
+ version: super::version,
+ private_data,
+ library_marker_id: crate::get_library_marker_id,
+ }
+}
+
+impl Drop for FFI_QueryPlanner {
+ fn drop(&mut self) {
+ unsafe { (self.release)(self) }
+ }
+}
+
+impl Clone for FFI_QueryPlanner {
+ fn clone(&self) -> Self {
+ unsafe { (self.clone)(self) }
+ }
+}
+
+impl FFI_QueryPlanner {
+ /// Creates an [`FFI_QueryPlanner`] with native extension codecs.
+ ///
+ /// Missing codecs use DataFusion's defaults. `runtime` and
+ /// `task_ctx_provider` support codec callbacks across the FFI boundary.
+ pub fn new(
+ planner: Arc<dyn QueryPlanner + Send + Sync>,
+ runtime: Option<Handle>,
+ task_ctx_provider: impl Into<FFI_TaskContextProvider>,
+ logical_codec: Option<Arc<dyn LogicalExtensionCodec>>,
+ physical_codec: Option<Arc<dyn PhysicalExtensionCodec + Send>>,
+ ) -> Self {
+ let task_ctx_provider = task_ctx_provider.into();
+ let logical_codec =
+ logical_codec.unwrap_or_else(||
Arc::new(DefaultLogicalExtensionCodec {}));
+ let logical_codec = FFI_LogicalExtensionCodec::new(
+ logical_codec,
+ runtime.clone(),
+ task_ctx_provider.clone(),
+ );
+ let physical_codec =
+ physical_codec.unwrap_or_else(||
Arc::new(DefaultPhysicalExtensionCodec {}));
+ let physical_codec =
+ FFI_PhysicalExtensionCodec::new(physical_codec, runtime,
task_ctx_provider);
+ Self::new_with_ffi_codecs(planner, logical_codec, physical_codec)
Review Comment:
If `None` is passed to `logical_codec` or `physical_codec`, those arguments
would default to `DefaultLogicalExtensionCodec` and
`DefaultPhysicalExtensionCodec`, replacing any previously registered codecs
even if users provide `None` as arguments.
:thinking: there seems to be a couple of footguns with the
`FFI_QueryPlanner::new` and `FFI_QueryPlanner::new_with_ffi_codecs` public
APIs, maybe there's a cleaner constructor method we can expose publicly that
makes these things irrepresentable?
--
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: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
