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

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

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
4121cdd
Fix datediff array length mismatch for dictionary-backed timestamps
vigneshsiva11 Jan 26, 2026
08f67c0
fix(spark-expr): remove unused import in datediff
vigneshsiva11 Jan 26, 2026
c9b27fb
test(spark): add regression test for datediff on dictionary-encoded t…
vigneshsiva11 Jan 29, 2026
395f863
fix: add missing aliases() method and improve regression test
vigneshsiva11 Mar 11, 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
15 changes: 11 additions & 4 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 @@ -16,6 +16,7 @@
// under the License.

use arrow::array::{Array, Date32Array, Int32Array};
use arrow::compute::cast;
use arrow::compute::kernels::arity::binary;
use arrow::datatypes::DataType;
use datafusion::common::{utils::take_function_args, DataFusionError, Result};
Expand Down Expand Up @@ -68,6 +69,10 @@ impl ScalarUDFImpl for SparkDateDiff {
Ok(DataType::Int32)
}

fn aliases(&self) -> &[String] {
&self.aliases
}

fn invoke_with_args(&self, args: ScalarFunctionArgs) -> Result<ColumnarValue> {
let [end_date, start_date] = take_function_args(self.name(), args.args)?;

Expand All @@ -84,6 +89,12 @@ impl ScalarUDFImpl for SparkDateDiff {
let end_arr = end_date.into_array(num_rows)?;
let start_arr = start_date.into_array(num_rows)?;

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

let end_date_array = end_arr
.as_any()
.downcast_ref::<Date32Array>()
Expand All @@ -106,8 +117,4 @@ impl ScalarUDFImpl for SparkDateDiff {

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

fn aliases(&self) -> &[String] {
&self.aliases
}
Comment on lines -110 to -112

@andygrove andygrove Jan 28, 2026

Copy link
Copy Markdown
Member

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 117 to 120

Copilot AI Jan 26, 2026

Copy link

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.
21 changes: 21 additions & 0 deletions spark/src/test/scala/org/apache/comet/CometTemporalExpressionSuite.scala
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -257,6 +257,27 @@ class CometTemporalExpressionSuite extends CometTestBase with AdaptiveSparkPlanH
}
}

test("datediff works with dictionary-encoded timestamp columns") {
withTempDir { path =>
withSQLConf(
CometConf.COMET_NATIVE_SCAN_IMPL.key -> CometConf.SCAN_NATIVE_COMET,
"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(SaveMode.Overwrite).parquet(path.toString)
spark.read.parquet(path.toString).createOrReplaceTempView("ts_tbl")

checkSparkAnswerAndOperator(
"SELECT id, datediff(DATE('2024-01-10'), ts) AS diff FROM ts_tbl ORDER BY id")
}
}
}

test("date_format with timestamp column") {
// Filter out formats with embedded quotes that need special handling
val supportedFormats = CometDateFormat.supportedFormats.keys.toSeq
Expand Down
Loading