arrow_data/equal/
variable_size.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 super::utils::equal_len;
6use crate::data::{ArrayData, contains_nulls};
7use crate::equal::list::lengths_equal;
8use arrow_buffer::ArrowNativeType;
9use num_integer::Integer;
10
11fn offset_value_equal<T: ArrowNativeType + Integer>(
12    lhs_values: &[u8],
13    rhs_values: &[u8],
14    lhs_offsets: &[T],
15    rhs_offsets: &[T],
16    lhs_pos: usize,
17    rhs_pos: usize,
18    len: usize,
19) -> bool {
20    let lhs_start = lhs_offsets[lhs_pos].as_usize();
21    let rhs_start = rhs_offsets[rhs_pos].as_usize();
22    let lhs_len = (lhs_offsets[lhs_pos + len] - lhs_offsets[lhs_pos])
23        .to_usize()
24        .unwrap();
25    let rhs_len = (rhs_offsets[rhs_pos + len] - rhs_offsets[rhs_pos])
26        .to_usize()
27        .unwrap();
28
29    if lhs_len == 0 && rhs_len == 0 {
30        return true;
31    }
32
33    lhs_len == rhs_len && equal_len(lhs_values, rhs_values, lhs_start, rhs_start, lhs_len)
34}
35
36pub(super) fn variable_sized_equal<T: ArrowNativeType + Integer>(
37    lhs: &ArrayData,
38    rhs: &ArrayData,
39    lhs_start: usize,
40    rhs_start: usize,
41    len: usize,
42) -> bool {
43    let lhs_offsets = lhs.buffer::<T>(0);
44    let rhs_offsets = rhs.buffer::<T>(0);
45
46    // the offsets of the `ArrayData` are ignored as they are only applied to the offset buffer.
47    let lhs_values = lhs.buffers()[1].as_slice();
48    let rhs_values = rhs.buffers()[1].as_slice();
49
50    // Only checking one null mask here because by the time the control flow reaches
51    // this point, the equality of the two masks would have already been verified.
52    if !contains_nulls(lhs.nulls(), lhs_start, len) {
53        let lhs_offsets_slice = &lhs_offsets[lhs_start..lhs_start + len + 1];
54        let rhs_offsets_slice = &rhs_offsets[rhs_start..rhs_start + len + 1];
55        lengths_equal(lhs_offsets_slice, rhs_offsets_slice)
56            && offset_value_equal(
57                lhs_values,
58                rhs_values,
59                lhs_offsets,
60                rhs_offsets,
61                lhs_start,
62                rhs_start,
63                len,
64            )
65    } else {
66        (0..len).all(|i| {
67            let lhs_pos = lhs_start + i;
68            let rhs_pos = rhs_start + i;
69
70            // the null bits can still be `None`, indicating that the value is valid.
71            let lhs_is_null = lhs.nulls().is_some_and(|v| v.is_null(lhs_pos));
72            let rhs_is_null = rhs.nulls().is_some_and(|v| v.is_null(rhs_pos));
73
74            lhs_is_null
75                || (lhs_is_null == rhs_is_null)
76                    && offset_value_equal(
77                        lhs_values,
78                        rhs_values,
79                        lhs_offsets,
80                        rhs_offsets,
81                        lhs_pos,
82                        rhs_pos,
83                        1,
84                    )
85        })
86    }
87}
88
89#[cfg(test)]
90mod tests {
91    use crate::ArrayData;
92    use crate::equal::variable_size::variable_sized_equal;
93    use arrow_buffer::Buffer;
94    use arrow_schema::DataType;
95
96    #[test]
97    fn test_variable_sized_equal_diff_offsets() {
98        let a = ArrayData::builder(DataType::Utf8)
99            .buffers(vec![
100                Buffer::from_vec(vec![0_i32, 3, 6]),
101                Buffer::from_slice_ref(b"foobar"),
102            ])
103            .null_bit_buffer(Some(Buffer::from_slice_ref([0b01_u8])))
104            .len(2)
105            .build()
106            .unwrap();
107        let b = ArrayData::builder(DataType::Utf8)
108            .buffers(vec![
109                Buffer::from_vec(vec![0_i32, 2, 6]),
110                Buffer::from_slice_ref(b"foobar"),
111            ])
112            .null_bit_buffer(Some(Buffer::from_slice_ref([0b01_u8])))
113            .len(2)
114            .build()
115            .unwrap();
116
117        assert!(!variable_sized_equal::<i32>(&a, &b, 0, 0, 2));
118    }
119}
arrow_data/equal/variable_size.rs

arrow_data/equal/
variable_size.rs
