Skip to content

fix(spark-expr): handle array length mismatch in datediff for dictionary-backed timestamps #3278

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
vigneshsiva11 wants to merge 4 commits into
base: main
Choose a base branch
from
Open

fix(spark-expr): handle array length mismatch in datediff for dictionary-backed timestamps #3278

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
4 commits
Select commit Hold shift + click to select a range
bb7ddc8
Fix datediff array length mismatch for dictionary-backed timestamps
vigneshsiva11 Jan 26, 2026
3783308
fix(spark-expr): remove unused import in datediff
vigneshsiva11 Jan 26, 2026
5da8d77
test(spark): add regression test for datediff on dictionary-encoded t…
vigneshsiva11 Jan 29, 2026
1d0fab5
test(spark): add regression test for datediff with dictionary-encoded…
vigneshsiva11 Feb 3, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
23 changes: 16 additions & 7 deletions native/spark-expr/src/datetime_funcs/date_diff.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -71,9 +71,22 @@ impl ScalarUDFImpl for SparkDateDiff {
fn invoke_with_args(&self, args: ScalarFunctionArgs) -> Result<ColumnarValue> {
let [end_date, start_date] = take_function_args(self.name(), args.args)?;

// Convert scalars to arrays for uniform processing
let end_arr = end_date.into_array(1)?;
let start_arr = start_date.into_array(1)?;
// Determine target length (broadcast scalars to column length)
let len = match (&end_date, &start_date) {
(ColumnarValue::Array(a), _) => a.len(),
(_, ColumnarValue::Array(a)) => a.len(),
_ => 1,
};

// Convert both arguments to arrays of the same length
let end_arr = end_date.into_array(len)?;
let start_arr = start_date.into_array(len)?;

// Normalize dictionary arrays (important for Iceberg)
let end_arr = arrow::compute::cast(&end_arr, &DataType::Date32)
.map_err(|e| DataFusionError::Execution(e.to_string()))?;
let start_arr = arrow::compute::cast(&start_arr, &DataType::Date32)
.map_err(|e| DataFusionError::Execution(e.to_string()))?;

let end_date_array = end_arr
.as_any()
Expand All @@ -97,8 +110,4 @@ impl ScalarUDFImpl for SparkDateDiff {

Ok(ColumnarValue::Array(Arc::new(result)))
}

fn aliases(&self) -> &[String] {
&self.aliases
}
Comment on lines -101 to -103
Copy link
Member

@andygrove andygrove Jan 28, 2026

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Why is this removed?

}
Comment on lines 110 to 113
Copy link

Copilot AI Jan 26, 2026

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

The aliases field still includes "datediff", but the fn aliases(&self) -> &[String] implementation was removed. If ScalarUDFImpl's default aliases() returns an empty slice (as used by other UDF impls in this crate), this will drop the datediff alias and can break Spark SQL/function resolution that relies on the alias rather than the primary name date_diff. Re-introduce aliases() (returning &self.aliases) or remove aliases entirely and ensure the UDF is registered under the intended name(s).

Copilot uses AI. Check for mistakes.
27 changes: 27 additions & 0 deletions spark/src/test/scala/org/apache/spark/sql/comet/ParquetDatetimeRebaseSuite.scala
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -116,6 +116,33 @@ abstract class ParquetDatetimeRebaseSuite extends CometTestBase {
}
}

test("datediff works with dictionary-encoded timestamp columns") {
withTempPath { path =>
withSQLConf(
CometConf.COMET_NATIVE_SCAN_IMPL.key -> CometConf.SCAN_NATIVE_COMET,
CometConf.COMET_ENABLED.key -> "true",
"spark.sql.parquet.enableDictionary" -> "true") {
val df = spark
.createDataFrame(
Seq(
("a", java.sql.Timestamp.valueOf("2024-01-02 10:00:00")),
("b", java.sql.Timestamp.valueOf("2024-01-03 11:00:00"))))
.toDF("id", "ts")

df.write.mode("overwrite").parquet(path.getAbsolutePath)

val readDf = spark.read.parquet(path.getAbsolutePath)

val result = readDf
.selectExpr("datediff(current_date(), ts) as diff")
.collect()

// Just verify it executes correctly (no CometNativeException)
assert(result.length == 2)
Copy link
Contributor

@coderfender coderfender Feb 3, 2026

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

We might want to leverage checkSparkAnswerAndOperator which also verifies the plan.

Copy link
Author

@vigneshsiva11 vigneshsiva11 Feb 3, 2026

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Thanks for the suggestion! For this regression test I focused on reproducing the original failure and ensuring "datediff" executes correctly without a CometNativeException. I can switch to "checkSparkAnswerAndOperator" if plan verification is preferred.

}
}
}

private def checkSparkNoRebaseAnswer(df: => DataFrame): Unit = {
var expected: Array[Row] = Array.empty

Expand Down