update

Author: xiedeyantu

datafusion/functions-table/src/generate_series.rs

@@ -263,6 +263,7 @@ impl GenerateSeriesTable {
                 batch_size,
                 include_end: *include_end,
                 name,
+                finished: false,
             })),
             GenSeriesArgs::TimestampArgs {
                 start,
@@ -303,6 +304,7 @@ impl GenerateSeriesTable {
                     batch_size,
                     include_end: *include_end,
                     name,
+                    finished: false,
                 }))
             }
             GenSeriesArgs::DateArgs {
@@ -332,6 +334,7 @@ impl GenerateSeriesTable {
                 batch_size,
                 include_end: *include_end,
                 name,
+                finished: false,
             })),
         };
 
@@ -349,6 +352,8 @@ pub struct GenericSeriesState<T: SeriesValue> {
     current: T,
     include_end: bool,
     name: &'static str,
+    /// Set to true when advancing would overflow, so the next call returns None.
+    finished: bool,
 }
 
 impl<T: SeriesValue> GenericSeriesState<T> {
@@ -387,6 +392,10 @@ impl<T: SeriesValue> LazyBatchGenerator for GenericSeriesState<T> {
     }
 
     fn generate_next_batch(&mut self) -> Result<Option<RecordBatch>> {
+        if self.finished {
+            return Ok(None);
+        }
+
         let mut buf = Vec::with_capacity(self.batch_size);
 
         while buf.len() < self.batch_size
@@ -399,6 +408,7 @@ impl<T: SeriesValue> LazyBatchGenerator for GenericSeriesState<T> {
             // is i64::MAX and step is 1), the series is exhausted – stop here
             // rather than returning an error, matching PostgreSQL/DuckDB behavior.
             if self.current.advance(&self.step).is_err() {
+                self.finished = true;
                 break;
             }
         }
@@ -415,6 +425,7 @@ impl<T: SeriesValue> LazyBatchGenerator for GenericSeriesState<T> {
     fn reset_state(&self) -> Arc<RwLock<dyn LazyBatchGenerator>> {
         let mut new = self.clone();
         new.current = new.start.clone();
+        new.finished = false;
         Arc::new(RwLock::new(new))
     }
 }
@@ -787,6 +798,7 @@ mod generate_series_tests {
             batch_size: 8192,
             include_end: true,
             name: "test",
+            finished: false,
         };
         let batch = state.generate_next_batch()?.expect("missing batch");