arrow_data/equal/
list.rs

1// SPDX-License-Identifier: Apache-2.0
2// SPDX-FileCopyrightText: See the NOTICE file distributed with this work for additional information regarding copyright ownership
3// SPDX-FileContributor: Licensed to the Apache Software Foundation (ASF) under one or more contributor license agreements
4
5use crate::data::{ArrayData, count_nulls};
6use arrow_buffer::ArrowNativeType;
7use num_integer::Integer;
8
9use super::equal_range;
10
11pub(super) fn lengths_equal<T: ArrowNativeType + Integer>(lhs: &[T], rhs: &[T]) -> bool {
12    // invariant from `base_equal`
13    debug_assert_eq!(lhs.len(), rhs.len());
14
15    if lhs.is_empty() {
16        return true;
17    }
18
19    if lhs[0] == T::zero() && rhs[0] == T::zero() {
20        return lhs == rhs;
21    }
22
23    // The expensive case, e.g.
24    // [0, 2, 4, 6, 9] == [4, 6, 8, 10, 13]
25    lhs.windows(2)
26        .zip(rhs.windows(2))
27        .all(|(lhs_offsets, rhs_offsets)| {
28            // length of left == length of right
29            (lhs_offsets[1] - lhs_offsets[0]) == (rhs_offsets[1] - rhs_offsets[0])
30        })
31}
32
33pub(super) fn list_equal<T: ArrowNativeType + Integer>(
34    lhs: &ArrayData,
35    rhs: &ArrayData,
36    lhs_start: usize,
37    rhs_start: usize,
38    len: usize,
39) -> bool {
40    let lhs_offsets = lhs.buffer::<T>(0);
41    let rhs_offsets = rhs.buffer::<T>(0);
42
43    // There is an edge-case where a n-length list that has 0 children, results in panics.
44    // For example; an array with offsets [0, 0, 0, 0, 0] has 4 slots, but will have
45    // no valid children.
46    // Under logical equality, the child null bitmap will be an empty buffer, as there are
47    // no child values. This causes panics when trying to count set bits.
48    //
49    // We caught this by chance from an accidental test-case, but due to the nature of this
50    // crash only occurring on list equality checks, we are adding a check here, instead of
51    // on the buffer/bitmap utilities, as a length check would incur a penalty for almost all
52    // other use-cases.
53    //
54    // The solution is to check the number of child values from offsets, and return `true` if
55    // they = 0. Empty arrays are equal, so this is correct.
56    //
57    // It's unlikely that one would create a n-length list array with no values, where n > 0,
58    // however, one is more likely to slice into a list array and get a region that has 0
59    // child values.
60    // The test that triggered this behaviour had [4, 4] as a slice of 1 value slot.
61    // For the edge case that zero length list arrays are always equal.
62    if len == 0 {
63        return true;
64    }
65
66    let lhs_child_length = lhs_offsets[lhs_start + len].to_usize().unwrap()
67        - lhs_offsets[lhs_start].to_usize().unwrap();
68
69    let rhs_child_length = rhs_offsets[rhs_start + len].to_usize().unwrap()
70        - rhs_offsets[rhs_start].to_usize().unwrap();
71
72    if lhs_child_length == 0 && lhs_child_length == rhs_child_length {
73        return true;
74    }
75
76    let lhs_values = &lhs.child_data()[0];
77    let rhs_values = &rhs.child_data()[0];
78
79    let lhs_null_count = count_nulls(lhs.nulls(), lhs_start, len);
80    let rhs_null_count = count_nulls(rhs.nulls(), rhs_start, len);
81
82    if lhs_null_count != rhs_null_count {
83        return false;
84    }
85
86    if lhs_null_count == 0 && rhs_null_count == 0 {
87        lhs_child_length == rhs_child_length
88            && lengths_equal(
89                &lhs_offsets[lhs_start..lhs_start + len],
90                &rhs_offsets[rhs_start..rhs_start + len],
91            )
92            && equal_range(
93                lhs_values,
94                rhs_values,
95                lhs_offsets[lhs_start].to_usize().unwrap(),
96                rhs_offsets[rhs_start].to_usize().unwrap(),
97                lhs_child_length,
98            )
99    } else {
100        // get a ref of the parent null buffer bytes, to use in testing for nullness
101        let lhs_nulls = lhs.nulls().unwrap();
102        let rhs_nulls = rhs.nulls().unwrap();
103
104        // with nulls, we need to compare item by item whenever it is not null
105        // TODO: Could potentially compare runs of not NULL values
106        (0..len).all(|i| {
107            let lhs_pos = lhs_start + i;
108            let rhs_pos = rhs_start + i;
109
110            let lhs_is_null = lhs_nulls.is_null(lhs_pos);
111            let rhs_is_null = rhs_nulls.is_null(rhs_pos);
112
113            if lhs_is_null != rhs_is_null {
114                return false;
115            }
116
117            let lhs_offset_start = lhs_offsets[lhs_pos].to_usize().unwrap();
118            let lhs_offset_end = lhs_offsets[lhs_pos + 1].to_usize().unwrap();
119            let rhs_offset_start = rhs_offsets[rhs_pos].to_usize().unwrap();
120            let rhs_offset_end = rhs_offsets[rhs_pos + 1].to_usize().unwrap();
121
122            let lhs_len = lhs_offset_end - lhs_offset_start;
123            let rhs_len = rhs_offset_end - rhs_offset_start;
124
125            lhs_is_null
126                || (lhs_len == rhs_len
127                    && equal_range(
128                        lhs_values,
129                        rhs_values,
130                        lhs_offset_start,
131                        rhs_offset_start,
132                        lhs_len,
133                    ))
134        })
135    }
136}
arrow_data/equal/list.rs

arrow_data/equal/
list.rs
