datafusion_functions_table/

generate_series.rs

// Licensed to the Apache Software Foundation (ASF) under one
// or more contributor license agreements.  See the NOTICE file
// distributed with this work for additional information
// regarding copyright ownership.  The ASF licenses this file
// to you under the Apache License, Version 2.0 (the
// "License"); you may not use this file except in compliance
// with the License.  You may obtain a copy of the License at
//
//   http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing,
// software distributed under the License is distributed on an
// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
// KIND, either express or implied.  See the License for the
// specific language governing permissions and limitations
// under the License.

use arrow::array::timezone::Tz;
use arrow::array::types::TimestampNanosecondType;
use arrow::array::{ArrayRef, Int64Array, TimestampNanosecondArray};
use arrow::compute::SortOptions;
use arrow::datatypes::{
    DataType, Field, IntervalMonthDayNano, Schema, SchemaRef, TimeUnit,
};
use arrow::record_batch::RecordBatch;
use async_trait::async_trait;
use datafusion_catalog::TableFunctionImpl;
use datafusion_catalog::TableProvider;
use datafusion_catalog::{Session, TableFunctionArgs};
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;
use datafusion_physical_plan::ExecutionPlan;
use datafusion_physical_plan::memory::{LazyBatchGenerator, LazyMemoryExec};
use parking_lot::RwLock;
use std::any::Any;
use std::fmt;
use std::str::FromStr;
use std::sync::Arc;

/// Empty generator that produces no rows - used when series arguments contain null values
#[derive(Debug, Clone)]
pub struct Empty {
    name: &'static str,
}

impl Empty {
    pub fn name(&self) -> &'static str {
        self.name
    }
}

impl LazyBatchGenerator for Empty {
    fn as_any(&self) -> &dyn Any {
        self
    }

    fn generate_next_batch(&mut self) -> Result<Option<RecordBatch>> {
        Ok(None)
    }

    fn reset_state(&self) -> Arc<RwLock<dyn LazyBatchGenerator>> {
        Arc::new(RwLock::new(Empty { name: self.name }))
    }
}

impl fmt::Display for Empty {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        write!(f, "{}: empty", self.name)
    }
}

/// Trait for values that can be generated in a series
pub trait SeriesValue: fmt::Debug + Clone + Send + Sync + 'static {
    type StepType: fmt::Debug + Clone + Send + Sync;
    type ValueType: fmt::Debug + Clone + Send + Sync;

    /// Check if we've reached the end of the series
    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<()>;

    /// Advance to the next value, adjusting the end of the series if needed.
    ///
    /// The default implementation preserves the behavior of [`Self::advance`].
    /// Implementations can override this method when they need to handle an
    /// overflow by terminating the series after the current value.
    fn advance_with_end(&mut self, _end: &mut Self, step: &Self::StepType) -> Result<()> {
        self.advance(step)
    }

    /// Create an Arrow array from a vector of values
    fn create_array(&self, values: Vec<Self::ValueType>) -> Result<ArrayRef>;

    /// Convert self to ValueType for array creation
    fn to_value_type(&self) -> Self::ValueType;

    /// Display the value for debugging
    fn display_value(&self) -> String;
}

impl SeriesValue for i64 {
    type StepType = i64;
    type ValueType = i64;

    fn should_stop(&self, end: Self, step: &Self::StepType, include_end: bool) -> bool {
        reach_end_int64(*self, end, *step, include_end)
    }

    fn advance(&mut self, step: &Self::StepType) -> Result<()> {
        *self += step;
        Ok(())
    }

    fn advance_with_end(&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(())
    }

    fn create_array(&self, values: Vec<Self::ValueType>) -> Result<ArrayRef> {
        Ok(Arc::new(Int64Array::from(values)))
    }

    fn to_value_type(&self) -> Self::ValueType {
        *self
    }

    fn display_value(&self) -> String {
        self.to_string()
    }
}

#[derive(Debug, Clone)]
pub struct TimestampValue {
    value: i64,
    parsed_tz: Option<Tz>,
    tz_str: Option<Arc<str>>,
}

impl TimestampValue {
    pub fn value(&self) -> i64 {
        self.value
    }

    pub fn tz_str(&self) -> Option<&Arc<str>> {
        self.tz_str.as_ref()
    }
}

impl SeriesValue for TimestampValue {
    type StepType = IntervalMonthDayNano;
    type ValueType = i64;

    fn should_stop(&self, end: Self, step: &Self::StepType, include_end: bool) -> bool {
        let step_negative = step.months < 0 || step.days < 0 || step.nanoseconds < 0;

        if include_end {
            if step_negative {
                self.value < end.value
            } else {
                self.value > end.value
            }
        } else if step_negative {
            self.value <= end.value
        } else {
            self.value >= end.value
        }
    }

    fn advance(&mut self, step: &Self::StepType) -> Result<()> {
        let tz = self
            .parsed_tz
            .unwrap_or_else(|| Tz::from_str("+00:00").unwrap());
        let Some(next_ts) =
            TimestampNanosecondType::add_month_day_nano(self.value, *step, tz)
        else {
            return plan_err!(
                "Failed to add interval {:?} to timestamp {}",
                step,
                self.value
            );
        };
        self.value = next_ts;
        Ok(())
    }

    fn advance_with_end(&mut self, end: &mut Self, step: &Self::StepType) -> Result<()> {
        let tz = self
            .parsed_tz
            .unwrap_or_else(|| Tz::from_str("+00:00").unwrap());
        if let Some(next_ts) =
            TimestampNanosecondType::add_month_day_nano(self.value, *step, tz)
        {
            self.value = next_ts;
        } else {
            // Advancing would exceed the timestamp range. Clamp `end` so the
            // series terminates after the current (last reachable) value.
            let step_negative = step.months < 0 || step.days < 0 || step.nanoseconds < 0;
            end.value = if step_negative {
                self.value.saturating_add(1)
            } else {
                self.value.saturating_sub(1)
            };
        }
        Ok(())
    }

    fn create_array(&self, values: Vec<Self::ValueType>) -> Result<ArrayRef> {
        let array = TimestampNanosecondArray::from(values);

        // Use timezone from self (now we have access to tz through &self)
        let array = match self.tz_str.as_ref() {
            Some(tz_str) => array.with_timezone(Arc::clone(tz_str)),
            None => array,
        };

        Ok(Arc::new(array))
    }

    fn to_value_type(&self) -> Self::ValueType {
        self.value
    }

    fn display_value(&self) -> String {
        self.value.to_string()
    }
}

/// Indicates the arguments used for generating a series.
#[derive(Debug, Clone)]
pub enum GenSeriesArgs {
    /// ContainsNull signifies that at least one argument(start, end, step) was null, thus no series will be generated.
    ContainsNull { name: &'static str },
    /// Int64Args holds the start, end, and step values for generating integer series when all arguments are not null.
    Int64Args {
        start: i64,
        end: i64,
        step: i64,
        /// Indicates whether the end value should be included in the series.
        include_end: bool,
        name: &'static str,
    },
    /// TimestampArgs holds the start, end, and step values for generating timestamp series when all arguments are not null.
    TimestampArgs {
        start: i64,
        end: i64,
        step: IntervalMonthDayNano,
        tz: Option<Arc<str>>,
        /// Indicates whether the end value should be included in the series.
        include_end: bool,
        name: &'static str,
    },
    /// DateArgs holds the start, end, and step values for generating date series when all arguments are not null.
    /// Internally, dates are converted to timestamps and use the timestamp logic.
    DateArgs {
        start: i64,
        end: i64,
        step: IntervalMonthDayNano,
        /// Indicates whether the end value should be included in the series.
        include_end: bool,
        name: &'static str,
    },
}

/// Table that generates a series of integers/timestamps from `start`(inclusive) to `end`, incrementing by step
#[derive(Debug, Clone)]
pub struct GenerateSeriesTable {
    schema: SchemaRef,
    args: GenSeriesArgs,
}

impl GenerateSeriesTable {
    pub fn new(schema: SchemaRef, args: GenSeriesArgs) -> Self {
        Self { schema, args }
    }

    pub fn as_generator(
        &self,
        batch_size: usize,
    ) -> Result<Arc<RwLock<dyn LazyBatchGenerator>>> {
        let generator: Arc<RwLock<dyn LazyBatchGenerator>> = match &self.args {
            GenSeriesArgs::ContainsNull { name } => Arc::new(RwLock::new(Empty { name })),
            GenSeriesArgs::Int64Args {
                start,
                end,
                step,
                include_end,
                name,
            } => Arc::new(RwLock::new(GenericSeriesState {
                schema: self.schema(),
                start: *start,
                end: *end,
                step: *step,
                current: *start,
                finished: false,
                batch_size,
                include_end: *include_end,
                name,
            })),
            GenSeriesArgs::TimestampArgs {
                start,
                end,
                step,
                tz,
                include_end,
                name,
            } => {
                let parsed_tz = tz
                    .as_ref()
                    .map(|s| Tz::from_str(s.as_ref()))
                    .transpose()
                    .map_err(|e| {
                        datafusion_common::internal_datafusion_err!(
                            "Failed to parse timezone: {e}"
                        )
                    })?
                    .unwrap_or_else(|| Tz::from_str("+00:00").unwrap());
                Arc::new(RwLock::new(GenericSeriesState {
                    schema: self.schema(),
                    start: TimestampValue {
                        value: *start,
                        parsed_tz: Some(parsed_tz),
                        tz_str: tz.clone(),
                    },
                    end: TimestampValue {
                        value: *end,
                        parsed_tz: Some(parsed_tz),
                        tz_str: tz.clone(),
                    },
                    step: *step,
                    current: TimestampValue {
                        value: *start,
                        parsed_tz: Some(parsed_tz),
                        tz_str: tz.clone(),
                    },
                    finished: false,
                    batch_size,
                    include_end: *include_end,
                    name,
                }))
            }
            GenSeriesArgs::DateArgs {
                start,
                end,
                step,
                include_end,
                name,
            } => Arc::new(RwLock::new(GenericSeriesState {
                schema: self.schema(),
                start: TimestampValue {
                    value: *start,
                    parsed_tz: None,
                    tz_str: None,
                },
                end: TimestampValue {
                    value: *end,
                    parsed_tz: None,
                    tz_str: None,
                },
                step: *step,
                current: TimestampValue {
                    value: *start,
                    parsed_tz: None,
                    tz_str: None,
                },
                finished: false,
                batch_size,
                include_end: *include_end,
                name,
            })),
        };

        Ok(generator)
    }

    /// Detects output sort order to potentially remove `SortExec`.
    /// Only the `Int64` argument type is currently supported.
    fn output_ordering(&self, schema: &Schema) -> Option<PhysicalSortExpr> {
        let step = match &self.args {
            GenSeriesArgs::Int64Args { step, .. } => *step,
            _ => return None,
        };

        if schema.fields().is_empty() {
            return None;
        }

        let descending = step < 0;
        Some(PhysicalSortExpr::new(
            Arc::new(Column::new(schema.field(0).name(), 0)),
            SortOptions {
                descending,
                // this table function won't output nulls, so either is fine
                nulls_first: false,
            },
        ))
    }
}

#[derive(Debug, Clone)]
pub struct GenericSeriesState<T: SeriesValue> {
    schema: SchemaRef,
    start: T,
    end: T,
    step: T::StepType,
    batch_size: usize,
    current: T,
    finished: bool,
    include_end: bool,
    name: &'static str,
}

impl<T: SeriesValue> GenericSeriesState<T> {
    pub fn name(&self) -> &'static str {
        self.name
    }

    pub fn batch_size(&self) -> usize {
        self.batch_size
    }

    pub fn include_end(&self) -> bool {
        self.include_end
    }

    pub fn start(&self) -> &T {
        &self.start
    }

    pub fn end(&self) -> &T {
        &self.end
    }

    pub fn step(&self) -> &T::StepType {
        &self.step
    }

    pub fn current(&self) -> &T {
        &self.current
    }
}

impl<T: SeriesValue> LazyBatchGenerator for GenericSeriesState<T> {
    fn as_any(&self) -> &dyn Any {
        self
    }

    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
            && !self
                .current
                .should_stop(self.end.clone(), &self.step, self.include_end)
        {
            buf.push(self.current.to_value_type());
            if self
                .current
                .should_stop(self.end.clone(), &self.step, false)
            {
                self.finished = true;
                break;
            }

            let original_end = self.end.clone();
            self.current.advance_with_end(&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() {
            return Ok(None);
        }

        let array = self.current.create_array(buf)?;
        let batch = RecordBatch::try_new(Arc::clone(&self.schema), vec![array])?;
        Ok(Some(batch))
    }

    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))
    }
}

impl<T: SeriesValue> fmt::Display for GenericSeriesState<T> {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        write!(
            f,
            "{}: start={}, end={}, batch_size={}",
            self.name,
            self.start.display_value(),
            self.end.display_value(),
            self.batch_size
        )
    }
}

fn reach_end_int64(val: i64, end: i64, step: i64, include_end: bool) -> bool {
    if step > 0 {
        if include_end { val > end } else { val >= end }
    } else if include_end {
        val < end
    } else {
        val <= end
    }
}

fn validate_interval_step(step: IntervalMonthDayNano) -> Result<()> {
    if step.months == 0 && step.days == 0 && step.nanoseconds == 0 {
        return plan_err!("Step interval cannot be zero");
    }

    Ok(())
}

#[async_trait]
impl TableProvider for GenerateSeriesTable {
    fn schema(&self) -> SchemaRef {
        Arc::clone(&self.schema)
    }

    fn table_type(&self) -> TableType {
        TableType::Base
    }

    async fn scan(
        &self,
        state: &dyn Session,
        projection: Option<&Vec<usize>>,
        _filters: &[Expr],
        _limit: Option<usize>,
    ) -> Result<Arc<dyn ExecutionPlan>> {
        let batch_size = state.config_options().execution.batch_size.get();
        let generator = self.as_generator(batch_size)?;
        let mut exec = LazyMemoryExec::try_new(self.schema(), vec![generator])?
            .with_projection(projection.cloned());

        if let Some(ordering) = self.output_ordering(exec.schema().as_ref()) {
            exec.add_ordering([ordering]);
        }

        Ok(Arc::new(exec))
    }
}

#[derive(Debug)]
struct GenerateSeriesFuncImpl {
    name: &'static str,
    include_end: bool,
}

impl TableFunctionImpl for GenerateSeriesFuncImpl {
    fn call_with_args(&self, args: TableFunctionArgs) -> Result<Arc<dyn TableProvider>> {
        let exprs = args.exprs();
        if exprs.is_empty() || exprs.len() > 3 {
            return plan_err!("{} function requires 1 to 3 arguments", self.name);
        }

        // Determine the data type from the first argument
        match &exprs[0] {
            Expr::Literal(
                // Default to int64 for null
                ScalarValue::Null | ScalarValue::Int64(_),
                _,
            ) => self.call_int64(exprs),
            Expr::Literal(s, _) if matches!(s.data_type(), DataType::Timestamp(_, _)) => {
                self.call_timestamp(exprs)
            }
            Expr::Literal(s, _) if matches!(s.data_type(), DataType::Date32) => {
                self.call_date(exprs)
            }
            Expr::Literal(scalar, _) => {
                plan_err!(
                    "Argument #1 must be an INTEGER, TIMESTAMP, DATE or NULL, got {:?}",
                    scalar.data_type()
                )
            }
            _ => plan_err!("Arguments must be literals"),
        }
    }
}

impl GenerateSeriesFuncImpl {
    fn call_int64(&self, exprs: &[Expr]) -> Result<Arc<dyn TableProvider>> {
        let mut normalize_args = Vec::new();
        for (expr_index, expr) in exprs.iter().enumerate() {
            match expr {
                Expr::Literal(ScalarValue::Null, _) => {}
                Expr::Literal(ScalarValue::Int64(Some(n)), _) => normalize_args.push(*n),
                other => {
                    return plan_err!(
                        "Argument #{} must be an INTEGER or NULL, got {:?}",
                        expr_index + 1,
                        other
                    );
                }
            };
        }

        let schema = Arc::new(Schema::new(vec![Field::new(
            "value",
            DataType::Int64,
            false,
        )]));

        if normalize_args.len() != exprs.len() {
            // contain null
            return Ok(Arc::new(GenerateSeriesTable {
                schema,
                args: GenSeriesArgs::ContainsNull { name: self.name },
            }));
        }

        let (start, end, step) = match &normalize_args[..] {
            [end] => (0, *end, 1),
            [start, end] => (*start, *end, 1),
            [start, end, step] => (*start, *end, *step),
            _ => {
                return plan_err!("{} function requires 1 to 3 arguments", self.name);
            }
        };

        if step == 0 {
            return plan_err!("Step cannot be zero");
        }

        Ok(Arc::new(GenerateSeriesTable {
            schema,
            args: GenSeriesArgs::Int64Args {
                start,
                end,
                step,
                include_end: self.include_end,
                name: self.name,
            },
        }))
    }

    fn call_timestamp(&self, exprs: &[Expr]) -> Result<Arc<dyn TableProvider>> {
        if exprs.len() != 3 {
            return plan_err!(
                "{} function with timestamps requires exactly 3 arguments",
                self.name
            );
        }

        // Parse start timestamp
        let (start_ts, tz) = match &exprs[0] {
            Expr::Literal(ScalarValue::TimestampNanosecond(ts, tz), _) => {
                (*ts, tz.clone())
            }
            other => {
                return plan_err!(
                    "First argument must be a timestamp or NULL, got {:?}",
                    other
                );
            }
        };

        // Parse end timestamp
        let end_ts = match &exprs[1] {
            Expr::Literal(ScalarValue::Null, _) => None,
            Expr::Literal(ScalarValue::TimestampNanosecond(ts, _), _) => *ts,
            other => {
                return plan_err!(
                    "Second argument must be a timestamp or NULL, got {:?}",
                    other
                );
            }
        };

        // Parse step interval
        let step_interval = match &exprs[2] {
            Expr::Literal(ScalarValue::Null, _) => None,
            Expr::Literal(ScalarValue::IntervalMonthDayNano(interval), _) => *interval,
            other => {
                return plan_err!(
                    "Third argument must be an interval or NULL, got {:?}",
                    other
                );
            }
        };

        let schema = Arc::new(Schema::new(vec![Field::new(
            "value",
            DataType::Timestamp(TimeUnit::Nanosecond, tz.clone()),
            false,
        )]));

        // Check if any argument is null
        let (Some(start), Some(end), Some(step)) = (start_ts, end_ts, step_interval)
        else {
            return Ok(Arc::new(GenerateSeriesTable {
                schema,
                args: GenSeriesArgs::ContainsNull { name: self.name },
            }));
        };

        // Validate step interval
        validate_interval_step(step)?;

        Ok(Arc::new(GenerateSeriesTable {
            schema,
            args: GenSeriesArgs::TimestampArgs {
                start,
                end,
                step,
                tz,
                include_end: self.include_end,
                name: self.name,
            },
        }))
    }

    fn call_date(&self, exprs: &[Expr]) -> Result<Arc<dyn TableProvider>> {
        if exprs.len() != 3 {
            return plan_err!(
                "{} function with dates requires exactly 3 arguments",
                self.name
            );
        }

        let schema = Arc::new(Schema::new(vec![Field::new(
            "value",
            DataType::Timestamp(TimeUnit::Nanosecond, None),
            false,
        )]));

        // Parse start date
        let start_date = match &exprs[0] {
            Expr::Literal(ScalarValue::Date32(Some(date)), _) => *date,
            Expr::Literal(ScalarValue::Date32(None), _)
            | Expr::Literal(ScalarValue::Null, _) => {
                return Ok(Arc::new(GenerateSeriesTable {
                    schema,
                    args: GenSeriesArgs::ContainsNull { name: self.name },
                }));
            }
            other => {
                return plan_err!(
                    "First argument must be a date or NULL, got {:?}",
                    other
                );
            }
        };

        // Parse end date
        let end_date = match &exprs[1] {
            Expr::Literal(ScalarValue::Date32(Some(date)), _) => *date,
            Expr::Literal(ScalarValue::Date32(None), _)
            | Expr::Literal(ScalarValue::Null, _) => {
                return Ok(Arc::new(GenerateSeriesTable {
                    schema,
                    args: GenSeriesArgs::ContainsNull { name: self.name },
                }));
            }
            other => {
                return plan_err!(
                    "Second argument must be a date or NULL, got {:?}",
                    other
                );
            }
        };

        // Parse step interval
        let step_interval = match &exprs[2] {
            Expr::Literal(ScalarValue::IntervalMonthDayNano(Some(interval)), _) => {
                *interval
            }
            Expr::Literal(ScalarValue::IntervalMonthDayNano(None), _)
            | Expr::Literal(ScalarValue::Null, _) => {
                return Ok(Arc::new(GenerateSeriesTable {
                    schema,
                    args: GenSeriesArgs::ContainsNull { name: self.name },
                }));
            }
            other => {
                return plan_err!(
                    "Third argument must be an interval or NULL, got {:?}",
                    other
                );
            }
        };

        // Convert Date32 (days since epoch) to timestamp nanoseconds (nanoseconds since epoch)
        // 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;

        // 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)?;

        Ok(Arc::new(GenerateSeriesTable {
            schema,
            args: GenSeriesArgs::DateArgs {
                start: start_ts,
                end: end_ts,
                step: step_interval,
                include_end: self.include_end,
                name: self.name,
            },
        }))
    }
}

#[derive(Debug)]
pub struct GenerateSeriesFunc {}

impl TableFunctionImpl for GenerateSeriesFunc {
    fn call_with_args(&self, args: TableFunctionArgs) -> Result<Arc<dyn TableProvider>> {
        let impl_func = GenerateSeriesFuncImpl {
            name: "generate_series",
            include_end: true,
        };
        impl_func.call_with_args(args)
    }
}

#[derive(Debug)]
pub struct RangeFunc {}

impl TableFunctionImpl for RangeFunc {
    fn call_with_args(&self, args: TableFunctionArgs) -> Result<Arc<dyn TableProvider>> {
        let impl_func = GenerateSeriesFuncImpl {
            name: "range",
            include_end: false,
        };
        impl_func.call_with_args(args)
    }
}

#[cfg(test)]
mod generate_series_tests {
    use std::sync::Arc;

    use arrow::datatypes::{DataType, Field, Schema};
    use datafusion_common::Result;
    use datafusion_physical_plan::memory::LazyBatchGenerator;

    use crate::generate_series::GenericSeriesState;

    #[test]
    fn test_generic_series_state_reset() -> Result<()> {
        let schema = Arc::new(Schema::new(vec![Field::new("a", DataType::Int64, false)]));
        let mut state = GenericSeriesState::<i64> {
            schema,
            start: 1,
            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::<i64> {
            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
            .write()
            .generate_next_batch()?
            .expect("missing reset batch");

        assert_eq!(batch, reset_batch);

        Ok(())
    }
}