Skip to content

Commit 0b498d2

Browse files
committed
fix(spark-expr): handle array length mismatch in datediff
1 parent 94dd587 commit 0b498d2

1 file changed

Lines changed: 16 additions & 3 deletions

File tree

native/spark-expr/src/datetime_funcs/date_diff.rs

Lines changed: 16 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -71,9 +71,22 @@ impl ScalarUDFImpl for SparkDateDiff {
7171
fn invoke_with_args(&self, args: ScalarFunctionArgs) -> Result<ColumnarValue> {
7272
let [end_date, start_date] = take_function_args(self.name(), args.args)?;
7373

74-
// Convert scalars to arrays for uniform processing
75-
let end_arr = end_date.into_array(1)?;
76-
let start_arr = start_date.into_array(1)?;
74+
// Determine target length (broadcast scalars to column length)
75+
let len = match (&end_date, &start_date) {
76+
(ColumnarValue::Array(a), _) => a.len(),
77+
(_, ColumnarValue::Array(a)) => a.len(),
78+
_ => 1,
79+
};
80+
81+
// Convert both arguments to arrays of the same length
82+
let end_arr = end_date.into_array(len)?;
83+
let start_arr = start_date.into_array(len)?;
84+
85+
// Normalize dictionary-backed arrays (important for Parquet / Iceberg)
86+
let end_arr = arrow::compute::cast(&end_arr, &DataType::Date32)
87+
.map_err(|e| DataFusionError::Execution(e.to_string()))?;
88+
let start_arr = arrow::compute::cast(&start_arr, &DataType::Date32)
89+
.map_err(|e| DataFusionError::Execution(e.to_string()))?;
7790

7891
let end_date_array = end_arr
7992
.as_any()

0 commit comments

Comments
 (0)