fixsizelist cmp fix

Signed-off-by: jayzhan211 <jayzhan211@gmail.com>

Author: jayzhan211

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_list_array;
+use crate::utils::{array_into_children_array_vec, array_into_list_array};
 use arrow::buffer::{NullBuffer, OffsetBuffer};
 use arrow::compute::kernels::numeric::*;
 use arrow::datatypes::{i256, Fields, SchemaBuilder};
@@ -312,31 +312,39 @@ impl PartialOrd for ScalarValue {
             (FixedSizeBinary(_, _), _) => None,
             (LargeBinary(v1), LargeBinary(v2)) => v1.partial_cmp(v2),
             (LargeBinary(_), _) => None,
-            (List(arr1), List(arr2)) | (FixedSizeList(arr1), FixedSizeList(arr2)) => {
-                if arr1.data_type() == arr2.data_type() {
-                    let list_arr1 = as_list_array(arr1);
-                    let list_arr2 = as_list_array(arr2);
+            (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() {
                         return None;
                     }
-                    for i in 0..list_arr1.len() {
-                        let arr1 = list_arr1.value(i);
-                        let arr2 = list_arr2.value(i);
-
-                        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);
-                            }
+
+                    // ScalarValue::List / ScalarValue::FixedSizeList should have only one list.
+                    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);
+
+                    // Single child data
+                    assert_eq!(arr1.len(), 1);
+                    assert_eq!(arr2.len(), 1);
+
+                    let arr1 = &arr1[0];
+                    let arr2 = &arr2[0];
+
+                    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
@@ -3060,6 +3068,7 @@ impl ScalarType<i64> for TimestampNanosecondType {
 }
 
 #[cfg(test)]
+#[cfg(feature = "parquet")]
 mod tests {
     use std::cmp::Ordering;
     use std::sync::Arc;
@@ -3534,24 +3543,6 @@ mod tests {
                 ])]),
             ));
         assert_eq!(a.partial_cmp(&b), Some(Ordering::Less));
-
-        let a =
-            ScalarValue::List(Arc::new(
-                ListArray::from_iter_primitive::<Int64Type, _, _>(vec![
-                    Some(vec![Some(10), Some(2), Some(3)]),
-                    None,
-                    Some(vec![Some(10), Some(2), Some(3)]),
-                ]),
-            ));
-        let b =
-            ScalarValue::List(Arc::new(
-                ListArray::from_iter_primitive::<Int64Type, _, _>(vec![
-                    Some(vec![Some(10), Some(2), Some(3)]),
-                    None,
-                    Some(vec![Some(10), Some(2), Some(3)]),
-                ]),
-            ));
-        assert_eq!(a.partial_cmp(&b), Some(Ordering::Equal));
     }
 
     #[test]

datafusion/common/src/utils.rs

@@ -390,6 +390,16 @@ pub fn arrays_into_list_array(
     ))
 }
 
+/// Get the child arrays from a `ListArray`.
+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).
 ///