arrow_data/equal/
run.rs

1// Licensed to the Apache Software Foundation (ASF) under one
2// or more contributor license agreements.  See the NOTICE file
3// distributed with this work for additional information
4// regarding copyright ownership.  The ASF licenses this file
5// to you under the Apache License, Version 2.0 (the
6// "License"); you may not use this file except in compliance
7// with the License.  You may obtain a copy of the License at
8//
9//   http://www.apache.org/licenses/LICENSE-2.0
10//
11// Unless required by applicable law or agreed to in writing,
12// software distributed under the License is distributed on an
13// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
14// KIND, either express or implied.  See the License for the
15// specific language governing permissions and limitations
16// under the License.
17
18use crate::data::ArrayData;
19use arrow_buffer::ArrowNativeType;
20use arrow_buffer::RunEndBuffer;
21use arrow_buffer::ScalarBuffer;
22use arrow_schema::DataType;
23use num_traits::ToPrimitive;
24
25use super::equal_range;
26
27/// Returns true if the two `RunEndEncoded` arrays are equal.
28///
29/// This provides a specialized implementation of equality for REE arrays that
30/// handles differences in run-encoding by iterating through the logical range.
31pub(super) fn run_equal(
32    lhs: &ArrayData,
33    rhs: &ArrayData,
34    lhs_start: usize,
35    rhs_start: usize,
36    len: usize,
37) -> bool {
38    let lhs_index_type = match lhs.data_type() {
39        DataType::RunEndEncoded(f, _) => f.data_type(),
40        _ => unreachable!(),
41    };
42
43    match lhs_index_type {
44        DataType::Int16 => run_equal_inner::<i16>(lhs, rhs, lhs_start, rhs_start, len),
45        DataType::Int32 => run_equal_inner::<i32>(lhs, rhs, lhs_start, rhs_start, len),
46        DataType::Int64 => run_equal_inner::<i64>(lhs, rhs, lhs_start, rhs_start, len),
47        _ => unreachable!(),
48    }
49}
50
51struct RunArrayData<'a, T: ArrowNativeType> {
52    run_ends: RunEndBuffer<T>,
53    values: &'a ArrayData,
54}
55
56impl<'a, T: ArrowNativeType + ToPrimitive> RunArrayData<'a, T> {
57    fn new(data: &'a ArrayData, start: usize, len: usize) -> Self {
58        debug_assert!(
59            data.child_data().len() == 2,
60            "RunEndEncoded arrays are guaranteed to have 2 children [run_ends, values]"
61        );
62        let run_ends_data = &data.child_data()[0];
63        let raw_run_ends_buffer = &run_ends_data.buffers()[0];
64        // SAFETY: we're reconstructing RunEndBuffer from a known valid RunArray
65        let run_ends = unsafe {
66            RunEndBuffer::<T>::new_unchecked(
67                ScalarBuffer::new(
68                    raw_run_ends_buffer.clone(),
69                    run_ends_data.offset(),
70                    run_ends_data.len(),
71                ),
72                data.offset() + start,
73                len,
74            )
75        };
76
77        let values = &data.child_data()[1];
78        Self { run_ends, values }
79    }
80
81    fn run_end(&self, index: usize) -> usize {
82        self.run_ends.values()[index].as_usize()
83    }
84
85    fn get_start_end_physical_indices(&self) -> (usize, usize) {
86        let start = self.run_ends.get_start_physical_index();
87        let end = self.run_ends.get_end_physical_index();
88        (start, end)
89    }
90}
91
92fn run_equal_inner<T: ArrowNativeType + ToPrimitive>(
93    lhs: &ArrayData,
94    rhs: &ArrayData,
95    lhs_start: usize,
96    rhs_start: usize,
97    len: usize,
98) -> bool {
99    if len == 0 {
100        return true;
101    }
102
103    let l_array = RunArrayData::<T>::new(lhs, lhs_start, len);
104    let r_array = RunArrayData::<T>::new(rhs, rhs_start, len);
105
106    let (l_start_phys, l_end_phys) = l_array.get_start_end_physical_indices();
107    let (r_start_phys, r_end_phys) = r_array.get_start_end_physical_indices();
108    let l_runs = l_end_phys - l_start_phys + 1;
109    let r_runs = r_end_phys - r_start_phys + 1;
110
111    if l_runs == r_runs {
112        // When the boundaries align perfectly, we don't need the complex stepping loop that calculates overlaps.
113        // Instead, we can simply treat the underlying values arrays as if they were standard primitive arrays.
114        let l_iter = l_array.run_ends.sliced_values();
115        let r_iter = r_array.run_ends.sliced_values();
116        let physical_match = l_iter.zip(r_iter).all(|(l_re, r_re)| l_re == r_re);
117
118        if physical_match {
119            // Both arrays are partitioned identically.
120            // We can just verify if the physical values in those partitions match.
121            return equal_range(
122                l_array.values,
123                r_array.values,
124                l_start_phys,
125                r_start_phys,
126                l_runs,
127            );
128        }
129    }
130
131    let mut l_phys = l_start_phys;
132    let mut r_phys = r_start_phys;
133    let mut processed = 0;
134    while processed < len {
135        if !equal_range(l_array.values, r_array.values, l_phys, r_phys, 1) {
136            return false;
137        }
138
139        let l_run_end = l_array.run_end(l_phys);
140        let r_run_end = r_array.run_end(r_phys);
141
142        //Calculate how many more logical elements are in the current run of the left and right array
143        let l_remaining_in_run = l_run_end - (l_array.run_ends.offset() + processed);
144        let r_remaining_in_run = r_run_end - (r_array.run_ends.offset() + processed);
145
146        //Calculate how many elements are left to compare in the requested range
147        let remaining_in_range = len - processed;
148
149        //Find the smallest of these three to determine our step size
150        //The goal is to move the logical cursor (processed) forward as far as possible without:
151        //Crossing the boundary of a run in the left or right array (where the value might change).
152        //Going past the total length we were asked to compare.
153        let step = l_remaining_in_run
154            .min(r_remaining_in_run)
155            .min(remaining_in_range);
156        processed += step;
157
158        if l_array.run_ends.offset() + processed == l_run_end {
159            l_phys += 1;
160        }
161        if r_array.run_ends.offset() + processed == r_run_end {
162            r_phys += 1;
163        }
164    }
165
166    true
167}
arrow_data/equal/run.rs

arrow_data/equal/
run.rs
