Skip to main content

arrow_data/equal/
utils.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::bit_chunk_iterator::BitChunks;
7use arrow_schema::DataType;
8
9// whether bits along the positions are equal
10// `lhs_start`, `rhs_start` and `len` are _measured in bits_.
11#[inline]
12pub(super) fn equal_bits(
13    lhs_values: &[u8],
14    rhs_values: &[u8],
15    lhs_start: usize,
16    rhs_start: usize,
17    len: usize,
18) -> bool {
19    let lhs = BitChunks::new(lhs_values, lhs_start, len).iter_padded();
20    let rhs = BitChunks::new(rhs_values, rhs_start, len).iter_padded();
21    lhs.zip(rhs).all(|(a, b)| a == b)
22}
23
24#[inline]
25pub(super) fn equal_nulls(
26    lhs: &ArrayData,
27    rhs: &ArrayData,
28    lhs_start: usize,
29    rhs_start: usize,
30    len: usize,
31) -> bool {
32    match (lhs.nulls(), rhs.nulls()) {
33        (Some(lhs), Some(rhs)) => equal_bits(
34            lhs.validity(),
35            rhs.validity(),
36            lhs.offset() + lhs_start,
37            rhs.offset() + rhs_start,
38            len,
39        ),
40        (Some(lhs), None) => !contains_nulls(Some(lhs), lhs_start, len),
41        (None, Some(rhs)) => !contains_nulls(Some(rhs), rhs_start, len),
42        (None, None) => true,
43    }
44}
45
46#[inline]
47pub(super) fn base_equal(lhs: &ArrayData, rhs: &ArrayData) -> bool {
48    let equal_type = match (lhs.data_type(), rhs.data_type()) {
49        (DataType::Union(l_fields, l_mode), DataType::Union(r_fields, r_mode)) => {
50            l_fields == r_fields && l_mode == r_mode
51        }
52        (DataType::Map(l_field, l_sorted), DataType::Map(r_field, r_sorted)) => {
53            let field_equal = match (l_field.data_type(), r_field.data_type()) {
54                (DataType::Struct(l_fields), DataType::Struct(r_fields))
55                    if l_fields.len() == 2 && r_fields.len() == 2 =>
56                {
57                    let l_key_field = &l_fields[0];
58                    let r_key_field = &r_fields[0];
59                    let l_value_field = &l_fields[1];
60                    let r_value_field = &r_fields[1];
61
62                    // We don't enforce the equality of field names
63                    let data_type_equal = l_key_field.data_type() == r_key_field.data_type()
64                        && l_value_field.data_type() == r_value_field.data_type();
65                    let nullability_equal = l_key_field.is_nullable() == r_key_field.is_nullable()
66                        && l_value_field.is_nullable() == r_value_field.is_nullable();
67                    let metadata_equal = l_key_field.metadata() == r_key_field.metadata()
68                        && l_value_field.metadata() == r_value_field.metadata();
69                    data_type_equal && nullability_equal && metadata_equal
70                }
71                _ => panic!("Map type should have 2 fields Struct in its field"),
72            };
73            field_equal && l_sorted == r_sorted
74        }
75        (l_data_type, r_data_type) => l_data_type == r_data_type,
76    };
77    equal_type && lhs.len() == rhs.len()
78}
79
80// whether the two memory regions are equal
81#[inline]
82pub(super) fn equal_len(
83    lhs_values: &[u8],
84    rhs_values: &[u8],
85    lhs_start: usize,
86    rhs_start: usize,
87    len: usize,
88) -> bool {
89    lhs_values[lhs_start..(lhs_start + len)] == rhs_values[rhs_start..(rhs_start + len)]
90}