diff --git a/datafusion/core/tests/user_defined/user_defined_scalar_functions.rs b/datafusion/core/tests/user_defined/user_defined_scalar_functions.rs index 4f8078542e4ac..36d094cd0360b 100644 --- a/datafusion/core/tests/user_defined/user_defined_scalar_functions.rs +++ b/datafusion/core/tests/user_defined/user_defined_scalar_functions.rs @@ -2255,3 +2255,99 @@ async fn test_extension_metadata_preserve_in_subquery() -> Result<()> { assert!(!df.collect().await?.is_empty()); Ok(()) } + +/// Uncorrelated scalar subqueries must keep Arrow field metadata on the +/// physical expression so UDFs that distinguish extension types still type-check. +/// https://github.com/apache/datafusion/issues/24933 +#[tokio::test] +async fn test_extension_metadata_preserve_in_uncorrelated_scalar_subquery() -> Result<()> +{ + #[derive(Debug, PartialEq, Eq, Hash)] + struct MetadataRequired { + signature: Signature, + } + + impl Default for MetadataRequired { + fn default() -> Self { + Self { + signature: Signature::user_defined(Volatility::Immutable), + } + } + } + + impl ScalarUDFImpl for MetadataRequired { + fn name(&self) -> &str { + "metadata_required" + } + + fn signature(&self) -> &Signature { + &self.signature + } + + fn coerce_types(&self, arg_types: &[DataType]) -> Result> { + Ok(arg_types.to_vec()) + } + + fn return_type(&self, _arg_types: &[DataType]) -> Result { + unreachable!("return_field_from_args is implemented") + } + + fn return_field_from_args(&self, args: ReturnFieldArgs) -> Result { + for (i, field) in args.arg_fields.iter().enumerate() { + if !field.metadata().contains_key("ARROW:extension:name") { + return exec_err!( + "argument {i} lost ARROW:extension:name; field={field:?} metadata={:?}", + field.metadata() + ); + } + } + + Ok(Field::new("metadata_required", DataType::Boolean, true).into()) + } + + fn invoke_with_args(&self, args: ScalarFunctionArgs) -> Result { + Ok(ColumnarValue::Scalar(ScalarValue::Boolean(Some( + args.arg_fields + .iter() + .all(|field| field.metadata().contains_key("ARROW:extension:name")), + )))) + } + } + + let schema = Schema::new(vec![ + Field::new("id", DataType::Int64, false), + Field::new("geometry", DataType::Utf8, false).with_metadata(HashMap::from([( + "ARROW:extension:name".to_string(), + "example.extension".to_string(), + )])), + ]); + + let batch = RecordBatch::try_new( + schema.clone().into(), + vec![ + create_array!(Int64, [1, 2]), + create_array!(Utf8, [Some("a"), Some("b")]), + ], + )?; + + let ctx = SessionContext::new(); + ctx.register_batch("l", batch.clone())?; + ctx.register_batch("r", batch)?; + ctx.register_udf(MetadataRequired::default().into()); + + let df = ctx + .sql( + " + SELECT id + FROM l + WHERE metadata_required( + l.geometry, + (SELECT r.geometry FROM r WHERE r.id = 1) + ) + ", + ) + .await?; + let batches = df.collect().await?; + assert!(!batches.is_empty()); + Ok(()) +} diff --git a/datafusion/physical-expr/src/planner.rs b/datafusion/physical-expr/src/planner.rs index 8c5347db742bb..c5b67cde8ec33 100644 --- a/datafusion/physical-expr/src/planner.rs +++ b/datafusion/physical-expr/src/planner.rs @@ -523,10 +523,12 @@ pub fn create_physical_expr( schema.fields().len() ); } - let dt = schema.field(0).data_type().clone(); - Ok(Arc::new(ScalarSubqueryExpr::new( - dt, + let output_field = schema.field(0); + let metadata = FieldMetadata::from(output_field.as_ref()); + Ok(Arc::new(ScalarSubqueryExpr::new_with_metadata( + output_field.data_type().clone(), e.nullable(input_dfschema)?, + (!metadata.is_empty()).then_some(metadata), index, planning_ctx.results().clone(), ))) @@ -827,6 +829,43 @@ mod tests { Ok(()) } + #[test] + fn scalar_subquery_preserves_output_field_metadata() -> Result<()> { + let metadata = FieldMetadata::from(HashMap::from([( + EXTENSION_TYPE_NAME_KEY.to_string(), + "example.extension".to_string(), + )])); + let plan = LogicalPlanBuilder::empty(true) + .project(vec![ + lit("a").alias_with_metadata("geometry", Some(metadata)), + ])? + .build()?; + let expr = scalar_subquery(Arc::new(plan)); + let Expr::ScalarSubquery(subquery) = &expr else { + unreachable!() + }; + + let index = SubqueryIndex::new(0); + let planning_ctx = PhysicalPlanningContext::new( + HashMap::from([(subquery.clone(), index)]), + ScalarSubqueryResults::new(1), + ); + let physical_expr = create_physical_expr( + &expr, + &DFSchema::empty(), + &ExecutionProps::new(), + &planning_ctx, + )?; + + let field = physical_expr.return_field(&Schema::empty())?; + assert_eq!( + field.metadata().get(EXTENSION_TYPE_NAME_KEY), + Some(&"example.extension".to_string()), + "scalar subquery physical expr must keep ARROW extension metadata: {field:?}" + ); + Ok(()) + } + #[test] fn test_cast_lowering_preserves_target_field_metadata() -> Result<()> { let schema = test_cast_schema(); diff --git a/datafusion/physical-expr/src/scalar_subquery.rs b/datafusion/physical-expr/src/scalar_subquery.rs index ac473ea9dd78d..389baf3505279 100644 --- a/datafusion/physical-expr/src/scalar_subquery.rs +++ b/datafusion/physical-expr/src/scalar_subquery.rs @@ -23,6 +23,7 @@ use std::sync::Arc; use arrow::datatypes::{DataType, Field, FieldRef, Schema}; use arrow::record_batch::RecordBatch; +use datafusion_common::metadata::FieldMetadata; use datafusion_common::{Result, internal_datafusion_err}; use datafusion_expr::physical_planning_context::{ScalarSubqueryResults, SubqueryIndex}; use datafusion_expr_common::columnar_value::ColumnarValue; @@ -36,8 +37,8 @@ use datafusion_physical_expr_common::physical_expr::PhysicalExpr; /// simply reads from that container at the appropriate index. #[derive(Debug)] pub struct ScalarSubqueryExpr { - data_type: DataType, - nullable: bool, + /// Output field of the scalar subquery, including Arrow extension metadata. + field: FieldRef, /// Index of this subquery in the shared results container. index: SubqueryIndex, /// Shared results container populated by `ScalarSubqueryExec`. @@ -51,9 +52,24 @@ impl ScalarSubqueryExpr { index: SubqueryIndex, results: ScalarSubqueryResults, ) -> Self { + Self::new_with_metadata(data_type, nullable, None, index, results) + } + + /// Create a scalar subquery expression, preserving optional field metadata + /// from the logical subquery output (for example Arrow extension type names). + pub fn new_with_metadata( + data_type: DataType, + nullable: bool, + metadata: Option, + index: SubqueryIndex, + results: ScalarSubqueryResults, + ) -> Self { + let mut field = Field::new("scalar_subquery", data_type, nullable); + if let Some(metadata) = metadata { + field = metadata.add_to_field(field); + } Self { - data_type, - nullable, + field: Arc::new(field), index, results, } @@ -68,7 +84,7 @@ impl ScalarSubqueryExpr { note = "was only used for proto serialization, which no longer needs it; use `return_field` for type/nullability. It will be removed in 61.0.0 or 6 months after 55.0.0 is released, whichever is longer." )] pub fn data_type(&self) -> &DataType { - &self.data_type + self.field.data_type() } #[deprecated( @@ -76,7 +92,7 @@ impl ScalarSubqueryExpr { note = "was only used for proto serialization, which no longer needs it; use `return_field` for type/nullability. It will be removed in 61.0.0 or 6 months after 55.0.0 is released, whichever is longer." )] pub fn nullable(&self) -> bool { - self.nullable + self.field.is_nullable() } /// Returns the index of this subquery in the shared results container. @@ -117,11 +133,7 @@ impl Eq for ScalarSubqueryExpr {} impl PhysicalExpr for ScalarSubqueryExpr { fn return_field(&self, _input_schema: &Schema) -> Result { - Ok(Arc::new(Field::new( - "scalar_subquery", - self.data_type.clone(), - self.nullable, - ))) + Ok(Arc::clone(&self.field)) } fn evaluate(&self, _batch: &RecordBatch) -> Result { @@ -163,13 +175,14 @@ impl PhysicalExpr for ScalarSubqueryExpr { expr_id: None, expr_type: Some(protobuf::physical_expr_node::ExprType::ScalarSubquery( protobuf::PhysicalScalarSubqueryExprNode { - data_type: Some((&self.data_type).try_into()?), - nullable: self.nullable, + data_type: Some(self.field.data_type().try_into()?), + nullable: self.field.is_nullable(), index: usize_to_wire( self.index.as_usize(), "ScalarSubqueryExpr", "index", )?, + metadata: self.field.metadata().clone(), }, )), })) @@ -207,9 +220,15 @@ impl ScalarSubqueryExpr { "data_type", )? .try_into()?; - Ok(Arc::new(ScalarSubqueryExpr::new( + let metadata = if sq.metadata.is_empty() { + None + } else { + Some(FieldMetadata::from(sq.metadata.clone())) + }; + Ok(Arc::new(ScalarSubqueryExpr::new_with_metadata( data_type, sq.nullable, + metadata, SubqueryIndex::new(sq.index as usize), results.clone(), ))) @@ -312,6 +331,31 @@ mod tests { ); assert_ne!(e1a, e3); } + + #[test] + fn return_field_preserves_extension_metadata() -> Result<()> { + let metadata = FieldMetadata::from(std::collections::HashMap::from([( + "ARROW:extension:name".to_string(), + "example.extension".to_string(), + )])); + let expr = ScalarSubqueryExpr::new_with_metadata( + DataType::Utf8, + true, + Some(metadata), + SubqueryIndex::new(0), + ScalarSubqueryResults::new(1), + ); + + let field = expr.return_field(&Schema::empty())?; + assert_eq!(field.name(), "scalar_subquery"); + assert_eq!(field.data_type(), &DataType::Utf8); + assert!(field.is_nullable()); + assert_eq!( + field.metadata().get("ARROW:extension:name"), + Some(&"example.extension".to_string()) + ); + Ok(()) + } } /// Tests for the `try_to_proto` / `try_from_proto` hooks. @@ -340,6 +384,7 @@ mod proto_tests { data_type, nullable, index, + metadata: Default::default(), }, )), } @@ -349,9 +394,13 @@ mod proto_tests { fn round_trips_through_proto() { // A three-slot results container so index 2 is meaningful. let results = ScalarSubqueryResults::new(3); - let expr = ScalarSubqueryExpr::new( + let expr = ScalarSubqueryExpr::new_with_metadata( DataType::Int32, true, + Some(FieldMetadata::from(std::collections::HashMap::from([( + "ARROW:extension:name".to_string(), + "example.extension".to_string(), + )]))), SubqueryIndex::new(2), results.clone(), ); @@ -378,6 +427,10 @@ mod proto_tests { .try_into() .unwrap(); assert_eq!(encoded_type, DataType::Int32); + assert_eq!( + sq.metadata.get("ARROW:extension:name").map(String::as_str), + Some("example.extension") + ); // Decode: reconstruct from the proto node, threading in the shared // results container the surrounding exec would provide. @@ -394,6 +447,10 @@ mod proto_tests { let field = decoded.return_field(&Schema::empty()).unwrap(); assert_eq!(field.data_type(), &DataType::Int32); assert!(field.is_nullable()); + assert_eq!( + field.metadata().get("ARROW:extension:name"), + Some(&"example.extension".to_string()) + ); // Same shared container + same index → equal to the original. assert_eq!(decoded, &expr); diff --git a/datafusion/proto-models/proto/datafusion.proto b/datafusion/proto-models/proto/datafusion.proto index e685b947bbaaf..eb0fa227df3a9 100644 --- a/datafusion/proto-models/proto/datafusion.proto +++ b/datafusion/proto-models/proto/datafusion.proto @@ -1738,4 +1738,6 @@ message PhysicalScalarSubqueryExprNode { datafusion_common.ArrowType data_type = 1; bool nullable = 2; uint32 index = 3; + // Serialized separately from data_type to keep older wire formats valid. + map metadata = 4; } diff --git a/datafusion/proto-models/src/generated/pbjson.rs b/datafusion/proto-models/src/generated/pbjson.rs index 4fcca074c8199..d5062bf3aba4c 100644 --- a/datafusion/proto-models/src/generated/pbjson.rs +++ b/datafusion/proto-models/src/generated/pbjson.rs @@ -21435,6 +21435,9 @@ impl serde::Serialize for PhysicalScalarSubqueryExprNode { if self.index != 0 { len += 1; } + if !self.metadata.is_empty() { + len += 1; + } let mut struct_ser = serializer.serialize_struct("datafusion.PhysicalScalarSubqueryExprNode", len)?; if let Some(v) = self.data_type.as_ref() { struct_ser.serialize_field("dataType", v)?; @@ -21445,6 +21448,9 @@ impl serde::Serialize for PhysicalScalarSubqueryExprNode { if self.index != 0 { struct_ser.serialize_field("index", &self.index)?; } + if !self.metadata.is_empty() { + struct_ser.serialize_field("metadata", &self.metadata)?; + } struct_ser.end() } } @@ -21459,6 +21465,7 @@ impl<'de> serde::Deserialize<'de> for PhysicalScalarSubqueryExprNode { "dataType", "nullable", "index", + "metadata", ]; #[allow(clippy::enum_variant_names)] @@ -21466,6 +21473,7 @@ impl<'de> serde::Deserialize<'de> for PhysicalScalarSubqueryExprNode { DataType, Nullable, Index, + Metadata, } impl<'de> serde::Deserialize<'de> for GeneratedField { fn deserialize(deserializer: D) -> std::result::Result @@ -21490,6 +21498,7 @@ impl<'de> serde::Deserialize<'de> for PhysicalScalarSubqueryExprNode { "dataType" | "data_type" => Ok(GeneratedField::DataType), "nullable" => Ok(GeneratedField::Nullable), "index" => Ok(GeneratedField::Index), + "metadata" => Ok(GeneratedField::Metadata), _ => Err(serde::de::Error::unknown_field(value, FIELDS)), } } @@ -21512,6 +21521,7 @@ impl<'de> serde::Deserialize<'de> for PhysicalScalarSubqueryExprNode { let mut data_type__ = None; let mut nullable__ = None; let mut index__ = None; + let mut metadata__ = None; while let Some(k) = map_.next_key()? { match k { GeneratedField::DataType => { @@ -21534,12 +21544,21 @@ impl<'de> serde::Deserialize<'de> for PhysicalScalarSubqueryExprNode { Some(map_.next_value::<::pbjson::private::NumberDeserialize<_>>()?.0) ; } + GeneratedField::Metadata => { + if metadata__.is_some() { + return Err(serde::de::Error::duplicate_field("metadata")); + } + metadata__ = Some( + map_.next_value::>()? + ); + } } } Ok(PhysicalScalarSubqueryExprNode { data_type: data_type__, nullable: nullable__.unwrap_or_default(), index: index__.unwrap_or_default(), + metadata: metadata__.unwrap_or_default(), }) } } diff --git a/datafusion/proto-models/src/generated/prost.rs b/datafusion/proto-models/src/generated/prost.rs index d14195ecc54a0..c5a4419b96eb0 100644 --- a/datafusion/proto-models/src/generated/prost.rs +++ b/datafusion/proto-models/src/generated/prost.rs @@ -2636,6 +2636,12 @@ pub struct PhysicalScalarSubqueryExprNode { pub nullable: bool, #[prost(uint32, tag = "3")] pub index: u32, + /// Serialized separately from data_type to keep older wire formats valid. + #[prost(map = "string, string", tag = "4")] + pub metadata: ::std::collections::HashMap< + ::prost::alloc::string::String, + ::prost::alloc::string::String, + >, } /// Identifies a built-in file format supported by DataFusion. /// Used by DefaultLogicalExtensionCodec to serialize/deserialize diff --git a/datafusion/sqllogictest/test_files/metadata.slt b/datafusion/sqllogictest/test_files/metadata.slt index 0fc74fa6cf602..68ee7a35412e8 100644 --- a/datafusion/sqllogictest/test_files/metadata.slt +++ b/datafusion/sqllogictest/test_files/metadata.slt @@ -518,6 +518,20 @@ from table_with_metadata limit 1; ---- NULL the id field +# Regression test: uncorrelated scalar subquery preserves field metadata +# https://github.com/apache/datafusion/issues/24933 +query IT +SELECT + (SELECT id FROM table_with_metadata WHERE id = 1), + arrow_metadata((SELECT id FROM table_with_metadata WHERE id = 1), 'metadata_key'); +---- +1 the id field + +query ? +select arrow_metadata((select id from table_with_metadata where id = 1)); +---- +{metadata_key: the id field} + statement ok drop table table_with_metadata;