diff --git a/datafusion/functions-table/src/generate_series.rs b/datafusion/functions-table/src/generate_series.rs index 0e8eca6bc2561..668eee4454f13 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()); @@ -259,6 +275,7 @@ impl GenerateSeriesTable { end: *end, step: *step, current: *start, + finished: false, batch_size, include_end: *include_end, name, @@ -299,6 +316,7 @@ impl GenerateSeriesTable { parsed_tz: Some(parsed_tz), tz_str: tz.clone(), }, + finished: false, batch_size, include_end: *include_end, name, @@ -328,6 +346,7 @@ impl GenerateSeriesTable { parsed_tz: None, tz_str: None, }, + finished: false, batch_size, include_end: *include_end, name, @@ -369,6 +388,7 @@ pub struct GenericSeriesState { step: T::StepType, batch_size: usize, current: T, + finished: bool, include_end: bool, name: &'static str, } @@ -409,6 +429,10 @@ impl LazyBatchGenerator for GenericSeriesState { } fn generate_next_batch(&mut self) -> Result> { + if self.finished { + return Ok(None); + } + let mut buf = Vec::with_capacity(self.batch_size); while buf.len() < self.batch_size @@ -417,7 +441,24 @@ 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.finished = true; + break; + } + + let original_end = self.end.clone(); + self.current.advance(&mut self.end, &self.step)?; + if self + .current + .should_stop(self.end.clone(), &self.step, self.include_end) + { + self.end = original_end; + self.finished = true; + break; + } } if buf.is_empty() { @@ -432,6 +473,7 @@ impl LazyBatchGenerator for GenericSeriesState { fn reset_state(&self) -> Arc> { let mut new = self.clone(); new.current = new.start.clone(); + new.finished = false; Arc::new(RwLock::new(new)) } } @@ -740,8 +782,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)?; @@ -804,11 +858,40 @@ mod generate_series_tests { end: 5, step: 1, current: 1, + finished: false, + batch_size: 8192, + include_end: true, + name: "test", + }; + let batch = state.generate_next_batch()?.expect("missing batch"); + + let state_reset = state.reset_state(); + let reset_batch = state_reset + .write() + .generate_next_batch()? + .expect("missing reset batch"); + + assert_eq!(batch, reset_batch); + + Ok(()) + } + + #[test] + fn test_generic_series_state_reset_after_overflow() -> Result<()> { + let schema = Arc::new(Schema::new(vec![Field::new("a", DataType::Int64, false)])); + let mut state = GenericSeriesState:: { + schema, + start: i64::MAX - 1, + end: i64::MAX, + step: 2, + current: i64::MAX - 1, + finished: false, batch_size: 8192, include_end: true, name: "test", }; let batch = state.generate_next_batch()?.expect("missing batch"); + assert!(state.generate_next_batch()?.is_none()); let state_reset = state.reset_state(); let reset_batch = state_reset diff --git a/datafusion/sqllogictest/test_files/table_functions.slt b/datafusion/sqllogictest/test_files/table_functions.slt index e1ab444d81044..93d55a3edb12f 100644 --- a/datafusion/sqllogictest/test_files/table_functions.slt +++ b/datafusion/sqllogictest/test_files/table_functions.slt @@ -197,6 +197,55 @@ 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) + +# Reaching the maximum representable date must not attempt to advance beyond it. +query P +SELECT * FROM generate_series(DATE '2262-04-11', DATE '2262-04-11', INTERVAL '1' DAY) +---- +2262-04-11T00:00:00 + +# Same for the maximum representable nanosecond timestamp. +query P +SELECT * FROM generate_series(TIMESTAMP '2262-04-11T23:47:16.854775807', TIMESTAMP '2262-04-11T23:47:16.854775807', INTERVAL '1' NANOSECOND) +---- +2262-04-11T23:47:16.854775807 + # 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)