Jefffrey commented on code in PR #19640:
URL: https://github.com/apache/datafusion/pull/19640#discussion_r2660680388
##########
datafusion/functions/src/datetime/date_trunc.rs:
##########
@@ -160,6 +183,7 @@ impl DateTruncFunc {
Self {
signature: Signature::one_of(
Review Comment:
We should probably refactor this to coercible API as it is quite bloated
##########
datafusion/functions/src/datetime/date_trunc.rs:
##########
@@ -339,16 +413,70 @@ impl ScalarUDFImpl for DateTruncFunc {
granularity,
tz_opt,
)?,
+ },
+ Time64(unit) => {
Review Comment:
We can fold the unit into the outer match so we don't nest as much:
```rust
Time64(Nanosecond) => { .. }
Time64(Microsecond) => { .. }
```
##########
datafusion/functions/src/datetime/date_trunc.rs:
##########
@@ -230,8 +262,10 @@ impl ScalarUDFImpl for DateTruncFunc {
Timestamp(Microsecond, tz_opt) => Ok(Timestamp(Microsecond,
tz_opt.clone())),
Timestamp(Millisecond, tz_opt) => Ok(Timestamp(Millisecond,
tz_opt.clone())),
Timestamp(Second, tz_opt) => Ok(Timestamp(Second, tz_opt.clone())),
+ Time64(tu) => Ok(Time64(*tu)),
+ Time32(tu) => Ok(Time32(*tu)),
_ => plan_err!(
- "The date_trunc function can only accept timestamp as the
second arg."
+ "The date_trunc function can only accept timestamp or time as
the second arg."
),
}
}
Review Comment:
```rust
fn return_type(&self, arg_types: &[DataType]) -> Result<DataType> {
if arg_types[1].is_null() {
Ok(Timestamp(Nanosecond, None))
} else {
Ok(arg_types[1].clone())
}
}
```
(Null handling technically wasn't necessary before, but if we use coercible
API it will be necessary)
##########
datafusion/functions/src/datetime/date_trunc.rs:
##########
@@ -315,10 +349,50 @@ impl ScalarUDFImpl for DateTruncFunc {
ColumnarValue::Scalar(ScalarValue::TimestampSecond(v, tz_opt)) => {
process_scalar::<TimestampSecondType>(v, granularity, tz_opt)?
}
+ ColumnarValue::Scalar(ScalarValue::Time64Nanosecond(v)) => {
+ if !granularity.valid_for_time() {
Review Comment:
This granularity check should be done in one place upfront instead of
repeated everywhere
##########
datafusion/functions/src/datetime/date_trunc.rs:
##########
@@ -136,11 +152,18 @@ impl DateTruncGranularity {
- second / SECOND
- millisecond / MILLISECOND
- microsecond / MICROSECOND
+
+ For Time types (hour, minute, second, millisecond, microsecond only):
+ - hour / HOUR
+ - minute / MINUTE
+ - second / SECOND
+ - millisecond / MILLISECOND
+ - microsecond / MICROSECOND
"#
),
argument(
name = "expression",
- description = "Time expression to operate on. Can be a constant,
column, or function."
+ description = "Timestamp or Time expression to operate on. Can be a
constant, column, or function."
Review Comment:
```suggestion
description = "Timestamp or time expression to operate on. Can be a
constant, column, or function."
```
##########
datafusion/functions/src/datetime/date_trunc.rs:
##########
@@ -374,6 +502,76 @@ impl ScalarUDFImpl for DateTruncFunc {
}
}
+const NANOS_PER_MICROSECOND: i64 = NANOSECONDS / MICROSECONDS;
+const NANOS_PER_MILLISECOND: i64 = NANOSECONDS / MILLISECONDS;
+const NANOS_PER_SECOND: i64 = NANOSECONDS;
+const NANOS_PER_MINUTE: i64 = 60 * NANOS_PER_SECOND;
+const NANOS_PER_HOUR: i64 = 60 * NANOS_PER_MINUTE;
+
+const MICROS_PER_MILLISECOND: i64 = MICROSECONDS / MILLISECONDS;
+const MICROS_PER_SECOND: i64 = MICROSECONDS;
+const MICROS_PER_MINUTE: i64 = 60 * MICROS_PER_SECOND;
+const MICROS_PER_HOUR: i64 = 60 * MICROS_PER_MINUTE;
+
+const MILLIS_PER_SECOND: i32 = MILLISECONDS as i32;
+const MILLIS_PER_MINUTE: i32 = 60 * MILLIS_PER_SECOND;
+const MILLIS_PER_HOUR: i32 = 60 * MILLIS_PER_MINUTE;
+
+const SECS_PER_MINUTE: i32 = 60;
+const SECS_PER_HOUR: i32 = 60 * SECS_PER_MINUTE;
+
+/// Truncate time in nanoseconds to the specified granularity
+fn truncate_time_nanos(value: i64, granularity: DateTruncGranularity) -> i64 {
Review Comment:
We might want to consider pulling this enum into a const generic so we don't
match at runtime each time (especially for arrays).
On a side note, are there potentially kernels in arrow-rs that we could
reuse? Or there are none suitable?
--
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]
