Fix array equal check (#571)

* pin this bug down & fix Signed-off-by: Ruihang Xia <[email protected]> * tidy Signed-off-by: Ruihang Xia <[email protected]> * add some test cases Signed-off-by: Ruihang Xia <[email protected]>
apache · Jul 20, 2021 · bd9f67a · bd9f67a
1 parent 6c82bdc
commit bd9f67a
Show file tree

Hide file tree

Showing 3 changed files with 74 additions and 2 deletions.
diff --git a/arrow/src/array/array_list.rs b/arrow/src/array/array_list.rs
@@ -1090,4 +1090,76 @@ mod tests {
             .build();
         ListArray::from(list_data);
     }
+
+    #[test]
+    fn list_array_equality() {
+        // test scaffold
+        fn do_comparison(
+            lhs_data: Vec<Option<Vec<Option<i32>>>>,
+            rhs_data: Vec<Option<Vec<Option<i32>>>>,
+            should_equal: bool,
+        ) {
+            let lhs = ListArray::from_iter_primitive::<Int32Type, _, _>(lhs_data.clone());
+            let rhs = ListArray::from_iter_primitive::<Int32Type, _, _>(rhs_data.clone());
+            assert_eq!(lhs == rhs, should_equal);
+
+            let lhs = LargeListArray::from_iter_primitive::<Int32Type, _, _>(lhs_data);
+            let rhs = LargeListArray::from_iter_primitive::<Int32Type, _, _>(rhs_data);
+            assert_eq!(lhs == rhs, should_equal);
+        }
+
+        do_comparison(
+            vec![
+                Some(vec![Some(0), Some(1), Some(2)]),
+                None,
+                Some(vec![Some(3), None, Some(5)]),
+                Some(vec![Some(6), Some(7)]),
+            ],
+            vec![
+                Some(vec![Some(0), Some(1), Some(2)]),
+                None,
+                Some(vec![Some(3), None, Some(5)]),
+                Some(vec![Some(6), Some(7)]),
+            ],
+            true,
+        );
+
+        do_comparison(
+            vec![
+                None,
+                None,
+                Some(vec![Some(3), None, Some(5)]),
+                Some(vec![Some(6), Some(7)]),
+            ],
+            vec![
+                Some(vec![Some(0), Some(1), Some(2)]),
+                None,
+                Some(vec![Some(3), None, Some(5)]),
+                Some(vec![Some(6), Some(7)]),
+            ],
+            false,
+        );
+
+        do_comparison(
+            vec![
+                None,
+                None,
+                Some(vec![Some(3), None, Some(5)]),
+                Some(vec![Some(6), Some(7)]),
+            ],
+            vec![
+                None,
+                None,
+                Some(vec![Some(3), None, Some(5)]),
+                Some(vec![Some(0), Some(0)]),
+            ],
+            false,
+        );
+
+        do_comparison(
+            vec![None, None, Some(vec![Some(1)])],
+            vec![None, None, Some(vec![Some(2)])],
+            false,
+        );
+    }
 }
diff --git a/arrow/src/array/equal/utils.rs b/arrow/src/array/equal/utils.rs
@@ -187,7 +187,7 @@ fn logical_list_bitmap<OffsetSize: OffsetSizeTrait>(
     offsets
         .windows(2)
         .enumerate()
-        .take(offset_len - offset_start)
+        .take(parent_data.len())
         .for_each(|(index, window)| {
             let start = window[0].to_usize().unwrap();
             let end = window[1].to_usize().unwrap();

diff --git a/arrow/src/compute/kernels/sort.rs b/arrow/src/compute/kernels/sort.rs
@@ -2347,7 +2347,7 @@ mod tests {
                 nulls_first: true,
             }),
             Some(3),
-            vec![None, None, Some(vec![Some(2)])],
+            vec![None, None, Some(vec![Some(1)])],
             None,
         );