You are viewing a plain text version of this content. The canonical link for it is here.
Posted to commits@datafusion.apache.org by gi...@apache.org on 2026/09/07 04:02:40 UTC
(datafusion) branch gh-readonly-queue/main/pr-24939-ca5d10b75ff1d1c7a222ebff0205e5132049fe14 created (now e1ca94fb11)
This is an automated email from the ASF dual-hosted git repository.
github-merge-queue[bot] pushed a change to branch gh-readonly-queue/main/pr-24939-ca5d10b75ff1d1c7a222ebff0205e5132049fe14
in repository https://gitbox.apache.org/repos/asf/datafusion.git
at e1ca94fb11 fix: preserve Arrow field metadata on scalar subquery expressions (#24939)
This branch includes the following new commits:
new e1ca94fb11 fix: preserve Arrow field metadata on scalar subquery expressions (#24939)
The 1 revisions listed above as "new" are entirely new to this
repository and will be described in separate emails. The revisions
listed as "add" were already present in the repository and have only
been added to this reference.
---------------------------------------------------------------------
To unsubscribe, e-mail: commits-unsubscribe@datafusion.apache.org
For additional commands, e-mail: commits-help@datafusion.apache.org
(datafusion) 01/01: fix: preserve Arrow field metadata on scalar subquery expressions (#24939)
Posted by gi...@apache.org.
This is an automated email from the ASF dual-hosted git repository.
github-merge-queue[bot] pushed a commit to branch gh-readonly-queue/main/pr-24939-ca5d10b75ff1d1c7a222ebff0205e5132049fe14
in repository https://gitbox.apache.org/repos/asf/datafusion.git
commit e1ca94fb11ab0f2ead9bddd2b5be6666e5d75c62
Author: Marcelo Tesla <ma...@gmail.com>
AuthorDate: Mon Sep 7 04:02:17 2026 +0000
fix: preserve Arrow field metadata on scalar subquery expressions (#24939)
## Which issue does this PR close?
- Closes #24933
Related:
[apache/sedona-db#1231](https://github.com/apache/sedona-db/issues/1231),
[apache/sedona-db#1226](https://github.com/apache/sedona-db/pull/1226).
This is the uncorrelated physical-plan counterpart of the earlier
outer-reference metadata fix in #17524 / #17422.
## Rationale for this change
UDFs that distinguish Arrow extension types from their storage types
(for example spatial predicates such as `ST_Intersects`) need
`ARROW:extension:name` on every argument. Uncorrelated scalar subqueries
kept that metadata in the logical plan, but physical planning built a
`ScalarSubqueryExpr` from only the data type and nullability. The
synthesized physical field had empty metadata, so queries like `WHERE
udf(col, (SELECT geometry FROM t WHERE id = 1))` failed even though the
equivalent join form worked.
## What changes are included in this PR?
- `ScalarSubqueryExpr` now stores the output `FieldRef` (name
`scalar_subquery`, original type/nullability, and metadata) via
`new_with_metadata`. The existing `new` constructor is unchanged and
still produces a field with empty metadata.
- Physical lowering copies metadata from the logical subquery output
field while still using `Expr::nullable` so zero-row subqueries remain
nullable.
- Protobuf encoding adds an additive `metadata` map on
`PhysicalScalarSubqueryExprNode` so plan round-trips keep extension
metadata.
## What is the testing strategy for this PR?
- Unit test `scalar_subquery_preserves_output_field_metadata` in
`planner.rs` reproduces the drop during physical lowering.
- Unit test `return_field_preserves_extension_metadata` and an updated
proto round-trip in `scalar_subquery.rs`.
- End-to-end regression
`test_extension_metadata_preserve_in_uncorrelated_scalar_subquery` in
`user_defined_scalar_functions.rs`, based on the issue reproducer. The
existing EXISTS-subquery metadata test still passes.
## Are there any user-facing changes?
Additive only: `ScalarSubqueryExpr::new_with_metadata` and an optional
protobuf `metadata` map (older payloads decode as empty metadata).
Existing `new(data_type, nullable, ...)` keeps working. Queries whose
UDFs inspect argument field metadata now see the subquery's original
extension metadata in the physical plan.
---------
Co-authored-by: Marcelo Tesla <90...@users.noreply.github.com>
---
.../user_defined/user_defined_scalar_functions.rs | 96 ++++++++++++++++++++++
datafusion/physical-expr/src/planner.rs | 45 +++++++++-
datafusion/physical-expr/src/scalar_subquery.rs | 87 ++++++++++++++++----
datafusion/proto-models/proto/datafusion.proto | 2 +
datafusion/proto-models/src/generated/pbjson.rs | 19 +++++
datafusion/proto-models/src/generated/prost.rs | 6 ++
datafusion/sqllogictest/test_files/metadata.slt | 14 ++++
7 files changed, 251 insertions(+), 18 deletions(-)
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 4f8078542e..36d094cd03 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<Vec<DataType>> {
+ Ok(arg_types.to_vec())
+ }
+
+ fn return_type(&self, _arg_types: &[DataType]) -> Result<DataType> {
+ unreachable!("return_field_from_args is implemented")
+ }
+
+ fn return_field_from_args(&self, args: ReturnFieldArgs) -> Result<FieldRef> {
+ 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<ColumnarValue> {
+ 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 8c5347db74..c5b67cde8e 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 ac473ea9dd..389baf3505 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<FieldMetadata>,
+ 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<FieldRef> {
- Ok(Arc::new(Field::new(
- "scalar_subquery",
- self.data_type.clone(),
- self.nullable,
- )))
+ Ok(Arc::clone(&self.field))
}
fn evaluate(&self, _batch: &RecordBatch) -> Result<ColumnarValue> {
@@ -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 8486410d07..e4607a9b28 100644
--- a/datafusion/proto-models/proto/datafusion.proto
+++ b/datafusion/proto-models/proto/datafusion.proto
@@ -1742,4 +1742,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<string, string> metadata = 4;
}
diff --git a/datafusion/proto-models/src/generated/pbjson.rs b/datafusion/proto-models/src/generated/pbjson.rs
index f568a6d853..0df2a4bd07 100644
--- a/datafusion/proto-models/src/generated/pbjson.rs
+++ b/datafusion/proto-models/src/generated/pbjson.rs
@@ -21471,6 +21471,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)?;
@@ -21481,6 +21484,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()
}
}
@@ -21495,6 +21501,7 @@ impl<'de> serde::Deserialize<'de> for PhysicalScalarSubqueryExprNode {
"dataType",
"nullable",
"index",
+ "metadata",
];
#[allow(clippy::enum_variant_names)]
@@ -21502,6 +21509,7 @@ impl<'de> serde::Deserialize<'de> for PhysicalScalarSubqueryExprNode {
DataType,
Nullable,
Index,
+ Metadata,
}
impl<'de> serde::Deserialize<'de> for GeneratedField {
fn deserialize<D>(deserializer: D) -> std::result::Result<GeneratedField, D::Error>
@@ -21526,6 +21534,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)),
}
}
@@ -21548,6 +21557,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 => {
@@ -21570,12 +21580,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::<std::collections::HashMap<_, _>>()?
+ );
+ }
}
}
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 f5c7c01e71..1af60e27a0 100644
--- a/datafusion/proto-models/src/generated/prost.rs
+++ b/datafusion/proto-models/src/generated/prost.rs
@@ -2640,6 +2640,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 0fc74fa6cf..68ee7a3541 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;
---------------------------------------------------------------------
To unsubscribe, e-mail: commits-unsubscribe@datafusion.apache.org
For additional commands, e-mail: commits-help@datafusion.apache.org