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

arrow_data/equal/
run.rs
