paleolimbot opened a new issue, #24933:
URL: https://github.com/apache/datafusion/issues/24933
### Describe the bug
Found when updating SedonaDB to DataFusion 54.1, but still present on main,
and fires during some fairly routine queries like:
```sql
SELECT * FROM some_geometry_table
WHERE ST_Intersects(
l.geometry,
(SELECT some_other_geometry_table.geometry FROM
some_other_geometry_table WHERE r.id = 1)
)
```
Previously these were optimized into a join but are now handled
independently.
From codex:
<details>
An uncorrelated scalar subquery preserves its output field metadata in the
logical plan, but DataFusion drops that metadata when it creates the physical
`ScalarSubqueryExpr`. This breaks UDFs that use Arrow field metadata to
distinguish extension types from their storage types. For example, a spatial
UDF receives `geometry, binary` instead of `geometry, geometry` when its
second
argument is produced by a scalar subquery.
</details>
### To Reproduce
Codex kindly put together a self-contained reproducer (also see the
conceptual SQL reproducer above).
<details>
## Self-contained reproducer
Create a new Cargo project alongside the DataFusion checkout:
```text
parent/
├── datafusion/
└── scalar-subquery-metadata-repro/
```
### `Cargo.toml`
```toml
[package]
name = "df-scalar-metadata-repro"
version = "0.1.0"
edition = "2024"
[dependencies]
datafusion = { path = "../datafusion/datafusion/core" }
tokio = { version = "1", features = ["macros", "rt-multi-thread"] }
```
### `src/main.rs`
```rust
use std::{collections::HashMap, sync::Arc};
use datafusion::{
arrow::{
array::{Int64Array, RecordBatch, StringArray},
datatypes::{DataType, Field, FieldRef, Schema},
},
common::{exec_err, Result, ScalarValue},
logical_expr::{
ColumnarValue, ReturnFieldArgs, ScalarFunctionArgs, ScalarUDFImpl,
Signature,
Volatility,
},
prelude::SessionContext,
};
const EXTENSION_KEY: &str = "ARROW:extension:name";
#[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(EXTENSION_KEY) {
return exec_err!(
"argument {i} lost {EXTENSION_KEY}; 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(EXTENSION_KEY)),
))))
}
}
fn batch() -> Result<RecordBatch> {
let metadata = HashMap::from([(
EXTENSION_KEY.to_string(),
"example.extension".to_string(),
)]);
let schema = Arc::new(Schema::new(vec![
Field::new("id", DataType::Int64, false),
Field::new("geometry", DataType::Utf8,
false).with_metadata(metadata),
]));
Ok(RecordBatch::try_new(
schema,
vec![
Arc::new(Int64Array::from(vec![1, 2])),
Arc::new(StringArray::from(vec!["a", "b"])),
],
)?)
}
#[tokio::main]
async fn main() -> Result<()> {
let ctx = SessionContext::new();
ctx.register_batch("l", batch()?)?;
ctx.register_batch("r", batch()?)?;
ctx.register_udf(MetadataRequired::default().into());
let sql = "
SELECT id
FROM l
WHERE metadata_required(
l.geometry,
(SELECT r.geometry FROM r WHERE r.id = 1)
)
";
let df = ctx.sql(sql).await?;
println!("{}", df.logical_plan().display_indent_schema());
let batches = df.collect().await?;
println!("query succeeded with {} batch(es)", batches.len());
Ok(())
}
```
Run it with:
```shell
cargo run
```
## Actual result
The UDF accepts both arguments during logical planning, demonstrating that
both
logical fields still contain `ARROW:extension:name`. Physical planning then
calls the same metadata-sensitive return-field logic with a newly created
field
for argument 1 that has no metadata:
```text
Projection: l.id [id:Int64]
Filter: metadata_required(l.geometry, (<subquery>)) [id:Int64,
geometry:Utf8]
Subquery: [geometry:Utf8]
Projection: r.geometry [geometry:Utf8]
Filter: r.id = Int64(1) [id:Int64, geometry:Utf8]
TableScan: r [id:Int64, geometry:Utf8]
TableScan: l [id:Int64, geometry:Utf8]
Error: Execution("argument 1 lost ARROW:extension:name; field=Field { name:
\"scalar_subquery\", data_type: Utf8, nullable: true } metadata={}")
```
## Expected result
The physical scalar-subquery expression should expose the subquery's original
output field metadata. The query should succeed and print:
```text
query succeeded with 1 batch(es)
```
</details>
### Expected behavior
The FieldRef of a value populated by a scalar subquery should be the same in
Logical and Physical plans.
### Additional context
I believe the issue is in `datafusion/physical-expr/src/planner.rs`...a
classic data type and nullability propagation without metadata.
Again, from our friend Codex:
<details>
## Drop point
In `datafusion/physical-expr/src/planner.rs`, the scalar-subquery planner
reads
the subquery output field but copies only its datatype:
```rust
let schema = sq.subquery.schema();
let dt = schema.field(0).data_type().clone();
Ok(Arc::new(ScalarSubqueryExpr::new(
dt,
e.nullable(input_dfschema)?,
index,
planning_ctx.results().clone(),
)))
```
`ScalarSubqueryExpr` consequently stores only `DataType` and nullability. Its
`return_field()` implementation in
`datafusion/physical-expr/src/scalar_subquery.rs` synthesizes a field with
empty
metadata:
```rust
fn return_field(&self, _input_schema: &Schema) -> Result<FieldRef> {
Ok(Arc::new(Field::new(
"scalar_subquery",
self.data_type.clone(),
self.nullable,
)))
}
```
One possible fix is for `ScalarSubqueryExpr` to retain the complete output
`FieldRef` (or otherwise carry its metadata) and return that information from
`return_field()`.
</details>
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
