Skip to content

Commit 5b22857

Browse files
feichai0017feichai0017
authored
fix: avoid extraneous casts for equivalent nested types (#20945)
## Summary This PR avoids inserting extraneous casts during function argument coercion when two nested types are structurally equivalent but differ only in nested field names or metadata. Specifically, it: - treats equivalent nested `DataType`s as matching during UDF argument coercion - avoids rewriting such arguments with unnecessary `CAST`s - adds regression coverage in both `datafusion-expr` and `datafusion-optimizer` Closes #19943. ## Tests - `cargo test -p datafusion-optimizer` - `./dev/rust_lint.sh` - `cargo test -p datafusion-expr` currently fails on an existing snapshot mismatch in `logical_plan::plan::tests::test_display_pg_json` on the current `main` baseline, unrelated to this change
1 parent 467d2c3 commit 5b22857

2 files changed

Lines changed: 128 additions & 13 deletions

File tree

‎datafusion/expr/src/type_coercion/functions.rs‎

Lines changed: 76 additions & 13 deletions
Original file line numberDiff line numberDiff line change
@@ -586,6 +586,52 @@ fn get_valid_types(
586586
arguments: &[ArrayFunctionArgument],
587587
array_coercion: Option<&ListCoercion>,
588588
) -> Result<Vec<Vec<DataType>>> {
589+
fn rebuild_array_type(
590+
current_type: &DataType,
591+
element_type: &DataType,
592+
nullable: bool,
593+
large_list: bool,
594+
fixed_size: Option<i32>,
595+
) -> DataType {
596+
// Preserve the original list field when possible so field name or
597+
// metadata differences do not introduce otherwise unnecessary casts.
598+
let field = match current_type {
599+
DataType::List(field)
600+
| DataType::LargeList(field)
601+
| DataType::FixedSizeList(field, _) => Some(Arc::new(
602+
field
603+
.as_ref()
604+
.clone()
605+
.with_data_type(element_type.clone())
606+
.with_nullable(nullable),
607+
)),
608+
_ => None,
609+
};
610+
611+
if large_list {
612+
field.map_or_else(
613+
|| DataType::new_large_list(element_type.clone(), nullable),
614+
DataType::LargeList,
615+
)
616+
} else if let Some(size) = fixed_size {
617+
field.map_or_else(
618+
|| {
619+
DataType::new_fixed_size_list(
620+
element_type.clone(),
621+
size,
622+
nullable,
623+
)
624+
},
625+
|field| DataType::FixedSizeList(field, size),
626+
)
627+
} else {
628+
field.map_or_else(
629+
|| DataType::new_list(element_type.clone(), nullable),
630+
DataType::List,
631+
)
632+
}
633+
}
634+
589635
if current_types.len() != arguments.len() {
590636
return Ok(vec![vec![]]);
591637
}
@@ -657,21 +703,13 @@ fn get_valid_types(
657703
ArrayFunctionArgument::Array => {
658704
if current_type.is_null() {
659705
DataType::Null
660-
} else if large_list {
661-
DataType::new_large_list(
662-
element_type.clone(),
663-
is_nested_item_nullable.unwrap_or(true),
664-
)
665-
} else if let Some(size) = list_sizes.next() {
666-
DataType::new_fixed_size_list(
667-
element_type.clone(),
668-
size,
669-
is_nested_item_nullable.unwrap_or(true),
670-
)
671706
} else {
672-
DataType::new_list(
673-
element_type.clone(),
707+
rebuild_array_type(
708+
current_type,
709+
&element_type,
674710
is_nested_item_nullable.unwrap_or(true),
711+
large_list,
712+
list_sizes.next(),
675713
)
676714
}
677715
}
@@ -1664,6 +1702,31 @@ mod tests {
16641702
Ok(())
16651703
}
16661704

1705+
#[test]
1706+
fn test_get_valid_types_array_and_index_preserves_list_field_name() -> Result<()> {
1707+
let struct_fields = vec![
1708+
Field::new("id", DataType::Utf8, true),
1709+
Field::new("prim", DataType::Boolean, true),
1710+
];
1711+
let current_type = DataType::List(Arc::new(Field::new(
1712+
"element",
1713+
DataType::Struct(struct_fields.into()),
1714+
true,
1715+
)));
1716+
let signature = Signature::array_and_index(Volatility::Immutable);
1717+
1718+
assert_eq!(
1719+
get_valid_types(
1720+
"array_element",
1721+
&signature.type_signature,
1722+
&[current_type.clone(), DataType::Int64],
1723+
)?,
1724+
vec![vec![current_type, DataType::Int64]]
1725+
);
1726+
1727+
Ok(())
1728+
}
1729+
16671730
#[test]
16681731
fn test_get_valid_types_element_and_array() -> Result<()> {
16691732
let function = "element_and_array";

‎datafusion/optimizer/src/analyzer/type_coercion.rs‎

Lines changed: 52 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1772,6 +1772,31 @@ mod test {
17721772
}
17731773
}
17741774

1775+
#[derive(Debug, Hash, PartialEq, Eq)]
1776+
struct TestArrayElementUDF;
1777+
1778+
impl ScalarUDFImpl for TestArrayElementUDF {
1779+
fn name(&self) -> &str {
1780+
"TestArrayElementUDF"
1781+
}
1782+
1783+
fn signature(&self) -> &Signature {
1784+
static SIGNATURE: std::sync::LazyLock<Signature> =
1785+
std::sync::LazyLock::new(|| {
1786+
Signature::array_and_index(Volatility::Immutable)
1787+
});
1788+
&SIGNATURE
1789+
}
1790+
1791+
fn return_type(&self, _args: &[DataType]) -> Result<DataType> {
1792+
Ok(Utf8)
1793+
}
1794+
1795+
fn invoke_with_args(&self, _args: ScalarFunctionArgs) -> Result<ColumnarValue> {
1796+
Ok(ColumnarValue::Scalar(ScalarValue::from("a")))
1797+
}
1798+
}
1799+
17751800
#[test]
17761801
fn scalar_udf() -> Result<()> {
17771802
let empty = empty();
@@ -2669,6 +2694,33 @@ mod test {
26692694
)
26702695
}
26712696

2697+
#[test]
2698+
fn array_element_preserves_parquet_list_field_name() -> Result<()> {
2699+
let list_type = DataType::List(Arc::new(Field::new(
2700+
"element",
2701+
DataType::Struct(
2702+
vec![
2703+
Field::new("id", Utf8, true),
2704+
Field::new("prim", DataType::Boolean, true),
2705+
]
2706+
.into(),
2707+
),
2708+
true,
2709+
)));
2710+
2711+
let expr = ScalarUDF::from(TestArrayElementUDF).call(vec![col("a"), lit(1_i64)]);
2712+
let empty = empty_with_type(list_type);
2713+
let plan = LogicalPlan::Projection(Projection::try_new(vec![expr], empty)?);
2714+
2715+
assert_analyzed_plan_eq!(
2716+
plan,
2717+
@r#"
2718+
Projection: TestArrayElementUDF(a, Int64(1))
2719+
EmptyRelation: rows=0
2720+
"#
2721+
)
2722+
}
2723+
26722724
#[test]
26732725
fn interval_plus_timestamp() -> Result<()> {
26742726
// SELECT INTERVAL '1' YEAR + '2000-01-01T00:00:00'::timestamp;

0 commit comments

Comments
 (0)