From cf5431e26d8b4eef145949cdfb554903a5cced87 Mon Sep 17 00:00:00 2001 From: naman Date: Tue, 15 Sep 2026 13:38:20 +0530 Subject: [PATCH 1/7] feat: give PhysicalExtensionCodec the full decode context PhysicalExtensionCodec::try_decode only receives the TaskContext, so a codec that decodes its own nested expressions has to build a fresh PhysicalPlanDecodeContext. That context has no scalar subquery results, so a ScalarSubqueryExpr inside an extension plan fails to deserialize even under the ScalarSubqueryExec that owns the results. Add try_decode_with_ctx, which receives the full PhysicalPlanDecodeContext and by default calls try_decode. Decode extension plans through it, forward it in ComposedPhysicalExtensionCodec, and call it with a root context on the FFI provider side. --- .../ffi/src/proto/physical_extension_codec.rs | 69 +++- datafusion/proto/src/physical_plan/mod.rs | 39 ++- .../tests/cases/plans/scalar_subquery.rs | 315 ++++++++++++++++++ 3 files changed, 417 insertions(+), 6 deletions(-) diff --git a/datafusion/ffi/src/proto/physical_extension_codec.rs b/datafusion/ffi/src/proto/physical_extension_codec.rs index 47a6f1e95cafb..a64b04e93bcb9 100644 --- a/datafusion/ffi/src/proto/physical_extension_codec.rs +++ b/datafusion/ffi/src/proto/physical_extension_codec.rs @@ -26,7 +26,7 @@ use datafusion_expr::{ }; use datafusion_physical_plan::ExecutionPlan; use datafusion_proto::physical_plan::{ - DefaultPhysicalProtoConverter, PhysicalExtensionCodec, + DefaultPhysicalProtoConverter, PhysicalExtensionCodec, PhysicalPlanDecodeContext, PhysicalProtoConverterExtension, }; @@ -148,10 +148,13 @@ unsafe extern "C" fn try_decode_fn_wrapper( .collect::>>(); let inputs = sresult_return!(inputs); - let plan = sresult_return!(codec.try_decode( + // The caller's decode context cannot cross the FFI boundary, so decode + // with a root context for this side's codec. + let decode_ctx = PhysicalPlanDecodeContext::new(task_ctx.as_ref(), codec.as_ref()); + let plan = sresult_return!(codec.try_decode_with_ctx( buf.as_ref(), &inputs, - task_ctx.as_ref(), + &decode_ctx, &DefaultPhysicalProtoConverter {}, )); @@ -458,7 +461,7 @@ pub(crate) mod tests { use datafusion_functions_window::rank::{Rank, RankType}; use datafusion_physical_plan::ExecutionPlan; use datafusion_proto::physical_plan::{ - DefaultPhysicalProtoConverter, PhysicalExtensionCodec, + DefaultPhysicalProtoConverter, PhysicalExtensionCodec, PhysicalPlanDecodeContext, PhysicalProtoConverterExtension, }; @@ -643,6 +646,64 @@ pub(crate) mod tests { Ok(()) } + /// A codec whose `try_decode` fails and that only decodes through + /// `try_decode_with_ctx`. + #[derive(Debug)] + struct ContextOnlyCodec; + + impl PhysicalExtensionCodec for ContextOnlyCodec { + fn try_decode( + &self, + _buf: &[u8], + _inputs: &[Arc], + _ctx: &TaskContext, + _proto_converter: &dyn PhysicalProtoConverterExtension, + ) -> Result> { + exec_err!("ContextOnlyCodec decodes through try_decode_with_ctx") + } + + fn try_decode_with_ctx( + &self, + _buf: &[u8], + _inputs: &[Arc], + _ctx: &PhysicalPlanDecodeContext<'_>, + _proto_converter: &dyn PhysicalProtoConverterExtension, + ) -> Result> { + Ok(create_test_exec()) + } + + fn try_encode( + &self, + _node: Arc, + _buf: &mut Vec, + _proto_converter: &dyn PhysicalProtoConverterExtension, + ) -> Result<()> { + Ok(()) + } + } + + #[test] + fn ffi_physical_extension_codec_decodes_through_try_decode_with_ctx() -> Result<()> { + let codec = Arc::new(ContextOnlyCodec); + let (ctx, task_ctx_provider) = crate::util::tests::test_session_and_ctx(); + + let mut ffi_codec = + FFI_PhysicalExtensionCodec::new(codec, None, task_ctx_provider); + ffi_codec.library_marker_id = crate::mock_foreign_marker_id; + let foreign_codec: Arc = (&ffi_codec).into(); + + let returned_exec = foreign_codec.try_decode( + &[], + &[create_test_exec()], + ctx.task_ctx().as_ref(), + &DefaultPhysicalProtoConverter {}, + )?; + + assert!(returned_exec.is::()); + + Ok(()) + } + #[test] fn roundtrip_ffi_physical_extension_codec_udf() -> Result<()> { let codec = Arc::new(TestExtensionCodec {}); diff --git a/datafusion/proto/src/physical_plan/mod.rs b/datafusion/proto/src/physical_plan/mod.rs index 940c732f81318..8d228088fc340 100644 --- a/datafusion/proto/src/physical_plan/mod.rs +++ b/datafusion/proto/src/physical_plan/mod.rs @@ -1399,10 +1399,10 @@ pub trait PhysicalPlanNodeExt: Sized { .map(|i| proto_converter.proto_to_execution_plan(i, ctx)) .collect::>()?; - let extension_node = ctx.codec().try_decode( + let extension_node = ctx.codec().try_decode_with_ctx( extension.node.as_slice(), &inputs, - ctx.task_ctx(), + ctx, proto_converter, )?; @@ -1665,6 +1665,29 @@ pub trait PhysicalExtensionCodec: Debug + Send + Sync + Any { proto_converter: &dyn PhysicalProtoConverterExtension, ) -> Result>; + /// Decode a custom extension plan from `buf`, with access to the full + /// decode context. + /// + /// Implement this instead of relying on [`Self::try_decode`] when the codec + /// decodes nested expressions or plans itself, for example through + /// [`PhysicalProtoConverterExtension::proto_to_physical_expr`], and pass + /// `ctx` on to those calls. A context built with + /// [`PhysicalPlanDecodeContext::new`] starts empty, so it lacks state that + /// only exists while decoding the surrounding plan, such as the results of an + /// enclosing `ScalarSubqueryExec` that a `ScalarSubqueryExpr` is bound to. + /// + /// The default implementation calls [`Self::try_decode`] with + /// `ctx.task_ctx()`, so codecs that only implement `try_decode` keep working. + fn try_decode_with_ctx( + &self, + buf: &[u8], + inputs: &[Arc], + ctx: &PhysicalPlanDecodeContext<'_>, + proto_converter: &dyn PhysicalProtoConverterExtension, + ) -> Result> { + self.try_decode(buf, inputs, ctx.task_ctx(), proto_converter) + } + fn try_encode( &self, node: Arc, @@ -2108,6 +2131,18 @@ impl PhysicalExtensionCodec for ComposedPhysicalExtensionCodec { }) } + fn try_decode_with_ctx( + &self, + buf: &[u8], + inputs: &[Arc], + ctx: &PhysicalPlanDecodeContext<'_>, + proto_converter: &dyn PhysicalProtoConverterExtension, + ) -> Result> { + self.decode_protobuf(buf, |codec, data| { + codec.try_decode_with_ctx(data, inputs, ctx, proto_converter) + }) + } + fn try_encode( &self, node: Arc, diff --git a/datafusion/proto/tests/cases/plans/scalar_subquery.rs b/datafusion/proto/tests/cases/plans/scalar_subquery.rs index 34d30aa03dece..8fef57b79271d 100644 --- a/datafusion/proto/tests/cases/plans/scalar_subquery.rs +++ b/datafusion/proto/tests/cases/plans/scalar_subquery.rs @@ -17,7 +17,9 @@ //! `ScalarSubqueryExec` and the results it scopes to its subtree. +use datafusion::arrow::datatypes::SchemaRef; use datafusion::arrow::datatypes::{DataType, Field, Schema}; +use datafusion::execution::TaskContext; use datafusion::logical_expr::Operator; use datafusion::physical_plan::ExecutionPlan; use datafusion::physical_plan::empty::EmptyExec; @@ -26,17 +28,30 @@ use datafusion::physical_plan::filter::FilterExec; use datafusion::physical_plan::scalar_subquery::{ ScalarSubqueryExec, ScalarSubqueryLink, }; +use datafusion::physical_plan::{ + ChildrenPropertiesMode, DisplayAs, DisplayFormatType, PhysicalExpr, PlanProperties, + ReplaceChildrenOptions, SendableRecordBatchStream, +}; use datafusion::prelude::SessionContext; use datafusion_common::Result; +use datafusion_common::tree_node::TreeNodeRecursion; +use datafusion_common::{internal_datafusion_err, internal_err}; use datafusion_expr::physical_planning_context::{ScalarSubqueryResults, SubqueryIndex}; use datafusion_physical_expr::scalar_subquery::ScalarSubqueryExpr; use datafusion_proto::bytes::{ physical_plan_from_bytes_with_proto_converter, physical_plan_to_bytes_with_proto_converter, }; +use datafusion_proto::physical_plan::{ + ComposedPhysicalExtensionCodec, DefaultPhysicalProtoConverter, + PhysicalExtensionCodec, PhysicalPlanDecodeContext, PhysicalProtoConverterExtension, +}; use datafusion_proto::physical_plan::{ DeduplicatingProtoConverter, DefaultPhysicalExtensionCodec, }; +use datafusion_proto::protobuf::PhysicalExprNode; +use prost::Message; +use std::fmt::Formatter; use std::sync::Arc; use std::vec; @@ -276,3 +291,303 @@ async fn roundtrip_scalar_subquery_exec_with_default_converter_executes() -> Res Ok(()) } + +/// An extension plan that carries physical expressions in its own payload, so +/// its codec has to decode them. +#[derive(Debug)] +struct ExprsExec { + exprs: Vec>, + child: Arc, +} + +impl DisplayAs for ExprsExec { + fn fmt_as(&self, _t: DisplayFormatType, f: &mut Formatter) -> std::fmt::Result { + write!(f, "ExprsExec") + } +} + +impl ExecutionPlan for ExprsExec { + fn name(&self) -> &str { + "ExprsExec" + } + + fn schema(&self) -> SchemaRef { + self.child.schema() + } + + fn properties(&self) -> &Arc { + self.child.properties() + } + + fn children(&self) -> Vec<&Arc> { + vec![&self.child] + } + + fn apply_expressions( + &self, + f: &mut dyn FnMut(&Arc) -> Result, + ) -> Result { + datafusion_physical_plan::apply_expression_roots(&self.exprs, f) + } + + fn replace_children( + self: Arc, + _: Vec>, + _: ReplaceChildrenOptions, + ) -> Result> { + unreachable!() + } + + fn with_new_children( + self: Arc, + children: Vec>, + ) -> Result> { + self.replace_children( + children, + ReplaceChildrenOptions::new(ChildrenPropertiesMode::Recompute), + ) + } + + fn execute( + &self, + _partition: usize, + _context: Arc, + ) -> Result { + unreachable!() + } +} + +#[derive(Clone, PartialEq, Message)] +struct ExprsExecProto { + #[prost(message, repeated, tag = "1")] + exprs: Vec, +} + +/// Decodes [`ExprsExec`] through `try_decode_with_ctx`, passing the decode +/// context on to its expressions. `try_decode` fails so the tests catch a +/// caller that drops the context. +#[derive(Debug)] +struct ExprsExecCodec; + +impl PhysicalExtensionCodec for ExprsExecCodec { + fn try_decode( + &self, + _buf: &[u8], + _inputs: &[Arc], + _ctx: &TaskContext, + _proto_converter: &dyn PhysicalProtoConverterExtension, + ) -> Result> { + internal_err!("ExprsExecCodec decodes through try_decode_with_ctx") + } + + fn try_decode_with_ctx( + &self, + buf: &[u8], + inputs: &[Arc], + ctx: &PhysicalPlanDecodeContext<'_>, + proto_converter: &dyn PhysicalProtoConverterExtension, + ) -> Result> { + let proto = ExprsExecProto::decode(buf) + .map_err(|e| internal_datafusion_err!("failed to decode ExprsExec: {e}"))?; + let schema = inputs[0].schema(); + let exprs = proto + .exprs + .iter() + .map(|e| proto_converter.proto_to_physical_expr(e, &schema, ctx)) + .collect::>()?; + Ok(Arc::new(ExprsExec { + exprs, + child: Arc::clone(&inputs[0]), + })) + } + + fn try_encode( + &self, + node: Arc, + buf: &mut Vec, + proto_converter: &dyn PhysicalProtoConverterExtension, + ) -> Result<()> { + let exec = node + .downcast_ref::() + .ok_or_else(|| internal_datafusion_err!("expected ExprsExec"))?; + let proto = ExprsExecProto { + exprs: exec + .exprs + .iter() + .map(|e| proto_converter.physical_expr_to_proto(e, self)) + .collect::>()?, + }; + proto + .encode(buf) + .map_err(|e| internal_datafusion_err!("failed to encode ExprsExec: {e}")) + } +} + +fn scalar_subquery_expr(results: &ScalarSubqueryResults) -> Arc { + Arc::new(ScalarSubqueryExpr::new( + DataType::Int64, + true, + SubqueryIndex::new(0), + results.clone(), + )) +} + +fn exprs_exec( + expr: Arc, + child: Arc, +) -> Arc { + Arc::new(ExprsExec { + exprs: vec![expr], + child, + }) +} + +fn scalar_subquery_exec( + input: Arc, + results: ScalarSubqueryResults, +) -> Arc { + let subquery_schema = + Arc::new(Schema::new(vec![Field::new("x", DataType::Int64, true)])); + Arc::new(ScalarSubqueryExec::new( + input, + vec![ScalarSubqueryLink { + plan: Arc::new(EmptyExec::new(subquery_schema)), + index: SubqueryIndex::new(0), + }], + results, + )) +} + +fn empty_input() -> Arc { + Arc::new(EmptyExec::new(Arc::new(Schema::new(vec![Field::new( + "a", + DataType::Int64, + false, + )])))) +} + +fn roundtrip_with( + plan: &Arc, + codec: &dyn PhysicalExtensionCodec, + converter: &dyn PhysicalProtoConverterExtension, +) -> Result> { + let bytes = + physical_plan_to_bytes_with_proto_converter(Arc::clone(plan), codec, converter)?; + let ctx = SessionContext::new(); + physical_plan_from_bytes_with_proto_converter( + bytes.as_ref(), + ctx.task_ctx().as_ref(), + codec, + converter, + ) +} + +/// Returns the results of `plan`, a `ScalarSubqueryExec`, and of the scalar +/// subquery expression in the `ExprsExec` directly below it. +fn exec_and_expr_results( + plan: &Arc, +) -> (&ScalarSubqueryResults, &ScalarSubqueryResults) { + let sq_exec = plan + .downcast_ref::() + .expect("expected ScalarSubqueryExec"); + let exprs_exec = sq_exec + .input() + .downcast_ref::() + .expect("expected ExprsExec"); + let sq_expr = exprs_exec.exprs[0] + .downcast_ref::() + .expect("expected ScalarSubqueryExpr"); + (sq_exec.results(), sq_expr.results()) +} + +/// A scalar subquery expression that an extension codec decodes itself is +/// wired to the results of the surrounding `ScalarSubqueryExec`. +#[test] +fn roundtrip_scalar_subquery_expr_in_extension_plan() -> Result<()> { + let results = ScalarSubqueryResults::new(1); + let plan = scalar_subquery_exec( + exprs_exec(scalar_subquery_expr(&results), empty_input()), + results, + ); + + let converters: [&dyn PhysicalProtoConverterExtension; 2] = [ + &DefaultPhysicalProtoConverter {}, + &DeduplicatingProtoConverter {}, + ]; + for converter in converters { + let deserialized = roundtrip_with(&plan, &ExprsExecCodec, converter)?; + let (exec_results, expr_results) = exec_and_expr_results(&deserialized); + assert!(ScalarSubqueryResults::ptr_eq(exec_results, expr_results)); + } + Ok(()) +} + +/// `ComposedPhysicalExtensionCodec` passes the decode context on to the codec +/// that decodes the node. +#[test] +fn roundtrip_scalar_subquery_expr_in_extension_plan_with_composed_codec() -> Result<()> { + let results = ScalarSubqueryResults::new(1); + let plan = scalar_subquery_exec( + exprs_exec(scalar_subquery_expr(&results), empty_input()), + results, + ); + let codec = ComposedPhysicalExtensionCodec::new(vec![ + Arc::new(DefaultPhysicalExtensionCodec {}), + Arc::new(ExprsExecCodec), + ]); + + let deserialized = roundtrip_with(&plan, &codec, &DefaultPhysicalProtoConverter {})?; + let (exec_results, expr_results) = exec_and_expr_results(&deserialized); + assert!(ScalarSubqueryResults::ptr_eq(exec_results, expr_results)); + Ok(()) +} + +/// With nested `ScalarSubqueryExec`s, each expression decoded by the extension +/// codec is wired to its own enclosing scope. +#[test] +fn roundtrip_nested_scalar_subquery_exprs_in_extension_plans() -> Result<()> { + let inner_results = ScalarSubqueryResults::new(1); + let inner = scalar_subquery_exec( + exprs_exec(scalar_subquery_expr(&inner_results), empty_input()), + inner_results, + ); + let outer_results = ScalarSubqueryResults::new(1); + let outer = scalar_subquery_exec( + exprs_exec(scalar_subquery_expr(&outer_results), inner), + outer_results, + ); + + let deserialized = + roundtrip_with(&outer, &ExprsExecCodec, &DefaultPhysicalProtoConverter {})?; + let (outer_exec, outer_expr) = exec_and_expr_results(&deserialized); + let inner_plan = deserialized + .downcast_ref::() + .expect("expected outer ScalarSubqueryExec") + .input() + .children()[0]; + let (inner_exec, inner_expr) = exec_and_expr_results(inner_plan); + + assert!(ScalarSubqueryResults::ptr_eq(outer_exec, outer_expr)); + assert!(ScalarSubqueryResults::ptr_eq(inner_exec, inner_expr)); + assert!(!ScalarSubqueryResults::ptr_eq(outer_exec, inner_exec)); + assert!(!ScalarSubqueryResults::ptr_eq(outer_exec, inner_expr)); + assert!(!ScalarSubqueryResults::ptr_eq(inner_exec, outer_expr)); + Ok(()) +} + +/// Without an enclosing `ScalarSubqueryExec` there are no results to wire the +/// expression to, so decoding still fails. +#[test] +fn scalar_subquery_expr_in_extension_plan_requires_enclosing_exec() { + let results = ScalarSubqueryResults::new(1); + let plan = exprs_exec(scalar_subquery_expr(&results), empty_input()); + + let err = roundtrip_with(&plan, &ExprsExecCodec, &DefaultPhysicalProtoConverter {}) + .unwrap_err(); + assert!( + err.to_string().contains( + "ScalarSubqueryExpr can only be deserialized as part of a surrounding ScalarSubqueryExec" + ), + "{err}" + ); +} From 00dfb1c71eb1a34c28c8a870b25b6b233eee6d56 Mon Sep 17 00:00:00 2001 From: naman Date: Fri, 2 Oct 2026 15:20:23 +0530 Subject: [PATCH 2/7] fix: preserve scalar subquery results scope across the FFI decode boundary ForeignPhysicalExtensionCodec now forwards the active ScalarSubqueryResults scope through an ABI-safe FFI_ScalarSubqueryResults handle instead of dropping it, so a ScalarSubqueryExpr decoded by a genuinely foreign codec shares the same results as the surrounding ScalarSubqueryExec. ScalarSubqueryResults gains an alternate backend representation (ScalarSubqueryResultsBackend) to support this proxy, with ScalarValues crossing the boundary as prost-encoded bytes like the rest of the codec's wire format. Adds a forced-foreign end-to-end ScalarSubqueryExec roundtrip test, and clarifies the legacy-decode test's name/intent. --- .../expr/src/physical_planning_context.rs | 139 ++++++-- datafusion/ffi/src/proto/mod.rs | 1 + .../ffi/src/proto/physical_extension_codec.rs | 324 +++++++++++++++++- .../ffi/src/proto/scalar_subquery_results.rs | 237 +++++++++++++ 4 files changed, 669 insertions(+), 32 deletions(-) create mode 100644 datafusion/ffi/src/proto/scalar_subquery_results.rs diff --git a/datafusion/expr/src/physical_planning_context.rs b/datafusion/expr/src/physical_planning_context.rs index b2e579ea3c7cb..133e2b03e58ff 100644 --- a/datafusion/expr/src/physical_planning_context.rs +++ b/datafusion/expr/src/physical_planning_context.rs @@ -123,70 +123,148 @@ impl SubqueryIndex { } } +/// An alternate backing store for a [`ScalarSubqueryResults`] container. +/// +/// A results container normally owns its slots directly (see +/// [`ScalarSubqueryResults::new`]). This trait lets a container instead +/// forward `get`/`set`/`clear` to some other shared container, which is how +/// `datafusion-ffi` bridges the results of a `ScalarSubqueryExec` across a +/// non-Rust-ABI boundary: the far side of the boundary gets a +/// [`ScalarSubqueryResults`] backed by a proxy that calls back into the +/// near side's real container for every operation, so both sides observe the +/// same populated values. +pub trait ScalarSubqueryResultsBackend: fmt::Debug + Send + Sync { + /// Returns the scalar value stored at `index`, if it has been populated. + fn get(&self, index: usize) -> Option; + + /// 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); +} + +#[derive(Clone)] +enum ScalarSubqueryResultsRepr { + Local(Arc>>>), + Remote(Arc), +} + /// 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>>>, + repr: ScalarSubqueryResultsRepr, +} + +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()), + repr: ScalarSubqueryResultsRepr::Local(Arc::new( + (0..n).map(|_| Mutex::new(None)).collect(), + )), + } + } + + /// Creates a results container that forwards every operation to + /// `backend`, for example a proxy that reaches across an FFI boundary to + /// a real container owned by the other side. + pub fn from_backend(backend: Arc) -> Self { + Self { + repr: ScalarSubqueryResultsRepr::Remote(backend), } } /// Returns the scalar value stored at `index`, if it has been populated. pub fn get(&self, index: SubqueryIndex) -> Option { - let slot = self.slots.get(index.as_usize())?; - slot.lock().unwrap().clone() + match &self.repr { + ScalarSubqueryResultsRepr::Local(slots) => { + let slot = slots.get(index.as_usize())?; + slot.lock().unwrap().clone() + } + ScalarSubqueryResultsRepr::Remote(backend) => 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() - ); + match &self.repr { + ScalarSubqueryResultsRepr::Local(slots) => { + let Some(slot) = 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(()) + } + ScalarSubqueryResultsRepr::Remote(backend) => { + backend.set(index.as_usize(), value) + } } - *slot = Some(value); - - Ok(()) } /// Clears all populated results so the container can be reused. pub fn clear(&self) { - for slot in self.slots.iter() { - *slot.lock().unwrap() = None; + match &self.repr { + ScalarSubqueryResultsRepr::Local(slots) => { + for slot in slots.iter() { + *slot.lock().unwrap() = None; + } + } + ScalarSubqueryResultsRepr::Remote(backend) => 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) + match (&this.repr, &other.repr) { + ( + ScalarSubqueryResultsRepr::Local(a), + ScalarSubqueryResultsRepr::Local(b), + ) => Arc::ptr_eq(a, b), + ( + ScalarSubqueryResultsRepr::Remote(a), + ScalarSubqueryResultsRepr::Remote(b), + ) => Arc::ptr_eq(a, b), + _ => false, + } } } 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() + match &self.repr { + ScalarSubqueryResultsRepr::Local(slots) => f + .debug_list() + .entries(slots.iter().map(|slot| slot.lock().unwrap().clone())) + .finish(), + ScalarSubqueryResultsRepr::Remote(backend) => f + .debug_tuple("ScalarSubqueryResults::Remote") + .field(backend) + .finish(), + } } } @@ -200,7 +278,12 @@ impl Eq for ScalarSubqueryResults {} impl Hash for ScalarSubqueryResults { fn hash(&self, state: &mut H) { - Arc::as_ptr(&self.slots).hash(state); + match &self.repr { + ScalarSubqueryResultsRepr::Local(slots) => Arc::as_ptr(slots).hash(state), + ScalarSubqueryResultsRepr::Remote(backend) => { + Arc::as_ptr(backend).hash(state) + } + } } } diff --git a/datafusion/ffi/src/proto/mod.rs b/datafusion/ffi/src/proto/mod.rs index ae76027ecb64e..1d580d0436c52 100644 --- a/datafusion/ffi/src/proto/mod.rs +++ b/datafusion/ffi/src/proto/mod.rs @@ -17,3 +17,4 @@ pub mod logical_extension_codec; pub mod physical_extension_codec; +pub mod scalar_subquery_results; diff --git a/datafusion/ffi/src/proto/physical_extension_codec.rs b/datafusion/ffi/src/proto/physical_extension_codec.rs index a64b04e93bcb9..c9e61b5e5c65e 100644 --- a/datafusion/ffi/src/proto/physical_extension_codec.rs +++ b/datafusion/ffi/src/proto/physical_extension_codec.rs @@ -37,10 +37,11 @@ use tokio::runtime::Handle; use crate::execution::FFI_TaskContextProvider; use crate::execution_plan::FFI_ExecutionPlan; +use crate::proto::scalar_subquery_results::FFI_ScalarSubqueryResults; use crate::udaf::FFI_AggregateUDF; use crate::udf::FFI_ScalarUDF; use crate::udwf::FFI_WindowUDF; -use crate::util::FFI_Result; +use crate::util::{FFI_Option, FFI_Result}; use crate::{df_result, sresult_return}; /// A stable struct for sharing [`PhysicalExtensionCodec`] across FFI boundaries. @@ -54,6 +55,17 @@ pub struct FFI_PhysicalExtensionCodec { inputs: SVec, ) -> FFI_Result, + /// Decode bytes into an execution plan, forwarding the active scalar + /// subquery results scope (if any) from the caller's decode context so a + /// `ScalarSubqueryExpr` decoded on this side of the boundary shares the + /// same populated results as the `ScalarSubqueryExec` that owns them. + try_decode_with_ctx: unsafe extern "C" fn( + &Self, + buf: SSlice, + inputs: SVec, + scalar_subquery_results: FFI_Option, + ) -> FFI_Result, + /// Encode an execution plan into bytes. try_encode: unsafe extern "C" fn(&Self, node: FFI_ExecutionPlan) -> FFI_Result>, @@ -161,6 +173,38 @@ unsafe extern "C" fn try_decode_fn_wrapper( FFI_Result::Ok(FFI_ExecutionPlan::new(plan, runtime)) } +unsafe extern "C" fn try_decode_with_ctx_fn_wrapper( + codec: &FFI_PhysicalExtensionCodec, + buf: SSlice, + inputs: SVec, + scalar_subquery_results: FFI_Option, +) -> FFI_Result { + let runtime = codec.runtime().cloned(); + let task_ctx: Arc = + sresult_return!((&codec.task_ctx_provider).try_into()); + let codec = codec.inner(); + let inputs = inputs + .into_iter() + .map(|plan| >::try_from(&plan)) + .collect::>>(); + let inputs = sresult_return!(inputs); + + let decode_ctx = PhysicalPlanDecodeContext::new(task_ctx.as_ref(), codec.as_ref()); + let decode_ctx = match scalar_subquery_results.into_option() { + Some(results) => decode_ctx.with_scalar_subquery_results(results.into()), + None => decode_ctx, + }; + + let plan = sresult_return!(codec.try_decode_with_ctx( + buf.as_ref(), + &inputs, + &decode_ctx, + &DefaultPhysicalProtoConverter {}, + )); + + FFI_Result::Ok(FFI_ExecutionPlan::new(plan, runtime)) +} + unsafe extern "C" fn try_encode_fn_wrapper( codec: &FFI_PhysicalExtensionCodec, node: FFI_ExecutionPlan, @@ -313,6 +357,7 @@ impl FFI_PhysicalExtensionCodec { Self { try_decode: try_decode_fn_wrapper, + try_decode_with_ctx: try_decode_with_ctx_fn_wrapper, try_encode: try_encode_fn_wrapper, try_decode_udf: try_decode_udf_fn_wrapper, try_encode_udf: try_encode_udf_fn_wrapper, @@ -377,6 +422,37 @@ impl PhysicalExtensionCodec for ForeignPhysicalExtensionCodec { Ok(plan) } + fn try_decode_with_ctx( + &self, + buf: &[u8], + inputs: &[Arc], + ctx: &PhysicalPlanDecodeContext<'_>, + _proto_converter: &dyn PhysicalProtoConverterExtension, + ) -> Result> { + let inputs = inputs + .iter() + .map(|plan| FFI_ExecutionPlan::new(Arc::clone(plan), None)) + .collect(); + + let scalar_subquery_results = ctx + .scalar_subquery_results() + .cloned() + .map(FFI_ScalarSubqueryResults::new) + .into(); + + let plan = df_result!(unsafe { + (self.0.try_decode_with_ctx)( + &self.0, + buf.into(), + inputs, + scalar_subquery_results, + ) + })?; + let plan: Arc = (&plan).try_into()?; + + Ok(plan) + } + fn try_encode( &self, node: Arc, @@ -452,18 +528,37 @@ pub(crate) mod tests { use std::sync::Arc; use arrow_schema::{DataType, Field, Schema}; - use datafusion_common::{Result, exec_err}; + use datafusion_common::tree_node::TreeNodeRecursion; + use datafusion_common::{Result, exec_err, internal_datafusion_err}; use datafusion_execution::TaskContext; + use datafusion_expr::physical_planning_context::{ + ScalarSubqueryResults, SubqueryIndex, + }; use datafusion_expr::ptr_eq::arc_ptr_eq; use datafusion_expr::{AggregateUDF, ScalarUDF, WindowUDF, WindowUDFImpl}; use datafusion_functions::math::abs::AbsFunc; use datafusion_functions_aggregate::sum::Sum; use datafusion_functions_window::rank::{Rank, RankType}; - use datafusion_physical_plan::ExecutionPlan; + use datafusion_physical_expr::scalar_subquery::ScalarSubqueryExpr; + use datafusion_physical_plan::empty::EmptyExec as RealEmptyExec; + use datafusion_physical_plan::scalar_subquery::{ + ScalarSubqueryExec, ScalarSubqueryLink, + }; + use datafusion_physical_plan::{ + ChildrenPropertiesMode, DisplayAs, DisplayFormatType, ExecutionPlan, + PhysicalExpr, PlanProperties, ReplaceChildrenOptions, SendableRecordBatchStream, + apply_expression_roots, + }; + use datafusion_proto::bytes::{ + physical_plan_from_bytes_with_proto_converter, + physical_plan_to_bytes_with_proto_converter, + }; use datafusion_proto::physical_plan::{ DefaultPhysicalProtoConverter, PhysicalExtensionCodec, PhysicalPlanDecodeContext, PhysicalProtoConverterExtension, }; + use datafusion_proto::protobuf::PhysicalExprNode; + use prost::Message; use crate::execution_plan::tests::EmptyExec; use crate::proto::physical_extension_codec::{ @@ -682,8 +777,14 @@ pub(crate) mod tests { } } + /// Calls the foreign adapter's legacy `try_decode`, which has no decode + /// context to forward, but the far side's own codec still reaches its + /// `try_decode_with_ctx` override (with a root context) because + /// `try_decode_fn_wrapper` always calls that method rather than + /// `try_decode` directly. #[test] - fn ffi_physical_extension_codec_decodes_through_try_decode_with_ctx() -> Result<()> { + fn ffi_physical_extension_codec_legacy_decode_uses_context_aware_default() + -> Result<()> { let codec = Arc::new(ContextOnlyCodec); let (ctx, task_ctx_provider) = crate::util::tests::test_session_and_ctx(); @@ -817,4 +918,219 @@ pub(crate) mod tests { .expect("rebound codec resolves"); assert_eq!(task_ctx.session_id(), ctx_b.task_ctx().session_id()); } + + /// An extension plan that carries a single physical expression, so its + /// codec has to decode that expression itself. + #[derive(Debug)] + struct ScalarSubqueryExprExec { + expr: Arc, + child: Arc, + } + + impl DisplayAs for ScalarSubqueryExprExec { + fn fmt_as( + &self, + _t: DisplayFormatType, + f: &mut std::fmt::Formatter, + ) -> std::fmt::Result { + write!(f, "ScalarSubqueryExprExec") + } + } + + impl ExecutionPlan for ScalarSubqueryExprExec { + fn name(&self) -> &str { + "ScalarSubqueryExprExec" + } + + fn properties(&self) -> &Arc { + self.child.properties() + } + + fn children(&self) -> Vec<&Arc> { + vec![&self.child] + } + + fn apply_expressions( + &self, + f: &mut dyn FnMut(&Arc) -> Result, + ) -> Result { + apply_expression_roots(std::slice::from_ref(&self.expr), f) + } + + fn replace_children( + self: Arc, + _: Vec>, + _: ReplaceChildrenOptions, + ) -> Result> { + unreachable!() + } + + fn with_new_children( + self: Arc, + children: Vec>, + ) -> Result> { + self.replace_children( + children, + ReplaceChildrenOptions::new(ChildrenPropertiesMode::Recompute), + ) + } + + fn execute( + &self, + _partition: usize, + _context: Arc, + ) -> Result { + unreachable!() + } + } + + #[derive(Clone, PartialEq, Message)] + struct ScalarSubqueryExprExecProto { + #[prost(message, optional, tag = "1")] + expr: Option, + } + + /// Decodes [`ScalarSubqueryExprExec`] through `try_decode_with_ctx`, + /// passing the decode context on to its expression. `try_decode` fails so + /// a caller that drops the context (by routing through the plain + /// `try_decode` FFI entry point instead of `try_decode_with_ctx`) is + /// caught by this test. + #[derive(Debug)] + struct ScalarSubqueryExprExecCodec; + + impl PhysicalExtensionCodec for ScalarSubqueryExprExecCodec { + fn try_decode( + &self, + _buf: &[u8], + _inputs: &[Arc], + _ctx: &TaskContext, + _proto_converter: &dyn PhysicalProtoConverterExtension, + ) -> Result> { + exec_err!("ScalarSubqueryExprExecCodec decodes through try_decode_with_ctx") + } + + fn try_decode_with_ctx( + &self, + buf: &[u8], + inputs: &[Arc], + ctx: &PhysicalPlanDecodeContext<'_>, + proto_converter: &dyn PhysicalProtoConverterExtension, + ) -> Result> { + let proto = ScalarSubqueryExprExecProto::decode(buf).map_err(|e| { + internal_datafusion_err!("failed to decode ScalarSubqueryExprExec: {e}") + })?; + let expr_proto = proto.expr.ok_or_else(|| { + internal_datafusion_err!("ScalarSubqueryExprExec is missing its expr") + })?; + let schema = inputs[0].schema(); + let expr = + proto_converter.proto_to_physical_expr(&expr_proto, &schema, ctx)?; + Ok(Arc::new(ScalarSubqueryExprExec { + expr, + child: Arc::clone(&inputs[0]), + })) + } + + fn try_encode( + &self, + node: Arc, + buf: &mut Vec, + proto_converter: &dyn PhysicalProtoConverterExtension, + ) -> Result<()> { + let exec = + node.downcast_ref::() + .ok_or_else(|| { + internal_datafusion_err!("expected ScalarSubqueryExprExec") + })?; + let proto = ScalarSubqueryExprExecProto { + expr: Some(proto_converter.physical_expr_to_proto(&exec.expr, self)?), + }; + proto.encode(buf).map_err(|e| { + internal_datafusion_err!("failed to encode ScalarSubqueryExprExec: {e}") + }) + } + } + + /// The decode context's active scalar subquery results scope must reach a + /// `ScalarSubqueryExpr` decoded by a codec that is forced foreign through + /// the FFI boundary, not just a local one. Without the fix, the far side + /// always decodes with a root context (no scope), so the embedded + /// `ScalarSubqueryExpr` fails to deserialize. + #[test] + fn ffi_physical_extension_codec_forced_foreign_scalar_subquery_roundtrip() + -> Result<()> { + let schema = Arc::new(Schema::new(vec![Field::new("a", DataType::Int64, false)])); + let subquery_schema = + Arc::new(Schema::new(vec![Field::new("x", DataType::Int64, true)])); + + let results = ScalarSubqueryResults::new(1); + let sq_expr: Arc = Arc::new(ScalarSubqueryExpr::new( + DataType::Int64, + true, + SubqueryIndex::new(0), + results.clone(), + )); + let extension_plan: Arc = Arc::new(ScalarSubqueryExprExec { + expr: sq_expr, + child: Arc::new(RealEmptyExec::new(Arc::clone(&schema))), + }); + let plan: Arc = Arc::new(ScalarSubqueryExec::new( + extension_plan, + vec![ScalarSubqueryLink { + plan: Arc::new(RealEmptyExec::new(subquery_schema)), + index: SubqueryIndex::new(0), + }], + results, + )); + + let bytes = physical_plan_to_bytes_with_proto_converter( + Arc::clone(&plan), + &ScalarSubqueryExprExecCodec, + &DefaultPhysicalProtoConverter {}, + )?; + + let (ctx, task_ctx_provider) = crate::util::tests::test_session_and_ctx(); + let mut ffi_codec = FFI_PhysicalExtensionCodec::new( + Arc::new(ScalarSubqueryExprExecCodec), + None, + task_ctx_provider, + ); + ffi_codec.library_marker_id = crate::mock_foreign_marker_id; + let foreign_codec: Arc = (&ffi_codec).into(); + + let deserialized = physical_plan_from_bytes_with_proto_converter( + bytes.as_ref(), + ctx.task_ctx().as_ref(), + foreign_codec.as_ref(), + &DefaultPhysicalProtoConverter {}, + )?; + + let sq_exec = deserialized + .downcast_ref::() + .expect("expected ScalarSubqueryExec"); + let decoded_extension = sq_exec + .input() + .downcast_ref::() + .expect( + "expected ScalarSubqueryExprExec decoded through the forced-foreign codec", + ); + let decoded_sq_expr = decoded_extension + .expr + .downcast_ref::() + .expect("expected ScalarSubqueryExpr"); + + // The expression decoded by the forced-foreign codec must observe + // values written to the host's ScalarSubqueryExec results container. + assert_eq!(decoded_sq_expr.results().get(SubqueryIndex::new(0)), None); + sq_exec.results().set( + SubqueryIndex::new(0), + datafusion_common::ScalarValue::Int64(Some(42)), + )?; + assert_eq!( + decoded_sq_expr.results().get(SubqueryIndex::new(0)), + Some(datafusion_common::ScalarValue::Int64(Some(42))) + ); + + Ok(()) + } } diff --git a/datafusion/ffi/src/proto/scalar_subquery_results.rs b/datafusion/ffi/src/proto/scalar_subquery_results.rs new file mode 100644 index 0000000000000..0d2277befbaec --- /dev/null +++ b/datafusion/ffi/src/proto/scalar_subquery_results.rs @@ -0,0 +1,237 @@ +// 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. + +use std::ffi::c_void; +use std::sync::Arc; + +use datafusion_common::error::Result; +use datafusion_common::{DataFusionError, ScalarValue}; +use datafusion_expr::physical_planning_context::{ + ScalarSubqueryResults, ScalarSubqueryResultsBackend, SubqueryIndex, +}; +use prost::Message; + +use stabby::vec::Vec as SVec; + +use crate::util::{FFI_Option, FFI_Result}; +use crate::{df_result, sresult, sresult_return}; + +/// A stable struct for sharing the shared results container of a +/// `ScalarSubqueryExec` across FFI boundaries. +/// +/// Unlike most FFI wrappers in this crate, which stand in for a whole +/// `dyn Trait` object, this one stands in for a single +/// [`ScalarSubqueryResults`] value: the active scope that a +/// `PhysicalExtensionCodec` on the near side of the boundary passes to +/// `try_decode_with_ctx` so a `ScalarSubqueryExpr` decoded on the far side +/// reads the same populated values as the `ScalarSubqueryExec` that owns +/// them. +/// +/// `ScalarValue`s cross the boundary as prost-encoded bytes, the same wire +/// format `PhysicalExtensionCodec::try_decode`/`try_encode` already use, so +/// this struct never assumes a stable memory layout for `ScalarValue` itself. +#[repr(C)] +#[derive(Debug)] +pub struct FFI_ScalarSubqueryResults { + get: unsafe extern "C" fn(&Self, index: u64) -> FFI_Result>>, + + set: unsafe extern "C" fn(&Self, index: u64, value: SVec) -> FFI_Result<()>, + + clear: unsafe extern "C" fn(&Self), + + /// Used to create a clone on the provider of the results container. This + /// should only need to be called by the receiver of the handle. + clone: unsafe extern "C" fn(&Self) -> Self, + + /// Release the memory of the private data when it is no longer being used. + release: unsafe extern "C" fn(&mut Self), + + private_data: *mut c_void, + + /// Utility to identify when FFI objects are accessed locally through + /// the foreign interface. + library_marker_id: extern "C" fn() -> usize, +} + +unsafe impl Send for FFI_ScalarSubqueryResults {} +unsafe impl Sync for FFI_ScalarSubqueryResults {} + +impl FFI_ScalarSubqueryResults { + fn inner(&self) -> &ScalarSubqueryResults { + let private_data = self.private_data as *const ScalarSubqueryResults; + unsafe { &*private_data } + } + + /// Creates a new [`FFI_ScalarSubqueryResults`] sharing `results`. + pub fn new(results: ScalarSubqueryResults) -> Self { + let private_data = Box::new(results); + + Self { + get: get_fn_wrapper, + set: set_fn_wrapper, + clear: clear_fn_wrapper, + clone: clone_fn_wrapper, + release: release_fn_wrapper, + private_data: Box::into_raw(private_data).cast::(), + library_marker_id: crate::get_library_marker_id, + } + } +} + +fn encode_scalar_value(value: ScalarValue) -> Result> { + let proto: datafusion_proto::protobuf::ScalarValue = + (&value).try_into().map_err(DataFusionError::from)?; + Ok(proto.encode_to_vec().into_iter().collect()) +} + +fn decode_scalar_value(bytes: &[u8]) -> Result { + let proto = datafusion_proto::protobuf::ScalarValue::decode(bytes) + .map_err(|e| DataFusionError::External(Box::new(e)))?; + ScalarValue::try_from(&proto).map_err(DataFusionError::from) +} + +unsafe extern "C" fn get_fn_wrapper( + results: &FFI_ScalarSubqueryResults, + index: u64, +) -> FFI_Result>> { + let value = results.inner().get(SubqueryIndex::new(index as usize)); + + let encoded = value.map(encode_scalar_value).transpose(); + + sresult!(encoded.map(FFI_Option::from)) +} + +unsafe extern "C" fn set_fn_wrapper( + results: &FFI_ScalarSubqueryResults, + index: u64, + value: SVec, +) -> FFI_Result<()> { + let value = sresult_return!(decode_scalar_value(value.as_ref())); + sresult!( + results + .inner() + .set(SubqueryIndex::new(index as usize), value) + ) +} + +unsafe extern "C" fn clear_fn_wrapper(results: &FFI_ScalarSubqueryResults) { + results.inner().clear(); +} + +unsafe extern "C" fn clone_fn_wrapper( + results: &FFI_ScalarSubqueryResults, +) -> FFI_ScalarSubqueryResults { + FFI_ScalarSubqueryResults::new(results.inner().clone()) +} + +unsafe extern "C" fn release_fn_wrapper(results: &mut FFI_ScalarSubqueryResults) { + unsafe { + drop(Box::from_raw( + results.private_data.cast::(), + )); + } +} + +impl Drop for FFI_ScalarSubqueryResults { + fn drop(&mut self) { + unsafe { (self.release)(self) } + } +} + +impl Clone for FFI_ScalarSubqueryResults { + fn clone(&self) -> Self { + unsafe { (self.clone)(self) } + } +} + +/// This wrapper struct exists on the receiver side of the FFI interface, so it +/// has no guarantees about being able to access the data behind +/// `private_data`. It forwards every [`ScalarSubqueryResultsBackend`] call +/// back across the boundary to the real container owned by the other side. +#[derive(Debug)] +struct ForeignScalarSubqueryResultsBackend(FFI_ScalarSubqueryResults); + +unsafe impl Send for ForeignScalarSubqueryResultsBackend {} +unsafe impl Sync for ForeignScalarSubqueryResultsBackend {} + +impl ScalarSubqueryResultsBackend for ForeignScalarSubqueryResultsBackend { + fn get(&self, index: usize) -> Option { + let bytes = match unsafe { (self.0.get)(&self.0, index as u64) } { + FFI_Result::Ok(bytes) => bytes.into_option()?, + FFI_Result::Err(_) => return None, + }; + decode_scalar_value(bytes.as_ref()).ok() + } + + fn set(&self, index: usize, value: ScalarValue) -> Result<()> { + let bytes = encode_scalar_value(value)?; + df_result!(unsafe { (self.0.set)(&self.0, index as u64, bytes) }) + } + + fn clear(&self) { + unsafe { (self.0.clear)(&self.0) } + } +} + +impl From for ScalarSubqueryResults { + fn from(results: FFI_ScalarSubqueryResults) -> Self { + if (results.library_marker_id)() == crate::get_library_marker_id() { + Self::clone(results.inner()) + } else { + Self::from_backend(Arc::new(ForeignScalarSubqueryResultsBackend(results))) + } + } +} + +#[cfg(test)] +mod tests { + use datafusion_common::Result; + use datafusion_common::ScalarValue; + use datafusion_expr::physical_planning_context::{ + ScalarSubqueryResults, SubqueryIndex, + }; + + use super::FFI_ScalarSubqueryResults; + + #[test] + fn roundtrip_ffi_scalar_subquery_results() -> Result<()> { + let results = ScalarSubqueryResults::new(1); + let mut ffi_results = FFI_ScalarSubqueryResults::new(results.clone()); + ffi_results.library_marker_id = crate::mock_foreign_marker_id; + + let foreign: ScalarSubqueryResults = ffi_results.into(); + + assert_eq!(foreign.get(SubqueryIndex::new(0)), None); + + results.set(SubqueryIndex::new(0), ScalarValue::Int64(Some(42)))?; + assert_eq!( + foreign.get(SubqueryIndex::new(0)), + Some(ScalarValue::Int64(Some(42))) + ); + + foreign.clear(); + assert_eq!(results.get(SubqueryIndex::new(0)), None); + + foreign.set(SubqueryIndex::new(0), ScalarValue::Int64(Some(7)))?; + assert_eq!( + results.get(SubqueryIndex::new(0)), + Some(ScalarValue::Int64(Some(7))) + ); + + Ok(()) + } +} From 880f3768a5e6a6e32eacd6613c0cd32c03d5672e Mon Sep 17 00:00:00 2001 From: naman Date: Fri, 2 Oct 2026 23:10:18 +0530 Subject: [PATCH 3/7] fix: resolve semver-check regressions from the scalar subquery FFI fix - Move the new FFI_PhysicalExtensionCodec::try_decode_with_ctx field to the end of the repr(C) struct instead of inserting it in the middle, so every pre-existing field keeps its original offset. - Add RefUnwindSafe as a supertrait bound on ScalarSubqueryResultsBackend so ScalarSubqueryResults keeps implementing UnwindSafe/RefUnwindSafe, with a regression test pinning it. --- .../expr/src/physical_planning_context.rs | 13 +++++++++- .../ffi/src/proto/physical_extension_codec.rs | 26 +++++++++++-------- 2 files changed, 27 insertions(+), 12 deletions(-) diff --git a/datafusion/expr/src/physical_planning_context.rs b/datafusion/expr/src/physical_planning_context.rs index 133e2b03e58ff..11a5639bd76b3 100644 --- a/datafusion/expr/src/physical_planning_context.rs +++ b/datafusion/expr/src/physical_planning_context.rs @@ -133,7 +133,9 @@ impl SubqueryIndex { /// [`ScalarSubqueryResults`] backed by a proxy that calls back into the /// near side's real container for every operation, so both sides observe the /// same populated values. -pub trait ScalarSubqueryResultsBackend: fmt::Debug + Send + Sync { +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; @@ -291,6 +293,15 @@ impl Hash for ScalarSubqueryResults { 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() {} + assert_bounds::(); + } + #[test] fn scalar_subquery_results_set_and_get() -> Result<()> { let results = ScalarSubqueryResults::new(1); diff --git a/datafusion/ffi/src/proto/physical_extension_codec.rs b/datafusion/ffi/src/proto/physical_extension_codec.rs index c9e61b5e5c65e..ff3a879a2c7d3 100644 --- a/datafusion/ffi/src/proto/physical_extension_codec.rs +++ b/datafusion/ffi/src/proto/physical_extension_codec.rs @@ -55,17 +55,6 @@ pub struct FFI_PhysicalExtensionCodec { inputs: SVec, ) -> FFI_Result, - /// Decode bytes into an execution plan, forwarding the active scalar - /// subquery results scope (if any) from the caller's decode context so a - /// `ScalarSubqueryExpr` decoded on this side of the boundary shares the - /// same populated results as the `ScalarSubqueryExec` that owns them. - try_decode_with_ctx: unsafe extern "C" fn( - &Self, - buf: SSlice, - inputs: SVec, - scalar_subquery_results: FFI_Option, - ) -> FFI_Result, - /// Encode an execution plan into bytes. try_encode: unsafe extern "C" fn(&Self, node: FFI_ExecutionPlan) -> FFI_Result>, @@ -123,6 +112,21 @@ pub struct FFI_PhysicalExtensionCodec { /// Utility to identify when FFI objects are accessed locally through /// the foreign interface. pub library_marker_id: extern "C" fn() -> usize, + + /// Decode bytes into an execution plan, forwarding the active scalar + /// subquery results scope (if any) from the caller's decode context so a + /// `ScalarSubqueryExpr` decoded on this side of the boundary shares the + /// same populated results as the `ScalarSubqueryExec` that owns them. + /// + /// Added after the original fields, at the end of the struct, so that + /// older code built against this `repr(C)` struct without this field + /// still sees every field it knows about at its original offset. + try_decode_with_ctx: unsafe extern "C" fn( + &Self, + buf: SSlice, + inputs: SVec, + scalar_subquery_results: FFI_Option, + ) -> FFI_Result, } unsafe impl Send for FFI_PhysicalExtensionCodec {} From dd38c1c6995fbec4c6ae15b85a36edcb8e0e4158 Mon Sep 17 00:00:00 2001 From: naman Date: Sun, 4 Oct 2026 19:21:27 +0530 Subject: [PATCH 4/7] fix: clippy needless_pass_by_value in encode_scalar_value CI's clippy (ci/scripts/rust_clippy.sh, not --all-features) flagged encode_scalar_value for taking ScalarValue by value while only borrowing it. Takes &ScalarValue instead. --- datafusion/ffi/src/proto/scalar_subquery_results.rs | 8 ++++---- 1 file changed, 4 insertions(+), 4 deletions(-) diff --git a/datafusion/ffi/src/proto/scalar_subquery_results.rs b/datafusion/ffi/src/proto/scalar_subquery_results.rs index 0d2277befbaec..a6c2c4320a6e0 100644 --- a/datafusion/ffi/src/proto/scalar_subquery_results.rs +++ b/datafusion/ffi/src/proto/scalar_subquery_results.rs @@ -92,9 +92,9 @@ impl FFI_ScalarSubqueryResults { } } -fn encode_scalar_value(value: ScalarValue) -> Result> { +fn encode_scalar_value(value: &ScalarValue) -> Result> { let proto: datafusion_proto::protobuf::ScalarValue = - (&value).try_into().map_err(DataFusionError::from)?; + value.try_into().map_err(DataFusionError::from)?; Ok(proto.encode_to_vec().into_iter().collect()) } @@ -110,7 +110,7 @@ unsafe extern "C" fn get_fn_wrapper( ) -> FFI_Result>> { let value = results.inner().get(SubqueryIndex::new(index as usize)); - let encoded = value.map(encode_scalar_value).transpose(); + let encoded = value.as_ref().map(encode_scalar_value).transpose(); sresult!(encoded.map(FFI_Option::from)) } @@ -178,7 +178,7 @@ impl ScalarSubqueryResultsBackend for ForeignScalarSubqueryResultsBackend { } fn set(&self, index: usize, value: ScalarValue) -> Result<()> { - let bytes = encode_scalar_value(value)?; + let bytes = encode_scalar_value(&value)?; df_result!(unsafe { (self.0.set)(&self.0, index as u64, bytes) }) } From b9b7162d99d7c5beae3ee3d03cb6f378fd429c36 Mon Sep 17 00:00:00 2001 From: naman Date: Tue, 6 Oct 2026 12:51:34 +0530 Subject: [PATCH 5/7] fix: preserve ScalarValue's exact DataType across the FFI scalar-subquery bridge datafusion_proto::protobuf::ScalarValue has no dedicated variant for every DataType: a non-null Float16 is encoded as Float32Value, so a foreign ScalarSubqueryExpr reading it back got a Float32 instead. Wrap the encoded bytes with the original DataType (as an ArrowType) and cast back to it on decode, so the value crossing FFI always matches the type it was declared with. Adds forced-foreign non-null Float16 get/set regression coverage. --- .../ffi/src/proto/scalar_subquery_results.rs | 89 ++++++++++++++++++- 1 file changed, 85 insertions(+), 4 deletions(-) diff --git a/datafusion/ffi/src/proto/scalar_subquery_results.rs b/datafusion/ffi/src/proto/scalar_subquery_results.rs index a6c2c4320a6e0..60f21346a55a8 100644 --- a/datafusion/ffi/src/proto/scalar_subquery_results.rs +++ b/datafusion/ffi/src/proto/scalar_subquery_results.rs @@ -19,10 +19,11 @@ use std::ffi::c_void; use std::sync::Arc; use datafusion_common::error::Result; -use datafusion_common::{DataFusionError, ScalarValue}; +use datafusion_common::{DataFusionError, ScalarValue, internal_datafusion_err}; use datafusion_expr::physical_planning_context::{ ScalarSubqueryResults, ScalarSubqueryResultsBackend, SubqueryIndex, }; +use datafusion_proto::protobuf::ArrowType; use prost::Message; use stabby::vec::Vec as SVec; @@ -92,16 +93,56 @@ impl FFI_ScalarSubqueryResults { } } +/// The wire form a [`ScalarValue`] crosses the FFI boundary in. +/// +/// `datafusion_proto::protobuf::ScalarValue` alone is not a type-preserving +/// transport: it has no dedicated variant for every `DataType` (`Float16`, +/// for example, is encoded as `Float32Value`), so decoding it back can yield +/// a `ScalarValue` of a different, widened type. Carrying the original +/// `DataType` alongside the encoded value lets [`decode_scalar_value`] cast +/// back down to it, so a foreign `ScalarSubqueryExpr` observes the same type +/// it was declared with. +#[derive(Clone, PartialEq, Message)] +struct ScalarValueEnvelope { + #[prost(message, optional, tag = "1")] + data_type: Option, + #[prost(bytes, tag = "2")] + value: Vec, +} + fn encode_scalar_value(value: &ScalarValue) -> Result> { + let data_type: ArrowType = (&value.data_type()) + .try_into() + .map_err(DataFusionError::from)?; let proto: datafusion_proto::protobuf::ScalarValue = value.try_into().map_err(DataFusionError::from)?; - Ok(proto.encode_to_vec().into_iter().collect()) + let envelope = ScalarValueEnvelope { + data_type: Some(data_type), + value: proto.encode_to_vec(), + }; + Ok(envelope.encode_to_vec().into_iter().collect()) } fn decode_scalar_value(bytes: &[u8]) -> Result { - let proto = datafusion_proto::protobuf::ScalarValue::decode(bytes) + let envelope = ScalarValueEnvelope::decode(bytes) .map_err(|e| DataFusionError::External(Box::new(e)))?; - ScalarValue::try_from(&proto).map_err(DataFusionError::from) + let data_type_proto = envelope.data_type.ok_or_else(|| { + internal_datafusion_err!("ScalarValueEnvelope is missing its data_type") + })?; + let data_type = (&data_type_proto) + .try_into() + .map_err(DataFusionError::from)?; + + let proto = + datafusion_proto::protobuf::ScalarValue::decode(envelope.value.as_slice()) + .map_err(|e| DataFusionError::External(Box::new(e)))?; + let value = ScalarValue::try_from(&proto).map_err(DataFusionError::from)?; + + if value.data_type() == data_type { + Ok(value) + } else { + value.cast_to(&data_type) + } } unsafe extern "C" fn get_fn_wrapper( @@ -234,4 +275,44 @@ mod tests { Ok(()) } + + /// `datafusion_proto::protobuf::ScalarValue` has no dedicated variant for + /// every `DataType`: a non-null `Float16` is encoded as `Float32Value`, + /// so decoding it back without the original type yields a `Float32` + /// instead. A value crossing a forced-foreign `FFI_ScalarSubqueryResults` + /// handle, in either direction, must still be a `Float16`, matching what + /// a `ScalarSubqueryExpr` declared it would return. + #[test] + fn ffi_scalar_subquery_results_preserves_non_null_float16() -> Result<()> { + use datafusion_common::arrow::datatypes::DataType; + + let float16 = ScalarValue::Float32(Some(1.5)).cast_to(&DataType::Float16)?; + + // get: the host sets a Float16, a forced-foreign reader must read a + // Float16 back, not the Float32 the wire format would otherwise widen + // it to. + let results = ScalarSubqueryResults::new(1); + let mut ffi_results = FFI_ScalarSubqueryResults::new(results.clone()); + ffi_results.library_marker_id = crate::mock_foreign_marker_id; + let foreign: ScalarSubqueryResults = ffi_results.into(); + + results.set(SubqueryIndex::new(0), float16.clone())?; + let read_back = foreign.get(SubqueryIndex::new(0)); + assert_eq!(read_back, Some(float16.clone())); + assert_eq!(read_back.unwrap().data_type(), DataType::Float16); + + // set: a forced-foreign writer sets a Float16, the host must read a + // Float16 back too. + let results = ScalarSubqueryResults::new(1); + let mut ffi_results = FFI_ScalarSubqueryResults::new(results.clone()); + ffi_results.library_marker_id = crate::mock_foreign_marker_id; + let foreign: ScalarSubqueryResults = ffi_results.into(); + + foreign.set(SubqueryIndex::new(0), float16.clone())?; + let read_back = results.get(SubqueryIndex::new(0)); + assert_eq!(read_back, Some(float16.clone())); + assert_eq!(read_back.unwrap().data_type(), DataType::Float16); + + Ok(()) + } } From 851d8d09b6ed0b56be78c14cd78b454c144f251f Mon Sep 17 00:00:00 2001 From: namanjain24-sudo Date: Wed, 7 Oct 2026 16:00:14 +0530 Subject: [PATCH 6/7] test: add cross-library coverage for the FFI scalar-subquery results bridge The existing unit test only forces the codec foreign through mock_foreign_marker_id, so the results handle still takes the local-bypass fast path and never exercises ForeignScalarSubqueryResultsBackend. Adds an integration-tests-gated test that decodes inside a separately dlopen'd copy of the cdylib, so the handle genuinely crosses the boundary. Shares the scalar-subquery fixtures between the unit test and the new integration test via a fixtures module gated on cfg(any(test, feature = "integration-tests")). Also flags the FFI ABI break from the appended try_decode_with_ctx field for a maintainer to label, since the struct's size changes even though existing try_decode-only implementations keep compiling. --- .../ffi/src/proto/physical_extension_codec.rs | 317 ++++++++++-------- datafusion/ffi/src/tests/mod.rs | 85 +++++ datafusion/ffi/src/tests/utils.rs | 53 +++ .../ffi/tests/ffi_physical_extension_codec.rs | 111 ++++++ 4 files changed, 425 insertions(+), 141 deletions(-) create mode 100644 datafusion/ffi/tests/ffi_physical_extension_codec.rs diff --git a/datafusion/ffi/src/proto/physical_extension_codec.rs b/datafusion/ffi/src/proto/physical_extension_codec.rs index ff3a879a2c7d3..fe056a654f662 100644 --- a/datafusion/ffi/src/proto/physical_extension_codec.rs +++ b/datafusion/ffi/src/proto/physical_extension_codec.rs @@ -527,13 +527,172 @@ impl PhysicalExtensionCodec for ForeignPhysicalExtensionCodec { } } +/// Fixtures shared by the in-process unit tests in this module and the +/// cross-library integration test in +/// `datafusion/ffi/tests/ffi_physical_extension_codec.rs`. +/// +/// Gated on `feature = "integration-tests"` (in addition to `test`) so the +/// external integration-test binary, which links this crate as an ordinary +/// dependency rather than compiling it with `cfg(test)`, can still reach +/// these types through `datafusion_ffi::proto::physical_extension_codec::fixtures`. +#[cfg(any(test, feature = "integration-tests"))] +pub mod fixtures { + use std::sync::Arc; + + use datafusion_common::tree_node::TreeNodeRecursion; + use datafusion_common::{Result, exec_err, internal_datafusion_err}; + use datafusion_execution::TaskContext; + use datafusion_physical_plan::{ + ChildrenPropertiesMode, DisplayAs, DisplayFormatType, ExecutionPlan, + PhysicalExpr, PlanProperties, ReplaceChildrenOptions, SendableRecordBatchStream, + apply_expression_roots, + }; + use datafusion_proto::physical_plan::{ + PhysicalExtensionCodec, PhysicalPlanDecodeContext, + PhysicalProtoConverterExtension, + }; + use datafusion_proto::protobuf::PhysicalExprNode; + use prost::Message; + + /// An extension plan that carries a single physical expression, so its + /// codec has to decode that expression itself. + #[derive(Debug)] + pub struct ScalarSubqueryExprExec { + pub expr: Arc, + pub child: Arc, + } + + impl DisplayAs for ScalarSubqueryExprExec { + fn fmt_as( + &self, + _t: DisplayFormatType, + f: &mut std::fmt::Formatter, + ) -> std::fmt::Result { + write!(f, "ScalarSubqueryExprExec") + } + } + + impl ExecutionPlan for ScalarSubqueryExprExec { + fn name(&self) -> &str { + "ScalarSubqueryExprExec" + } + + fn properties(&self) -> &Arc { + self.child.properties() + } + + fn children(&self) -> Vec<&Arc> { + vec![&self.child] + } + + fn apply_expressions( + &self, + f: &mut dyn FnMut(&Arc) -> Result, + ) -> Result { + apply_expression_roots(std::slice::from_ref(&self.expr), f) + } + + fn replace_children( + self: Arc, + _: Vec>, + _: ReplaceChildrenOptions, + ) -> Result> { + unreachable!() + } + + fn with_new_children( + self: Arc, + children: Vec>, + ) -> Result> { + self.replace_children( + children, + ReplaceChildrenOptions::new(ChildrenPropertiesMode::Recompute), + ) + } + + fn execute( + &self, + _partition: usize, + _context: Arc, + ) -> Result { + unreachable!() + } + } + + #[derive(Clone, PartialEq, Message)] + pub struct ScalarSubqueryExprExecProto { + #[prost(message, optional, tag = "1")] + expr: Option, + } + + /// Decodes [`ScalarSubqueryExprExec`] through `try_decode_with_ctx`, + /// passing the decode context on to its expression. `try_decode` fails so + /// a caller that drops the context (by routing through the plain + /// `try_decode` FFI entry point instead of `try_decode_with_ctx`) is + /// caught by callers of this codec. + #[derive(Debug)] + pub struct ScalarSubqueryExprExecCodec; + + impl PhysicalExtensionCodec for ScalarSubqueryExprExecCodec { + fn try_decode( + &self, + _buf: &[u8], + _inputs: &[Arc], + _ctx: &TaskContext, + _proto_converter: &dyn PhysicalProtoConverterExtension, + ) -> Result> { + exec_err!("ScalarSubqueryExprExecCodec decodes through try_decode_with_ctx") + } + + fn try_decode_with_ctx( + &self, + buf: &[u8], + inputs: &[Arc], + ctx: &PhysicalPlanDecodeContext<'_>, + proto_converter: &dyn PhysicalProtoConverterExtension, + ) -> Result> { + let proto = ScalarSubqueryExprExecProto::decode(buf).map_err(|e| { + internal_datafusion_err!("failed to decode ScalarSubqueryExprExec: {e}") + })?; + let expr_proto = proto.expr.ok_or_else(|| { + internal_datafusion_err!("ScalarSubqueryExprExec is missing its expr") + })?; + let schema = inputs[0].schema(); + let expr = + proto_converter.proto_to_physical_expr(&expr_proto, &schema, ctx)?; + Ok(Arc::new(ScalarSubqueryExprExec { + expr, + child: Arc::clone(&inputs[0]), + })) + } + + fn try_encode( + &self, + node: Arc, + buf: &mut Vec, + proto_converter: &dyn PhysicalProtoConverterExtension, + ) -> Result<()> { + let exec = + node.downcast_ref::() + .ok_or_else(|| { + internal_datafusion_err!("expected ScalarSubqueryExprExec") + })?; + let proto = ScalarSubqueryExprExecProto { + expr: Some(proto_converter.physical_expr_to_proto(&exec.expr, self)?), + }; + proto.encode(buf).map_err(|e| { + internal_datafusion_err!("failed to encode ScalarSubqueryExprExec: {e}") + }) + } + } +} + #[cfg(test)] pub(crate) mod tests { use std::sync::Arc; use arrow_schema::{DataType, Field, Schema}; - use datafusion_common::tree_node::TreeNodeRecursion; - use datafusion_common::{Result, exec_err, internal_datafusion_err}; + use datafusion_common::{Result, exec_err}; use datafusion_execution::TaskContext; use datafusion_expr::physical_planning_context::{ ScalarSubqueryResults, SubqueryIndex, @@ -548,11 +707,7 @@ pub(crate) mod tests { use datafusion_physical_plan::scalar_subquery::{ ScalarSubqueryExec, ScalarSubqueryLink, }; - use datafusion_physical_plan::{ - ChildrenPropertiesMode, DisplayAs, DisplayFormatType, ExecutionPlan, - PhysicalExpr, PlanProperties, ReplaceChildrenOptions, SendableRecordBatchStream, - apply_expression_roots, - }; + use datafusion_physical_plan::{ExecutionPlan, PhysicalExpr}; use datafusion_proto::bytes::{ physical_plan_from_bytes_with_proto_converter, physical_plan_to_bytes_with_proto_converter, @@ -561,10 +716,11 @@ pub(crate) mod tests { DefaultPhysicalProtoConverter, PhysicalExtensionCodec, PhysicalPlanDecodeContext, PhysicalProtoConverterExtension, }; - use datafusion_proto::protobuf::PhysicalExprNode; - use prost::Message; use crate::execution_plan::tests::EmptyExec; + use crate::proto::physical_extension_codec::fixtures::{ + ScalarSubqueryExprExec, ScalarSubqueryExprExecCodec, + }; use crate::proto::physical_extension_codec::{ FFI_PhysicalExtensionCodec, ForeignPhysicalExtensionCodec, }; @@ -923,143 +1079,22 @@ pub(crate) mod tests { assert_eq!(task_ctx.session_id(), ctx_b.task_ctx().session_id()); } - /// An extension plan that carries a single physical expression, so its - /// codec has to decode that expression itself. - #[derive(Debug)] - struct ScalarSubqueryExprExec { - expr: Arc, - child: Arc, - } - - impl DisplayAs for ScalarSubqueryExprExec { - fn fmt_as( - &self, - _t: DisplayFormatType, - f: &mut std::fmt::Formatter, - ) -> std::fmt::Result { - write!(f, "ScalarSubqueryExprExec") - } - } - - impl ExecutionPlan for ScalarSubqueryExprExec { - fn name(&self) -> &str { - "ScalarSubqueryExprExec" - } - - fn properties(&self) -> &Arc { - self.child.properties() - } - - fn children(&self) -> Vec<&Arc> { - vec![&self.child] - } - - fn apply_expressions( - &self, - f: &mut dyn FnMut(&Arc) -> Result, - ) -> Result { - apply_expression_roots(std::slice::from_ref(&self.expr), f) - } - - fn replace_children( - self: Arc, - _: Vec>, - _: ReplaceChildrenOptions, - ) -> Result> { - unreachable!() - } - - fn with_new_children( - self: Arc, - children: Vec>, - ) -> Result> { - self.replace_children( - children, - ReplaceChildrenOptions::new(ChildrenPropertiesMode::Recompute), - ) - } - - fn execute( - &self, - _partition: usize, - _context: Arc, - ) -> Result { - unreachable!() - } - } - - #[derive(Clone, PartialEq, Message)] - struct ScalarSubqueryExprExecProto { - #[prost(message, optional, tag = "1")] - expr: Option, - } - - /// Decodes [`ScalarSubqueryExprExec`] through `try_decode_with_ctx`, - /// passing the decode context on to its expression. `try_decode` fails so - /// a caller that drops the context (by routing through the plain - /// `try_decode` FFI entry point instead of `try_decode_with_ctx`) is - /// caught by this test. - #[derive(Debug)] - struct ScalarSubqueryExprExecCodec; - - impl PhysicalExtensionCodec for ScalarSubqueryExprExecCodec { - fn try_decode( - &self, - _buf: &[u8], - _inputs: &[Arc], - _ctx: &TaskContext, - _proto_converter: &dyn PhysicalProtoConverterExtension, - ) -> Result> { - exec_err!("ScalarSubqueryExprExecCodec decodes through try_decode_with_ctx") - } - - fn try_decode_with_ctx( - &self, - buf: &[u8], - inputs: &[Arc], - ctx: &PhysicalPlanDecodeContext<'_>, - proto_converter: &dyn PhysicalProtoConverterExtension, - ) -> Result> { - let proto = ScalarSubqueryExprExecProto::decode(buf).map_err(|e| { - internal_datafusion_err!("failed to decode ScalarSubqueryExprExec: {e}") - })?; - let expr_proto = proto.expr.ok_or_else(|| { - internal_datafusion_err!("ScalarSubqueryExprExec is missing its expr") - })?; - let schema = inputs[0].schema(); - let expr = - proto_converter.proto_to_physical_expr(&expr_proto, &schema, ctx)?; - Ok(Arc::new(ScalarSubqueryExprExec { - expr, - child: Arc::clone(&inputs[0]), - })) - } - - fn try_encode( - &self, - node: Arc, - buf: &mut Vec, - proto_converter: &dyn PhysicalProtoConverterExtension, - ) -> Result<()> { - let exec = - node.downcast_ref::() - .ok_or_else(|| { - internal_datafusion_err!("expected ScalarSubqueryExprExec") - })?; - let proto = ScalarSubqueryExprExecProto { - expr: Some(proto_converter.physical_expr_to_proto(&exec.expr, self)?), - }; - proto.encode(buf).map_err(|e| { - internal_datafusion_err!("failed to encode ScalarSubqueryExprExec: {e}") - }) - } - } - /// The decode context's active scalar subquery results scope must reach a /// `ScalarSubqueryExpr` decoded by a codec that is forced foreign through /// the FFI boundary, not just a local one. Without the fix, the far side /// always decodes with a root context (no scope), so the embedded /// `ScalarSubqueryExpr` fails to deserialize. + /// + /// This only mocks the foreign path in-process (see + /// `mock_foreign_marker_id`): both the codec and the results handle are + /// marked foreign by overwriting `library_marker_id`, but every call + /// still runs in this same compilation unit. The cross-library + /// integration test `test_ffi_physical_extension_codec_scalar_subquery_cross_library` + /// in `datafusion/ffi/tests/ffi_physical_extension_codec.rs` exercises the + /// same scenario through a genuinely separate `dlopen`'d copy of this + /// crate, so the results handle actually round-trips through + /// `ForeignScalarSubqueryResultsBackend` rather than taking the + /// local-bypass fast path. #[test] fn ffi_physical_extension_codec_forced_foreign_scalar_subquery_roundtrip() -> Result<()> { diff --git a/datafusion/ffi/src/tests/mod.rs b/datafusion/ffi/src/tests/mod.rs index 7e583b6c5d5bf..989ae914ea6cc 100644 --- a/datafusion/ffi/src/tests/mod.rs +++ b/datafusion/ffi/src/tests/mod.rs @@ -273,6 +273,91 @@ pub extern "C" fn datafusion_ffi_test_create_exec_with_byte_metrics() -> FFI_Exe FFI_ExecutionPlan::new(plan, None) } +/// Decodes a [`crate::proto::physical_extension_codec::fixtures::ScalarSubqueryExprExec`] +/// from `bytes` using +/// [`crate::proto::physical_extension_codec::fixtures::ScalarSubqueryExprExecCodec`] +/// *inside this image*, forwarding `scalar_subquery_results` (if present) +/// into the decode context exactly the way a real +/// `FFI_PhysicalExtensionCodec::try_decode_with_ctx` call would, then reads +/// back the embedded `ScalarSubqueryExpr`'s `Int64` result at `index`. +/// +/// There is no generic `FFI_PhysicalExpr` wrapper in this crate to carry an +/// arbitrary decoded expression back across the boundary, so the decoded +/// `ScalarSubqueryExprExec`/`ScalarSubqueryExpr` never leave this image as +/// Rust values - only the `i64` their results container reports does. That +/// is still the assertion that matters: when `scalar_subquery_results` is +/// built in a genuinely different image (see the integration test in +/// `datafusion/ffi/tests/ffi_physical_extension_codec.rs`), reading it back +/// here only succeeds if `get`/`set` actually round-tripped through +/// `ForeignScalarSubqueryResultsBackend` into the caller's own results +/// container - the cross-library coverage gap the in-process, marker-mocked +/// unit test in `datafusion/ffi/src/proto/physical_extension_codec.rs` +/// cannot close on its own. +/// +/// Exported as its own top-level symbol rather than a new field on +/// [`ForeignLibraryModule`] for the same reason as +/// [`datafusion_ffi_test_create_exec_with_byte_metrics`]: that struct is +/// public, `#[repr(C)]`, and has no private/gated constructor, so adding a +/// field changes its exhaustive-construction ABI surface even for this +/// test-only, `integration-tests`-gated struct. +#[unsafe(no_mangle)] +pub extern "C" fn datafusion_ffi_test_decode_scalar_subquery_expr_exec_result( + task_ctx_provider: crate::execution::FFI_TaskContextProvider, + bytes: stabby::slice::Slice, + scalar_subquery_results: FFI_Option< + crate::proto::scalar_subquery_results::FFI_ScalarSubqueryResults, + >, + index: u64, +) -> crate::util::FFI_Result> { + use datafusion_common::ScalarValue; + use datafusion_execution::TaskContext; + use datafusion_expr::physical_planning_context::SubqueryIndex; + use datafusion_physical_expr::scalar_subquery::ScalarSubqueryExpr; + use datafusion_proto::physical_plan::{ + DefaultPhysicalProtoConverter, PhysicalExtensionCodec, PhysicalPlanDecodeContext, + }; + + use crate::proto::physical_extension_codec::fixtures::{ + ScalarSubqueryExprExec, ScalarSubqueryExprExecCodec, + }; + use crate::sresult_return; + use crate::util::FFI_Result; + + let task_ctx: Arc = sresult_return!((&task_ctx_provider).try_into()); + + let codec = ScalarSubqueryExprExecCodec; + let decode_ctx = PhysicalPlanDecodeContext::new(task_ctx.as_ref(), &codec); + let decode_ctx = match scalar_subquery_results.into_option() { + Some(results) => decode_ctx.with_scalar_subquery_results(results.into()), + None => decode_ctx, + }; + + // `ScalarSubqueryExprExecCodec::try_decode_with_ctx` only uses this to + // read a schema; `ScalarSubqueryExpr` carries no column references. + let dummy_input: Arc = + Arc::new(EmptyExec::new(Arc::new(Schema::empty()))); + + let plan = sresult_return!(codec.try_decode_with_ctx( + bytes.as_ref(), + &[dummy_input], + &decode_ctx, + &DefaultPhysicalProtoConverter {}, + )); + + let Some(exec) = plan.downcast_ref::() else { + return FFI_Result::Err("expected ScalarSubqueryExprExec".into()); + }; + let Some(sq_expr) = exec.expr.downcast_ref::() else { + return FFI_Result::Err("expected ScalarSubqueryExpr".into()); + }; + + match sq_expr.results().get(SubqueryIndex::new(index as usize)) { + None => FFI_Result::Ok(FFI_Option::None), + Some(ScalarValue::Int64(v)) => FFI_Result::Ok(v.into()), + Some(_) => FFI_Result::Err("expected an Int64 ScalarValue".into()), + } +} + /// Thin wrapper that attaches a fixed [`Statistics`] snapshot to any inner /// [`TableProvider`] without changing its scan behaviour. #[derive(Debug)] diff --git a/datafusion/ffi/src/tests/utils.rs b/datafusion/ffi/src/tests/utils.rs index 2169cf04186d9..d7b9da10dfee5 100644 --- a/datafusion/ffi/src/tests/utils.rs +++ b/datafusion/ffi/src/tests/utils.rs @@ -120,6 +120,59 @@ pub fn get_byte_metrics_exec() -> Result, + index: u64, +) -> Result> { + let lib_path = find_library()?; + + let lib = unsafe { + libloading::Library::new(&lib_path) + .map_err(|e| DataFusionError::External(Box::new(e)))? + }; + + let decode: libloading::Symbol< + extern "C" fn( + crate::execution::FFI_TaskContextProvider, + stabby::slice::Slice, + crate::util::FFI_Option< + crate::proto::scalar_subquery_results::FFI_ScalarSubqueryResults, + >, + u64, + ) -> crate::util::FFI_Result>, + > = unsafe { + lib.get(b"datafusion_ffi_test_decode_scalar_subquery_expr_exec_result") + .map_err(|e| DataFusionError::External(Box::new(e)))? + }; + + let result = decode( + task_ctx_provider, + bytes.into(), + scalar_subquery_results, + index, + ); + + // Leak the library to keep it loaded for the duration of the test + #[expect(clippy::mem_forget)] + std::mem::forget(lib); + + match result { + crate::util::FFI_Result::Ok(value) => Ok(value.into_option()), + crate::util::FFI_Result::Err(e) => { + Err(DataFusionError::External(e.to_string().into())) + } + } +} + /// Load an independent copy of the integration-test cdylib. /// /// Copying to a unique path makes the dynamic loader create a separate image diff --git a/datafusion/ffi/tests/ffi_physical_extension_codec.rs b/datafusion/ffi/tests/ffi_physical_extension_codec.rs new file mode 100644 index 0000000000000..4a983bbe8e003 --- /dev/null +++ b/datafusion/ffi/tests/ffi_physical_extension_codec.rs @@ -0,0 +1,111 @@ +// 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. + +#[cfg(feature = "integration-tests")] +mod tests { + use std::sync::Arc; + + use arrow_schema::{DataType, 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::proto::physical_extension_codec::fixtures::{ + ScalarSubqueryExprExec, ScalarSubqueryExprExecCodec, + }; + use datafusion_ffi::proto::scalar_subquery_results::FFI_ScalarSubqueryResults; + use datafusion_ffi::tests::utils::decode_scalar_subquery_expr_exec_result; + use datafusion_ffi::util::FFI_Option; + 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, + }; + + /// The scalar-subquery results scope a `PhysicalExtensionCodec` receives + /// through `try_decode_with_ctx` must reach a `ScalarSubqueryExpr` + /// decoded by a codec that is foreign *for real*: the decode in this + /// test runs inside a separately `dlopen`'d copy of the + /// `datafusion-ffi` cdylib (see + /// [`decode_scalar_subquery_expr_exec_result`]), a genuinely different + /// compiled image from this test binary, which links the crate as an + /// ordinary (statically-linked) rlib dependency instead. + /// + /// That image difference means the `FFI_ScalarSubqueryResults` handle + /// this test builds and passes across is not recognized as local on the + /// far side, so decoding there must go through + /// `ForeignScalarSubqueryResultsBackend`, round-tripping every + /// `get`/`set` back across the boundary - not the local-bypass fast path + /// the in-process, marker-mocked unit test + /// `ffi_physical_extension_codec_forced_foreign_scalar_subquery_roundtrip` + /// (in `datafusion/ffi/src/proto/physical_extension_codec.rs`) takes. + #[test] + fn test_ffi_physical_extension_codec_scalar_subquery_cross_library() + -> Result<(), DataFusionError> { + let results = ScalarSubqueryResults::new(1); + let sq_expr: Arc = Arc::new(ScalarSubqueryExpr::new( + DataType::Int64, + true, + SubqueryIndex::new(0), + results.clone(), + )); + let extension_plan: Arc = Arc::new(ScalarSubqueryExprExec { + expr: sq_expr, + child: 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; + + // Nothing has been written to `results` yet. + let read = decode_scalar_subquery_expr_exec_result( + FFI_TaskContextProvider::from(&host_task_ctx_provider), + &bytes, + FFI_Option::Some(FFI_ScalarSubqueryResults::new(results.clone())), + 0, + )?; + assert_eq!(read, None); + + // The host writes a value into its own results container. A fresh + // decode inside the foreign image must observe it - this only + // happens if the handle genuinely round-tripped through + // `ForeignScalarSubqueryResultsBackend` into *this* container, + // rather than a disconnected one private to the foreign side. + results.set(SubqueryIndex::new(0), ScalarValue::Int64(Some(42)))?; + + let read = decode_scalar_subquery_expr_exec_result( + FFI_TaskContextProvider::from(&host_task_ctx_provider), + &bytes, + FFI_Option::Some(FFI_ScalarSubqueryResults::new(results.clone())), + 0, + )?; + assert_eq!(read, Some(42)); + + Ok(()) + } +} From 24854ee53b64b6aa78eadf71a9ebb7875e76d908 Mon Sep 17 00:00:00 2001 From: namanjain24-sudo <180642416+namanjain24-sudo@users.noreply.github.com> Date: Fri, 9 Oct 2026 10:42:39 +0530 Subject: [PATCH 7/7] test: exercise the production try_decode_with_ctx callback in the cross-library test The cross-library integration test decoded through a bespoke test-only entry point instead of the real FFI_PhysicalExtensionCodec callback a foreign codec actually goes through, so it never exercised try_decode_with_ctx_fn_wrapper or ForeignPhysicalExtensionCodec itself. Now it builds the codec inside the dlopen'd cdylib, converts it to Arc (confirming it is really foreign), and calls try_decode_with_ctx on that Arc directly - the same call any caller of a foreign codec makes. The decoded plan comes back as an opaque ForeignExecutionPlan, so the host-value assertion now goes through execute() instead of a downcast: ScalarSubqueryExprExec gets a real execute() that evaluates its expression and emits the result as a RecordBatch, giving the host-populated value a way to cross back out through the FFI_ExecutionPlan/FFI_RecordBatchStream machinery every ExecutionPlan already has, without a generic FFI_PhysicalExpr wrapper. ScalarSubqueryExprExec gains a dedicated single-column PlanProperties (no longer delegating to its child's) and a constructor that computes it, used by the fixture's own decode path and both tests. --- .../ffi/src/proto/physical_extension_codec.rs | 73 +++++++-- datafusion/ffi/src/tests/mod.rs | 103 ++++--------- datafusion/ffi/src/tests/utils.rs | 44 ++---- .../ffi/tests/ffi_physical_extension_codec.rs | 139 ++++++++++++------ 4 files changed, 195 insertions(+), 164 deletions(-) diff --git a/datafusion/ffi/src/proto/physical_extension_codec.rs b/datafusion/ffi/src/proto/physical_extension_codec.rs index fe056a654f662..ca9eec90f6f95 100644 --- a/datafusion/ffi/src/proto/physical_extension_codec.rs +++ b/datafusion/ffi/src/proto/physical_extension_codec.rs @@ -539,13 +539,18 @@ impl PhysicalExtensionCodec for ForeignPhysicalExtensionCodec { pub mod fixtures { use std::sync::Arc; + use arrow::array::RecordBatch; + use arrow_schema::{Field, Schema}; use datafusion_common::tree_node::TreeNodeRecursion; use datafusion_common::{Result, exec_err, internal_datafusion_err}; use datafusion_execution::TaskContext; + use datafusion_physical_expr::EquivalenceProperties; + use datafusion_physical_plan::execution_plan::{Boundedness, EmissionType}; + use datafusion_physical_plan::stream::RecordBatchStreamAdapter; use datafusion_physical_plan::{ ChildrenPropertiesMode, DisplayAs, DisplayFormatType, ExecutionPlan, - PhysicalExpr, PlanProperties, ReplaceChildrenOptions, SendableRecordBatchStream, - apply_expression_roots, + Partitioning, PhysicalExpr, PlanProperties, ReplaceChildrenOptions, + SendableRecordBatchStream, apply_expression_roots, }; use datafusion_proto::physical_plan::{ PhysicalExtensionCodec, PhysicalPlanDecodeContext, @@ -555,11 +560,42 @@ pub mod fixtures { use prost::Message; /// An extension plan that carries a single physical expression, so its - /// codec has to decode that expression itself. + /// codec has to decode that expression itself. Executing it evaluates + /// that expression and emits the result as the single row of a + /// single-column batch, which is how the cross-library integration test + /// observes a decoded `ScalarSubqueryExpr`'s value without a generic + /// `FFI_PhysicalExpr` wrapper: the value crosses back through the + /// `FFI_ExecutionPlan`/`FFI_RecordBatchStream` machinery every + /// `ExecutionPlan` already has, via the Arrow C Data Interface. #[derive(Debug)] pub struct ScalarSubqueryExprExec { pub expr: Arc, pub child: Arc, + properties: Arc, + } + + impl ScalarSubqueryExprExec { + pub fn new(expr: Arc, child: Arc) -> Self { + let field = expr + .return_field(child.schema().as_ref()) + .expect("ScalarSubqueryExpr always returns a field"); + let schema = Arc::new(Schema::new(vec![Field::new( + field.name(), + field.data_type().clone(), + field.is_nullable(), + )])); + let properties = Arc::new(PlanProperties::new( + EquivalenceProperties::new(schema), + Partitioning::UnknownPartitioning(1), + EmissionType::Incremental, + Boundedness::Bounded, + )); + Self { + expr, + child, + properties, + } + } } impl DisplayAs for ScalarSubqueryExprExec { @@ -578,7 +614,7 @@ pub mod fixtures { } fn properties(&self) -> &Arc { - self.child.properties() + &self.properties } fn children(&self) -> Vec<&Arc> { @@ -610,12 +646,24 @@ pub mod fixtures { ) } + /// `ScalarSubqueryExpr::evaluate` ignores its batch argument - it + /// only reads its own cached result - so this evaluates it directly + /// rather than pulling a batch from `child`, and always emits + /// exactly one row regardless of what `child` would have produced. fn execute( &self, _partition: usize, _context: Arc, ) -> Result { - unreachable!() + let schema = self.properties.eq_properties.schema(); + let empty_batch = RecordBatch::new_empty(self.child.schema()); + let value = self.expr.evaluate(&empty_batch)?; + let array = value.into_array(1)?; + let batch = RecordBatch::try_new(Arc::clone(schema), vec![array])?; + Ok(Box::pin(RecordBatchStreamAdapter::new( + Arc::clone(schema), + futures::stream::once(async move { Ok(batch) }), + ))) } } @@ -660,10 +708,10 @@ pub mod fixtures { let schema = inputs[0].schema(); let expr = proto_converter.proto_to_physical_expr(&expr_proto, &schema, ctx)?; - Ok(Arc::new(ScalarSubqueryExprExec { + Ok(Arc::new(ScalarSubqueryExprExec::new( expr, - child: Arc::clone(&inputs[0]), - })) + Arc::clone(&inputs[0]), + ))) } fn try_encode( @@ -1109,10 +1157,11 @@ pub(crate) mod tests { SubqueryIndex::new(0), results.clone(), )); - let extension_plan: Arc = Arc::new(ScalarSubqueryExprExec { - expr: sq_expr, - child: Arc::new(RealEmptyExec::new(Arc::clone(&schema))), - }); + let extension_plan: Arc = + Arc::new(ScalarSubqueryExprExec::new( + sq_expr, + Arc::new(RealEmptyExec::new(Arc::clone(&schema))), + )); let plan: Arc = Arc::new(ScalarSubqueryExec::new( extension_plan, vec![ScalarSubqueryLink { diff --git a/datafusion/ffi/src/tests/mod.rs b/datafusion/ffi/src/tests/mod.rs index 989ae914ea6cc..5181eb0682f17 100644 --- a/datafusion/ffi/src/tests/mod.rs +++ b/datafusion/ffi/src/tests/mod.rs @@ -273,26 +273,30 @@ pub extern "C" fn datafusion_ffi_test_create_exec_with_byte_metrics() -> FFI_Exe FFI_ExecutionPlan::new(plan, None) } -/// Decodes a [`crate::proto::physical_extension_codec::fixtures::ScalarSubqueryExprExec`] -/// from `bytes` using +/// Builds a [`FFI_PhysicalExtensionCodec`] wrapping /// [`crate::proto::physical_extension_codec::fixtures::ScalarSubqueryExprExecCodec`] -/// *inside this image*, forwarding `scalar_subquery_results` (if present) -/// into the decode context exactly the way a real -/// `FFI_PhysicalExtensionCodec::try_decode_with_ctx` call would, then reads -/// back the embedded `ScalarSubqueryExpr`'s `Int64` result at `index`. +/// whose function pointers - including `try_decode_with_ctx` - point at code +/// compiled into *this* image. /// -/// There is no generic `FFI_PhysicalExpr` wrapper in this crate to carry an -/// arbitrary decoded expression back across the boundary, so the decoded -/// `ScalarSubqueryExprExec`/`ScalarSubqueryExpr` never leave this image as -/// Rust values - only the `i64` their results container reports does. That -/// is still the assertion that matters: when `scalar_subquery_results` is -/// built in a genuinely different image (see the integration test in -/// `datafusion/ffi/tests/ffi_physical_extension_codec.rs`), reading it back -/// here only succeeds if `get`/`set` actually round-tripped through -/// `ForeignScalarSubqueryResultsBackend` into the caller's own results -/// container - the cross-library coverage gap the in-process, marker-mocked -/// unit test in `datafusion/ffi/src/proto/physical_extension_codec.rs` -/// cannot close on its own. +/// A consumer that loads this through `libloading` (see +/// [`crate::tests::utils::get_scalar_subquery_expr_exec_codec`]) and converts +/// the result into `Arc` takes the real +/// `ForeignPhysicalExtensionCodec` branch, because the two copies of this +/// crate's `LIBRARY_MARKER` static sit at different addresses. Calling the +/// real `try_decode_with_ctx` on that `Arc` - the same production entry +/// point `try_decode_with_ctx_fn_wrapper` dispatches through for any foreign +/// codec - then makes `ForeignPhysicalExtensionCodec::try_decode_with_ctx` +/// wrap the caller's live +/// [`datafusion_expr::physical_planning_context::ScalarSubqueryResults`] in a +/// genuinely foreign `FFI_ScalarSubqueryResults`, so the `ScalarSubqueryExpr` +/// decoded on this side reads through `ForeignScalarSubqueryResultsBackend` +/// instead of the local-bypass fast path - the cross-library coverage gap +/// the in-process, marker-mocked unit test in +/// `datafusion/ffi/src/proto/physical_extension_codec.rs` cannot close on +/// its own. The decoded plan's `execute` (also dispatched through the real +/// `FFI_ExecutionPlan`) then carries that value back out as a `RecordBatch` +/// over the Arrow C Data Interface, with no need for a generic +/// `FFI_PhysicalExpr` wrapper. /// /// Exported as its own top-level symbol rather than a new field on /// [`ForeignLibraryModule`] for the same reason as @@ -301,61 +305,16 @@ pub extern "C" fn datafusion_ffi_test_create_exec_with_byte_metrics() -> FFI_Exe /// field changes its exhaustive-construction ABI surface even for this /// test-only, `integration-tests`-gated struct. #[unsafe(no_mangle)] -pub extern "C" fn datafusion_ffi_test_decode_scalar_subquery_expr_exec_result( +pub extern "C" fn datafusion_ffi_test_create_scalar_subquery_expr_exec_codec( task_ctx_provider: crate::execution::FFI_TaskContextProvider, - bytes: stabby::slice::Slice, - scalar_subquery_results: FFI_Option< - crate::proto::scalar_subquery_results::FFI_ScalarSubqueryResults, - >, - index: u64, -) -> crate::util::FFI_Result> { - use datafusion_common::ScalarValue; - use datafusion_execution::TaskContext; - use datafusion_expr::physical_planning_context::SubqueryIndex; - use datafusion_physical_expr::scalar_subquery::ScalarSubqueryExpr; - use datafusion_proto::physical_plan::{ - DefaultPhysicalProtoConverter, PhysicalExtensionCodec, PhysicalPlanDecodeContext, - }; - - use crate::proto::physical_extension_codec::fixtures::{ - ScalarSubqueryExprExec, ScalarSubqueryExprExecCodec, - }; - use crate::sresult_return; - use crate::util::FFI_Result; - - let task_ctx: Arc = sresult_return!((&task_ctx_provider).try_into()); - - let codec = ScalarSubqueryExprExecCodec; - let decode_ctx = PhysicalPlanDecodeContext::new(task_ctx.as_ref(), &codec); - let decode_ctx = match scalar_subquery_results.into_option() { - Some(results) => decode_ctx.with_scalar_subquery_results(results.into()), - None => decode_ctx, - }; - - // `ScalarSubqueryExprExecCodec::try_decode_with_ctx` only uses this to - // read a schema; `ScalarSubqueryExpr` carries no column references. - let dummy_input: Arc = - Arc::new(EmptyExec::new(Arc::new(Schema::empty()))); - - let plan = sresult_return!(codec.try_decode_with_ctx( - bytes.as_ref(), - &[dummy_input], - &decode_ctx, - &DefaultPhysicalProtoConverter {}, - )); - - let Some(exec) = plan.downcast_ref::() else { - return FFI_Result::Err("expected ScalarSubqueryExprExec".into()); - }; - let Some(sq_expr) = exec.expr.downcast_ref::() else { - return FFI_Result::Err("expected ScalarSubqueryExpr".into()); - }; - - match sq_expr.results().get(SubqueryIndex::new(index as usize)) { - None => FFI_Result::Ok(FFI_Option::None), - Some(ScalarValue::Int64(v)) => FFI_Result::Ok(v.into()), - Some(_) => FFI_Result::Err("expected an Int64 ScalarValue".into()), - } +) -> FFI_PhysicalExtensionCodec { + use crate::proto::physical_extension_codec::fixtures::ScalarSubqueryExprExecCodec; + + FFI_PhysicalExtensionCodec::new( + Arc::new(ScalarSubqueryExprExecCodec), + None, + task_ctx_provider, + ) } /// Thin wrapper that attaches a fixed [`Statistics`] snapshot to any inner diff --git a/datafusion/ffi/src/tests/utils.rs b/datafusion/ffi/src/tests/utils.rs index d7b9da10dfee5..d751cc7fc30a2 100644 --- a/datafusion/ffi/src/tests/utils.rs +++ b/datafusion/ffi/src/tests/utils.rs @@ -120,19 +120,14 @@ pub fn get_byte_metrics_exec() -> Result, - index: u64, -) -> Result> { +) -> Result { let lib_path = find_library()?; let lib = unsafe { @@ -140,37 +135,22 @@ pub fn decode_scalar_subquery_expr_exec_result( .map_err(|e| DataFusionError::External(Box::new(e)))? }; - let decode: libloading::Symbol< + let create_codec: libloading::Symbol< extern "C" fn( crate::execution::FFI_TaskContextProvider, - stabby::slice::Slice, - crate::util::FFI_Option< - crate::proto::scalar_subquery_results::FFI_ScalarSubqueryResults, - >, - u64, - ) -> crate::util::FFI_Result>, + ) -> crate::proto::physical_extension_codec::FFI_PhysicalExtensionCodec, > = unsafe { - lib.get(b"datafusion_ffi_test_decode_scalar_subquery_expr_exec_result") + lib.get(b"datafusion_ffi_test_create_scalar_subquery_expr_exec_codec") .map_err(|e| DataFusionError::External(Box::new(e)))? }; - let result = decode( - task_ctx_provider, - bytes.into(), - scalar_subquery_results, - index, - ); + let codec = create_codec(task_ctx_provider); // Leak the library to keep it loaded for the duration of the test #[expect(clippy::mem_forget)] std::mem::forget(lib); - match result { - crate::util::FFI_Result::Ok(value) => Ok(value.into_option()), - crate::util::FFI_Result::Err(e) => { - Err(DataFusionError::External(e.to_string().into())) - } - } + Ok(codec) } /// Load an independent copy of the integration-test cdylib. diff --git a/datafusion/ffi/tests/ffi_physical_extension_codec.rs b/datafusion/ffi/tests/ffi_physical_extension_codec.rs index 4a983bbe8e003..41e51ba49b778 100644 --- a/datafusion/ffi/tests/ffi_physical_extension_codec.rs +++ b/datafusion/ffi/tests/ffi_physical_extension_codec.rs @@ -17,9 +17,11 @@ #[cfg(feature = "integration-tests")] mod tests { + use std::any::Any; use std::sync::Arc; - use arrow_schema::{DataType, Schema}; + use arrow::array::{AsArray, RecordBatch}; + use arrow::datatypes::{DataType, Int64Type, Schema}; use datafusion::prelude::SessionContext; use datafusion_common::{DataFusionError, ScalarValue}; use datafusion_execution::TaskContextProvider; @@ -27,50 +29,60 @@ mod tests { 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::proto::scalar_subquery_results::FFI_ScalarSubqueryResults; - use datafusion_ffi::tests::utils::decode_scalar_subquery_expr_exec_result; - use datafusion_ffi::util::FFI_Option; + 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, + DefaultPhysicalProtoConverter, PhysicalExtensionCodec, PhysicalPlanDecodeContext, }; - /// The scalar-subquery results scope a `PhysicalExtensionCodec` receives - /// through `try_decode_with_ctx` must reach a `ScalarSubqueryExpr` - /// decoded by a codec that is foreign *for real*: the decode in this - /// test runs inside a separately `dlopen`'d copy of the - /// `datafusion-ffi` cdylib (see - /// [`decode_scalar_subquery_expr_exec_result`]), a genuinely different - /// compiled image from this test binary, which links the crate as an - /// ordinary (statically-linked) rlib dependency instead. + /// The decode context's active scalar-subquery results scope must reach + /// a `ScalarSubqueryExpr` decoded by the real, production + /// `PhysicalExtensionCodec::try_decode_with_ctx` entry point when the + /// codec is foreign *for real*. /// - /// That image difference means the `FFI_ScalarSubqueryResults` handle - /// this test builds and passes across is not recognized as local on the - /// far side, so decoding there must go through - /// `ForeignScalarSubqueryResultsBackend`, round-tripping every - /// `get`/`set` back across the boundary - not the local-bypass fast path - /// the in-process, marker-mocked unit test - /// `ffi_physical_extension_codec_forced_foreign_scalar_subquery_roundtrip` - /// (in `datafusion/ffi/src/proto/physical_extension_codec.rs`) takes. - #[test] - fn test_ffi_physical_extension_codec_scalar_subquery_cross_library() + /// The codec here is built inside a separately `dlopen`'d copy of this + /// crate's cdylib (see [`get_scalar_subquery_expr_exec_codec`]), so + /// converting it to `Arc` takes the real + /// `ForeignPhysicalExtensionCodec` branch, and calling + /// `try_decode_with_ctx` on it dispatches through the same + /// `try_decode_with_ctx_fn_wrapper` any foreign codec goes through in + /// production - not a bespoke test-only entry point. That makes + /// `ForeignPhysicalExtensionCodec::try_decode_with_ctx` wrap this test's + /// live `ScalarSubqueryResults` in a genuinely foreign + /// `FFI_ScalarSubqueryResults`, so the decoded expression reads through + /// `ForeignScalarSubqueryResultsBackend`. + /// + /// The decoded plan comes back as a `ForeignExecutionPlan`, opaque to + /// this side - there is no generic `FFI_PhysicalExpr` wrapper to pull the + /// embedded `ScalarSubqueryExpr` back out directly. Executing it instead + /// (through the real `FFI_ExecutionPlan`/`FFI_RecordBatchStream` + /// machinery every `ExecutionPlan` already has) carries the value back + /// out as an ordinary `RecordBatch`, which is where the host-value + /// assertion happens. + #[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 = Arc::new(ScalarSubqueryExpr::new( DataType::Int64, true, SubqueryIndex::new(0), results.clone(), )); - let extension_plan: Arc = Arc::new(ScalarSubqueryExprExec { - expr: sq_expr, - child: Arc::new(EmptyExec::new(Arc::new(Schema::empty()))), - }); + let extension_plan: Arc = + Arc::new(ScalarSubqueryExprExec::new( + sq_expr, + Arc::new(EmptyExec::new(Arc::new(Schema::empty()))), + )); let mut bytes = Vec::new(); ScalarSubqueryExprExecCodec.try_encode( @@ -81,30 +93,61 @@ mod tests { let ctx = Arc::new(SessionContext::default()); let host_task_ctx_provider = Arc::clone(&ctx) as Arc; + let ffi_task_ctx_provider = + FFI_TaskContextProvider::from(&host_task_ctx_provider); - // Nothing has been written to `results` yet. - let read = decode_scalar_subquery_expr_exec_result( - FFI_TaskContextProvider::from(&host_task_ctx_provider), - &bytes, - FFI_Option::Some(FFI_ScalarSubqueryResults::new(results.clone())), - 0, - )?; - assert_eq!(read, None); - - // The host writes a value into its own results container. A fresh - // decode inside the foreign image must observe it - this only - // happens if the handle genuinely round-tripped through - // `ForeignScalarSubqueryResultsBackend` into *this* container, - // rather than a disconnected one private to the foreign side. - results.set(SubqueryIndex::new(0), ScalarValue::Int64(Some(42)))?; + // The codec itself is produced inside a genuinely separate copy of + // this crate's cdylib, loaded through `libloading`. + let foreign_ffi_codec = + get_scalar_subquery_expr_exec_codec(ffi_task_ctx_provider)?; + let foreign_codec: Arc = (&foreign_ffi_codec).into(); + assert!( + (Arc::clone(&foreign_codec) as Arc) + .downcast_ref::() + .is_some(), + "codec built inside the dlopen'd cdylib must be seen as foreign" + ); - let read = decode_scalar_subquery_expr_exec_result( - FFI_TaskContextProvider::from(&host_task_ctx_provider), - &bytes, - FFI_Option::Some(FFI_ScalarSubqueryResults::new(results.clone())), - 0, + // The production entry point: calling `try_decode_with_ctx` directly + // on the foreign `Arc` is exactly what + // `PhysicalPlanDecodeContext`-aware decoding does for any extension + // plan, foreign or not - there is no separate test-only path here. + 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 = + 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_eq!(read, Some(42)); + assert!( + decoded.downcast_ref::().is_none(), + "a plan decoded by a genuinely foreign codec must stay opaque, not downcast locally" + ); + assert!( + decoded.downcast_ref::().is_some(), + "expected the decoded plan to be a ForeignExecutionPlan" + ); + + // Executing the foreign plan crosses back into the cdylib, where the + // decoded ScalarSubqueryExpr reads through + // ForeignScalarSubqueryResultsBackend into this test's `results`, + // and the value crosses back out as a RecordBatch. + let stream = decoded.execute(0, task_ctx)?; + let batches: Vec = + 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::().value(0); + assert_eq!(value, 42); Ok(()) }