You are viewing a plain text version of this content. The canonical link for it is here.
Posted to commits@datafusion.apache.org by gi...@apache.org on 2026/09/29 09:24:40 UTC
(datafusion-comet) 01/01: fix: keep the input's timezone label on native date_trunc results (#6345)
This is an automated email from the ASF dual-hosted git repository.
github-merge-queue[bot] pushed a commit to branch gh-readonly-queue/main/pr-6345-4ca04b155bed7abe9e0e267f6438c31cca2c278d
in repository https://gitbox.apache.org/repos/asf/datafusion-comet.git
commit 4c2e4b901cc0a369f74af8e0c176d56ac88f47a2
Author: Andy Grove <ag...@apache.org>
AuthorDate: Tue Sep 29 00:05:32 2026 +0000
fix: keep the input's timezone label on native date_trunc results (#6345)
TimestampTruncExpr declared and emitted Timestamp(Microsecond, <session timezone>), while every other TimestampType value in a native plan is labelled UTC. Arrow's comparison kernels require identical types, so comparing the result with another timestamp failed, by default in Etc/UTC sessions. Truncate in the session timezone as before, but keep the input's label on the result.
Closes #6330.
---
.../src/datetime_funcs/timestamp_trunc.rs | 104 +++++++++++++++++----
.../expressions/datetime/trunc_timestamp_label.sql | 60 ++++++++++++
2 files changed, 146 insertions(+), 18 deletions(-)
diff --git a/native/spark-expr/src/datetime_funcs/timestamp_trunc.rs b/native/spark-expr/src/datetime_funcs/timestamp_trunc.rs
index 3a40aa79a2..91a83b77e5 100644
--- a/native/spark-expr/src/datetime_funcs/timestamp_trunc.rs
+++ b/native/spark-expr/src/datetime_funcs/timestamp_trunc.rs
@@ -17,6 +17,7 @@
use crate::utils::array_with_timezone;
use arrow::array::ArrayRef;
+use arrow::compute::cast;
use arrow::datatypes::{DataType, Schema, TimeUnit::Microsecond};
use arrow::record_batch::RecordBatch;
use datafusion::common::{DataFusionError, ScalarValue, ScalarValue::Utf8};
@@ -87,21 +88,7 @@ impl PhysicalExpr for TimestampTruncExpr {
}
fn data_type(&self, input_schema: &Schema) -> datafusion::common::Result<DataType> {
- // The kernel converts the input to the session timezone and emits an array stamped with
- // that timezone (NTZ stays NTZ). The declared `data_type` has to match or
- // shuffle/sort/`RowConverter` will fail with a schema-mismatch error.
- let session_tz = || DataType::Timestamp(Microsecond, Some(Arc::clone(&self.timezone)));
- match self.child.data_type(input_schema)? {
- DataType::Timestamp(_, None) => Ok(DataType::Timestamp(Microsecond, None)),
- DataType::Dictionary(key_type, inner) => {
- let inner_out = match inner.as_ref() {
- DataType::Timestamp(_, None) => DataType::Timestamp(Microsecond, None),
- _ => session_tz(),
- };
- Ok(DataType::Dictionary(key_type, Box::new(inner_out)))
- }
- _ => Ok(session_tz()),
- }
+ Ok(output_type(&self.child.data_type(input_schema)?))
}
fn nullable(&self, _: &Schema) -> datafusion::common::Result<bool> {
@@ -111,6 +98,7 @@ impl PhysicalExpr for TimestampTruncExpr {
fn evaluate(&self, batch: &RecordBatch) -> datafusion::common::Result<ColumnarValue> {
let timestamp = self.child.evaluate(batch)?;
let format = self.format.evaluate(batch)?;
+ let output_type = output_type(×tamp.data_type());
let tz = &self.timezone;
let resolve_tz = |ts: ArrayRef| -> datafusion::common::Result<ArrayRef> {
// For TimestampNTZ (Timestamp(Microsecond, None)), skip timezone conversion.
@@ -126,18 +114,23 @@ impl PhysicalExpr for TimestampTruncExpr {
)?)
}
};
+ // The kernels label their output with the session timezone they truncated in. Relabel it
+ // with the input's timezone, which leaves the values unchanged.
+ let relabel = |result: ArrayRef| -> datafusion::common::Result<ArrayRef> {
+ Ok(cast(&result, &output_type)?)
+ };
match (timestamp, format) {
(ColumnarValue::Array(ts), ColumnarValue::Scalar(Utf8(Some(format)))) => {
let result = timestamp_trunc_dyn(&resolve_tz(ts)?, format)?;
- Ok(ColumnarValue::Array(result))
+ Ok(ColumnarValue::Array(relabel(result)?))
}
(ColumnarValue::Array(ts), ColumnarValue::Array(formats)) => {
let result = timestamp_trunc_array_fmt_dyn(&resolve_tz(ts)?, &formats)?;
- Ok(ColumnarValue::Array(result))
+ Ok(ColumnarValue::Array(relabel(result)?))
}
(ColumnarValue::Scalar(ts_scalar), ColumnarValue::Scalar(Utf8(Some(format)))) => {
let result = timestamp_trunc_dyn(&resolve_tz(ts_scalar.to_array()?)?, format)?;
- let scalar = ScalarValue::try_from_array(&result, 0)?;
+ let scalar = ScalarValue::try_from_array(&relabel(result)?, 0)?;
Ok(ColumnarValue::Scalar(scalar))
}
_ => Err(DataFusionError::Execution(
@@ -163,3 +156,78 @@ impl PhysicalExpr for TimestampTruncExpr {
)))
}
}
+
+/// The result type for an input of type `input`. The kernels truncate in the session timezone, but
+/// the result keeps the input's timezone label. Every `TimestampType` value in a native plan is
+/// labelled "UTC", and a result labelled with the session timezone could not be compared with
+/// other timestamps or mixed with them in `CASE` and `coalesce`.
+fn output_type(input: &DataType) -> DataType {
+ match input {
+ DataType::Dictionary(key_type, value_type) => {
+ DataType::Dictionary(key_type.clone(), Box::new(output_type(value_type)))
+ }
+ DataType::Timestamp(_, tz) => DataType::Timestamp(Microsecond, tz.clone()),
+ other => other.clone(),
+ }
+}
+
+#[cfg(test)]
+mod tests {
+ use super::*;
+ use arrow::array::{Array, AsArray, DictionaryArray, Int32Array, TimestampMicrosecondArray};
+ use arrow::datatypes::{Field, Int32Type, TimestampMicrosecondType};
+ use datafusion::physical_expr::expressions::{Column, Literal};
+
+ /// 2024-01-15 18:30:45 UTC, which is 2024-01-16 00:00:45 in Asia/Kolkata (+05:30).
+ const MICROS: i64 = 1_705_343_445_000_000;
+ /// `MICROS` truncated to the hour in Asia/Kolkata: 2024-01-15 18:30:00 UTC.
+ const HOUR_IN_KOLKATA: i64 = 1_705_343_400_000_000;
+
+ fn utc_timestamp() -> DataType {
+ DataType::Timestamp(Microsecond, Some("UTC".into()))
+ }
+
+ /// Truncates `input` to the hour in Asia/Kolkata, returning the declared and actual types.
+ fn trunc_to_hour_in_kolkata(input: ArrayRef) -> (DataType, ArrayRef) {
+ let schema = Schema::new(vec![Field::new("ts", input.data_type().clone(), true)]);
+ let batch = RecordBatch::try_new(Arc::new(schema.clone()), vec![input]).unwrap();
+ let expr = TimestampTruncExpr::new(
+ Arc::new(Column::new("ts", 0)),
+ Arc::new(Literal::new(Utf8(Some("HOUR".to_string())))),
+ "Asia/Kolkata".to_string(),
+ );
+ let declared = expr.data_type(&schema).unwrap();
+ let ColumnarValue::Array(result) = expr.evaluate(&batch).unwrap() else {
+ panic!("expected an array");
+ };
+ (declared, result)
+ }
+
+ #[test]
+ fn result_keeps_the_input_label() {
+ let input = TimestampMicrosecondArray::from(vec![Some(MICROS), None]).with_timezone("UTC");
+ let (declared, result) = trunc_to_hour_in_kolkata(Arc::new(input));
+ assert_eq!(declared, utc_timestamp());
+ assert_eq!(result.data_type(), &utc_timestamp());
+ let result = result.as_primitive::<TimestampMicrosecondType>();
+ assert_eq!(result.value(0), HOUR_IN_KOLKATA);
+ assert!(result.is_null(1));
+ }
+
+ #[test]
+ fn dictionary_result_keeps_the_input_label() {
+ let values = TimestampMicrosecondArray::from(vec![MICROS]).with_timezone("UTC");
+ let input =
+ DictionaryArray::<Int32Type>::try_new(Int32Array::from(vec![0, 0]), Arc::new(values))
+ .unwrap();
+ let expected = DataType::Dictionary(Box::new(DataType::Int32), Box::new(utc_timestamp()));
+ let (declared, result) = trunc_to_hour_in_kolkata(Arc::new(input));
+ assert_eq!(declared, expected);
+ assert_eq!(result.data_type(), &expected);
+ let values = result.as_any_dictionary().values();
+ assert_eq!(
+ values.as_primitive::<TimestampMicrosecondType>().value(0),
+ HOUR_IN_KOLKATA
+ );
+ }
+}
diff --git a/spark/src/test/resources/sql-tests/expressions/datetime/trunc_timestamp_label.sql b/spark/src/test/resources/sql-tests/expressions/datetime/trunc_timestamp_label.sql
new file mode 100644
index 0000000000..8dabb1c27e
--- /dev/null
+++ b/spark/src/test/resources/sql-tests/expressions/datetime/trunc_timestamp_label.sql
@@ -0,0 +1,60 @@
+-- Licensed to the Apache Software Foundation (ASF) under one
+-- or more contributor license agreements. See the NOTICE file
+-- distributed with this work for additional information
+-- regarding copyright ownership. The ASF licenses this file
+-- to you under the Apache License, Version 2.0 (the
+-- "License"); you may not use this file except in compliance
+-- with the License. You may obtain a copy of the License at
+--
+-- http://www.apache.org/licenses/LICENSE-2.0
+--
+-- Unless required by applicable law or agreed to in writing,
+-- software distributed under the License is distributed on an
+-- "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+-- KIND, either express or implied. See the License for the
+-- specific language governing permissions and limitations
+-- under the License.
+
+-- date_trunc truncates in the session timezone, but its result must keep the UTC label that
+-- every other timestamp column has, so that it can be compared with them and mixed with them in
+-- conditionals. Etc/UTC takes the native path by default. The other zones take it through
+-- allowIncompatible, and have no DST transitions, which keeps these queries clear of
+-- https://github.com/apache/datafusion-comet/issues/5633.
+-- https://github.com/apache/datafusion-comet/issues/6330
+
+-- Config: spark.comet.expression.TruncTimestamp.allowIncompatible=true
+-- ConfigMatrix: spark.sql.session.timeZone=UTC,Etc/UTC,Asia/Tokyo,Asia/Kolkata
+
+statement
+CREATE TABLE test_trunc_ts_label(id int, ts timestamp, other timestamp) USING parquet
+
+statement
+INSERT INTO test_trunc_ts_label VALUES
+ (0, TIMESTAMP'2024-01-15 18:30:45Z', TIMESTAMP'2024-01-15 18:00:00Z'),
+ (1, TIMESTAMP'2024-06-30 23:30:00Z', TIMESTAMP'2024-06-01 00:00:00Z'),
+ (2, TIMESTAMP'1969-12-31 23:59:59.999999Z', TIMESTAMP'1969-12-31 23:00:00Z'),
+ (3, NULL, TIMESTAMP'2024-01-01 00:00:00Z')
+
+-- comparisons with a timestamp column and a timestamp literal
+query expect_native(date_trunc)
+SELECT id, date_trunc('HOUR', ts) = other, date_trunc('DAY', ts) >= other FROM test_trunc_ts_label
+
+query expect_native(date_trunc)
+SELECT id FROM test_trunc_ts_label WHERE date_trunc('DAY', ts) >= TIMESTAMP'2024-06-01 00:00:00Z'
+
+query expect_native(date_trunc)
+SELECT id, date_trunc('HOUR', ts) BETWEEN other AND ts, date_trunc('HOUR', ts) <=> other FROM test_trunc_ts_label
+
+query expect_native(date_trunc)
+SELECT id, nullif(date_trunc('HOUR', ts), other) FROM test_trunc_ts_label
+
+-- conditionals that mix the result with a timestamp column
+query expect_native(date_trunc)
+SELECT id, CASE WHEN id > 0 THEN date_trunc('HOUR', ts) ELSE other END FROM test_trunc_ts_label
+
+query expect_native(date_trunc)
+SELECT id, coalesce(date_trunc('HOUR', ts), other) FROM test_trunc_ts_label
+
+-- a join condition
+query expect_native(date_trunc)
+SELECT a.id FROM test_trunc_ts_label a JOIN test_trunc_ts_label b ON a.id = b.id AND date_trunc('DAY', a.ts) <= b.other
---------------------------------------------------------------------
To unsubscribe, e-mail: commits-unsubscribe@datafusion.apache.org
For additional commands, e-mail: commits-help@datafusion.apache.org