Skip to main content

arrow_data/equal/
primitive.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::bit_iterator::BitSliceIterator;
6use crate::contains_nulls;
7use std::mem::size_of;
8
9use crate::data::ArrayData;
10
11use super::utils::equal_len;
12
13pub(crate) const NULL_SLICES_SELECTIVITY_THRESHOLD: f64 = 0.4;
14
15pub(super) fn primitive_equal<T>(
16    lhs: &ArrayData,
17    rhs: &ArrayData,
18    lhs_start: usize,
19    rhs_start: usize,
20    len: usize,
21) -> bool {
22    let byte_width = size_of::<T>();
23    let lhs_values = &lhs.buffers()[0].as_slice()[lhs.offset() * byte_width..];
24    let rhs_values = &rhs.buffers()[0].as_slice()[rhs.offset() * byte_width..];
25
26    // Only checking one null mask here because by the time the control flow reaches
27    // this point, the equality of the two masks would have already been verified.
28    if !contains_nulls(lhs.nulls(), lhs_start, len) {
29        // without nulls, we just need to compare slices
30        equal_len(
31            lhs_values,
32            rhs_values,
33            lhs_start * byte_width,
34            rhs_start * byte_width,
35            len * byte_width,
36        )
37    } else {
38        let selectivity_frac = lhs.null_count() as f64 / lhs.len() as f64;
39
40        // get a ref of the null buffer bytes, to use in testing for nullness
41        let lhs_nulls = lhs.nulls().unwrap();
42        let rhs_nulls = rhs.nulls().unwrap();
43
44        if selectivity_frac >= NULL_SLICES_SELECTIVITY_THRESHOLD {
45            // with nulls, we need to compare item by item whenever it is not null
46            (0..len).all(|i| {
47                let lhs_pos = lhs_start + i;
48                let rhs_pos = rhs_start + i;
49                let lhs_is_null = lhs_nulls.is_null(lhs_pos);
50                let rhs_is_null = rhs_nulls.is_null(rhs_pos);
51
52                lhs_is_null
53                    || (lhs_is_null == rhs_is_null)
54                        && equal_len(
55                            lhs_values,
56                            rhs_values,
57                            lhs_pos * byte_width,
58                            rhs_pos * byte_width,
59                            byte_width, // 1 * byte_width since we are comparing a single entry
60                        )
61            })
62        } else {
63            let lhs_slices_iter =
64                BitSliceIterator::new(lhs_nulls.validity(), lhs_start + lhs_nulls.offset(), len);
65            let rhs_slices_iter =
66                BitSliceIterator::new(rhs_nulls.validity(), rhs_start + rhs_nulls.offset(), len);
67
68            lhs_slices_iter
69                .zip(rhs_slices_iter)
70                .all(|((l_start, l_end), (r_start, r_end))| {
71                    l_start == r_start
72                        && l_end == r_end
73                        && equal_len(
74                            lhs_values,
75                            rhs_values,
76                            (lhs_start + l_start) * byte_width,
77                            (rhs_start + r_start) * byte_width,
78                            (l_end - l_start) * byte_width,
79                        )
80                })
81        }
82    }
83}