alamb commented on code in PR #25173:
URL: https://github.com/apache/datafusion/pull/25173#discussion_r4018664743
##########
datafusion/functions-table/src/generate_series.rs:
##########
@@ -673,38 +735,42 @@ impl GenerateSeriesFuncImpl {
);
}
- // Parse start timestamp
- let (start_ts, tz) = match &exprs[0] {
- Expr::Literal(ScalarValue::TimestampNanosecond(ts, tz), _) => {
- (*ts, tz.clone())
- }
- other => {
- return plan_err!(
- "First argument must be a timestamp or NULL, got {:?}",
- other
- );
- }
- };
-
- // Parse end timestamp
- let end_ts = match &exprs[1] {
- Expr::Literal(ScalarValue::Null, _) => None,
- Expr::Literal(ScalarValue::TimestampNanosecond(ts, _), _) => *ts,
- other => {
- return plan_err!(
- "Second argument must be a timestamp or NULL, got {:?}",
- other
- );
- }
- };
+ // Parse the start and end timestamps.
+ //
+ // Both are widened to nanoseconds, so the two arguments do not have to
+ // agree on a `TimeUnit`: an Arrow timestamp denotes an instant
+ // regardless of the unit it happens to be stored in, and the output is
+ // nanoseconds either way (see the schema below).
+ let (start_ts, tz) =
+ timestamp_arg_to_nanos(&exprs[0], "First argument", self.name)?;
+ let (end_ts, _end_tz) =
+ timestamp_arg_to_nanos(&exprs[1], "Second argument", self.name)?;
+
+ // `_end_tz` is deliberately discarded: the output timezone comes from
+ // the start argument alone. A timezone on an Arrow timestamp does not
+ // change which instant it denotes, only how that instant is rendered,
+ // so a start and end carrying different timezones are still directly
+ // comparable once both are nanoseconds since the epoch -- there is
+ // nothing to reject. The start's zone is the one that is kept because
+ // it also anchors the calendar arithmetic that advances the series:
+ // month and day components of the step are applied in local time, so
+ // they follow that zone's DST rules.
// Parse step interval
let step_interval = match &exprs[2] {
Expr::Literal(ScalarValue::Null, _) => None,
Expr::Literal(ScalarValue::IntervalMonthDayNano(interval), _) =>
*interval,
+ Expr::Literal(scalar, _) => {
+ return plan_err!(
+ "Third argument for {} must be an INTERVAL or NULL, got
{:?}",
Review Comment:
This code is also looking for a literal (not just interval/null). Maybe we
could say that
```
"Third argument for {} must be an INTERVAL or NULL
constant, got {:?}",
```
##########
datafusion/functions-table/src/generate_series.rs:
##########
@@ -673,38 +735,42 @@ impl GenerateSeriesFuncImpl {
);
}
- // Parse start timestamp
- let (start_ts, tz) = match &exprs[0] {
- Expr::Literal(ScalarValue::TimestampNanosecond(ts, tz), _) => {
- (*ts, tz.clone())
- }
- other => {
- return plan_err!(
- "First argument must be a timestamp or NULL, got {:?}",
- other
- );
- }
- };
-
- // Parse end timestamp
- let end_ts = match &exprs[1] {
- Expr::Literal(ScalarValue::Null, _) => None,
- Expr::Literal(ScalarValue::TimestampNanosecond(ts, _), _) => *ts,
- other => {
- return plan_err!(
- "Second argument must be a timestamp or NULL, got {:?}",
- other
- );
- }
- };
+ // Parse the start and end timestamps.
+ //
+ // Both are widened to nanoseconds, so the two arguments do not have to
+ // agree on a `TimeUnit`: an Arrow timestamp denotes an instant
+ // regardless of the unit it happens to be stored in, and the output is
+ // nanoseconds either way (see the schema below).
+ let (start_ts, tz) =
+ timestamp_arg_to_nanos(&exprs[0], "First argument", self.name)?;
+ let (end_ts, _end_tz) =
+ timestamp_arg_to_nanos(&exprs[1], "Second argument", self.name)?;
+
+ // `_end_tz` is deliberately discarded: the output timezone comes from
Review Comment:
this isn't different behavior, right? It jus documents what currently
happens?
##########
datafusion/functions-table/src/generate_series.rs:
##########
@@ -673,38 +735,42 @@ impl GenerateSeriesFuncImpl {
);
}
- // Parse start timestamp
- let (start_ts, tz) = match &exprs[0] {
- Expr::Literal(ScalarValue::TimestampNanosecond(ts, tz), _) => {
- (*ts, tz.clone())
- }
- other => {
- return plan_err!(
- "First argument must be a timestamp or NULL, got {:?}",
- other
- );
- }
- };
-
- // Parse end timestamp
- let end_ts = match &exprs[1] {
- Expr::Literal(ScalarValue::Null, _) => None,
- Expr::Literal(ScalarValue::TimestampNanosecond(ts, _), _) => *ts,
- other => {
- return plan_err!(
- "Second argument must be a timestamp or NULL, got {:?}",
- other
- );
- }
- };
+ // Parse the start and end timestamps.
+ //
+ // Both are widened to nanoseconds, so the two arguments do not have to
+ // agree on a `TimeUnit`: an Arrow timestamp denotes an instant
+ // regardless of the unit it happens to be stored in, and the output is
+ // nanoseconds either way (see the schema below).
+ let (start_ts, tz) =
+ timestamp_arg_to_nanos(&exprs[0], "First argument", self.name)?;
+ let (end_ts, _end_tz) =
+ timestamp_arg_to_nanos(&exprs[1], "Second argument", self.name)?;
+
+ // `_end_tz` is deliberately discarded: the output timezone comes from
+ // the start argument alone. A timezone on an Arrow timestamp does not
+ // change which instant it denotes, only how that instant is rendered,
+ // so a start and end carrying different timezones are still directly
+ // comparable once both are nanoseconds since the epoch -- there is
+ // nothing to reject. The start's zone is the one that is kept because
+ // it also anchors the calendar arithmetic that advances the series:
+ // month and day components of the step are applied in local time, so
+ // they follow that zone's DST rules.
// Parse step interval
let step_interval = match &exprs[2] {
Expr::Literal(ScalarValue::Null, _) => None,
Expr::Literal(ScalarValue::IntervalMonthDayNano(interval), _) =>
*interval,
+ Expr::Literal(scalar, _) => {
+ return plan_err!(
+ "Third argument for {} must be an INTERVAL or NULL, got
{:?}",
Review Comment:
same comment below
##########
datafusion/functions-table/src/generate_series.rs:
##########
@@ -874,16 +964,347 @@ impl TableFunctionImpl for RangeFunc {
#[cfg(test)]
mod generate_series_tests {
+ use std::any::Any;
use std::sync::Arc;
- use arrow::datatypes::{DataType, Field, Schema};
- use datafusion_common::Result;
+ use arrow::datatypes::{
+ DataType, Field, IntervalMonthDayNano, Schema, SchemaRef, TimeUnit,
+ };
+ use datafusion_catalog::TableProvider;
+ use datafusion_common::{Result, ScalarValue};
+ use datafusion_expr::Expr;
use datafusion_physical_plan::memory::LazyBatchGenerator;
use crate::generate_series::{
- GenSeriesArgs, GenerateSeriesTable, GenericSeriesState,
+ GenSeriesArgs, GenerateSeriesFuncImpl, GenerateSeriesTable,
GenericSeriesState,
+ timestamp_arg_to_nanos,
};
+ /// Nanoseconds in a day, for readable expectations below.
+ const DAY_NANOS: i64 = 24 * 60 * 60 * 1_000_000_000;
+
+ fn lit(scalar: ScalarValue) -> Expr {
+ Expr::Literal(scalar, None)
+ }
+
+ /// `2024-01-01T00:00:00Z` in seconds since the epoch.
+ const JAN_1_2024_SECS: i64 = 1_704_067_200;
+
+ fn tz(s: &str) -> Option<Arc<str>> {
+ Some(Arc::from(s))
+ }
+
+ fn generate_series_impl() -> GenerateSeriesFuncImpl {
+ GenerateSeriesFuncImpl {
+ name: "generate_series",
+ include_end: true,
+ }
+ }
+
+ fn range_impl() -> GenerateSeriesFuncImpl {
+ GenerateSeriesFuncImpl {
+ name: "range",
+ include_end: false,
+ }
+ }
+
+ /// Run `call_timestamp` and return the resulting table's schema and args.
+ fn call_timestamp(
+ func: &GenerateSeriesFuncImpl,
+ start: ScalarValue,
+ end: ScalarValue,
+ step_days: i32,
+ ) -> Result<(SchemaRef, GenSeriesArgs)> {
+ let exprs = vec![
+ lit(start),
+ lit(end),
+ lit(ScalarValue::IntervalMonthDayNano(Some(
+ IntervalMonthDayNano::new(0, step_days, 0),
+ ))),
+ ];
+ let provider = func.call_timestamp(&exprs)?;
+ let table = (provider.as_ref() as &dyn Any)
+ .downcast_ref::<GenerateSeriesTable>()
+ .expect("call_timestamp returns a GenerateSeriesTable");
+ Ok((table.schema(), table.args.clone()))
+ }
+
+ fn timestamp_args(args: &GenSeriesArgs) -> (i64, i64, Option<Arc<str>>) {
+ match args {
+ GenSeriesArgs::TimestampArgs { start, end, tz, .. } => {
+ (*start, *end, tz.clone())
+ }
+ other => panic!("expected TimestampArgs, got {other:?}"),
+ }
+ }
+
+ /// Every `TimeUnit` is accepted and widened to nanoseconds, keeping its
Review Comment:
Can we please use slt for this test? It is more concise and doesn't require
compiling / linking
--
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]
