fix list

Signed-off-by: jayzhan211 <jayzhan211@gmail.com>
apache · Nov 17, 2023 · d17a682 · d17a682
1 parent 6b56441
commit d17a682
Showing 1 changed file with 34 additions and 4 deletions.
diff --git a/datafusion/common/src/scalar.rs b/datafusion/common/src/scalar.rs
@@ -311,10 +311,40 @@ impl PartialOrd for ScalarValue {
             (FixedSizeBinary(_, _), _) => None,
             (LargeBinary(v1), LargeBinary(v2)) => v1.partial_cmp(v2),
             (LargeBinary(_), _) => None,
-            (List(arr1), List(arr2)) | (FixedSizeList(arr1), FixedSizeList(arr2)) => {
+            (List(arr1), List(arr2)) => {
                 if arr1.data_type() == arr2.data_type() {
-                    let list_arr1 = as_list_array(arr1).unwrap();
-                    let list_arr2 = as_list_array(arr2).unwrap();
+                    let list_arr1 = arr1.as_list::<i32>();
+                    let list_arr2 = arr2.as_list::<i32>();
+                    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);
+                            }
+                        }
+                    }
+                    Some(Ordering::Equal)
+                } else {
+                    None
+                }
+            }
+            (FixedSizeList(arr1), FixedSizeList(arr2)) => {
+                if arr1.data_type() == arr2.data_type() {
+                    let list_arr1 = arr1.as_fixed_size_list();
+                    let list_arr2 = arr2.as_fixed_size_list();
                     if list_arr1.len() != list_arr2.len() {
                         return None;
                     }
@@ -3173,7 +3203,7 @@ mod tests {
         let actual_arr = sv
             .to_array_of_size(2)
             .expect("Failed to convert to array of size");
-        let actual_list_arr = as_list_array(&actual_arr);
+        let actual_list_arr = as_list_array(&actual_arr).unwrap();
 
         let arr = ListArray::from_iter_primitive::<Int32Type, _, _>(vec![
             Some(vec![Some(1), None, Some(2)]),