fix: wrong result of range function (#8313)

* fix: wrong result of range function * fix test * add ut * add ut * nit * nit --------- Co-authored-by: zhongjingxiong <[email protected]>
apache · Nov 27, 2023 · 9dfb224 · 9dfb224
1 parent 6d5a350
commit 9dfb224
Show file tree

Hide file tree

Showing 2 changed files with 73 additions and 5 deletions.
diff --git a/datafusion/physical-expr/src/array_expressions.rs b/datafusion/physical-expr/src/array_expressions.rs
@@ -746,8 +746,14 @@ pub fn gen_range(args: &[ArrayRef]) -> Result<ArrayRef> {
         if step == 0 {
             return exec_err!("step can't be 0 for function range(start [, stop, step]");
         }
-        let value = (start..stop).step_by(step as usize);
-        values.extend(value);
+        if step < 0 {
+            // Decreasing range
+            values.extend((stop + 1..start + 1).rev().step_by((-step) as usize));
+        } else {
+            // Increasing range
+            values.extend((start..stop).step_by(step as usize));
+        }
+
         offsets.push(values.len() as i32);
     }
     let arr = Arc::new(ListArray::try_new(
@@ -2514,6 +2520,67 @@ mod tests {
             .is_null(0));
     }
 
+    #[test]
+    fn test_array_range() {
+        // range(1, 5, 1) = [1, 2, 3, 4]
+        let args1 = Arc::new(Int64Array::from(vec![Some(1)])) as ArrayRef;
+        let args2 = Arc::new(Int64Array::from(vec![Some(5)])) as ArrayRef;
+        let args3 = Arc::new(Int64Array::from(vec![Some(1)])) as ArrayRef;
+        let arr = gen_range(&[args1, args2, args3]).unwrap();
+
+        let result = as_list_array(&arr).expect("failed to initialize function range");
+        assert_eq!(
+            &[1, 2, 3, 4],
+            result
+                .value(0)
+                .as_any()
+                .downcast_ref::<Int64Array>()
+                .unwrap()
+                .values()
+        );
+
+        // range(1, -5, -1) = [1, 0, -1, -2, -3, -4]
+        let args1 = Arc::new(Int64Array::from(vec![Some(1)])) as ArrayRef;
+        let args2 = Arc::new(Int64Array::from(vec![Some(-5)])) as ArrayRef;
+        let args3 = Arc::new(Int64Array::from(vec![Some(-1)])) as ArrayRef;
+        let arr = gen_range(&[args1, args2, args3]).unwrap();
+
+        let result = as_list_array(&arr).expect("failed to initialize function range");
+        assert_eq!(
+            &[1, 0, -1, -2, -3, -4],
+            result
+                .value(0)
+                .as_any()
+                .downcast_ref::<Int64Array>()
+                .unwrap()
+                .values()
+        );
+
+        // range(1, 5, -1) = []
+        let args1 = Arc::new(Int64Array::from(vec![Some(1)])) as ArrayRef;
+        let args2 = Arc::new(Int64Array::from(vec![Some(5)])) as ArrayRef;
+        let args3 = Arc::new(Int64Array::from(vec![Some(-1)])) as ArrayRef;
+        let arr = gen_range(&[args1, args2, args3]).unwrap();
+
+        let result = as_list_array(&arr).expect("failed to initialize function range");
+        assert_eq!(
+            &[],
+            result
+                .value(0)
+                .as_any()
+                .downcast_ref::<Int64Array>()
+                .unwrap()
+                .values()
+        );
+
+        // range(1, 5, 0) = []
+        let args1 = Arc::new(Int64Array::from(vec![Some(1)])) as ArrayRef;
+        let args2 = Arc::new(Int64Array::from(vec![Some(5)])) as ArrayRef;
+        let args3 = Arc::new(Int64Array::from(vec![Some(0)])) as ArrayRef;
+        let is_err = gen_range(&[args1, args2, args3]).is_err();
+        assert!(is_err)
+    }
+
     #[test]
     fn test_nested_array_slice() {
         // array_slice([[1, 2, 3, 4], [5, 6, 7, 8]], 1, 1) = [[1, 2, 3, 4]]

diff --git a/datafusion/sqllogictest/test_files/array.slt b/datafusion/sqllogictest/test_files/array.slt
@@ -2744,15 +2744,16 @@ from arrays_range;
 [0, 1, 2, 3, 4, 5, 6, 7, 8, 9] [3, 4, 5, 6, 7, 8, 9] [3, 5, 7, 9]
 [0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12] [4, 5, 6, 7, 8, 9, 10, 11, 12] [4, 7, 10]
 
-query ?????
+query ??????
 select range(5),
        range(2, 5),
        range(2, 10, 3),
        range(1, 5, -1),
-       range(1, -5, 1)
+       range(1, -5, 1),
+       range(1, -5, -1)
 ;
 ----
-[0, 1, 2, 3, 4] [2, 3, 4] [2, 5, 8] [1] []
+[0, 1, 2, 3, 4] [2, 3, 4] [2, 5, 8] [] [] [1, 0, -1, -2, -3, -4]
 
 query ???
 select generate_series(5),