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:39 UTC
(datafusion-comet) branch gh-readonly-queue/main/pr-6345-4ca04b155bed7abe9e0e267f6438c31cca2c278d created (now 4c2e4b901c)
This is an automated email from the ASF dual-hosted git repository.
github-merge-queue[bot] pushed a change to branch gh-readonly-queue/main/pr-6345-4ca04b155bed7abe9e0e267f6438c31cca2c278d
in repository https://gitbox.apache.org/repos/asf/datafusion-comet.git
at 4c2e4b901c fix: keep the input's timezone label on native date_trunc results (#6345)
This branch includes the following new commits:
new 4c2e4b901c fix: keep the input's timezone label on native date_trunc results (#6345)
The 1 revisions listed above as "new" are entirely new to this
repository and will be described in separate emails. The revisions
listed as "add" were already present in the repository and have only
been added to this reference.
---------------------------------------------------------------------
To unsubscribe, e-mail: commits-unsubscribe@datafusion.apache.org
For additional commands, e-mail: commits-help@datafusion.apache.org
(datafusion-comet) 01/01: fix: keep the input's timezone label on native date_trunc results (#6345)
Posted by gi...@apache.org.
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