diff --git a/Cargo.lock b/Cargo.lock index e3bfa7f5959b1..f940bf44de641 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -2269,6 +2269,7 @@ dependencies = [ "datafusion-execution", "datafusion-expr", "datafusion-ffi", + "datafusion-physical-expr", "datafusion-physical-expr-common", "datafusion-physical-optimizer", "datafusion-physical-plan", diff --git a/datafusion/expr/src/physical_planning_context.rs b/datafusion/expr/src/physical_planning_context.rs index 71bbdd4257bef..06db226422779 100644 --- a/datafusion/expr/src/physical_planning_context.rs +++ b/datafusion/expr/src/physical_planning_context.rs @@ -110,70 +110,119 @@ impl SubqueryIndex { } } +/// Where a [`ScalarSubqueryResults`] container keeps its slots. +/// +/// The slots are normally held in memory next to the plan. `datafusion-ffi` +/// implements this to forward each operation back over an FFI boundary, so a +/// plan decoded on the far side reads the values populated on the near side. +pub trait ScalarSubqueryResultsBackend: + fmt::Debug + Send + Sync + std::panic::RefUnwindSafe +{ + /// Returns the scalar value stored at `index`, if it has been populated. + fn get(&self, index: usize) -> Option; + + /// Stores `value` in the slot at `index`. + fn set(&self, index: usize, value: ScalarValue) -> Result<()>; + + /// Clears all populated results so the container can be reused. + fn clear(&self); +} + +/// The in-memory backend behind [`ScalarSubqueryResults::new`]. +struct LocalResults(Vec>>); + +impl fmt::Debug for LocalResults { + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + f.debug_list() + .entries(self.0.iter().map(|slot| slot.lock().unwrap().clone())) + .finish() + } +} + +impl ScalarSubqueryResultsBackend for LocalResults { + fn get(&self, index: usize) -> Option { + self.0.get(index)?.lock().unwrap().clone() + } + + fn set(&self, index: usize, value: ScalarValue) -> Result<()> { + let Some(slot) = self.0.get(index) else { + return internal_err!( + "ScalarSubqueryResults: result index {index} is out of bounds" + ); + }; + + let mut slot = slot.lock().unwrap(); + if slot.is_some() { + return internal_err!( + "ScalarSubqueryResults: result for index {index} was already populated" + ); + } + *slot = Some(value); + + Ok(()) + } + + fn clear(&self) { + for slot in &self.0 { + *slot.lock().unwrap() = None; + } + } +} + /// Shared results container for uncorrelated scalar subqueries. /// /// Each entry corresponds to one scalar subquery, identified by its index. /// Each slot is populated at execution time by `ScalarSubqueryExec`, read by /// `ScalarSubqueryExpr` instances that share this container, and cleared when /// the plan is reset for re-execution. -#[derive(Clone, Default)] +#[derive(Clone)] pub struct ScalarSubqueryResults { - slots: Arc>>>, + backend: Arc, +} + +impl Default for ScalarSubqueryResults { + fn default() -> Self { + Self::new(0) + } } impl ScalarSubqueryResults { /// Creates a new shared results container with `n` empty slots. pub fn new(n: usize) -> Self { - Self { - slots: Arc::new((0..n).map(|_| Mutex::new(None)).collect()), - } + Self::from_backend(Arc::new(LocalResults( + (0..n).map(|_| Mutex::new(None)).collect(), + ))) + } + + /// Creates a results container whose slots live in `backend`. + pub fn from_backend(backend: Arc) -> Self { + Self { 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() + self.backend.get(index.as_usize()) } /// Stores `value` in the slot at `index`. pub fn set(&self, index: SubqueryIndex, value: ScalarValue) -> Result<()> { - let Some(slot) = self.slots.get(index.as_usize()) else { - return internal_err!( - "ScalarSubqueryResults: result index {} is out of bounds", - index.as_usize() - ); - }; - - let mut slot = slot.lock().unwrap(); - if slot.is_some() { - return internal_err!( - "ScalarSubqueryResults: result for index {} was already populated", - index.as_usize() - ); - } - *slot = Some(value); - - Ok(()) + self.backend.set(index.as_usize(), value) } /// Clears all populated results so the container can be reused. pub fn clear(&self) { - for slot in self.slots.iter() { - *slot.lock().unwrap() = None; - } + self.backend.clear() } /// Returns true if `this` and `other` point to the same shared container. pub fn ptr_eq(this: &Self, other: &Self) -> bool { - Arc::ptr_eq(&this.slots, &other.slots) + Arc::ptr_eq(&this.backend, &other.backend) } } impl fmt::Debug for ScalarSubqueryResults { fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { - f.debug_list() - .entries(self.slots.iter().map(|slot| slot.lock().unwrap().clone())) - .finish() + self.backend.fmt(f) } } @@ -187,7 +236,7 @@ impl Eq for ScalarSubqueryResults {} impl Hash for ScalarSubqueryResults { fn hash(&self, state: &mut H) { - Arc::as_ptr(&self.slots).hash(state); + Arc::as_ptr(&self.backend).cast::<()>().hash(state); } } @@ -195,6 +244,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/integration-tests/Cargo.toml b/datafusion/ffi/integration-tests/Cargo.toml index 5c94315210497..5f32936ba1134 100644 --- a/datafusion/ffi/integration-tests/Cargo.toml +++ b/datafusion/ffi/integration-tests/Cargo.toml @@ -30,6 +30,7 @@ datafusion-catalog = { workspace = true } datafusion-common = { workspace = true } datafusion-execution = { workspace = true } datafusion-expr = { workspace = true } +datafusion-physical-expr = { workspace = true } datafusion-physical-expr-common = { workspace = true } datafusion-physical-optimizer = { workspace = true } datafusion-physical-plan = { workspace = true } diff --git a/datafusion/ffi/integration-tests/tests/cases/mod.rs b/datafusion/ffi/integration-tests/tests/cases/mod.rs index c1d47cc4eac4b..c167292f247df 100644 --- a/datafusion/ffi/integration-tests/tests/cases/mod.rs +++ b/datafusion/ffi/integration-tests/tests/cases/mod.rs @@ -5,6 +5,7 @@ mod catalog; mod config; mod execution_plan; +mod physical_extension_codec; mod physical_optimizer; mod query_planner; mod udaf; diff --git a/datafusion/ffi/integration-tests/tests/cases/physical_extension_codec.rs b/datafusion/ffi/integration-tests/tests/cases/physical_extension_codec.rs new file mode 100644 index 0000000000000..303f24bbc188e --- /dev/null +++ b/datafusion/ffi/integration-tests/tests/cases/physical_extension_codec.rs @@ -0,0 +1,112 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: See the NOTICE file distributed with this work for additional information regarding copyright ownership +// SPDX-FileContributor: Licensed to the Apache Software Foundation (ASF) under one or more contributor license agreements + +mod tests { + use std::any::Any; + use std::sync::Arc; + + use arrow::array::{AsArray, RecordBatch}; + use arrow::datatypes::{DataType, Int64Type, Schema}; + use datafusion::prelude::SessionContext; + use datafusion_common::{DataFusionError, ScalarValue}; + use datafusion_execution::TaskContextProvider; + use datafusion_expr::physical_planning_context::{ + ScalarSubqueryResults, SubqueryIndex, + }; + use datafusion_ffi::execution::FFI_TaskContextProvider; + use datafusion_ffi::execution_plan::ForeignExecutionPlan; + use datafusion_ffi::proto::physical_extension_codec::ForeignPhysicalExtensionCodec; + use datafusion_ffi::proto::physical_extension_codec::fixtures::{ + ScalarSubqueryExprExec, ScalarSubqueryExprExecCodec, + }; + use datafusion_ffi::tests::utils::get_scalar_subquery_expr_exec_codec; + use datafusion_physical_expr::scalar_subquery::ScalarSubqueryExpr; + use datafusion_physical_plan::empty::EmptyExec; + use datafusion_physical_plan::{ExecutionPlan, PhysicalExpr}; + use datafusion_proto::physical_plan::{ + DefaultPhysicalProtoConverter, PhysicalExtensionCodec, PhysicalPlanDecodeContext, + }; + + /// A plan decoded by a genuinely foreign codec must read the scalar + /// subquery results the host populated. + /// + /// The codec comes from a separately `dlopen`'d copy of this crate, so + /// the decode takes the same path a real foreign codec takes rather than + /// a test-only one. + #[tokio::test] + async fn test_ffi_physical_extension_codec_scalar_subquery_cross_library() + -> Result<(), DataFusionError> { + let results = ScalarSubqueryResults::new(1); + results.set(SubqueryIndex::new(0), ScalarValue::Int64(Some(42)))?; + + let sq_expr: Arc = 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); + + 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 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" + ); + + 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/ffi/src/proto/mod.rs b/datafusion/ffi/src/proto/mod.rs index d3d473ca89e37..2996d61b27c49 100644 --- a/datafusion/ffi/src/proto/mod.rs +++ b/datafusion/ffi/src/proto/mod.rs @@ -4,3 +4,4 @@ pub mod logical_extension_codec; pub mod physical_extension_codec; +pub mod scalar_subquery_results; diff --git a/datafusion/ffi/src/proto/physical_extension_codec.rs b/datafusion/ffi/src/proto/physical_extension_codec.rs index 2a9ec86488c21..1e6a97bbf59a1 100644 --- a/datafusion/ffi/src/proto/physical_extension_codec.rs +++ b/datafusion/ffi/src/proto/physical_extension_codec.rs @@ -13,7 +13,7 @@ use datafusion_expr::{ }; use datafusion_physical_plan::ExecutionPlan; use datafusion_proto::physical_plan::{ - DefaultPhysicalProtoConverter, PhysicalExtensionCodec, + DefaultPhysicalProtoConverter, PhysicalExtensionCodec, PhysicalPlanDecodeContext, PhysicalProtoConverterExtension, }; @@ -24,10 +24,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. @@ -98,6 +99,16 @@ 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 caller's scalar + /// subquery results so a `ScalarSubqueryExpr` decoded on this side reads + /// the same values 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, } unsafe impl Send for FFI_PhysicalExtensionCodec {} @@ -135,10 +146,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, + &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, - task_ctx.as_ref(), + &decode_ctx, &DefaultPhysicalProtoConverter {}, )); @@ -297,6 +343,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, @@ -361,6 +408,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, @@ -431,6 +509,207 @@ impl PhysicalExtensionCodec for ForeignPhysicalExtensionCodec { } } +/// An extension plan and codec for the cross-library integration test. +/// +/// That test loads them from `datafusion-ffi-integration-tests`'s cdylib, +/// which links this crate, so they have to be reachable from here. The feature +/// keeps them out of ordinary builds. +#[cfg(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 carrying a single physical expression, so its codec + /// has to decode that expression itself. Executing it emits the + /// expression's value as a one-row batch, which is how the test reads it + /// back. + #[derive(Debug)] + pub struct ScalarSubqueryExprExec { + expr: Arc, + 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)] + 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; @@ -445,7 +724,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, }; @@ -630,6 +909,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 {}); 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..9190445de0401 --- /dev/null +++ b/datafusion/ffi/src/proto/scalar_subquery_results.rs @@ -0,0 +1,177 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: See the NOTICE file distributed with this work for additional information regarding copyright ownership +// SPDX-FileContributor: Licensed to the Apache Software Foundation (ASF) under one or more contributor license agreements + +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}; + +/// Shares a [`ScalarSubqueryResults`] container across an FFI boundary, so a +/// `ScalarSubqueryExpr` decoded on the far side reads the values the +/// `ScalarSubqueryExec` on the near side populates. +/// +/// `ScalarValue`s cross as encoded bytes, so this does not rely on their +/// memory layout being stable across the boundary. +#[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.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))) + } + } +} diff --git a/datafusion/ffi/src/tests/mod.rs b/datafusion/ffi/src/tests/mod.rs index 32ac7aceefb21..62208dceeef0b 100644 --- a/datafusion/ffi/src/tests/mod.rs +++ b/datafusion/ffi/src/tests/mod.rs @@ -281,6 +281,26 @@ pub extern "C" fn datafusion_ffi_test_create_exec_with_byte_metrics() -> FFI_Exe FFI_ExecutionPlan::new(plan, None) } +/// Builds an [`FFI_PhysicalExtensionCodec`] inside this library, so the +/// cross-library test gets a codec that is genuinely foreign to the caller. +/// +/// Exported as its own symbol rather than a [`ForeignLibraryModule`] field, +/// for the same reason as [`datafusion_ffi_test_create_exec_with_byte_metrics`]: +/// that struct is `#[repr(C)]` with no gated constructor, so adding a field +/// changes its ABI. +#[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 0efca296183e9..5e3750c43a9b6 100644 --- a/datafusion/ffi/src/tests/utils.rs +++ b/datafusion/ffi/src/tests/utils.rs @@ -107,6 +107,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/proto/src/physical_plan/mod.rs b/datafusion/proto/src/physical_plan/mod.rs index 253735a51814d..81e3ff79230ef 100644 --- a/datafusion/proto/src/physical_plan/mod.rs +++ b/datafusion/proto/src/physical_plan/mod.rs @@ -1455,12 +1455,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 @@ -1513,6 +1508,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, @@ -1965,6 +1983,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 11f329025c266..e9377921fc112 100644 --- a/datafusion/proto/tests/cases/plans/scalar_subquery.rs +++ b/datafusion/proto/tests/cases/plans/scalar_subquery.rs @@ -4,7 +4,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; @@ -13,17 +15,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; @@ -263,3 +278,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}" + ); +} diff --git a/docs/source/library-user-guide/upgrading/57.0.0.md b/docs/source/library-user-guide/upgrading/57.0.0.md index 2f63a9d39a392..90e1753d55298 100644 --- a/docs/source/library-user-guide/upgrading/57.0.0.md +++ b/docs/source/library-user-guide/upgrading/57.0.0.md @@ -38,3 +38,18 @@ parents and siblings unchanged. The values themselves are shared through `Arc`. Update logical-plan snapshots if needed: generated lambda parameter names are local to lexical scopes, so sibling lambdas may reuse names. + +### `PhysicalExtensionCodec` can decode with the full decode context + +`PhysicalExtensionCodec::try_decode` takes a `TaskContext`, which does not carry +the state that only exists while the surrounding plan is being decoded. The +trait gains `try_decode_with_ctx`, which takes the `PhysicalPlanDecodeContext` +that decode is running under. See [issue #25089] for details. + +**Who is affected:** + +- No changes are required. +- If you need the additional context, implement `try_decode_with_ctx` and pass + `ctx` on to any nested decode. + +[issue #25089]: https://github.com/apache/datafusion/issues/25089