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..1565b25d9 --- /dev/null +++ b/examples/datafusion-ffi-example/src/foreign_plan_workaround.rs @@ -0,0 +1,116 @@ +// 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`. Fixing either one +//! retires this module. +//! +//! **Copying this will:** claim every other library's plan nodes, and produce +//! payloads that decode only in the writing process, exactly once each. +//! +//! `DataSourceExec` is this library's own node. The `ForeignExecutionPlan` +//! arm is a workaround, and it is load-bearing: a host physical optimizer +//! rule that runs during a foreign planner's `create_physical_plan` -- +//! `EnsureCooperative` always does -- hands the library back a +//! `ForeignExecutionPlan` wrapping the host's `CooperativeExec`. That type has +//! no reachable `try_to_proto`, so nothing can encode it natively and +//! `FFI_QueryPlanner` must serialize the plan it returns. Claiming the opaque +//! wrapper here and parking it in a process-local registry, rather than +//! serializing it, is what lets those plans round-trip at all. +//! +//! The cost is that this codec also claims every *other* library's nodes, +//! since that is the type any node arrives as once it has crossed the +//! boundary -- see `extension_codec_order`. Narrowing [`claims`] to +//! `DataSourceExec` alone makes 31 tests in +//! `datafusion-ffi-query-planner-example` fail because the wrapped plan can +//! no longer be encoded. +//! +//! A library whose planner controls its own physical optimizer rules never +//! sees a foreign node and needs no such arm. +//! +//! Both halves are upstream defects, tracked together in the issue above: +//! `FFI_PlanProperties` carries no `scheduling_type`, so `EnsureCooperative` +//! reads every foreign leaf as non-cooperative and wraps it, and the resulting +//! `ForeignExecutionPlan` then has no way to serialize itself. +//! +//! 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. +//! +//! The registry is the 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. + +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 810e090de..c8964eb76 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, @@ -69,9 +49,9 @@ pub(crate) struct PhysicalCallCounters { /// owning cdylib can restore their concrete Rust type after the plan travels /// through the independent query-planner and datafusion-python libraries. /// -/// See [`execution_plans`] for the token lifecycle, which is narrower than it -/// looks: a decode consumes its token, so the same encoded plan cannot be -/// decoded twice. +/// See [`crate::foreign_plan_workaround`] for the token lifecycle, which is +/// narrower than it looks: a decode consumes its token, so the same encoded +/// plan cannot be decoded twice. struct CountingPhysicalExtensionCodec { inner: DefaultPhysicalExtensionCodec, counters: Arc, @@ -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,42 +94,13 @@ impl PhysicalExtensionCodec for CountingPhysicalExtensionCodec { buf: &mut Vec, proto_converter: &dyn PhysicalProtoConverterExtension, ) -> Result<()> { - // `DataSourceExec` is this library's own node. The `ForeignExecutionPlan` - // arm is a workaround, not a pattern to copy, and it is load-bearing: - // a host physical optimizer rule that runs during a foreign planner's - // `create_physical_plan` -- `EnsureCooperative` always does -- hands the - // library back a `ForeignExecutionPlan` wrapping the host's - // `CooperativeExec`. That type has no reachable `try_to_proto`, so - // nothing can encode it natively and `FFI_QueryPlanner` must serialize - // the plan it returns. Claiming it here is what lets those plans - // round-trip at all. - // - // The cost is that this codec also claims every *other* library's - // nodes, since that is the type any node arrives as once it has crossed - // the boundary -- see `extension_codec_order`. Narrowing this to - // `DataSourceExec` alone makes 31 tests in - // `datafusion-ffi-query-planner-example` fail with the error above. - // - // A library whose planner controls its own physical optimizer rules - // never sees a foreign node and needs no such arm. - // - // Both halves are upstream defects, tracked together in - // https://github.com/apache/datafusion/issues/25152: `FFI_PlanProperties` - // carries no `scheduling_type`, so `EnsureCooperative` reads every - // foreign leaf as non-cooperative and wraps it, and the resulting - // `ForeignExecutionPlan` then has no way to serialize itself. Fixing - // either one retires this arm. - if node.is::() || node.is::() { + // See [`crate::foreign_plan_workaround`] for why the + // `ForeignExecutionPlan` arm exists and what retires it. + 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)