mbutrovich commented on code in PR #4915:
URL: https://github.com/apache/datafusion-comet/pull/4915#discussion_r3591090356
##########
native/spark-expr/src/kernels/temporal.rs:
##########
@@ -321,21 +322,48 @@ where
}
}
+/// Truncates a date expressed as days since the epoch, returning `None` if it
is out of range.
+type DateTruncFn = fn(i32) -> Option<i32>;
+
+/// The `date_trunc` formats Spark accepts, and the truncation each one
selects.
+const DATE_TRUNC_FORMATS: [(&str, DateTruncFn); 8] = [
+ ("YEAR", trunc_days_to_year),
+ ("YYYY", trunc_days_to_year),
+ ("YY", trunc_days_to_year),
+ ("QUARTER", trunc_days_to_quarter),
+ ("MONTH", trunc_days_to_month),
+ ("MON", trunc_days_to_month),
+ ("MM", trunc_days_to_month),
+ ("WEEK", trunc_days_to_week),
+];
+
+/// Resolve a `date_trunc` format string to the corresponding truncation
function.
+///
+/// The format is matched case-insensitively. Every supported format is ASCII,
so ASCII input can
+/// be compared in place, leaving the per-row format column path
allocation-free. Non-ASCII input
+/// still needs full Unicode case folding to fold the same way Spark does.
+fn date_trunc_fn_for_format(format: &str) -> Result<DateTruncFn, SparkError> {
+ let key: Cow<str> = if format.is_ascii() {
+ Cow::Borrowed(format)
+ } else {
+ Cow::Owned(format.to_uppercase())
+ };
+ DATE_TRUNC_FORMATS
+ .iter()
+ .find(|(name, _)| name.eq_ignore_ascii_case(&key))
+ .map(|(_, trunc_fn)| *trunc_fn)
+ .ok_or_else(|| {
+ SparkError::Internal(format!(
+ "Unsupported format: {format:?} for function 'date_trunc'"
+ ))
+ })
+}
Review Comment:
```rust
let key: Cow<str> = if format.is_ascii() {
Cow::Borrowed(format)
} else {
Cow::Owned(format.to_uppercase())
};
DATE_TRUNC_FORMATS
.iter()
.find(|(name, _)| name.eq_ignore_ascii_case(&key))
```
Every entry in `DATE_TRUNC_FORMATS` is ASCII. `eq_ignore_ascii_case` only
folds ASCII bytes, so a non-ASCII `key` can never equal an ASCII key regardless
of whether it was `to_uppercase`d first. The `else` branch allocates a `String`
and then cannot change the outcome: a non-ASCII format always falls through to
the error. The comment claiming non-ASCII "still needs full Unicode case
folding to fold the same way Spark does" is incorrect for this table, because
Spark also only matches ASCII literals after folding, so a non-ASCII input is
invalid in Spark too.
Suggested change: drop the `Cow` and the `is_ascii()` split entirely and
match on the borrowed `&str`:
```rust
fn date_trunc_fn_for_format(format: &str) -> Result<DateTruncFn, SparkError>
{
DATE_TRUNC_FORMATS
.iter()
.find(|(name, _)| name.eq_ignore_ascii_case(format))
.map(|(_, trunc_fn)| *trunc_fn)
.ok_or_else(|| {
SparkError::Internal(format!(
"Unsupported format: {format:?} for function 'date_trunc'"
))
})
}
```
This removes the `use std::borrow::Cow;` import, removes an allocation on
the non-ASCII path, and is exactly as Spark-compatible.
##########
native/spark-expr/src/kernels/temporal.rs:
##########
Review Comment:
`native/spark-expr/src/kernels/temporal.rs:591` (`ntz_trunc_fn_for_format`
uses `format.to_uppercase().as_str()`) and `832` (the tz array-format helper
still calls `$formats.value(index).to_uppercase().as_str()` per row). These are
the `timestamp_trunc` siblings of the code this PR just optimized, and they
carry the identical per-row `String` allocation the PR is removing. Leaving
them behind means the "resolve once into a table plus memo" pattern this PR
introduces is applied to one of three structurally identical loops, so the
improvement is narrower than the general problem.
Suggested change: either extend this PR to build a `TIMESTAMP_TRUNC_FORMATS`
table plus the same last-format memo for the NTZ and tz array paths, or file a
follow-up issue and reference it in this PR so the remaining allocations are
tracked. I would fold at least the NTZ path in here since it is a direct analog
of `DATE_TRUNC_FORMATS`.
--
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]
