feat: Add TimestampNTZType support for casts and unix_timestamp

Author: andygrove

docs/source/user-guide/latest/compatibility.md

@@ -88,6 +88,22 @@ the [Comet Supported Expressions Guide](expressions.md) for more information on
 - **StructsToJson (to_json)**: Does not support `+Infinity` and `-Infinity` for numeric types (float, double).
   [#3016](https://github.com/apache/datafusion-comet/issues/3016)
 
+## Date and Time Functions
+
+Comet's native implementation of date and time functions may produce different results than Spark for dates
+far in the future (approximately beyond year 2100). This is because Comet uses the chrono-tz library for
+timezone calculations, which has limited support for Daylight Saving Time (DST) rules beyond the IANA
+time zone database's explicit transitions.
+
+For dates within a reasonable range (approximately 1970-2100), Comet's date and time functions are compatible
+with Spark. For dates beyond this range, functions that involve timezone-aware calculations (such as
+`date_trunc` with timezone-aware timestamps) may produce results with incorrect DST offsets.
+
+If you need to process dates far in the future with accurate timezone handling, consider:
+
+- Using timezone-naive types (`timestamp_ntz`) when timezone conversion is not required
+- Falling back to Spark for these specific operations
+
 ## Regular Expressions
 
 Comet uses the Rust regexp crate for evaluating regular expressions, and this has different behavior from Java's

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

@@ -78,6 +78,27 @@ impl ScalarUDFImpl for SparkUnixTimestamp {
 
         match args {
             [ColumnarValue::Array(array)] => match array.data_type() {
+                DataType::Timestamp(Microsecond, None) => {
+                    // TimestampNTZ: No timezone conversion needed - simply divide microseconds
+                    // by MICROS_PER_SECOND. TimestampNTZ stores local time without timezone.
+                    let timestamp_array =
+                        array.as_primitive::<arrow::datatypes::TimestampMicrosecondType>();
+
+                    let result: PrimitiveArray<Int64Type> = if timestamp_array.null_count() == 0 {
+                        timestamp_array
+                            .values()
+                            .iter()
+                            .map(|&micros| micros / MICROS_PER_SECOND)
+                            .collect()
+                    } else {
+                        timestamp_array
+                            .iter()
+                            .map(|v| v.map(|micros| div_floor(micros, MICROS_PER_SECOND)))
+                            .collect()
+                    };
+
+                    Ok(ColumnarValue::Array(Arc::new(result)))
+                }
                 DataType::Timestamp(_, _) => {
                     let is_utc = self.timezone == "UTC";
                     let array = if is_utc

native/spark-expr/src/kernels/temporal.rs

@@ -17,7 +17,9 @@
 
 //! temporal kernels
 
-use chrono::{DateTime, Datelike, Duration, NaiveDate, Timelike, Utc};
+use chrono::{
+    DateTime, Datelike, Duration, LocalResult, NaiveDate, Offset, TimeZone, Timelike, Utc,
+};
 
 use std::sync::Arc;
 
@@ -153,10 +155,30 @@ where
     Ok(())
 }
 
-// Apply the Tz to the Naive Date Time,,convert to UTC, and return as microseconds in Unix epoch
+// Apply the Tz to the Naive Date Time, convert to UTC, and return as microseconds in Unix epoch.
+// This function re-interprets the local datetime in the timezone to ensure the correct DST offset
+// is used for the target date (not the original date's offset). This is important when truncation
+// changes the date to a different DST period (e.g., from December/PST to October/PDT).
+//
+// Note: For far-future dates (approximately beyond year 2100), chrono-tz may not accurately
+// calculate DST transitions, which can result in incorrect offsets. See the compatibility
+// guide for more information.
 #[inline]
 fn as_micros_from_unix_epoch_utc(dt: Option<DateTime<Tz>>) -> i64 {
-    dt.unwrap().with_timezone(&Utc).timestamp_micros()
+    let dt = dt.unwrap();
+    let naive = dt.naive_local();
+    let tz = dt.timezone();
+
+    // Re-interpret the local time in the timezone to get the correct DST offset
+    // for the truncated date. Use noon to avoid DST gaps that occur around midnight.
+    let noon = naive.date().and_hms_opt(12, 0, 0).unwrap_or(naive);
+
+    let offset = match tz.offset_from_local_datetime(&noon) {
+        LocalResult::Single(off) | LocalResult::Ambiguous(off, _) => off.fix(),
+        LocalResult::None => return dt.with_timezone(&Utc).timestamp_micros(),
+    };
+
+    (naive - offset).and_utc().timestamp_micros()
 }
 
 #[inline]

spark/src/main/scala/org/apache/comet/serde/datetime.scala

@@ -295,11 +295,9 @@ object CometSecond extends CometExpressionSerde[Second] {
 object CometUnixTimestamp extends CometExpressionSerde[UnixTimestamp] {
 
   private def isSupportedInputType(expr: UnixTimestamp): Boolean = {
-    // Note: TimestampNTZType is not supported because Comet incorrectly applies
-    // timezone conversion to TimestampNTZ values. TimestampNTZ stores local time
-    // without timezone, so no conversion should be applied.
     expr.children.head.dataType match {
       case TimestampType | DateType => true
+      case dt if dt.typeName == "timestamp_ntz" => true
       case _ => false
     }
   }

spark/src/test/scala/org/apache/comet/CometTemporalExpressionSuite.scala

@@ -139,17 +139,55 @@ class CometTemporalExpressionSuite extends CometTestBase with AdaptiveSparkPlanH
     }
   }
 
-  test("unix_timestamp - timestamp_ntz input falls back to Spark") {
-    // TimestampNTZ is not supported because Comet incorrectly applies timezone
-    // conversion. TimestampNTZ stores local time without timezone, so the unix
-    // timestamp should just be the value divided by microseconds per second.
+  test("unix_timestamp - timestamp_ntz input") {
     val r = new Random(42)
     val ntzSchema = StructType(Seq(StructField("ts_ntz", DataTypes.TimestampNTZType, true)))
     val ntzDF = FuzzDataGenerator.generateDataFrame(r, spark, ntzSchema, 100, DataGenOptions())
     ntzDF.createOrReplaceTempView("ntz_tbl")
-    checkSparkAnswerAndFallbackReason(
-      "SELECT ts_ntz, unix_timestamp(ts_ntz) from ntz_tbl order by ts_ntz",
-      "unix_timestamp does not support input type: TimestampNTZType")
+    checkSparkAnswerAndOperator(
+      "SELECT ts_ntz, unix_timestamp(ts_ntz) from ntz_tbl order by ts_ntz")
+  }
+
+  test("hour/minute/second - timestamp_ntz input") {
+    val r = new Random(42)
+    val ntzSchema = StructType(Seq(StructField("ts_ntz", DataTypes.TimestampNTZType, true)))
+    val ntzDF = FuzzDataGenerator.generateDataFrame(r, spark, ntzSchema, 100, DataGenOptions())
+    ntzDF.createOrReplaceTempView("ntz_tbl")
+    checkSparkAnswerAndOperator(
+      "SELECT ts_ntz, hour(ts_ntz), minute(ts_ntz), second(ts_ntz) from ntz_tbl order by ts_ntz")
+  }
+
+  test("date_trunc - timestamp_ntz input") {
+    val r = new Random(42)
+    val ntzSchema = StructType(Seq(StructField("ts_ntz", DataTypes.TimestampNTZType, true)))
+    val reasonableBaseDate =
+      new java.text.SimpleDateFormat("yyyy-MM-dd HH:mm:ss").parse("2024-06-15 12:00:00").getTime
+    val ntzDF = FuzzDataGenerator.generateDataFrame(
+      r,
+      spark,
+      ntzSchema,
+      100,
+      DataGenOptions(baseDate = reasonableBaseDate))
+    ntzDF.createOrReplaceTempView("ntz_tbl")
+    for (format <- CometTruncTimestamp.supportedFormats) {
+      checkSparkAnswerAndOperator(
+        s"SELECT ts_ntz, date_trunc('$format', ts_ntz) from ntz_tbl order by ts_ntz")
+    }
+  }
+
+  test("date_format - timestamp_ntz input") {
+    withSQLConf(SQLConf.SESSION_LOCAL_TIMEZONE.key -> "UTC") {
+      val r = new Random(42)
+      val ntzSchema = StructType(Seq(StructField("ts_ntz", DataTypes.TimestampNTZType, true)))
+      val ntzDF = FuzzDataGenerator.generateDataFrame(r, spark, ntzSchema, 100, DataGenOptions())
+      ntzDF.createOrReplaceTempView("ntz_tbl")
+      val supportedFormats =
+        CometDateFormat.supportedFormats.keys.toSeq.filterNot(_.contains("'"))
+      for (format <- supportedFormats) {
+        checkSparkAnswerAndOperator(
+          s"SELECT ts_ntz, date_format(ts_ntz, '$format') from ntz_tbl order by ts_ntz")
+      }
+    }
   }
 
   test("unix_timestamp - string input falls back to Spark") {