fix: truncate per-row date_trunc formats with the literal-format kernel

docs/source/contributor-guide/expression-audits/datetime_funcs.md

@@ -35,7 +35,7 @@
 
 ## date_trunc
 
-- Spark 4.2.0 raises `long overflow` when SECOND or MILLISECOND truncation falls below the smallest timestamp, like every other unit (SPARK-56663). Earlier versions wrap the Long subtraction. The serde tells the native kernel which Spark it runs against, so on 4.2 those two units also keep dictionary key masking for values near the lower bound.
+- Spark 4.2.0 and later raise `long overflow` when SECOND or MILLISECOND truncation falls below the smallest timestamp, like every other unit (SPARK-56663). Earlier versions wrap the Long subtraction. The serde tells the native kernel which Spark it runs against, so on 4.2 those two units also keep dictionary key masking for values near the lower bound.
 - Performance (tuned 2026-10-06, PR [#5956](https://github.com/apache/datafusion-comet/pull/5956)): scalar timestamp truncation skips dictionary key masking for infallible fine units and safe UTC/NTZ values with substantial repetition. A values-only range check retains masking near the lower timestamp bound; non-UTC calendar paths and high-cardinality coarse batches retain masking to avoid errors and unnecessary calendar work on unused entries. About 19x faster for MICROSECOND with 65,536 keys and 32 distinct values. Matched Criterion coverage includes sparse/dense/all-NULL keys, high-cardinality YEAR batches, and used/unused overflow values. Benchmark: `benches/timestamp_trunc.rs`.
 
 ## dayname

native/proto/src/proto/expr.proto

@@ -598,8 +598,8 @@ message TruncTimestamp {
   Expr child = 2;
   string timezone = 3;
   // Whether SECOND and MILLISECOND truncation wraps when the result would fall below the
-  // smallest timestamp. Spark before 4.2.0 wraps the Long subtraction; 4.2.0 (SPARK-56663)
-  // raises `long overflow` like every other unit. This tracks the Spark version Comet is
+  // smallest timestamp. Spark before 4.2.0 wraps the Long subtraction; 4.2.0 and later
+  // (SPARK-56663) raise `long overflow` like every other unit. This tracks the Spark version Comet is
   // running against.
   bool wrap_second_millisecond_overflow = 4;
 }

native/spark-expr/benches/timestamp_trunc.rs

@@ -15,7 +15,9 @@
 // specific language governing permissions and limitations
 // under the License.
 
-use arrow::array::{Array, DictionaryArray, Int32Array, RecordBatch, TimestampMicrosecondArray};
+use arrow::array::{
+    Array, DictionaryArray, Int32Array, RecordBatch, StringArray, TimestampMicrosecondArray,
+};
 use arrow::datatypes::{DataType, Field, Int32Type, Schema, TimeUnit};
 use criterion::{criterion_group, criterion_main, BenchmarkId, Criterion};
 use datafusion::physical_expr::expressions::{lit, Column};
@@ -153,6 +155,63 @@ fn criterion_benchmark(c: &mut Criterion) {
         }
     }
     group.finish();
+
+    // Formats from a column: one format for every row, and every unit mixed across the rows.
+    // Daytime instants (16:00 to 22:00 UTC) avoid DST transitions, so older kernels can be
+    // measured on the same input.
+    let mut group = c.benchmark_group("timestamp_trunc_format_column");
+    let rows = 8_192;
+    let units = [
+        "YEAR",
+        "QUARTER",
+        "MONTH",
+        "WEEK",
+        "DAY",
+        "HOUR",
+        "MINUTE",
+        "SECOND",
+        "MILLISECOND",
+        "MICROSECOND",
+    ];
+    for timezone in ["UTC", "America/Los_Angeles"] {
+        let schema = Arc::new(Schema::new(vec![
+            Field::new(
+                "a",
+                DataType::Timestamp(TimeUnit::Microsecond, Some(timezone.into())),
+                true,
+            ),
+            Field::new("fmt", DataType::Utf8, true),
+        ]));
+        let timestamps = TimestampMicrosecondArray::from_iter_values((0..rows).map(|i| {
+            BASE_MICROS
+                + (i % 366) as i64 * MICROS_PER_DAY
+                + 57_600_000_000
+                + (i as i64 * 1_234_567) % 21_600_000_000
+        }))
+        .with_timezone(timezone);
+        for (tag, unit) in [("YEAR", Some(0)), ("HOUR", Some(5)), ("mixed", None)] {
+            let formats = StringArray::from_iter_values(
+                (0..rows).map(|i| units[unit.unwrap_or(i % units.len())]),
+            );
+            let batch = RecordBatch::try_new(
+                Arc::clone(&schema),
+                vec![Arc::new(timestamps.clone()), Arc::new(formats)],
+            )
+            .unwrap();
+            let expr = TimestampTruncExpr::new(
+                Arc::new(Column::new("a", 0)),
+                Arc::new(Column::new("fmt", 1)),
+                timezone.to_string(),
+                true,
+            );
+            group.bench_with_input(
+                BenchmarkId::from_parameter(format!("{timezone}/{tag}/{rows}")),
+                &batch,
+                |b, batch| b.iter(|| black_box(expr.evaluate(black_box(batch)).unwrap())),
+            );
+        }
+    }
+    group.finish();
 }
 
 criterion_group!(benches, criterion_benchmark);

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

@@ -16,7 +16,7 @@
 // under the License.
 
 use crate::utils::array_with_timezone;
-use arrow::array::ArrayRef;
+use arrow::array::{new_null_array, ArrayRef};
 use arrow::compute::cast;
 use arrow::datatypes::{DataType, Schema, TimeUnit::Microsecond};
 use arrow::record_batch::RecordBatch;
@@ -42,7 +42,7 @@ pub struct TimestampTruncExpr {
     /// reallocating, and parsed once into a `chrono::TimeZone` per batch.
     timezone: Arc<str>,
     /// Whether SECOND and MILLISECOND truncation wraps below the smallest timestamp, as Spark
-    /// did before 4.2.0, instead of raising `long overflow` as 4.2.0 does (SPARK-56663).
+    /// did before 4.2.0, instead of raising `long overflow` as 4.2.0 and later do (SPARK-56663).
     wrap_second_millisecond_overflow: bool,
 }
 
@@ -133,7 +133,7 @@ impl PhysicalExpr for TimestampTruncExpr {
                 Ok(ColumnarValue::Array(relabel(result)?))
             }
             (ColumnarValue::Array(ts), ColumnarValue::Array(formats)) => {
-                let result = timestamp_trunc_array_fmt_dyn(&resolve_tz(ts)?, &formats)?;
+                let result = timestamp_trunc_array_fmt_dyn(&resolve_tz(ts)?, &formats, wrap)?;
                 Ok(ColumnarValue::Array(relabel(result)?))
             }
             (ColumnarValue::Scalar(ts_scalar), ColumnarValue::Scalar(Utf8(Some(format)))) => {
@@ -142,6 +142,18 @@ impl PhysicalExpr for TimestampTruncExpr {
                 let scalar = ScalarValue::try_from_array(&relabel(result)?, 0)?;
                 Ok(ColumnarValue::Scalar(scalar))
             }
+            (ColumnarValue::Scalar(ts_scalar), ColumnarValue::Array(formats)) => {
+                let ts = ts_scalar.to_array_of_size(formats.len())?;
+                let result = timestamp_trunc_array_fmt_dyn(&resolve_tz(ts)?, &formats, wrap)?;
+                Ok(ColumnarValue::Array(relabel(result)?))
+            }
+            // A NULL format gives NULL, as in Spark.
+            (ColumnarValue::Array(ts), ColumnarValue::Scalar(Utf8(None))) => {
+                Ok(ColumnarValue::Array(new_null_array(&output_type, ts.len())))
+            }
+            (ColumnarValue::Scalar(_), ColumnarValue::Scalar(Utf8(None))) => {
+                Ok(ColumnarValue::Scalar(ScalarValue::try_from(&output_type)?))
+            }
             _ => Err(DataFusionError::Execution(
                 "Invalid input to function TimestampTrunc. \
                     Expected (Timestamp, Utf8)"
@@ -184,7 +196,9 @@ fn output_type(input: &DataType) -> DataType {
 #[cfg(test)]
 mod tests {
     use super::*;
-    use arrow::array::{Array, AsArray, DictionaryArray, Int32Array, TimestampMicrosecondArray};
+    use arrow::array::{
+        Array, AsArray, DictionaryArray, Int32Array, StringArray, TimestampMicrosecondArray,
+    };
     use arrow::datatypes::{Field, Int32Type, TimestampMicrosecondType};
     use datafusion::physical_expr::expressions::{Column, Literal};
 
@@ -241,4 +255,68 @@ mod tests {
             HOUR_IN_KOLKATA
         );
     }
+
+    /// Evaluates `date_trunc` in America/Los_Angeles with the given timestamp and format inputs.
+    fn trunc_in_los_angeles(
+        timestamp: Arc<dyn PhysicalExpr>,
+        format: Arc<dyn PhysicalExpr>,
+        batch: &RecordBatch,
+    ) -> ColumnarValue {
+        TimestampTruncExpr::new(timestamp, format, "America/Los_Angeles".to_string(), false)
+            .evaluate(batch)
+            .unwrap()
+    }
+
+    #[test]
+    fn literal_timestamp_with_a_format_column() {
+        // 2024-11-03 01:30 PDT, inside the overlap. DAY keeps the input's offset, YEAR does not
+        // depend on it, and a NULL format gives NULL.
+        let micros = 1_730_622_600_000_000;
+        let formats = StringArray::from(vec![Some("DAY"), Some("YEAR"), None]);
+        let schema = Schema::new(vec![Field::new("fmt", DataType::Utf8, true)]);
+        let batch = RecordBatch::try_new(Arc::new(schema), vec![Arc::new(formats)]).unwrap();
+        let timestamp = Literal::new(ScalarValue::TimestampMicrosecond(
+            Some(micros),
+            Some("UTC".into()),
+        ));
+        let ColumnarValue::Array(result) =
+            trunc_in_los_angeles(Arc::new(timestamp), Arc::new(Column::new("fmt", 0)), &batch)
+        else {
+            panic!("expected an array");
+        };
+        assert_eq!(result.data_type(), &utc_timestamp());
+        let result = result.as_primitive::<TimestampMicrosecondType>();
+        // 2024-11-03 00:00 PDT and 2024-01-01 00:00 PST.
+        assert_eq!(result.value(0), 1_730_617_200_000_000);
+        assert_eq!(result.value(1), 1_704_096_000_000_000);
+        assert!(result.is_null(2));
+    }
+
+    #[test]
+    fn null_literal_format_gives_null() {
+        let input = TimestampMicrosecondArray::from(vec![Some(MICROS), None]).with_timezone("UTC");
+        let schema = Schema::new(vec![Field::new("ts", utc_timestamp(), true)]);
+        let batch = RecordBatch::try_new(Arc::new(schema), vec![Arc::new(input)]).unwrap();
+        let null_format = || Arc::new(Literal::new(Utf8(None)));
+        let ColumnarValue::Array(result) =
+            trunc_in_los_angeles(Arc::new(Column::new("ts", 0)), null_format(), &batch)
+        else {
+            panic!("expected an array");
+        };
+        assert_eq!(result.data_type(), &utc_timestamp());
+        assert_eq!(result.null_count(), 2);
+
+        let timestamp = Arc::new(Literal::new(ScalarValue::TimestampMicrosecond(
+            Some(MICROS),
+            Some("UTC".into()),
+        )));
+        let ColumnarValue::Scalar(result) = trunc_in_los_angeles(timestamp, null_format(), &batch)
+        else {
+            panic!("expected a scalar");
+        };
+        assert_eq!(
+            result,
+            ScalarValue::TimestampMicrosecond(None, Some("UTC".into()))
+        );
+    }
 }