sunchao commented on code in PR #5956:
URL: https://github.com/apache/datafusion-comet/pull/5956#discussion_r4019031293
##########
native/spark-expr/src/kernels/temporal.rs:
##########
@@ -353,19 +409,110 @@ fn date_trunc_fn_for_format(format: &str) ->
Result<DateTruncFn, SparkError> {
})
}
-/// Optimized date truncation for Date32 arrays
-/// Works directly with days since epoch instead of converting to/from
NaiveDateTime
+/// Normalize a Spark `trunc` format without exposing additional DataFusion
granularities.
+fn normalize_date_trunc_format(format: &str) -> Result<&'static str,
SparkError> {
+ DATE_TRUNC_ALIASES
+ .iter()
+ .find(|(name, _)| name.eq_ignore_ascii_case(format))
+ .map(|(_, granularity)| *granularity)
+ .ok_or_else(|| {
+ SparkError::Internal(format!(
+ "Unsupported format: {format:?} for function 'date_trunc'"
+ ))
+ })
+}
+
+const MICROS_PER_DAY: i64 = 86_400_000_000;
+
+#[inline]
+fn fits_timestamp_nanosecond(micros: i64) -> bool {
+ micros.checked_mul(1_000).is_some()
+}
+
+/// DataFusion's coarse truncation first converts the input to nanoseconds.
Although an input near
+/// the lower TimestampNanosecond bound can itself be represented, truncating
it may move the
+/// result before that bound: for example, `1677-09-22` truncated to YEAR
becomes `1677-01-01`.
+/// The result can move backward by 365 days (366 when starting from December
31 in a leap year),
+/// and timezone gap handling can shift it by a few more hours. Round that
worst case up to 370
+/// days so both DataFusion's input and coarse-truncation result remain
representable. The
+/// effective microsecond interval is approximately
`1678-09-26T00:12:43.145225Z` through
+/// `2262-04-11T23:47:16.854775Z`; because Date32 values are UTC midnight, its
first upstream date
+/// is 1678-09-27 and its last is 2262-04-11.
+#[inline]
+fn fits_datafusion_coarse_trunc_range(micros: i64) -> bool {
+ const LOWER_NANOSECOND_MICROS: i64 = i64::MIN / 1_000;
+ const COARSE_TRUNC_MARGIN_MICROS: i64 = 370 * MICROS_PER_DAY;
+
+ fits_timestamp_nanosecond(micros)
+ && micros >= LOWER_NANOSECOND_MICROS + COARSE_TRUNC_MARGIN_MICROS
+}
+
+#[inline]
+fn date32_to_utc_midnight_micros(days: i32) -> Option<i64> {
+ i64::from(days).checked_mul(MICROS_PER_DAY)
+}
+
+#[inline]
+fn date32_fits_upstream(days: i32) -> bool {
+
date32_to_utc_midnight_micros(days).is_some_and(fits_datafusion_coarse_trunc_range)
+}
+
+/// Truncate scalar-format Date32 values through DataFusion's physical
`date_trunc`.
+///
+/// DataFusion 55.1 scales coarse timestamp granularities to nanoseconds
internally. Spark Date
+/// supports approximately years 0001 through 9999, while TimestampNanosecond
only spans roughly
+/// 1677 through 2262. Values outside the guarded TimestampNanosecond range
therefore retain the
+/// established Date32 calculation; values inside it use the upstream cast
sandwich.
fn date_trunc_date32(array: &Date32Array, format: String) ->
Result<Date32Array, SparkError> {
- // Select the truncation function based on format
+ let granularity = normalize_date_trunc_format(&format)?;
let trunc_fn = date_trunc_fn_for_format(&format)?;
-
- // Apply truncation to each element
- let result: Date32Array = array
+ let mut has_wide_value = false;
+ let upstream_input: Date32Array = array
.iter()
- .map(|opt_days| opt_days.and_then(trunc_fn))
+ .map(|value| {
+ value.and_then(|days| {
+ if date32_fits_upstream(days) {
+ Some(days)
+ } else {
+ has_wide_value = true;
+ None
+ }
+ })
+ })
.collect();
- Ok(result)
+ if upstream_input.null_count() == array.len() {
+ return Ok(array.iter().map(|value|
value.and_then(trunc_fn)).collect());
+ }
+
+ let timestamps = cast(
+ &upstream_input,
+ &DataType::Timestamp(TimeUnit::Microsecond, None),
+ )?;
+ let truncated = datafusion_date_trunc(timestamps, granularity)?;
+ let truncated = cast(truncated.as_ref(), &DataType::Date32)?;
Review Comment:
### Performance
[P2] Could you include matched parent/head microbenchmark results before
switching the default scalar Date32 path? Even an ordinary all-modern batch now
builds a filtered Date32 array, casts the whole batch to microseconds, runs the
upstream kernel, and casts the result back to Date32. The previous
implementation produced its output in one pass, and mixed modern/year-3333
batches add another merge pass. The existing
`native/spark-expr/benches/date_trunc.rs` already exercises
YEAR/QUARTER/MONTH/WEEK over 10,000 modern dates. Please compare that benchmark
on the parent and this head, add a representative mixed-range/null case, and
use the results to justify this dispatch or retain the existing Date32 kernel
where the conversions cause a material regression. The correctness tests do not
establish the performance tradeoff, and there are no benchmark results in the
PR yet.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
