Skip to main content

arrow_data/equal/
dictionary.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, contains_nulls};
6use arrow_buffer::ArrowNativeType;
7
8use super::equal_range;
9
10pub(super) fn dictionary_equal<T: ArrowNativeType>(
11    lhs: &ArrayData,
12    rhs: &ArrayData,
13    lhs_start: usize,
14    rhs_start: usize,
15    len: usize,
16) -> bool {
17    let lhs_keys = lhs.buffer::<T>(0);
18    let rhs_keys = rhs.buffer::<T>(0);
19
20    let lhs_values = &lhs.child_data()[0];
21    let rhs_values = &rhs.child_data()[0];
22
23    // Only checking one null mask here because by the time the control flow reaches
24    // this point, the equality of the two masks would have already been verified.
25    if !contains_nulls(lhs.nulls(), lhs_start, len) {
26        (0..len).all(|i| {
27            let lhs_pos = lhs_start + i;
28            let rhs_pos = rhs_start + i;
29
30            equal_range(
31                lhs_values,
32                rhs_values,
33                lhs_keys[lhs_pos].to_usize().unwrap(),
34                rhs_keys[rhs_pos].to_usize().unwrap(),
35                1,
36            )
37        })
38    } else {
39        // get a ref of the null buffer bytes, to use in testing for nullness
40        let lhs_nulls = lhs.nulls().unwrap();
41        let rhs_nulls = rhs.nulls().unwrap();
42        (0..len).all(|i| {
43            let lhs_pos = lhs_start + i;
44            let rhs_pos = rhs_start + i;
45
46            let lhs_is_null = lhs_nulls.is_null(lhs_pos);
47            let rhs_is_null = rhs_nulls.is_null(rhs_pos);
48
49            lhs_is_null
50                || (lhs_is_null == rhs_is_null)
51                    && equal_range(
52                        lhs_values,
53                        rhs_values,
54                        lhs_keys[lhs_pos].to_usize().unwrap(),
55                        rhs_keys[rhs_pos].to_usize().unwrap(),
56                        1,
57                    )
58        })
59    }
60}