add more tests; debug some of the struct problems

apache · Jul 16, 2021 · 16db65a · 16db65a
1 parent bb4d526
commit 16db65a
Show file tree

Hide file tree

Showing 3 changed files with 71 additions and 8 deletions.
diff --git a/arrow/src/array/equal/mod.rs b/arrow/src/array/equal/mod.rs
@@ -1084,6 +1084,71 @@ mod tests {
         test_equal(a.data(), b.data(), true);
     }
 
+    #[test]
+    fn test_struct_equal_large_aligned_slice() {
+        let a = make_struct(vec![
+            None,
+            None,
+            None,
+            None,
+            None,
+            None,
+            None,
+            None,
+            Some((Some("joe"), Some(1))),
+            Some((None, Some(2))),
+            Some((None, None)),
+            Some((Some("mark"), Some(4))),
+            Some((Some("doe"), Some(5))),
+        ]);
+        let a = a.slice(8, 5);
+        let a = a.as_any().downcast_ref::<StructArray>().unwrap();
+
+        let b = make_struct(vec![
+            Some((Some("joe"), Some(1))),
+            Some((None, Some(2))),
+            Some((None, None)),
+            Some((Some("mark"), Some(4))),
+            Some((Some("doe"), Some(5))),
+        ]);
+        assert_eq!(a, &b);
+
+        test_equal(a.data(), b.data(), true);
+    }
+
+    #[test]
+    fn test_struct_equal_large_unaligned_slice() {
+        let a = make_struct(vec![
+            None,
+            None,
+            None,
+            None,
+            None,
+            None,
+            None,
+            None,
+            None,
+            Some((Some("joe"), Some(1))),
+            Some((None, Some(2))),
+            Some((None, None)),
+            Some((Some("mark"), Some(4))),
+            Some((Some("doe"), Some(5))),
+        ]);
+        let a = a.slice(9, 5);
+        let a = a.as_any().downcast_ref::<StructArray>().unwrap();
+
+        let b = make_struct(vec![
+            Some((Some("joe"), Some(1))),
+            Some((None, Some(2))),
+            Some((None, None)),
+            Some((Some("mark"), Some(4))),
+            Some((Some("doe"), Some(5))),
+        ]);
+        assert_eq!(a, &b);
+
+        test_equal(a.data(), b.data(), true);
+    }
+
     #[test]
     fn test_struct_equal_null() {
         let strings: ArrayRef = Arc::new(StringArray::from(vec![

diff --git a/arrow/src/array/equal/structure.rs b/arrow/src/array/equal/structure.rs
@@ -66,8 +66,8 @@ pub(super) fn struct_equal(
     len: usize,
 ) -> bool {
     // we have to recalculate null counts from the null buffers
-    let lhs_null_count = count_nulls(lhs_nulls, lhs_start, len);
-    let rhs_null_count = count_nulls(rhs_nulls, rhs_start, len);
+    let lhs_null_count = count_nulls(lhs_nulls, lhs.offset() + lhs_start, len);
+    let rhs_null_count = count_nulls(rhs_nulls, rhs.offset() + rhs_start, len);
     if lhs_null_count == 0 && rhs_null_count == 0 {
         equal_values(lhs, rhs, lhs_nulls, rhs_nulls, lhs_start, rhs_start, len)
     } else {

diff --git a/arrow/src/array/equal/utils.rs b/arrow/src/array/equal/utils.rs
@@ -47,16 +47,14 @@ pub(super) fn equal_nulls(
     rhs_start: usize,
     len: usize,
 ) -> bool {
-    let lhs_null_count = count_nulls(lhs_nulls, lhs_start, len);
-    let rhs_null_count = count_nulls(rhs_nulls, rhs_start, len);
+    let lhs_null_count = count_nulls(lhs_nulls, lhs.offset() + lhs_start, len);
+    let rhs_null_count = count_nulls(rhs_nulls, rhs.offset() + rhs_start, len);
     if lhs_null_count > 0 || rhs_null_count > 0 {
         let lhs_values = lhs_nulls.unwrap().as_slice();
         let rhs_values = rhs_nulls.unwrap().as_slice();
         equal_bits(
-            lhs_values,
-            rhs_values,
-            lhs_start + lhs.offset(),
-            rhs_start + rhs.offset(),
+            lhs_values, rhs_values, lhs_start, // + lhs.offset(),
+            rhs_start, //  + rhs.offset(),
             len,
         )
     } else {