Skip to content

Commit 754f39c

Browse files
committed
split into two functions
Signed-off-by: jayzhan211 <jayzhan211@gmail.com>
1 parent 629e3b7 commit 754f39c

File tree

2 files changed

+48
-22
lines changed

2 files changed

+48
-22
lines changed

datafusion/common/src/scalar.rs

Lines changed: 48 additions & 12 deletions
Original file line numberDiff line numberDiff line change
@@ -30,7 +30,7 @@ use crate::cast::{
3030
};
3131
use crate::error::{DataFusionError, Result, _internal_err, _not_impl_err};
3232
use crate::hash_utils::create_hashes;
33-
use crate::utils::{array_into_children_array_vec, array_into_list_array};
33+
use crate::utils::array_into_list_array;
3434
use arrow::buffer::{NullBuffer, OffsetBuffer};
3535
use arrow::compute::kernels::numeric::*;
3636
use arrow::datatypes::{i256, Fields, SchemaBuilder};
@@ -312,29 +312,65 @@ impl PartialOrd for ScalarValue {
312312
(FixedSizeBinary(_, _), _) => None,
313313
(LargeBinary(v1), LargeBinary(v2)) => v1.partial_cmp(v2),
314314
(LargeBinary(_), _) => None,
315-
(List(list_arr1), List(list_arr2))
316-
| (FixedSizeList(list_arr1), FixedSizeList(list_arr2)) => {
317-
if list_arr1.data_type() == list_arr2.data_type() {
318-
if list_arr1.len() != list_arr2.len() {
315+
(List(arr1), List(arr2)) => {
316+
if arr1.data_type() == arr2.data_type() {
317+
if arr1.len() != arr2.len() {
319318
return None;
320319
}
321320

322321
// ScalarValue::List / ScalarValue::FixedSizeList should have only one list.
322+
assert_eq!(arr1.len(), 1);
323+
assert_eq!(arr2.len(), 1);
324+
325+
let list_arr1 = arr1.as_list::<i32>();
326+
let list_arr2 = arr2.as_list::<i32>();
327+
328+
// Single child data
323329
assert_eq!(list_arr1.len(), 1);
324330
assert_eq!(list_arr2.len(), 1);
325331

326-
let arr1 = array_into_children_array_vec(list_arr1);
327-
let arr2 = array_into_children_array_vec(list_arr2);
332+
let arr1 = list_arr1.value(0);
333+
let arr2 = list_arr2.value(0);
328334

329-
// Single child data
335+
let lt_res = arrow::compute::kernels::cmp::lt(&arr1, &arr2).ok()?;
336+
let eq_res = arrow::compute::kernels::cmp::eq(&arr1, &arr2).ok()?;
337+
338+
for j in 0..lt_res.len() {
339+
if lt_res.is_valid(j) && lt_res.value(j) {
340+
return Some(Ordering::Less);
341+
}
342+
if eq_res.is_valid(j) && !eq_res.value(j) {
343+
return Some(Ordering::Greater);
344+
}
345+
}
346+
347+
Some(Ordering::Equal)
348+
} else {
349+
None
350+
}
351+
}
352+
(FixedSizeList(arr1), FixedSizeList(arr2)) => {
353+
if arr1.data_type() == arr2.data_type() {
354+
if arr1.len() != arr2.len() {
355+
return None;
356+
}
357+
358+
// ScalarValue::List / ScalarValue::FixedSizeList should have only one list.
330359
assert_eq!(arr1.len(), 1);
331360
assert_eq!(arr2.len(), 1);
332361

333-
let arr1 = &arr1[0];
334-
let arr2 = &arr2[0];
362+
let list_arr1 = arr1.as_fixed_size_list();
363+
let list_arr2 = arr2.as_fixed_size_list();
364+
365+
// Single child data
366+
assert_eq!(list_arr1.len(), 1);
367+
assert_eq!(list_arr2.len(), 1);
368+
369+
let arr1 = list_arr1.value(0);
370+
let arr2 = list_arr2.value(0);
335371

336-
let lt_res = arrow::compute::kernels::cmp::lt(arr1, arr2).ok()?;
337-
let eq_res = arrow::compute::kernels::cmp::eq(arr1, arr2).ok()?;
372+
let lt_res = arrow::compute::kernels::cmp::lt(&arr1, &arr2).ok()?;
373+
let eq_res = arrow::compute::kernels::cmp::eq(&arr1, &arr2).ok()?;
338374

339375
for j in 0..lt_res.len() {
340376
if lt_res.is_valid(j) && lt_res.value(j) {

datafusion/common/src/utils.rs

Lines changed: 0 additions & 10 deletions
Original file line numberDiff line numberDiff line change
@@ -390,16 +390,6 @@ pub fn arrays_into_list_array(
390390
))
391391
}
392392

393-
/// Get the child arrays from a `ArrayRef`.
394-
pub fn array_into_children_array_vec(list_arr: &ArrayRef) -> Vec<ArrayRef> {
395-
let data = list_arr.to_data();
396-
let children = data.child_data();
397-
children
398-
.iter()
399-
.map(|x| arrow_array::make_array(x.to_owned()))
400-
.collect::<Vec<_>>()
401-
}
402-
403393
/// An extension trait for smart pointers. Provides an interface to get a
404394
/// raw pointer to the data (with metadata stripped away).
405395
///

0 commit comments

Comments
 (0)