feat: give PhysicalExtensionCodec the full decode context

datafusion/expr/src/physical_planning_context.rs

@@ -110,70 +110,119 @@ impl SubqueryIndex {
     }
 }
 
+/// Where a [`ScalarSubqueryResults`] container keeps its slots.
+///
+/// The slots are normally held in memory next to the plan. `datafusion-ffi`
+/// implements this to forward each operation back over an FFI boundary, so a
+/// plan decoded on the far side reads the values populated on the near side.
+pub trait ScalarSubqueryResultsBackend:
+    fmt::Debug + Send + Sync + std::panic::RefUnwindSafe
+{
+    /// Returns the scalar value stored at `index`, if it has been populated.
+    fn get(&self, index: usize) -> Option<ScalarValue>;
+
+    /// Stores `value` in the slot at `index`.
+    fn set(&self, index: usize, value: ScalarValue) -> Result<()>;
+
+    /// Clears all populated results so the container can be reused.
+    fn clear(&self);
+}
+
+/// The in-memory backend behind [`ScalarSubqueryResults::new`].
+struct LocalResults(Vec<Mutex<Option<ScalarValue>>>);
+
+impl fmt::Debug for LocalResults {
+    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
+        f.debug_list()
+            .entries(self.0.iter().map(|slot| slot.lock().unwrap().clone()))
+            .finish()
+    }
+}
+
+impl ScalarSubqueryResultsBackend for LocalResults {
+    fn get(&self, index: usize) -> Option<ScalarValue> {
+        self.0.get(index)?.lock().unwrap().clone()
+    }
+
+    fn set(&self, index: usize, value: ScalarValue) -> Result<()> {
+        let Some(slot) = self.0.get(index) else {
+            return internal_err!(
+                "ScalarSubqueryResults: result index {index} is out of bounds"
+            );
+        };
+
+        let mut slot = slot.lock().unwrap();
+        if slot.is_some() {
+            return internal_err!(
+                "ScalarSubqueryResults: result for index {index} was already populated"
+            );
+        }
+        *slot = Some(value);
+
+        Ok(())
+    }
+
+    fn clear(&self) {
+        for slot in &self.0 {
+            *slot.lock().unwrap() = None;
+        }
+    }
+}
+
 /// Shared results container for uncorrelated scalar subqueries.
 ///
 /// Each entry corresponds to one scalar subquery, identified by its index.
 /// Each slot is populated at execution time by `ScalarSubqueryExec`, read by
 /// `ScalarSubqueryExpr` instances that share this container, and cleared when
 /// the plan is reset for re-execution.
-#[derive(Clone, Default)]
+#[derive(Clone)]
 pub struct ScalarSubqueryResults {
-    slots: Arc<Vec<Mutex<Option<ScalarValue>>>>,
+    backend: Arc<dyn ScalarSubqueryResultsBackend>,
+}
+
+impl Default for ScalarSubqueryResults {
+    fn default() -> Self {
+        Self::new(0)
+    }
 }
 
 impl ScalarSubqueryResults {
     /// Creates a new shared results container with `n` empty slots.
     pub fn new(n: usize) -> Self {
-        Self {
-            slots: Arc::new((0..n).map(|_| Mutex::new(None)).collect()),
-        }
+        Self::from_backend(Arc::new(LocalResults(
+            (0..n).map(|_| Mutex::new(None)).collect(),
+        )))
+    }
+
+    /// Creates a results container whose slots live in `backend`.
+    pub fn from_backend(backend: Arc<dyn ScalarSubqueryResultsBackend>) -> Self {
+        Self { backend }
     }
 
     /// Returns the scalar value stored at `index`, if it has been populated.
     pub fn get(&self, index: SubqueryIndex) -> Option<ScalarValue> {
-        let slot = self.slots.get(index.as_usize())?;
-        slot.lock().unwrap().clone()
+        self.backend.get(index.as_usize())
     }
 
     /// Stores `value` in the slot at `index`.
     pub fn set(&self, index: SubqueryIndex, value: ScalarValue) -> Result<()> {
-        let Some(slot) = self.slots.get(index.as_usize()) else {
-            return internal_err!(
-                "ScalarSubqueryResults: result index {} is out of bounds",
-                index.as_usize()
-            );
-        };
-
-        let mut slot = slot.lock().unwrap();
-        if slot.is_some() {
-            return internal_err!(
-                "ScalarSubqueryResults: result for index {} was already populated",
-                index.as_usize()
-            );
-        }
-        *slot = Some(value);
-
-        Ok(())
+        self.backend.set(index.as_usize(), value)
     }
 
     /// Clears all populated results so the container can be reused.
     pub fn clear(&self) {
-        for slot in self.slots.iter() {
-            *slot.lock().unwrap() = None;
-        }
+        self.backend.clear()
     }
 
     /// Returns true if `this` and `other` point to the same shared container.
     pub fn ptr_eq(this: &Self, other: &Self) -> bool {
-        Arc::ptr_eq(&this.slots, &other.slots)
+        Arc::ptr_eq(&this.backend, &other.backend)
     }
 }
 
 impl fmt::Debug for ScalarSubqueryResults {
     fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
-        f.debug_list()
-            .entries(self.slots.iter().map(|slot| slot.lock().unwrap().clone()))
-            .finish()
+        self.backend.fmt(f)
     }
 }
 
@@ -187,14 +236,23 @@ impl Eq for ScalarSubqueryResults {}
 
 impl Hash for ScalarSubqueryResults {
     fn hash<H: Hasher>(&self, state: &mut H) {
-        Arc::as_ptr(&self.slots).hash(state);
+        Arc::as_ptr(&self.backend).cast::<()>().hash(state);
     }
 }
 
 #[cfg(test)]
 mod tests {
     use super::*;
 
+    /// `ScalarSubqueryResults` must stay `UnwindSafe`/`RefUnwindSafe` so a
+    /// `ScalarSubqueryResultsBackend` implementation (e.g. `datafusion-ffi`'s
+    /// FFI proxy) doesn't regress these auto traits for callers.
+    #[test]
+    fn scalar_subquery_results_is_unwind_safe() {
+        fn assert_bounds<T: std::panic::UnwindSafe + std::panic::RefUnwindSafe>() {}
+        assert_bounds::<ScalarSubqueryResults>();
+    }
+
     #[test]
     fn scalar_subquery_results_set_and_get() -> Result<()> {
         let results = ScalarSubqueryResults::new(1);

datafusion/ffi/integration-tests/Cargo.toml

@@ -30,6 +30,7 @@ datafusion-catalog = { workspace = true }
 datafusion-common = { workspace = true }
 datafusion-execution = { workspace = true }
 datafusion-expr = { workspace = true }
+datafusion-physical-expr = { workspace = true }
 datafusion-physical-expr-common = { workspace = true }
 datafusion-physical-optimizer = { workspace = true }
 datafusion-physical-plan = { workspace = true }

datafusion/ffi/integration-tests/tests/cases/mod.rs

@@ -5,6 +5,7 @@
 mod catalog;
 mod config;
 mod execution_plan;
+mod physical_extension_codec;
 mod physical_optimizer;
 mod query_planner;
 mod udaf;

datafusion/ffi/integration-tests/tests/cases/physical_extension_codec.rs

@@ -0,0 +1,112 @@
+// SPDX-License-Identifier: Apache-2.0
+// SPDX-FileCopyrightText: See the NOTICE file distributed with this work for additional information regarding copyright ownership
+// SPDX-FileContributor: Licensed to the Apache Software Foundation (ASF) under one or more contributor license agreements
+
+mod tests {
+    use std::any::Any;
+    use std::sync::Arc;
+
+    use arrow::array::{AsArray, RecordBatch};
+    use arrow::datatypes::{DataType, Int64Type, Schema};
+    use datafusion::prelude::SessionContext;
+    use datafusion_common::{DataFusionError, ScalarValue};
+    use datafusion_execution::TaskContextProvider;
+    use datafusion_expr::physical_planning_context::{
+        ScalarSubqueryResults, SubqueryIndex,
+    };
+    use datafusion_ffi::execution::FFI_TaskContextProvider;
+    use datafusion_ffi::execution_plan::ForeignExecutionPlan;
+    use datafusion_ffi::proto::physical_extension_codec::ForeignPhysicalExtensionCodec;
+    use datafusion_ffi::proto::physical_extension_codec::fixtures::{
+        ScalarSubqueryExprExec, ScalarSubqueryExprExecCodec,
+    };
+    use datafusion_ffi::tests::utils::get_scalar_subquery_expr_exec_codec;
+    use datafusion_physical_expr::scalar_subquery::ScalarSubqueryExpr;
+    use datafusion_physical_plan::empty::EmptyExec;
+    use datafusion_physical_plan::{ExecutionPlan, PhysicalExpr};
+    use datafusion_proto::physical_plan::{
+        DefaultPhysicalProtoConverter, PhysicalExtensionCodec, PhysicalPlanDecodeContext,
+    };
+
+    /// A plan decoded by a genuinely foreign codec must read the scalar
+    /// subquery results the host populated.
+    ///
+    /// The codec comes from a separately `dlopen`'d copy of this crate, so
+    /// the decode takes the same path a real foreign codec takes rather than
+    /// a test-only one.
+    #[tokio::test]
+    async fn test_ffi_physical_extension_codec_scalar_subquery_cross_library()
+    -> Result<(), DataFusionError> {
+        let results = ScalarSubqueryResults::new(1);
+        results.set(SubqueryIndex::new(0), ScalarValue::Int64(Some(42)))?;
+
+        let sq_expr: Arc<dyn PhysicalExpr> = Arc::new(ScalarSubqueryExpr::new(
+            DataType::Int64,
+            true,
+            SubqueryIndex::new(0),
+            results.clone(),
+        ));
+        let extension_plan: Arc<dyn ExecutionPlan> =
+            Arc::new(ScalarSubqueryExprExec::new(
+                sq_expr,
+                Arc::new(EmptyExec::new(Arc::new(Schema::empty()))),
+            ));
+
+        let mut bytes = Vec::new();
+        ScalarSubqueryExprExecCodec.try_encode(
+            Arc::clone(&extension_plan),
+            &mut bytes,
+            &DefaultPhysicalProtoConverter {},
+        )?;
+
+        let ctx = Arc::new(SessionContext::default());
+        let host_task_ctx_provider = Arc::clone(&ctx) as Arc<dyn TaskContextProvider>;
+        let ffi_task_ctx_provider =
+            FFI_TaskContextProvider::from(&host_task_ctx_provider);
+
+        let foreign_ffi_codec =
+            get_scalar_subquery_expr_exec_codec(ffi_task_ctx_provider)?;
+        let foreign_codec: Arc<dyn PhysicalExtensionCodec> = (&foreign_ffi_codec).into();
+        assert!(
+            (Arc::clone(&foreign_codec) as Arc<dyn Any>)
+                .downcast_ref::<ForeignPhysicalExtensionCodec>()
+                .is_some(),
+            "codec built inside the dlopen'd cdylib must be seen as foreign"
+        );
+
+        let task_ctx = ctx.task_ctx();
+        let decode_ctx =
+            PhysicalPlanDecodeContext::new(task_ctx.as_ref(), foreign_codec.as_ref())
+                .with_scalar_subquery_results(results.clone());
+        // `ScalarSubqueryExprExecCodec::try_decode_with_ctx` only uses this
+        // to read a schema; `ScalarSubqueryExpr` carries no column
+        // references of its own.
+        let dummy_input: Arc<dyn ExecutionPlan> =
+            Arc::new(EmptyExec::new(Arc::new(Schema::empty())));
+        let decoded = foreign_codec.try_decode_with_ctx(
+            bytes.as_ref(),
+            &[dummy_input],
+            &decode_ctx,
+            &DefaultPhysicalProtoConverter {},
+        )?;
+        assert!(
+            decoded.downcast_ref::<ScalarSubqueryExprExec>().is_none(),
+            "a plan decoded by a genuinely foreign codec must stay opaque, not downcast locally"
+        );
+        assert!(
+            decoded.downcast_ref::<ForeignExecutionPlan>().is_some(),
+            "expected the decoded plan to be a ForeignExecutionPlan"
+        );
+
+        let stream = decoded.execute(0, task_ctx)?;
+        let batches: Vec<RecordBatch> =
+            datafusion_physical_plan::common::collect(stream).await?;
+        assert_eq!(batches.len(), 1);
+        let batch = &batches[0];
+        assert_eq!(batch.num_rows(), 1);
+        let value = batch.column(0).as_primitive::<Int64Type>().value(0);
+        assert_eq!(value, 42);
+
+        Ok(())
+    }
+}

datafusion/ffi/src/proto/mod.rs

@@ -4,3 +4,4 @@
 
 pub mod logical_extension_codec;
 pub mod physical_extension_codec;
+pub mod scalar_subquery_results;