diff --git a/examples/datafusion-ffi-example/src/foreign_plan_workaround.rs b/examples/datafusion-ffi-example/src/foreign_plan_workaround.rs new file mode 100644 index 000000000..cdca98150 --- /dev/null +++ b/examples/datafusion-ffi-example/src/foreign_plan_workaround.rs @@ -0,0 +1,89 @@ +// 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. + +//! # NOT A PATTERN +//! +//! **Blocked on:** +//! +//! **Delete when:** `FFI_PlanProperties` carries `scheduling_type`, or +//! `ForeignExecutionPlan` gains a reachable `try_to_proto`. +//! +//! **Copying this will:** claim every other library's plan nodes, and produce +//! payloads that decode only in the writing process, exactly once each. +//! +//! A provider owns `DataSourceExec`, but a host-added execution decorator can +//! wrap that scan in `ForeignExecutionPlan`. This workaround claims the opaque +//! wrapper because the native encoder cannot reach its underlying plan. It +//! parks the plan in a process-local registry rather than serializing it. +//! +//! The `DataSourceExec` arm could use durable metadata and does not, because +//! the registry must exist for the `ForeignExecutionPlan` arm regardless. +//! Splitting the two arms across wire formats costs real code and removes +//! nothing; this module keeps the temporary compromise obvious. + +use std::collections::HashMap; +use std::sync::atomic::{AtomicU64, Ordering}; +use std::sync::{Arc, Mutex, OnceLock}; + +use datafusion::common::{DataFusionError, Result}; +use datafusion::datasource::source::DataSourceExec; +use datafusion::physical_plan::ExecutionPlan; +use datafusion_ffi::execution_plan::ForeignExecutionPlan; + +const EXECUTION_PLAN_TOKEN: &[u8] = b"DFPYEXEP"; +static NEXT_EXECUTION_PLAN_ID: AtomicU64 = AtomicU64::new(1); +static EXECUTION_PLANS: OnceLock>>> = OnceLock::new(); + +fn execution_plans() -> &'static Mutex>> { + EXECUTION_PLANS.get_or_init(|| Mutex::new(HashMap::new())) +} + +fn token_id(buf: &[u8]) -> Option { + let id: [u8; 8] = buf.strip_prefix(EXECUTION_PLAN_TOKEN)?.try_into().ok()?; + Some(u64::from_le_bytes(id)) +} + +pub(crate) fn claims(node: &Arc) -> bool { + node.is::() || node.is::() +} + +pub(crate) fn park(node: Arc, buf: &mut Vec) -> Result<()> { + let id = NEXT_EXECUTION_PLAN_ID.fetch_add(1, Ordering::SeqCst); + execution_plans() + .lock() + .map_err(|err| DataFusionError::Internal(err.to_string()))? + .insert(id, node); + buf.extend_from_slice(EXECUTION_PLAN_TOKEN); + buf.extend_from_slice(&id.to_le_bytes()); + Ok(()) +} + +pub(crate) fn take(buf: &[u8]) -> Result>> { + let Some(id) = token_id(buf) else { + return Ok(None); + }; + let plan = execution_plans() + .lock() + .map_err(|err| DataFusionError::Internal(err.to_string()))? + .remove(&id) + .ok_or_else(|| { + DataFusionError::Internal(format!( + "Unknown datafusion-ffi-example execution plan token {id}" + )) + })?; + Ok(Some(plan)) +} diff --git a/examples/datafusion-ffi-example/src/lib.rs b/examples/datafusion-ffi-example/src/lib.rs index 92fccb1e2..2eba7e041 100644 --- a/examples/datafusion-ffi-example/src/lib.rs +++ b/examples/datafusion-ffi-example/src/lib.rs @@ -33,6 +33,7 @@ use crate::window_udf::MyRankUDF; pub(crate) mod aggregate_udf; pub(crate) mod catalog_provider; pub(crate) mod config; +pub(crate) mod foreign_plan_workaround; pub(crate) mod logical_extension_codec; pub(crate) mod name_only_codec; pub(crate) mod physical_extension_codec; diff --git a/examples/datafusion-ffi-example/src/physical_extension_codec.rs b/examples/datafusion-ffi-example/src/physical_extension_codec.rs index f9e96382e..905a77e45 100644 --- a/examples/datafusion-ffi-example/src/physical_extension_codec.rs +++ b/examples/datafusion-ffi-example/src/physical_extension_codec.rs @@ -15,17 +15,14 @@ // specific language governing permissions and limitations // under the License. -use std::collections::HashMap; use std::fmt; -use std::sync::atomic::{AtomicU64, AtomicUsize, Ordering}; -use std::sync::{Arc, Mutex, OnceLock}; +use std::sync::Arc; +use std::sync::atomic::{AtomicUsize, Ordering}; -use datafusion::common::{DataFusionError, Result}; -use datafusion::datasource::source::DataSourceExec; +use datafusion::common::Result; use datafusion::execution::TaskContext; use datafusion::logical_expr::ScalarUDF; use datafusion::physical_plan::ExecutionPlan; -use datafusion_ffi::execution_plan::ForeignExecutionPlan; use datafusion_ffi::proto::physical_extension_codec::FFI_PhysicalExtensionCodec; use datafusion_proto::physical_plan::{ DefaultPhysicalExtensionCodec, PhysicalExtensionCodec, PhysicalProtoConverterExtension, @@ -34,26 +31,9 @@ use datafusion_python_util::{ffi_task_context_provider_from_pycapsule, get_tokio use pyo3::prelude::*; use pyo3::types::PyCapsule; +use crate::foreign_plan_workaround; use crate::required_udf::{TaskContextProbe, resolve_required_udf}; -const EXECUTION_PLAN_TOKEN: &[u8] = b"DFPYEXEP"; -static NEXT_EXECUTION_PLAN_ID: AtomicU64 = AtomicU64::new(1); -static EXECUTION_PLANS: OnceLock>>> = OnceLock::new(); - -/// Execution-plan counterpart of the logical codec's provider registry, with -/// the same lifecycle: encoding inserts, decoding removes, so a decode -/// consumes its token and an encode that is never decoded leaks. See -/// [`crate::logical_extension_codec`] for why that is acceptable here and not -/// in a real codec. -fn execution_plans() -> &'static Mutex>> { - EXECUTION_PLANS.get_or_init(|| Mutex::new(HashMap::new())) -} - -fn token_id(buf: &[u8]) -> Option { - let id: [u8; 8] = buf.strip_prefix(EXECUTION_PLAN_TOKEN)?.try_into().ok()?; - Some(u64::from_le_bytes(id)) -} - #[derive(Debug, Default)] pub(crate) struct PhysicalCallCounters { pub encode_udf: AtomicUsize, @@ -99,19 +79,11 @@ impl PhysicalExtensionCodec for CountingPhysicalExtensionCodec { proto_converter: &dyn PhysicalProtoConverterExtension, ) -> Result> { resolve_required_udf(self.required_udf.as_deref(), ctx, &self.counters.task_ctx)?; - if let Some(id) = token_id(buf) { + if let Some(plan) = foreign_plan_workaround::take(buf)? { self.counters .decode_execution_plan .fetch_add(1, Ordering::SeqCst); - return execution_plans() - .lock() - .map_err(|err| DataFusionError::Internal(err.to_string()))? - .remove(&id) - .ok_or_else(|| { - DataFusionError::Internal(format!( - "Unknown datafusion-ffi-example execution plan token {id}" - )) - }); + return Ok(plan); } self.inner.try_decode(buf, inputs, ctx, proto_converter) } @@ -122,20 +94,11 @@ impl PhysicalExtensionCodec for CountingPhysicalExtensionCodec { buf: &mut Vec, proto_converter: &dyn PhysicalProtoConverterExtension, ) -> Result<()> { - // The provider owns DataSourceExec. A ForeignExecutionPlan can wrap a - // host-added execution decorator around that scan; retaining the opaque - // wrapper preserves its original library identity without downcasting it. - if node.is::() || node.is::() { + if foreign_plan_workaround::claims(&node) { self.counters .encode_execution_plan .fetch_add(1, Ordering::SeqCst); - let id = NEXT_EXECUTION_PLAN_ID.fetch_add(1, Ordering::SeqCst); - execution_plans() - .lock() - .map_err(|err| DataFusionError::Internal(err.to_string()))? - .insert(id, node); - buf.extend_from_slice(EXECUTION_PLAN_TOKEN); - buf.extend_from_slice(&id.to_le_bytes()); + foreign_plan_workaround::park(node, buf)?; return Ok(()); } self.inner.try_encode(node, buf, proto_converter)