fn gen_range_timestamp()

in datafusion/functions-nested/src/range.rs [511:609]
81 lines of code
20 McCabe index (conditional complexity)

fn gen_range_timestamp(args: &[ArrayRef], include_upper_bound: bool) -> Result<ArrayRef> {
    let func_name = if include_upper_bound {
        "GENERATE_SERIES"
    } else {
        "RANGE"
    };
    let [start, stop, step] = take_function_args(func_name, args)?;

    // coerce_types fn should coerce all types to Timestamp(Nanosecond, tz)
    let (start_arr, start_tz_opt) = cast_timestamp_arg(start, include_upper_bound)?;
    let (stop_arr, stop_tz_opt) = cast_timestamp_arg(stop, include_upper_bound)?;
    let step_arr = as_interval_mdn_array(step)?;
    let start_tz = parse_tz(start_tz_opt)?;
    let stop_tz = parse_tz(stop_tz_opt)?;

    // values are timestamps
    let values_builder = start_tz_opt
        .clone()
        .map_or_else(TimestampNanosecondBuilder::new, |start_tz_str| {
            TimestampNanosecondBuilder::new().with_timezone(start_tz_str)
        });
    let mut list_builder = ListBuilder::new(values_builder);

    for idx in 0..start_arr.len() {
        if start_arr.is_null(idx) || stop_arr.is_null(idx) || step_arr.is_null(idx) {
            list_builder.append_null();
            continue;
        }

        let start = start_arr.value(idx);
        let stop = stop_arr.value(idx);
        let step = step_arr.value(idx);

        let (months, days, ns) = IntervalMonthDayNanoType::to_parts(step);
        if months == 0 && days == 0 && ns == 0 {
            return exec_err!(
                "Interval argument to {} must not be 0",
                if include_upper_bound {
                    "GENERATE_SERIES"
                } else {
                    "RANGE"
                }
            );
        }

        let neg = TSNT::add_month_day_nano(start, step, start_tz)
            .ok_or(exec_datafusion_err!(
                "Cannot generate timestamp range where start + step overflows"
            ))?
            .cmp(&start)
            == Ordering::Less;

        let stop_dt = as_datetime_with_timezone::<TSNT>(stop, stop_tz).ok_or(
            exec_datafusion_err!(
                "Cannot generate timestamp for stop: {}: {:?}",
                stop,
                stop_tz
            ),
        )?;

        let mut current = start;
        let mut current_dt = as_datetime_with_timezone::<TSNT>(current, start_tz).ok_or(
            exec_datafusion_err!(
                "Cannot generate timestamp for start: {}: {:?}",
                current,
                start_tz
            ),
        )?;

        let values = from_fn(|| {
            if (include_upper_bound
                && ((neg && current_dt < stop_dt) || (!neg && current_dt > stop_dt)))
                || (!include_upper_bound
                    && ((neg && current_dt <= stop_dt)
                        || (!neg && current_dt >= stop_dt)))
            {
                return None;
            }

            let prev_current = current;

            if let Some(ts) = TSNT::add_month_day_nano(current, step, start_tz) {
                current = ts;
                current_dt = as_datetime_with_timezone::<TSNT>(current, start_tz)?;

                Some(Some(prev_current))
            } else {
                // we failed to parse the timestamp here so terminate the series
                None
            }
        });

        list_builder.append_value(values);
    }

    let arr = Arc::new(list_builder.finish());

    Ok(arr)
}