From bd9f67aeca45513fcc15d0378bf8e6637acfcefc Mon Sep 17 00:00:00 2001
From: Ruihang Xia <waynestxia@gmail.com>
Date: Wed, 21 Jul 2021 02:01:44 +0800
Subject: [PATCH] Fix array equal check (#571)

* pin this bug down & fix

Signed-off-by: Ruihang Xia <waynestxia@gmail.com>

* tidy

Signed-off-by: Ruihang Xia <waynestxia@gmail.com>

* add some test cases

Signed-off-by: Ruihang Xia <waynestxia@gmail.com>
---
 arrow/src/array/array_list.rs     | 72 +++++++++++++++++++++++++++++++
 arrow/src/array/equal/utils.rs    |  2 +-
 arrow/src/compute/kernels/sort.rs |  2 +-
 3 files changed, 74 insertions(+), 2 deletions(-)

diff --git a/arrow/src/array/array_list.rs b/arrow/src/array/array_list.rs
index 0efb65c99d5a..990ec8aae1c6 100644
--- 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
index d0108d236491..2a1ce88d1474 100644
--- 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
index 54529facd650..ee723da0cc43 100644
--- 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,
         );