docs(example): quarantine foreign plan workaround

Move the temporary execution-plan token registry into a deliberately named module and document its process-local behavior and removal criteria.

The rationale that #1720 added to the try_encode comment (why the ForeignExecutionPlan arm is load-bearing, what it costs, and what retires it) now lives in the module docs, and the codec points there instead.

Signed-off-by: Akash Kumar <116457960+akashchamp@users.noreply.github.com>

Author: akashchamp

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:** <https://github.com/apache/datafusion/issues/25152>
+//!
+//! **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<Mutex<HashMap<u64, Arc<dyn ExecutionPlan>>>> = OnceLock::new();
+
+fn execution_plans() -> &'static Mutex<HashMap<u64, Arc<dyn ExecutionPlan>>> {
+    EXECUTION_PLANS.get_or_init(|| Mutex::new(HashMap::new()))
+}
+
+fn token_id(buf: &[u8]) -> Option<u64> {
+    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<dyn ExecutionPlan>) -> bool {
+    node.is::<DataSourceExec>() || node.is::<ForeignExecutionPlan>()
+}
+
+pub(crate) fn park(node: Arc<dyn ExecutionPlan>, buf: &mut Vec<u8>) -> 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<Option<Arc<dyn ExecutionPlan>>> {
+    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))
+}

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;

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<Mutex<HashMap<u64, Arc<dyn ExecutionPlan>>>> = 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<HashMap<u64, Arc<dyn ExecutionPlan>>> {
-    EXECUTION_PLANS.get_or_init(|| Mutex::new(HashMap::new()))
-}
-
-fn token_id(buf: &[u8]) -> Option<u64> {
-    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<PhysicalCallCounters>,
@@ -99,19 +79,11 @@ impl PhysicalExtensionCodec for CountingPhysicalExtensionCodec {
         proto_converter: &dyn PhysicalProtoConverterExtension,
     ) -> Result<Arc<dyn ExecutionPlan>> {
         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<u8>,
         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::<DataSourceExec>() || node.is::<ForeignExecutionPlan>() {
+        // 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)