diff --git a/datafusion/physical-expr/src/array_expressions.rs b/datafusion/physical-expr/src/array_expressions.rs index 8968bcf2ea4e..6b7bef8e6a36 100644 --- 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 { 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::() + .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::() + .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::() + .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 index d33555509e6c..db657ff22bd5 100644 --- 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),