diff --git a/datafusion/functions-table/src/generate_series.rs b/datafusion/functions-table/src/generate_series.rs index 0e8eca6bc2561..3f669fd94999d 100644 --- a/datafusion/functions-table/src/generate_series.rs +++ b/datafusion/functions-table/src/generate_series.rs @@ -27,7 +27,7 @@ use async_trait::async_trait; use datafusion_catalog::TableFunctionImpl; use datafusion_catalog::TableProvider; use datafusion_catalog::{Session, TableFunctionArgs}; -use datafusion_common::{Result, ScalarValue, plan_err}; +use datafusion_common::{Result, ScalarValue, plan_datafusion_err, plan_err}; use datafusion_expr::{Expr, TableType}; use datafusion_physical_expr::PhysicalSortExpr; use datafusion_physical_expr::expressions::Column; @@ -80,7 +80,12 @@ pub trait SeriesValue: fmt::Debug + Clone + Send + Sync + 'static { fn should_stop(&self, end: Self, step: &Self::StepType, include_end: bool) -> bool; /// Advance to the next value in the series - fn advance(&mut self, step: &Self::StepType) -> Result<()>; + /// + /// If advancing would overflow the value range, `end` is updated so that + /// the series terminates after the current value (matching the behavior + /// of PostgreSQL and DuckDB, which return the reachable values instead of + /// erroring). + fn advance(&mut self, end: &mut Self, step: &Self::StepType) -> Result<()>; /// Create an Arrow array from a vector of values fn create_array(&self, values: Vec) -> Result; @@ -100,8 +105,19 @@ impl SeriesValue for i64 { reach_end_int64(*self, end, *step, include_end) } - fn advance(&mut self, step: &Self::StepType) -> Result<()> { - *self += step; + fn advance(&mut self, end: &mut Self, step: &Self::StepType) -> Result<()> { + if let Some(next) = self.checked_add(*step) { + *self = next; + } else { + // Advancing would overflow: clamp `end` so the series stops after + // the current (last reachable) value instead of panicking or + // wrapping around. + *end = if *step > 0 { + self.saturating_sub(1) + } else { + self.saturating_add(1) + }; + } Ok(()) } @@ -155,7 +171,7 @@ impl SeriesValue for TimestampValue { } } - fn advance(&mut self, step: &Self::StepType) -> Result<()> { + fn advance(&mut self, _end: &mut Self, step: &Self::StepType) -> Result<()> { let tz = self .parsed_tz .unwrap_or_else(|| Tz::from_str("+00:00").unwrap()); @@ -417,7 +433,15 @@ impl LazyBatchGenerator for GenericSeriesState { .should_stop(self.end.clone(), &self.step, self.include_end) { buf.push(self.current.to_value_type()); - self.current.advance(&self.step)?; + if self + .current + .should_stop(self.end.clone(), &self.step, false) + { + self.current.advance(&mut self.end, &self.step)?; + break; + } + + self.current.advance(&mut self.end, &self.step)?; } if buf.is_empty() { @@ -740,8 +764,20 @@ impl GenerateSeriesFuncImpl { // Date32 is days since 1970-01-01, so multiply by nanoseconds per day const NANOS_PER_DAY: i64 = 24 * 60 * 60 * 1_000_000_000; - let start_ts = start_date as i64 * NANOS_PER_DAY; - let end_ts = end_date as i64 * NANOS_PER_DAY; + // Dates outside the nanosecond timestamp range (1677-09-21 to + // 2262-04-11) cannot be represented; return an error instead of + // panicking (debug) or silently wrapping (release). + let date_to_ts_nanos = |date: i32, arg: &str| { + (date as i64).checked_mul(NANOS_PER_DAY).ok_or_else(|| { + plan_datafusion_err!( + "{arg} for {} is out of range of nanosecond timestamps", + self.name + ) + }) + }; + + let start_ts = date_to_ts_nanos(start_date, "First argument")?; + let end_ts = date_to_ts_nanos(end_date, "Second argument")?; // Validate step interval validate_interval_step(step_interval)?; diff --git a/datafusion/sqllogictest/test_files/table_functions.slt b/datafusion/sqllogictest/test_files/table_functions.slt index e1ab444d81044..c2c8ce24fa081 100644 --- a/datafusion/sqllogictest/test_files/table_functions.slt +++ b/datafusion/sqllogictest/test_files/table_functions.slt @@ -197,6 +197,43 @@ SELECT * FROM generate_series(1, 2, 3, 4) statement error DataFusion error: Error during planning: Argument \#1 must be an INTEGER, TIMESTAMP, DATE or NULL, got Utf8 SELECT * FROM generate_series('foo', 'bar') +# Regression test for https://github.com/apache/datafusion/issues/22208 +# A step that would overflow i64 after the last reachable value must return the +# reachable values instead of panicking, matching PostgreSQL/DuckDB behavior. +query I +SELECT * FROM generate_series(9223372036854775806, 9223372036854775807, 2) +---- +9223372036854775806 + +# Same, in the descending direction +query I +SELECT * FROM generate_series(-9223372036854775806, -9223372036854775808, -2) +---- +-9223372036854775806 +-9223372036854775808 + +# Landing exactly on i64::MAX must include it +query I +SELECT * FROM generate_series(9223372036854775805, 9223372036854775807, 2) +---- +9223372036854775805 +9223372036854775807 + +# Same overflow behavior for `range` (end exclusive) +query I +SELECT * FROM range(9223372036854775806, 9223372036854775807, 2) +---- +9223372036854775806 + +# Regression test for https://github.com/apache/datafusion/issues/22193 +# Dates outside the nanosecond timestamp range must produce a clean planning +# error instead of panicking (debug) or silently wrapping (release). +statement error DataFusion error: Error during planning: First argument for generate_series is out of range of nanosecond timestamps +SELECT * FROM generate_series(DATE '0001-01-01', DATE '2000-01-01', INTERVAL '1' DAY) + +statement error DataFusion error: Error during planning: Second argument for generate_series is out of range of nanosecond timestamps +SELECT * FROM generate_series(DATE '2000-01-01', DATE '3000-01-01', INTERVAL '1' DAY) + # UDF and UDTF `generate_series` can be used simultaneously query ? rowsort SELECT generate_series(1, t1.end) FROM generate_series(3, 5) as t1(end)