Skip to content

Commit 8b6d60f

Browse files
committed
cleanup three type of list
Signed-off-by: jayzhan211 <jayzhan211@gmail.com>
1 parent 2ee3ed5 commit 8b6d60f

File tree

1 file changed

+31
-97
lines changed

1 file changed

+31
-97
lines changed

datafusion/common/src/scalar.rs

Lines changed: 31 additions & 97 deletions
Original file line numberDiff line numberDiff line change
@@ -317,113 +317,47 @@ impl PartialOrd for ScalarValue {
317317
(FixedSizeBinary(_, _), _) => None,
318318
(LargeBinary(v1), LargeBinary(v2)) => v1.partial_cmp(v2),
319319
(LargeBinary(_), _) => None,
320-
(List(arr1), List(arr2)) => {
321-
if arr1.data_type() == arr2.data_type() {
322-
if arr1.len() != arr2.len() {
323-
return None;
324-
}
325-
326-
// ScalarValue::List / ScalarValue::FixedSizeList should have only one list.
327-
assert_eq!(arr1.len(), 1);
328-
assert_eq!(arr2.len(), 1);
329-
330-
let list_arr1 = arr1.as_list::<i32>();
331-
let list_arr2 = arr2.as_list::<i32>();
332-
333-
// Single child data
334-
assert_eq!(list_arr1.len(), 1);
335-
assert_eq!(list_arr2.len(), 1);
336-
337-
let arr1 = list_arr1.value(0);
338-
let arr2 = list_arr2.value(0);
339-
340-
let lt_res = arrow::compute::kernels::cmp::lt(&arr1, &arr2).ok()?;
341-
let eq_res = arrow::compute::kernels::cmp::eq(&arr1, &arr2).ok()?;
320+
(List(arr1), List(arr2))
321+
| (FixedSizeList(arr1), FixedSizeList(arr2))
322+
| (LargeList(arr1), LargeList(arr2)) => {
323+
// ScalarValue::List / ScalarValue::FixedSizeList / ScalarValue::LargeList are ensure to have length 1
324+
assert_eq!(arr1.len(), 1);
325+
assert_eq!(arr2.len(), 1);
326+
327+
if arr1.data_type() != arr2.data_type() {
328+
return None;
329+
}
342330

343-
for j in 0..lt_res.len() {
344-
if lt_res.is_valid(j) && lt_res.value(j) {
345-
return Some(Ordering::Less);
346-
}
347-
if eq_res.is_valid(j) && !eq_res.value(j) {
348-
return Some(Ordering::Greater);
349-
}
331+
fn first_array_for_list(arr: &ArrayRef) -> ArrayRef {
332+
if let Some(arr) = arr.as_list_opt::<i32>() {
333+
arr.value(0)
334+
} else if let Some(arr) = arr.as_list_opt::<i64>() {
335+
arr.value(0)
336+
} else if let Some(arr) = arr.as_fixed_size_list_opt() {
337+
arr.value(0)
338+
} else {
339+
unreachable!("Since only List / LargeList / FixedSizeList are supported, this should never happen")
350340
}
351-
352-
Some(Ordering::Equal)
353-
} else {
354-
None
355341
}
356-
}
357-
(FixedSizeList(arr1), FixedSizeList(arr2)) => {
358-
if arr1.data_type() == arr2.data_type() {
359-
if arr1.len() != arr2.len() {
360-
return None;
361-
}
362-
363-
// ScalarValue::List / ScalarValue::FixedSizeList should have only one list.
364-
assert_eq!(arr1.len(), 1);
365-
assert_eq!(arr2.len(), 1);
366342

367-
let list_arr1 = arr1.as_fixed_size_list();
368-
let list_arr2 = arr2.as_fixed_size_list();
343+
let arr1 = first_array_for_list(arr1);
344+
let arr2 = first_array_for_list(arr2);
369345

370-
// Single child data
371-
assert_eq!(list_arr1.len(), 1);
372-
assert_eq!(list_arr2.len(), 1);
346+
let lt_res = arrow::compute::kernels::cmp::lt(&arr1, &arr2).ok()?;
347+
let eq_res = arrow::compute::kernels::cmp::eq(&arr1, &arr2).ok()?;
373348

374-
let arr1 = list_arr1.value(0);
375-
let arr2 = list_arr2.value(0);
376-
377-
let lt_res = arrow::compute::kernels::cmp::lt(&arr1, &arr2).ok()?;
378-
let eq_res = arrow::compute::kernels::cmp::eq(&arr1, &arr2).ok()?;
379-
380-
for j in 0..lt_res.len() {
381-
if lt_res.is_valid(j) && lt_res.value(j) {
382-
return Some(Ordering::Less);
383-
}
384-
if eq_res.is_valid(j) && !eq_res.value(j) {
385-
return Some(Ordering::Greater);
386-
}
387-
}
388-
389-
Some(Ordering::Equal)
390-
} else {
391-
None
392-
}
393-
}
394-
(LargeList(arr1), LargeList(arr2)) => {
395-
if arr1.data_type() == arr2.data_type() {
396-
let list_arr1 = as_large_list_array(arr1);
397-
let list_arr2 = as_large_list_array(arr2);
398-
if list_arr1.len() != list_arr2.len() {
399-
return None;
349+
for j in 0..lt_res.len() {
350+
if lt_res.is_valid(j) && lt_res.value(j) {
351+
return Some(Ordering::Less);
400352
}
401-
for i in 0..list_arr1.len() {
402-
let arr1 = list_arr1.value(i);
403-
let arr2 = list_arr2.value(i);
404-
405-
let lt_res =
406-
arrow::compute::kernels::cmp::lt(&arr1, &arr2).ok()?;
407-
let eq_res =
408-
arrow::compute::kernels::cmp::eq(&arr1, &arr2).ok()?;
409-
410-
for j in 0..lt_res.len() {
411-
if lt_res.is_valid(j) && lt_res.value(j) {
412-
return Some(Ordering::Less);
413-
}
414-
if eq_res.is_valid(j) && !eq_res.value(j) {
415-
return Some(Ordering::Greater);
416-
}
417-
}
353+
if eq_res.is_valid(j) && !eq_res.value(j) {
354+
return Some(Ordering::Greater);
418355
}
419-
Some(Ordering::Equal)
420-
} else {
421-
None
422356
}
357+
358+
Some(Ordering::Equal)
423359
}
424-
(List(_), _) => None,
425-
(LargeList(_), _) => None,
426-
(FixedSizeList(_), _) => None,
360+
(List(_), _) | (LargeList(_), _) | (FixedSizeList(_), _) => None,
427361
(Date32(v1), Date32(v2)) => v1.partial_cmp(v2),
428362
(Date32(_), _) => None,
429363
(Date64(v1), Date64(v2)) => v1.partial_cmp(v2),

0 commit comments

Comments
 (0)