Method gen_range_inner

datafusion/functions-nested/src/range.rs:298–350

Generates an array of integers from start to stop with a given step. This function takes 1 to 3 ArrayRefs as arguments, representing start, stop, and step values. It returns a `Result ` representing the resulting ListArray after the operation. # Arguments `args` - An array of 1 to 3 ArrayRefs representing start, stop, and step (step value can not be zero) values. # Examples gen_range

(&self, args: &[ArrayRef])

    /// gen_range(1, 4) => [1, 2, 3]
    /// gen_range(1, 7, 2) => [1, 3, 5]
    fn gen_range_inner(&self, args: &[ArrayRef]) -> Result<ArrayRef> {
        let (start_array, stop_array, step_array) = match args {
            [stop_array] => (None, as_int64_array(stop_array)?, None),
            [start_array, stop_array] => (
                Some(as_int64_array(start_array)?),
                as_int64_array(stop_array)?,
                None,
            ),
            [start_array, stop_array, step_array] => (
                Some(as_int64_array(start_array)?),
                as_int64_array(stop_array)?,
                Some(as_int64_array(step_array)?),
            ),
            _ => return internal_err!("{} expects 1 to 3 arguments", self.name()),
        };

        let mut values = vec![];
        let mut offsets = vec![0];
        let mut valid = NullBufferBuilder::new(stop_array.len());
        for (idx, stop) in stop_array.iter().enumerate() {
            match retrieve_range_args(start_array, stop, step_array, idx) {
                Some((_, _, 0)) => {
                    return exec_err!(
                        "step can't be 0 for function {}(start [, stop, step])",
                        self.name()
                    );
                }
                Some((start, stop, step)) => {
                    generate_range_values(
                        start,
                        stop,
                        step,
                        self.include_upper_bound,
                        &mut values,
                    )?;
                    offsets.push(values.len() as i32);
                    valid.append_non_null();
                }
                // If any of the arguments is NULL, append a NULL value to the result.
                None => {
                    offsets.push(values.len() as i32);
                    valid.append_null();
                }
            };
        }
        let arr = Arc::new(ListArray::try_new(
            Arc::new(Field::new_list_field(DataType::Int64, true)),
            OffsetBuffer::new(offsets.into()),
            Arc::new(Int64Array::from(values)),
            valid.finish(),
        )?);
        Ok(arr)
    }

    fn gen_range_date(&self, args: &[ArrayRef]) -> Result<ArrayRef> {
        let [start, stop, step] = take_function_args(self.name(), args)?;