Apply refactoring to fix array.slt test failures

Author: erenavsarogullari

datafusion/functions-nested/src/range.rs

@@ -297,7 +297,7 @@ impl Range {
     ///
     /// # Arguments
     ///
-    /// * `args` - An array of 1 to 3 ArrayRefs representing start, stop, and step(step value can not be zero) values.
+    /// * `args` - An array of 1 to 3 ArrayRefs representing start, stop, and step(step value can not be zero.) values.
     ///
     /// # Examples
     ///
@@ -338,23 +338,10 @@ impl Range {
                         usize::try_from(step.unsigned_abs()).map_err(|_| {
                             not_impl_datafusion_err!("step {} can't fit into usize", step)
                         })?;
-                    if start < stop {
-                        values.extend(
-                            gen_range_iter(
-                                start,
-                                stop,
-                                step < 0,
-                                self.include_upper_bound,
-                            )
+                    values.extend(
+                        gen_range_iter(start, stop, step < 0, self.include_upper_bound)
                             .step_by(step_abs),
-                        )
-                    } else {
-                        values.extend(
-                            gen_range_iter(start, stop, true, self.include_upper_bound)
-                                .step_by(step_abs),
-                        )
-                    };
-
+                    );
                     offsets.push(values.len() as i32);
                     valid.append_non_null();
                 }

datafusion/spark/src/function/array/sequence.rs

@@ -16,7 +16,9 @@
 // under the License.
 
 use crate::function::functions_nested_utils::make_scalar_function;
+use arrow::array::{Array, Int64Builder};
 use arrow::datatypes::{DataType, Field, FieldRef, IntervalMonthDayNano};
+use datafusion_common::cast::as_int64_array;
 use datafusion_common::internal_err;
 use datafusion_common::{DataFusionError, Result, ScalarValue, exec_err};
 use datafusion_expr::{
@@ -146,12 +148,21 @@ impl ScalarUDFImpl for SparkSequence {
         if args.iter().any(|arg| arg.data_type().is_null()) {
             return Ok(ColumnarValue::Scalar(ScalarValue::Null));
         }
+
         match args[0].data_type() {
-            DataType::Int64 => make_scalar_function(|args| {
-                Range::generate_series().gen_range_inner(args)
-            })(args),
+            DataType::Int64 => {
+                validate_int64_sequence_step(args)?;
+                let optional_new_args = add_step_argument_if_not_exists(args)?;
+                let new_args = match optional_new_args {
+                    Some(new_args) => &new_args.to_owned(),
+                    None => args,
+                };
+                make_scalar_function(|args| {
+                    Range::generate_series().gen_range_inner(args)
+                })(new_args)
+            }
             DataType::Date32 | DataType::Date64 => {
-                let optional_new_args = add_interval_if_not_exists(args);
+                let optional_new_args = add_interval_argument_if_not_exists(args);
                 let new_args = match optional_new_args {
                     Some(new_args) => &new_args.to_owned(),
                     None => args,
@@ -161,7 +172,7 @@ impl ScalarUDFImpl for SparkSequence {
                 )
             }
             DataType::Timestamp(_, _) => {
-                let optional_new_args = add_interval_if_not_exists(args);
+                let optional_new_args = add_interval_argument_if_not_exists(args);
                 let new_args = match optional_new_args {
                     Some(new_args) => &new_args.to_owned(),
                     None => args,
@@ -180,6 +191,62 @@ impl ScalarUDFImpl for SparkSequence {
     }
 }
 
+/// Validates explicit `step` for 3-argument integer `sequence` (Spark semantics).
+fn validate_int64_sequence_step(args: &[ColumnarValue]) -> Result<()> {
+    if args.len() != 3 {
+        return Ok(());
+    }
+    let arrays = ColumnarValue::values_to_arrays(args)?;
+    let start = as_int64_array(&arrays[0])?;
+    let stop = as_int64_array(&arrays[1])?;
+    let step = as_int64_array(&arrays[2])?;
+    for i in 0..start.len() {
+        if start.is_null(i) || stop.is_null(i) || step.is_null(i) {
+            continue;
+        }
+        let s = start.value(i);
+        let e = stop.value(i);
+        let st = step.value(i);
+        if st == 0 {
+            return exec_err!("Step cannot be 0 for sequence");
+        }
+        if s < e && st <= 0 {
+            return exec_err!("When start < stop, step must be positive");
+        }
+        if s > e && st >= 0 {
+            return exec_err!("When start > stop, step must be negative");
+        }
+    }
+    Ok(())
+}
+
+/// When only start and stop are given, Spark picks step `1` if start ≤ stop and `-1` if start > stop.
+fn add_step_argument_if_not_exists(args: &[ColumnarValue]) -> Result<Option<Vec<ColumnarValue>>> {
+    if args.len() != 2 {
+        return Ok(None);
+    }
+    let arrays = ColumnarValue::values_to_arrays(args)?;
+    let start = as_int64_array(&arrays[0])?;
+    let stop = as_int64_array(&arrays[1])?;
+    let len = start.len();
+    let mut step = Int64Builder::with_capacity(len);
+    for i in 0..len {
+        if start.is_null(i) || stop.is_null(i) {
+            step.append_null();
+        } else if start.value(i) > stop.value(i) {
+            step.append_value(-1);
+        } else {
+            step.append_value(1);
+        }
+    }
+    let step = step.finish();
+    Ok(Some(vec![
+        args[0].clone(),
+        args[1].clone(),
+        ColumnarValue::Array(Arc::new(step)),
+    ]))
+}
+
 fn check_type(
     data_type: DataType,
     param_name: &str,
@@ -242,7 +309,7 @@ fn check_interval_type_by_first_type(
     }
 }
 
-fn add_interval_if_not_exists(args: &[ColumnarValue]) -> Option<Vec<ColumnarValue>> {
+fn add_interval_argument_if_not_exists(args: &[ColumnarValue]) -> Option<Vec<ColumnarValue>> {
     if args.len() == 2 {
         let mut new_args = args.to_owned();
         new_args.push(ColumnarValue::Scalar(ScalarValue::IntervalMonthDayNano(

datafusion/sqllogictest/test_files/spark/array/sequence.slt

@@ -39,11 +39,6 @@ SELECT sequence(0, 5, 1);
 ----
 [0, 1, 2, 3, 4, 5]
 
-query ?
-SELECT sequence(5, 0, 1);
-----
-[5, 4, 3, 2, 1, 0]
-
 query ?
 SELECT sequence(-3::int, 3::int);
 ----
@@ -290,3 +285,39 @@ SELECT sequence(1);
 DataFusion error: Error during planning: Execution error: Function 'sequence' user-defined coercion failed with: Execution error: num of input parameters should be 2 or 3. No function matches the given name and argument types 'sequence(Int64)'. You might need to add explicit type casts.
 	Candidate functions:
 	sequence(UserDefined)
+
+
+query ?
+SELECT sequence(2, 2);
+----
+[2]
+
+query ?
+SELECT sequence(-2, -2);
+----
+[-2]
+
+query error
+SELECT sequence(5, 0, 1);
+----
+DataFusion error: Execution error: When start > stop, step must be negative
+
+
+query error
+SELECT sequence(0, 5, -1);
+----
+DataFusion error: Execution error: When start < stop, step must be positive
+
+
+query error
+SELECT sequence(1, 5, 0);
+----
+DataFusion error: Execution error: Step cannot be 0 for sequence
+
+
+query error
+SELECT sequence(5, 1, 0);
+----
+DataFusion error: Execution error: Step cannot be 0 for sequence
+
+