Skip to content
Open
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
52 changes: 44 additions & 8 deletions datafusion/functions-table/src/generate_series.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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;
Expand Down Expand Up @@ -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<Self::ValueType>) -> Result<ArrayRef>;
Expand All @@ -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(())
}

Expand Down Expand Up @@ -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());
Expand Down Expand Up @@ -417,7 +433,15 @@ impl<T: SeriesValue> LazyBatchGenerator for GenericSeriesState<T> {
.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() {
Expand Down Expand Up @@ -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)?;
Expand Down
37 changes: 37 additions & 0 deletions datafusion/sqllogictest/test_files/table_functions.slt
Original file line number Diff line number Diff line change
Expand Up @@ -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)
Expand Down
Loading