diff --git a/datafusion/expr/src/physical_planning_context.rs b/datafusion/expr/src/physical_planning_context.rs index b2e579ea3c7cb..11a5639bd76b3 100644 --- a/datafusion/expr/src/physical_planning_context.rs +++ b/datafusion/expr/src/physical_planning_context.rs @@ -123,70 +123,150 @@ 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 + std::panic::RefUnwindSafe +{ + /// 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 +280,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) + } + } } } @@ -208,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/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 47a6f1e95cafb..ca9eec90f6f95 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, }; @@ -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. @@ -111,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 {} @@ -148,10 +164,45 @@ 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 {}, + )); + + 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 {}, )); @@ -310,6 +361,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, @@ -374,6 +426,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, @@ -444,6 +527,214 @@ 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 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, + Partitioning, 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. 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 { + 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.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), + ) + } + + /// `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 { + 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) }), + ))) + } + } + + #[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::new( + expr, + 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; @@ -451,18 +742,33 @@ pub(crate) mod tests { use arrow_schema::{DataType, Field, Schema}; use datafusion_common::{Result, exec_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::{ExecutionPlan, PhysicalExpr}; + 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, + DefaultPhysicalProtoConverter, PhysicalExtensionCodec, PhysicalPlanDecodeContext, PhysicalProtoConverterExtension, }; use crate::execution_plan::tests::EmptyExec; + use crate::proto::physical_extension_codec::fixtures::{ + ScalarSubqueryExprExec, ScalarSubqueryExprExecCodec, + }; use crate::proto::physical_extension_codec::{ FFI_PhysicalExtensionCodec, ForeignPhysicalExtensionCodec, }; @@ -643,6 +949,70 @@ 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(()) + } + } + + /// 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_legacy_decode_uses_context_aware_default() + -> 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 {}); @@ -756,4 +1126,99 @@ pub(crate) mod tests { .expect("rebound codec resolves"); assert_eq!(task_ctx.session_id(), ctx_b.task_ctx().session_id()); } + + /// 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<()> { + 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::new( + sq_expr, + 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..60f21346a55a8 --- /dev/null +++ b/datafusion/ffi/src/proto/scalar_subquery_results.rs @@ -0,0 +1,318 @@ +// 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, 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; + +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, + } + } +} + +/// 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)?; + 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 envelope = ScalarValueEnvelope::decode(bytes) + .map_err(|e| DataFusionError::External(Box::new(e)))?; + 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( + results: &FFI_ScalarSubqueryResults, + index: u64, +) -> FFI_Result>> { + let value = results.inner().get(SubqueryIndex::new(index as usize)); + + let encoded = value.as_ref().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(()) + } + + /// `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(()) + } +} diff --git a/datafusion/ffi/src/tests/mod.rs b/datafusion/ffi/src/tests/mod.rs index a334bd6d21f43..cfaf5d435c3f2 100644 --- a/datafusion/ffi/src/tests/mod.rs +++ b/datafusion/ffi/src/tests/mod.rs @@ -291,6 +291,50 @@ pub extern "C" fn datafusion_ffi_test_create_exec_with_byte_metrics() -> FFI_Exe FFI_ExecutionPlan::new(plan, None) } +/// Builds a [`FFI_PhysicalExtensionCodec`] wrapping +/// [`crate::proto::physical_extension_codec::fixtures::ScalarSubqueryExprExecCodec`] +/// whose function pointers - including `try_decode_with_ctx` - point at code +/// compiled into *this* image. +/// +/// 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 +/// [`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_create_scalar_subquery_expr_exec_codec( + task_ctx_provider: crate::execution::FFI_TaskContextProvider, +) -> 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 /// [`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..d751cc7fc30a2 100644 --- a/datafusion/ffi/src/tests/utils.rs +++ b/datafusion/ffi/src/tests/utils.rs @@ -120,6 +120,39 @@ pub fn get_byte_metrics_exec() -> Result Result { + let lib_path = find_library()?; + + let lib = unsafe { + libloading::Library::new(&lib_path) + .map_err(|e| DataFusionError::External(Box::new(e)))? + }; + + let create_codec: libloading::Symbol< + extern "C" fn( + crate::execution::FFI_TaskContextProvider, + ) -> crate::proto::physical_extension_codec::FFI_PhysicalExtensionCodec, + > = unsafe { + lib.get(b"datafusion_ffi_test_create_scalar_subquery_expr_exec_codec") + .map_err(|e| DataFusionError::External(Box::new(e)))? + }; + + 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); + + Ok(codec) +} + /// 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..41e51ba49b778 --- /dev/null +++ b/datafusion/ffi/tests/ffi_physical_extension_codec.rs @@ -0,0 +1,154 @@ +// 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::any::Any; + use std::sync::Arc; + + use arrow::array::{AsArray, RecordBatch}; + use arrow::datatypes::{DataType, Int64Type, Schema}; + use datafusion::prelude::SessionContext; + use datafusion_common::{DataFusionError, ScalarValue}; + use datafusion_execution::TaskContextProvider; + use datafusion_expr::physical_planning_context::{ + ScalarSubqueryResults, SubqueryIndex, + }; + use datafusion_ffi::execution::FFI_TaskContextProvider; + use datafusion_ffi::execution_plan::ForeignExecutionPlan; + use datafusion_ffi::proto::physical_extension_codec::ForeignPhysicalExtensionCodec; + use datafusion_ffi::proto::physical_extension_codec::fixtures::{ + ScalarSubqueryExprExec, ScalarSubqueryExprExecCodec, + }; + use datafusion_ffi::tests::utils::get_scalar_subquery_expr_exec_codec; + use datafusion_physical_expr::scalar_subquery::ScalarSubqueryExpr; + use datafusion_physical_plan::empty::EmptyExec; + use datafusion_physical_plan::{ExecutionPlan, PhysicalExpr}; + use datafusion_proto::physical_plan::{ + DefaultPhysicalProtoConverter, PhysicalExtensionCodec, PhysicalPlanDecodeContext, + }; + + /// 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*. + /// + /// 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::new( + sq_expr, + Arc::new(EmptyExec::new(Arc::new(Schema::empty()))), + )); + + let mut bytes = Vec::new(); + ScalarSubqueryExprExecCodec.try_encode( + Arc::clone(&extension_plan), + &mut bytes, + &DefaultPhysicalProtoConverter {}, + )?; + + let ctx = Arc::new(SessionContext::default()); + let host_task_ctx_provider = Arc::clone(&ctx) as Arc; + let ffi_task_ctx_provider = + FFI_TaskContextProvider::from(&host_task_ctx_provider); + + // 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" + ); + + // 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!( + 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(()) + } +} diff --git a/datafusion/proto/src/physical_plan/mod.rs b/datafusion/proto/src/physical_plan/mod.rs index e9b7fd647b302..2fbfb59fc79ab 100644 --- a/datafusion/proto/src/physical_plan/mod.rs +++ b/datafusion/proto/src/physical_plan/mod.rs @@ -1440,12 +1440,7 @@ pub trait PhysicalPlanNodeExt: Sized { .collect::>()?; ctx.codec() - .try_decode( - extension.node.as_slice(), - &inputs, - ctx.task_ctx(), - proto_converter, - ) + .try_decode_with_ctx(extension.node.as_slice(), &inputs, ctx, proto_converter) .map_err(|e| match extension.plan_name.as_deref() { // The writer named the plan but this session has no decoder for // it. Add that as context, rather than leaving only the codec's @@ -1713,6 +1708,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, @@ -2156,6 +2174,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}" + ); +}