Skip to content

Commit 9dfb224

Browse files
smallzhongfengzhongjingxiong
andauthored
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 <zhongjingxiong@bytedance.com>
1 parent 6d5a350 commit 9dfb224

File tree

2 files changed

+73
-5
lines changed

2 files changed

+73
-5
lines changed

datafusion/physical-expr/src/array_expressions.rs

Lines changed: 69 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -746,8 +746,14 @@ pub fn gen_range(args: &[ArrayRef]) -> Result<ArrayRef> {
746746
if step == 0 {
747747
return exec_err!("step can't be 0 for function range(start [, stop, step]");
748748
}
749-
let value = (start..stop).step_by(step as usize);
750-
values.extend(value);
749+
if step < 0 {
750+
// Decreasing range
751+
values.extend((stop + 1..start + 1).rev().step_by((-step) as usize));
752+
} else {
753+
// Increasing range
754+
values.extend((start..stop).step_by(step as usize));
755+
}
756+
751757
offsets.push(values.len() as i32);
752758
}
753759
let arr = Arc::new(ListArray::try_new(
@@ -2514,6 +2520,67 @@ mod tests {
25142520
.is_null(0));
25152521
}
25162522

2523+
#[test]
2524+
fn test_array_range() {
2525+
// range(1, 5, 1) = [1, 2, 3, 4]
2526+
let args1 = Arc::new(Int64Array::from(vec![Some(1)])) as ArrayRef;
2527+
let args2 = Arc::new(Int64Array::from(vec![Some(5)])) as ArrayRef;
2528+
let args3 = Arc::new(Int64Array::from(vec![Some(1)])) as ArrayRef;
2529+
let arr = gen_range(&[args1, args2, args3]).unwrap();
2530+
2531+
let result = as_list_array(&arr).expect("failed to initialize function range");
2532+
assert_eq!(
2533+
&[1, 2, 3, 4],
2534+
result
2535+
.value(0)
2536+
.as_any()
2537+
.downcast_ref::<Int64Array>()
2538+
.unwrap()
2539+
.values()
2540+
);
2541+
2542+
// range(1, -5, -1) = [1, 0, -1, -2, -3, -4]
2543+
let args1 = Arc::new(Int64Array::from(vec![Some(1)])) as ArrayRef;
2544+
let args2 = Arc::new(Int64Array::from(vec![Some(-5)])) as ArrayRef;
2545+
let args3 = Arc::new(Int64Array::from(vec![Some(-1)])) as ArrayRef;
2546+
let arr = gen_range(&[args1, args2, args3]).unwrap();
2547+
2548+
let result = as_list_array(&arr).expect("failed to initialize function range");
2549+
assert_eq!(
2550+
&[1, 0, -1, -2, -3, -4],
2551+
result
2552+
.value(0)
2553+
.as_any()
2554+
.downcast_ref::<Int64Array>()
2555+
.unwrap()
2556+
.values()
2557+
);
2558+
2559+
// range(1, 5, -1) = []
2560+
let args1 = Arc::new(Int64Array::from(vec![Some(1)])) as ArrayRef;
2561+
let args2 = Arc::new(Int64Array::from(vec![Some(5)])) as ArrayRef;
2562+
let args3 = Arc::new(Int64Array::from(vec![Some(-1)])) as ArrayRef;
2563+
let arr = gen_range(&[args1, args2, args3]).unwrap();
2564+
2565+
let result = as_list_array(&arr).expect("failed to initialize function range");
2566+
assert_eq!(
2567+
&[],
2568+
result
2569+
.value(0)
2570+
.as_any()
2571+
.downcast_ref::<Int64Array>()
2572+
.unwrap()
2573+
.values()
2574+
);
2575+
2576+
// range(1, 5, 0) = []
2577+
let args1 = Arc::new(Int64Array::from(vec![Some(1)])) as ArrayRef;
2578+
let args2 = Arc::new(Int64Array::from(vec![Some(5)])) as ArrayRef;
2579+
let args3 = Arc::new(Int64Array::from(vec![Some(0)])) as ArrayRef;
2580+
let is_err = gen_range(&[args1, args2, args3]).is_err();
2581+
assert!(is_err)
2582+
}
2583+
25172584
#[test]
25182585
fn test_nested_array_slice() {
25192586
// array_slice([[1, 2, 3, 4], [5, 6, 7, 8]], 1, 1) = [[1, 2, 3, 4]]

datafusion/sqllogictest/test_files/array.slt

Lines changed: 4 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -2744,15 +2744,16 @@ from arrays_range;
27442744
[0, 1, 2, 3, 4, 5, 6, 7, 8, 9] [3, 4, 5, 6, 7, 8, 9] [3, 5, 7, 9]
27452745
[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]
27462746

2747-
query ?????
2747+
query ??????
27482748
select range(5),
27492749
range(2, 5),
27502750
range(2, 10, 3),
27512751
range(1, 5, -1),
2752-
range(1, -5, 1)
2752+
range(1, -5, 1),
2753+
range(1, -5, -1)
27532754
;
27542755
----
2755-
[0, 1, 2, 3, 4] [2, 3, 4] [2, 5, 8] [1] []
2756+
[0, 1, 2, 3, 4] [2, 3, 4] [2, 5, 8] [] [] [1, 0, -1, -2, -3, -4]
27562757

27572758
query ???
27582759
select generate_series(5),

0 commit comments

Comments
 (0)