datafusion_functions_table/
generate_series.rs

1// Licensed to the Apache Software Foundation (ASF) under one
2// or more contributor license agreements.  See the NOTICE file
3// distributed with this work for additional information
4// regarding copyright ownership.  The ASF licenses this file
5// to you under the Apache License, Version 2.0 (the
6// "License"); you may not use this file except in compliance
7// with the License.  You may obtain a copy of the License at
8//
9//   http://www.apache.org/licenses/LICENSE-2.0
10//
11// Unless required by applicable law or agreed to in writing,
12// software distributed under the License is distributed on an
13// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
14// KIND, either express or implied.  See the License for the
15// specific language governing permissions and limitations
16// under the License.
17
18use arrow::array::timezone::Tz;
19use arrow::array::types::TimestampNanosecondType;
20use arrow::array::{ArrayRef, Int64Array, TimestampNanosecondArray};
21use arrow::compute::SortOptions;
22use arrow::datatypes::{
23    DataType, Field, IntervalMonthDayNano, Schema, SchemaRef, TimeUnit,
24};
25use arrow::record_batch::RecordBatch;
26use async_trait::async_trait;
27use datafusion_catalog::TableFunctionImpl;
28use datafusion_catalog::TableProvider;
29use datafusion_catalog::{Session, TableFunctionArgs};
30use datafusion_common::{Result, ScalarValue, plan_err};
31use datafusion_expr::{Expr, TableType};
32use datafusion_physical_expr::PhysicalSortExpr;
33use datafusion_physical_expr::expressions::Column;
34use datafusion_physical_plan::ExecutionPlan;
35use datafusion_physical_plan::memory::{LazyBatchGenerator, LazyMemoryExec};
36use parking_lot::RwLock;
37use std::any::Any;
38use std::fmt;
39use std::str::FromStr;
40use std::sync::Arc;
41
42/// Empty generator that produces no rows - used when series arguments contain null values
43#[derive(Debug, Clone)]
44pub struct Empty {
45    name: &'static str,
46}
47
48impl Empty {
49    pub fn name(&self) -> &'static str {
50        self.name
51    }
52}
53
54impl LazyBatchGenerator for Empty {
55    fn as_any(&self) -> &dyn Any {
56        self
57    }
58
59    fn generate_next_batch(&mut self) -> Result<Option<RecordBatch>> {
60        Ok(None)
61    }
62
63    fn reset_state(&self) -> Arc<RwLock<dyn LazyBatchGenerator>> {
64        Arc::new(RwLock::new(Empty { name: self.name }))
65    }
66}
67
68impl fmt::Display for Empty {
69    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
70        write!(f, "{}: empty", self.name)
71    }
72}
73
74/// Trait for values that can be generated in a series
75pub trait SeriesValue: fmt::Debug + Clone + Send + Sync + 'static {
76    type StepType: fmt::Debug + Clone + Send + Sync;
77    type ValueType: fmt::Debug + Clone + Send + Sync;
78
79    /// Check if we've reached the end of the series
80    fn should_stop(&self, end: Self, step: &Self::StepType, include_end: bool) -> bool;
81
82    /// Advance to the next value in the series
83    fn advance(&mut self, step: &Self::StepType) -> Result<()>;
84
85    /// Create an Arrow array from a vector of values
86    fn create_array(&self, values: Vec<Self::ValueType>) -> Result<ArrayRef>;
87
88    /// Convert self to ValueType for array creation
89    fn to_value_type(&self) -> Self::ValueType;
90
91    /// Display the value for debugging
92    fn display_value(&self) -> String;
93}
94
95impl SeriesValue for i64 {
96    type StepType = i64;
97    type ValueType = i64;
98
99    fn should_stop(&self, end: Self, step: &Self::StepType, include_end: bool) -> bool {
100        reach_end_int64(*self, end, *step, include_end)
101    }
102
103    fn advance(&mut self, step: &Self::StepType) -> Result<()> {
104        *self += step;
105        Ok(())
106    }
107
108    fn create_array(&self, values: Vec<Self::ValueType>) -> Result<ArrayRef> {
109        Ok(Arc::new(Int64Array::from(values)))
110    }
111
112    fn to_value_type(&self) -> Self::ValueType {
113        *self
114    }
115
116    fn display_value(&self) -> String {
117        self.to_string()
118    }
119}
120
121#[derive(Debug, Clone)]
122pub struct TimestampValue {
123    value: i64,
124    parsed_tz: Option<Tz>,
125    tz_str: Option<Arc<str>>,
126}
127
128impl TimestampValue {
129    pub fn value(&self) -> i64 {
130        self.value
131    }
132
133    pub fn tz_str(&self) -> Option<&Arc<str>> {
134        self.tz_str.as_ref()
135    }
136}
137
138impl SeriesValue for TimestampValue {
139    type StepType = IntervalMonthDayNano;
140    type ValueType = i64;
141
142    fn should_stop(&self, end: Self, step: &Self::StepType, include_end: bool) -> bool {
143        let step_negative = step.months < 0 || step.days < 0 || step.nanoseconds < 0;
144
145        if include_end {
146            if step_negative {
147                self.value < end.value
148            } else {
149                self.value > end.value
150            }
151        } else if step_negative {
152            self.value <= end.value
153        } else {
154            self.value >= end.value
155        }
156    }
157
158    fn advance(&mut self, step: &Self::StepType) -> Result<()> {
159        let tz = self
160            .parsed_tz
161            .unwrap_or_else(|| Tz::from_str("+00:00").unwrap());
162        let Some(next_ts) =
163            TimestampNanosecondType::add_month_day_nano(self.value, *step, tz)
164        else {
165            return plan_err!(
166                "Failed to add interval {:?} to timestamp {}",
167                step,
168                self.value
169            );
170        };
171        self.value = next_ts;
172        Ok(())
173    }
174
175    fn create_array(&self, values: Vec<Self::ValueType>) -> Result<ArrayRef> {
176        let array = TimestampNanosecondArray::from(values);
177
178        // Use timezone from self (now we have access to tz through &self)
179        let array = match self.tz_str.as_ref() {
180            Some(tz_str) => array.with_timezone(Arc::clone(tz_str)),
181            None => array,
182        };
183
184        Ok(Arc::new(array))
185    }
186
187    fn to_value_type(&self) -> Self::ValueType {
188        self.value
189    }
190
191    fn display_value(&self) -> String {
192        self.value.to_string()
193    }
194}
195
196/// Indicates the arguments used for generating a series.
197#[derive(Debug, Clone)]
198pub enum GenSeriesArgs {
199    /// ContainsNull signifies that at least one argument(start, end, step) was null, thus no series will be generated.
200    ContainsNull { name: &'static str },
201    /// Int64Args holds the start, end, and step values for generating integer series when all arguments are not null.
202    Int64Args {
203        start: i64,
204        end: i64,
205        step: i64,
206        /// Indicates whether the end value should be included in the series.
207        include_end: bool,
208        name: &'static str,
209    },
210    /// TimestampArgs holds the start, end, and step values for generating timestamp series when all arguments are not null.
211    TimestampArgs {
212        start: i64,
213        end: i64,
214        step: IntervalMonthDayNano,
215        tz: Option<Arc<str>>,
216        /// Indicates whether the end value should be included in the series.
217        include_end: bool,
218        name: &'static str,
219    },
220    /// DateArgs holds the start, end, and step values for generating date series when all arguments are not null.
221    /// Internally, dates are converted to timestamps and use the timestamp logic.
222    DateArgs {
223        start: i64,
224        end: i64,
225        step: IntervalMonthDayNano,
226        /// Indicates whether the end value should be included in the series.
227        include_end: bool,
228        name: &'static str,
229    },
230}
231
232/// Table that generates a series of integers/timestamps from `start`(inclusive) to `end`, incrementing by step
233#[derive(Debug, Clone)]
234pub struct GenerateSeriesTable {
235    schema: SchemaRef,
236    args: GenSeriesArgs,
237}
238
239impl GenerateSeriesTable {
240    pub fn new(schema: SchemaRef, args: GenSeriesArgs) -> Self {
241        Self { schema, args }
242    }
243
244    pub fn as_generator(
245        &self,
246        batch_size: usize,
247    ) -> Result<Arc<RwLock<dyn LazyBatchGenerator>>> {
248        let generator: Arc<RwLock<dyn LazyBatchGenerator>> = match &self.args {
249            GenSeriesArgs::ContainsNull { name } => Arc::new(RwLock::new(Empty { name })),
250            GenSeriesArgs::Int64Args {
251                start,
252                end,
253                step,
254                include_end,
255                name,
256            } => Arc::new(RwLock::new(GenericSeriesState {
257                schema: self.schema(),
258                start: *start,
259                end: *end,
260                step: *step,
261                current: *start,
262                batch_size,
263                include_end: *include_end,
264                name,
265            })),
266            GenSeriesArgs::TimestampArgs {
267                start,
268                end,
269                step,
270                tz,
271                include_end,
272                name,
273            } => {
274                let parsed_tz = tz
275                    .as_ref()
276                    .map(|s| Tz::from_str(s.as_ref()))
277                    .transpose()
278                    .map_err(|e| {
279                        datafusion_common::internal_datafusion_err!(
280                            "Failed to parse timezone: {e}"
281                        )
282                    })?
283                    .unwrap_or_else(|| Tz::from_str("+00:00").unwrap());
284                Arc::new(RwLock::new(GenericSeriesState {
285                    schema: self.schema(),
286                    start: TimestampValue {
287                        value: *start,
288                        parsed_tz: Some(parsed_tz),
289                        tz_str: tz.clone(),
290                    },
291                    end: TimestampValue {
292                        value: *end,
293                        parsed_tz: Some(parsed_tz),
294                        tz_str: tz.clone(),
295                    },
296                    step: *step,
297                    current: TimestampValue {
298                        value: *start,
299                        parsed_tz: Some(parsed_tz),
300                        tz_str: tz.clone(),
301                    },
302                    batch_size,
303                    include_end: *include_end,
304                    name,
305                }))
306            }
307            GenSeriesArgs::DateArgs {
308                start,
309                end,
310                step,
311                include_end,
312                name,
313            } => Arc::new(RwLock::new(GenericSeriesState {
314                schema: self.schema(),
315                start: TimestampValue {
316                    value: *start,
317                    parsed_tz: None,
318                    tz_str: None,
319                },
320                end: TimestampValue {
321                    value: *end,
322                    parsed_tz: None,
323                    tz_str: None,
324                },
325                step: *step,
326                current: TimestampValue {
327                    value: *start,
328                    parsed_tz: None,
329                    tz_str: None,
330                },
331                batch_size,
332                include_end: *include_end,
333                name,
334            })),
335        };
336
337        Ok(generator)
338    }
339
340    /// Detects output sort order to potentially remove `SortExec`.
341    /// Only the `Int64` argument type is currently supported.
342    fn output_ordering(&self, schema: &Schema) -> Option<PhysicalSortExpr> {
343        let step = match &self.args {
344            GenSeriesArgs::Int64Args { step, .. } => *step,
345            _ => return None,
346        };
347
348        if schema.fields().is_empty() {
349            return None;
350        }
351
352        let descending = step < 0;
353        Some(PhysicalSortExpr::new(
354            Arc::new(Column::new(schema.field(0).name(), 0)),
355            SortOptions {
356                descending,
357                // this table function won't output nulls, so either is fine
358                nulls_first: false,
359            },
360        ))
361    }
362}
363
364#[derive(Debug, Clone)]
365pub struct GenericSeriesState<T: SeriesValue> {
366    schema: SchemaRef,
367    start: T,
368    end: T,
369    step: T::StepType,
370    batch_size: usize,
371    current: T,
372    include_end: bool,
373    name: &'static str,
374}
375
376impl<T: SeriesValue> GenericSeriesState<T> {
377    pub fn name(&self) -> &'static str {
378        self.name
379    }
380
381    pub fn batch_size(&self) -> usize {
382        self.batch_size
383    }
384
385    pub fn include_end(&self) -> bool {
386        self.include_end
387    }
388
389    pub fn start(&self) -> &T {
390        &self.start
391    }
392
393    pub fn end(&self) -> &T {
394        &self.end
395    }
396
397    pub fn step(&self) -> &T::StepType {
398        &self.step
399    }
400
401    pub fn current(&self) -> &T {
402        &self.current
403    }
404}
405
406impl<T: SeriesValue> LazyBatchGenerator for GenericSeriesState<T> {
407    fn as_any(&self) -> &dyn Any {
408        self
409    }
410
411    fn generate_next_batch(&mut self) -> Result<Option<RecordBatch>> {
412        let mut buf = Vec::with_capacity(self.batch_size);
413
414        while buf.len() < self.batch_size
415            && !self
416                .current
417                .should_stop(self.end.clone(), &self.step, self.include_end)
418        {
419            buf.push(self.current.to_value_type());
420            self.current.advance(&self.step)?;
421        }
422
423        if buf.is_empty() {
424            return Ok(None);
425        }
426
427        let array = self.current.create_array(buf)?;
428        let batch = RecordBatch::try_new(Arc::clone(&self.schema), vec![array])?;
429        Ok(Some(batch))
430    }
431
432    fn reset_state(&self) -> Arc<RwLock<dyn LazyBatchGenerator>> {
433        let mut new = self.clone();
434        new.current = new.start.clone();
435        Arc::new(RwLock::new(new))
436    }
437}
438
439impl<T: SeriesValue> fmt::Display for GenericSeriesState<T> {
440    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
441        write!(
442            f,
443            "{}: start={}, end={}, batch_size={}",
444            self.name,
445            self.start.display_value(),
446            self.end.display_value(),
447            self.batch_size
448        )
449    }
450}
451
452fn reach_end_int64(val: i64, end: i64, step: i64, include_end: bool) -> bool {
453    if step > 0 {
454        if include_end { val > end } else { val >= end }
455    } else if include_end {
456        val < end
457    } else {
458        val <= end
459    }
460}
461
462fn validate_interval_step(step: IntervalMonthDayNano) -> Result<()> {
463    if step.months == 0 && step.days == 0 && step.nanoseconds == 0 {
464        return plan_err!("Step interval cannot be zero");
465    }
466
467    Ok(())
468}
469
470#[async_trait]
471impl TableProvider for GenerateSeriesTable {
472    fn schema(&self) -> SchemaRef {
473        Arc::clone(&self.schema)
474    }
475
476    fn table_type(&self) -> TableType {
477        TableType::Base
478    }
479
480    async fn scan(
481        &self,
482        state: &dyn Session,
483        projection: Option<&Vec<usize>>,
484        _filters: &[Expr],
485        _limit: Option<usize>,
486    ) -> Result<Arc<dyn ExecutionPlan>> {
487        let batch_size = state.config_options().execution.batch_size;
488        let generator = self.as_generator(batch_size)?;
489        let mut exec = LazyMemoryExec::try_new(self.schema(), vec![generator])?
490            .with_projection(projection.cloned());
491
492        if let Some(ordering) = self.output_ordering(exec.schema().as_ref()) {
493            exec.add_ordering([ordering]);
494        }
495
496        Ok(Arc::new(exec))
497    }
498}
499
500#[derive(Debug)]
501struct GenerateSeriesFuncImpl {
502    name: &'static str,
503    include_end: bool,
504}
505
506impl TableFunctionImpl for GenerateSeriesFuncImpl {
507    fn call_with_args(&self, args: TableFunctionArgs) -> Result<Arc<dyn TableProvider>> {
508        let exprs = args.exprs();
509        if exprs.is_empty() || exprs.len() > 3 {
510            return plan_err!("{} function requires 1 to 3 arguments", self.name);
511        }
512
513        // Determine the data type from the first argument
514        match &exprs[0] {
515            Expr::Literal(
516                // Default to int64 for null
517                ScalarValue::Null | ScalarValue::Int64(_),
518                _,
519            ) => self.call_int64(exprs),
520            Expr::Literal(s, _) if matches!(s.data_type(), DataType::Timestamp(_, _)) => {
521                self.call_timestamp(exprs)
522            }
523            Expr::Literal(s, _) if matches!(s.data_type(), DataType::Date32) => {
524                self.call_date(exprs)
525            }
526            Expr::Literal(scalar, _) => {
527                plan_err!(
528                    "Argument #1 must be an INTEGER, TIMESTAMP, DATE or NULL, got {:?}",
529                    scalar.data_type()
530                )
531            }
532            _ => plan_err!("Arguments must be literals"),
533        }
534    }
535}
536
537impl GenerateSeriesFuncImpl {
538    fn call_int64(&self, exprs: &[Expr]) -> Result<Arc<dyn TableProvider>> {
539        let mut normalize_args = Vec::new();
540        for (expr_index, expr) in exprs.iter().enumerate() {
541            match expr {
542                Expr::Literal(ScalarValue::Null, _) => {}
543                Expr::Literal(ScalarValue::Int64(Some(n)), _) => normalize_args.push(*n),
544                other => {
545                    return plan_err!(
546                        "Argument #{} must be an INTEGER or NULL, got {:?}",
547                        expr_index + 1,
548                        other
549                    );
550                }
551            };
552        }
553
554        let schema = Arc::new(Schema::new(vec![Field::new(
555            "value",
556            DataType::Int64,
557            false,
558        )]));
559
560        if normalize_args.len() != exprs.len() {
561            // contain null
562            return Ok(Arc::new(GenerateSeriesTable {
563                schema,
564                args: GenSeriesArgs::ContainsNull { name: self.name },
565            }));
566        }
567
568        let (start, end, step) = match &normalize_args[..] {
569            [end] => (0, *end, 1),
570            [start, end] => (*start, *end, 1),
571            [start, end, step] => (*start, *end, *step),
572            _ => {
573                return plan_err!("{} function requires 1 to 3 arguments", self.name);
574            }
575        };
576
577        if step == 0 {
578            return plan_err!("Step cannot be zero");
579        }
580
581        Ok(Arc::new(GenerateSeriesTable {
582            schema,
583            args: GenSeriesArgs::Int64Args {
584                start,
585                end,
586                step,
587                include_end: self.include_end,
588                name: self.name,
589            },
590        }))
591    }
592
593    fn call_timestamp(&self, exprs: &[Expr]) -> Result<Arc<dyn TableProvider>> {
594        if exprs.len() != 3 {
595            return plan_err!(
596                "{} function with timestamps requires exactly 3 arguments",
597                self.name
598            );
599        }
600
601        // Parse start timestamp
602        let (start_ts, tz) = match &exprs[0] {
603            Expr::Literal(ScalarValue::TimestampNanosecond(ts, tz), _) => {
604                (*ts, tz.clone())
605            }
606            other => {
607                return plan_err!(
608                    "First argument must be a timestamp or NULL, got {:?}",
609                    other
610                );
611            }
612        };
613
614        // Parse end timestamp
615        let end_ts = match &exprs[1] {
616            Expr::Literal(ScalarValue::Null, _) => None,
617            Expr::Literal(ScalarValue::TimestampNanosecond(ts, _), _) => *ts,
618            other => {
619                return plan_err!(
620                    "Second argument must be a timestamp or NULL, got {:?}",
621                    other
622                );
623            }
624        };
625
626        // Parse step interval
627        let step_interval = match &exprs[2] {
628            Expr::Literal(ScalarValue::Null, _) => None,
629            Expr::Literal(ScalarValue::IntervalMonthDayNano(interval), _) => *interval,
630            other => {
631                return plan_err!(
632                    "Third argument must be an interval or NULL, got {:?}",
633                    other
634                );
635            }
636        };
637
638        let schema = Arc::new(Schema::new(vec![Field::new(
639            "value",
640            DataType::Timestamp(TimeUnit::Nanosecond, tz.clone()),
641            false,
642        )]));
643
644        // Check if any argument is null
645        let (Some(start), Some(end), Some(step)) = (start_ts, end_ts, step_interval)
646        else {
647            return Ok(Arc::new(GenerateSeriesTable {
648                schema,
649                args: GenSeriesArgs::ContainsNull { name: self.name },
650            }));
651        };
652
653        // Validate step interval
654        validate_interval_step(step)?;
655
656        Ok(Arc::new(GenerateSeriesTable {
657            schema,
658            args: GenSeriesArgs::TimestampArgs {
659                start,
660                end,
661                step,
662                tz,
663                include_end: self.include_end,
664                name: self.name,
665            },
666        }))
667    }
668
669    fn call_date(&self, exprs: &[Expr]) -> Result<Arc<dyn TableProvider>> {
670        if exprs.len() != 3 {
671            return plan_err!(
672                "{} function with dates requires exactly 3 arguments",
673                self.name
674            );
675        }
676
677        let schema = Arc::new(Schema::new(vec![Field::new(
678            "value",
679            DataType::Timestamp(TimeUnit::Nanosecond, None),
680            false,
681        )]));
682
683        // Parse start date
684        let start_date = match &exprs[0] {
685            Expr::Literal(ScalarValue::Date32(Some(date)), _) => *date,
686            Expr::Literal(ScalarValue::Date32(None), _)
687            | Expr::Literal(ScalarValue::Null, _) => {
688                return Ok(Arc::new(GenerateSeriesTable {
689                    schema,
690                    args: GenSeriesArgs::ContainsNull { name: self.name },
691                }));
692            }
693            other => {
694                return plan_err!(
695                    "First argument must be a date or NULL, got {:?}",
696                    other
697                );
698            }
699        };
700
701        // Parse end date
702        let end_date = match &exprs[1] {
703            Expr::Literal(ScalarValue::Date32(Some(date)), _) => *date,
704            Expr::Literal(ScalarValue::Date32(None), _)
705            | Expr::Literal(ScalarValue::Null, _) => {
706                return Ok(Arc::new(GenerateSeriesTable {
707                    schema,
708                    args: GenSeriesArgs::ContainsNull { name: self.name },
709                }));
710            }
711            other => {
712                return plan_err!(
713                    "Second argument must be a date or NULL, got {:?}",
714                    other
715                );
716            }
717        };
718
719        // Parse step interval
720        let step_interval = match &exprs[2] {
721            Expr::Literal(ScalarValue::IntervalMonthDayNano(Some(interval)), _) => {
722                *interval
723            }
724            Expr::Literal(ScalarValue::IntervalMonthDayNano(None), _)
725            | Expr::Literal(ScalarValue::Null, _) => {
726                return Ok(Arc::new(GenerateSeriesTable {
727                    schema,
728                    args: GenSeriesArgs::ContainsNull { name: self.name },
729                }));
730            }
731            other => {
732                return plan_err!(
733                    "Third argument must be an interval or NULL, got {:?}",
734                    other
735                );
736            }
737        };
738
739        // Convert Date32 (days since epoch) to timestamp nanoseconds (nanoseconds since epoch)
740        // Date32 is days since 1970-01-01, so multiply by nanoseconds per day
741        const NANOS_PER_DAY: i64 = 24 * 60 * 60 * 1_000_000_000;
742
743        let start_ts = start_date as i64 * NANOS_PER_DAY;
744        let end_ts = end_date as i64 * NANOS_PER_DAY;
745
746        // Validate step interval
747        validate_interval_step(step_interval)?;
748
749        Ok(Arc::new(GenerateSeriesTable {
750            schema,
751            args: GenSeriesArgs::DateArgs {
752                start: start_ts,
753                end: end_ts,
754                step: step_interval,
755                include_end: self.include_end,
756                name: self.name,
757            },
758        }))
759    }
760}
761
762#[derive(Debug)]
763pub struct GenerateSeriesFunc {}
764
765impl TableFunctionImpl for GenerateSeriesFunc {
766    fn call_with_args(&self, args: TableFunctionArgs) -> Result<Arc<dyn TableProvider>> {
767        let impl_func = GenerateSeriesFuncImpl {
768            name: "generate_series",
769            include_end: true,
770        };
771        impl_func.call_with_args(args)
772    }
773}
774
775#[derive(Debug)]
776pub struct RangeFunc {}
777
778impl TableFunctionImpl for RangeFunc {
779    fn call_with_args(&self, args: TableFunctionArgs) -> Result<Arc<dyn TableProvider>> {
780        let impl_func = GenerateSeriesFuncImpl {
781            name: "range",
782            include_end: false,
783        };
784        impl_func.call_with_args(args)
785    }
786}
787
788#[cfg(test)]
789mod generate_series_tests {
790    use std::sync::Arc;
791
792    use arrow::datatypes::{DataType, Field, Schema};
793    use datafusion_common::Result;
794    use datafusion_physical_plan::memory::LazyBatchGenerator;
795
796    use crate::generate_series::GenericSeriesState;
797
798    #[test]
799    fn test_generic_series_state_reset() -> Result<()> {
800        let schema = Arc::new(Schema::new(vec![Field::new("a", DataType::Int64, false)]));
801        let mut state = GenericSeriesState::<i64> {
802            schema,
803            start: 1,
804            end: 5,
805            step: 1,
806            current: 1,
807            batch_size: 8192,
808            include_end: true,
809            name: "test",
810        };
811        let batch = state.generate_next_batch()?.expect("missing batch");
812
813        let state_reset = state.reset_state();
814        let reset_batch = state_reset
815            .write()
816            .generate_next_batch()?
817            .expect("missing reset batch");
818
819        assert_eq!(batch, reset_batch);
820
821        Ok(())
822    }
823}
datafusion_functions_table/generate_series.rs

datafusion_functions_table/
generate_series.rs
