split into two functions

Signed-off-by: jayzhan211 <jayzhan211@gmail.com>
apache · Nov 18, 2023 · 754f39c · 754f39c
1 parent 629e3b7
commit 754f39c
Show file tree

Hide file tree

Showing 2 changed files with 48 additions and 22 deletions.
diff --git a/datafusion/common/src/scalar.rs b/datafusion/common/src/scalar.rs
@@ -30,7 +30,7 @@ use crate::cast::{
 };
 use crate::error::{DataFusionError, Result, _internal_err, _not_impl_err};
 use crate::hash_utils::create_hashes;
-use crate::utils::{array_into_children_array_vec, array_into_list_array};
+use crate::utils::array_into_list_array;
 use arrow::buffer::{NullBuffer, OffsetBuffer};
 use arrow::compute::kernels::numeric::*;
 use arrow::datatypes::{i256, Fields, SchemaBuilder};
@@ -312,29 +312,65 @@ impl PartialOrd for ScalarValue {
             (FixedSizeBinary(_, _), _) => None,
             (LargeBinary(v1), LargeBinary(v2)) => v1.partial_cmp(v2),
             (LargeBinary(_), _) => None,
-            (List(list_arr1), List(list_arr2))
-            | (FixedSizeList(list_arr1), FixedSizeList(list_arr2)) => {
-                if list_arr1.data_type() == list_arr2.data_type() {
-                    if list_arr1.len() != list_arr2.len() {
+            (List(arr1), List(arr2)) => {
+                if arr1.data_type() == arr2.data_type() {
+                    if arr1.len() != arr2.len() {
                         return None;
                     }
 
                     // ScalarValue::List / ScalarValue::FixedSizeList should have only one list.
+                    assert_eq!(arr1.len(), 1);
+                    assert_eq!(arr2.len(), 1);
+
+                    let list_arr1 = arr1.as_list::<i32>();
+                    let list_arr2 = arr2.as_list::<i32>();
+
+                    // Single child data
                     assert_eq!(list_arr1.len(), 1);
                     assert_eq!(list_arr2.len(), 1);
 
-                    let arr1 = array_into_children_array_vec(list_arr1);
-                    let arr2 = array_into_children_array_vec(list_arr2);
+                    let arr1 = list_arr1.value(0);
+                    let arr2 = list_arr2.value(0);
 
-                    // Single child data
+                    let lt_res = arrow::compute::kernels::cmp::lt(&arr1, &arr2).ok()?;
+                    let eq_res = arrow::compute::kernels::cmp::eq(&arr1, &arr2).ok()?;
+
+                    for j in 0..lt_res.len() {
+                        if lt_res.is_valid(j) && lt_res.value(j) {
+                            return Some(Ordering::Less);
+                        }
+                        if eq_res.is_valid(j) && !eq_res.value(j) {
+                            return Some(Ordering::Greater);
+                        }
+                    }
+
+                    Some(Ordering::Equal)
+                } else {
+                    None
+                }
+            }
+            (FixedSizeList(arr1), FixedSizeList(arr2)) => {
+                if arr1.data_type() == arr2.data_type() {
+                    if arr1.len() != arr2.len() {
+                        return None;
+                    }
+
+                    // ScalarValue::List / ScalarValue::FixedSizeList should have only one list.
                     assert_eq!(arr1.len(), 1);
                     assert_eq!(arr2.len(), 1);
 
-                    let arr1 = &arr1[0];
-                    let arr2 = &arr2[0];
+                    let list_arr1 = arr1.as_fixed_size_list();
+                    let list_arr2 = arr2.as_fixed_size_list();
+
+                    // Single child data
+                    assert_eq!(list_arr1.len(), 1);
+                    assert_eq!(list_arr2.len(), 1);
+
+                    let arr1 = list_arr1.value(0);
+                    let arr2 = list_arr2.value(0);
 
-                    let lt_res = arrow::compute::kernels::cmp::lt(arr1, arr2).ok()?;
-                    let eq_res = arrow::compute::kernels::cmp::eq(arr1, arr2).ok()?;
+                    let lt_res = arrow::compute::kernels::cmp::lt(&arr1, &arr2).ok()?;
+                    let eq_res = arrow::compute::kernels::cmp::eq(&arr1, &arr2).ok()?;
 
                     for j in 0..lt_res.len() {
                         if lt_res.is_valid(j) && lt_res.value(j) {

diff --git a/datafusion/common/src/utils.rs b/datafusion/common/src/utils.rs
@@ -390,16 +390,6 @@ pub fn arrays_into_list_array(
     ))
 }
 
-/// Get the child arrays from a `ArrayRef`.
-pub fn array_into_children_array_vec(list_arr: &ArrayRef) -> Vec<ArrayRef> {
-    let data = list_arr.to_data();
-    let children = data.child_data();
-    children
-        .iter()
-        .map(|x| arrow_array::make_array(x.to_owned()))
-        .collect::<Vec<_>>()
-}
-
 /// An extension trait for smart pointers. Provides an interface to get a
 /// raw pointer to the data (with metadata stripped away).
 ///