arrow_data/equal/
utils.rs

1// Licensed to the Apache Software Foundation (ASF) under one
2// or more contributor license agreements.  See the NOTICE file
3// distributed with this work for additional information
4// regarding copyright ownership.  The ASF licenses this file
5// to you under the Apache License, Version 2.0 (the
6// "License"); you may not use this file except in compliance
7// with the License.  You may obtain a copy of the License at
8//
9//   http://www.apache.org/licenses/LICENSE-2.0
10//
11// Unless required by applicable law or agreed to in writing,
12// software distributed under the License is distributed on an
13// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
14// KIND, either express or implied.  See the License for the
15// specific language governing permissions and limitations
16// under the License.
17
18use crate::data::{contains_nulls, ArrayData};
19use arrow_buffer::bit_chunk_iterator::BitChunks;
20use arrow_schema::DataType;
21
22// whether bits along the positions are equal
23// `lhs_start`, `rhs_start` and `len` are _measured in bits_.
24#[inline]
25pub(super) fn equal_bits(
26    lhs_values: &[u8],
27    rhs_values: &[u8],
28    lhs_start: usize,
29    rhs_start: usize,
30    len: usize,
31) -> bool {
32    let lhs = BitChunks::new(lhs_values, lhs_start, len).iter_padded();
33    let rhs = BitChunks::new(rhs_values, rhs_start, len).iter_padded();
34    lhs.zip(rhs).all(|(a, b)| a == b)
35}
36
37#[inline]
38pub(super) fn equal_nulls(
39    lhs: &ArrayData,
40    rhs: &ArrayData,
41    lhs_start: usize,
42    rhs_start: usize,
43    len: usize,
44) -> bool {
45    match (lhs.nulls(), rhs.nulls()) {
46        (Some(lhs), Some(rhs)) => equal_bits(
47            lhs.validity(),
48            rhs.validity(),
49            lhs.offset() + lhs_start,
50            rhs.offset() + rhs_start,
51            len,
52        ),
53        (Some(lhs), None) => !contains_nulls(Some(lhs), lhs_start, len),
54        (None, Some(rhs)) => !contains_nulls(Some(rhs), rhs_start, len),
55        (None, None) => true,
56    }
57}
58
59#[inline]
60pub(super) fn base_equal(lhs: &ArrayData, rhs: &ArrayData) -> bool {
61    let equal_type = match (lhs.data_type(), rhs.data_type()) {
62        (DataType::Union(l_fields, l_mode), DataType::Union(r_fields, r_mode)) => {
63            l_fields == r_fields && l_mode == r_mode
64        }
65        (DataType::Map(l_field, l_sorted), DataType::Map(r_field, r_sorted)) => {
66            let field_equal = match (l_field.data_type(), r_field.data_type()) {
67                (DataType::Struct(l_fields), DataType::Struct(r_fields))
68                    if l_fields.len() == 2 && r_fields.len() == 2 =>
69                {
70                    let l_key_field = &l_fields[0];
71                    let r_key_field = &r_fields[0];
72                    let l_value_field = &l_fields[1];
73                    let r_value_field = &r_fields[1];
74
75                    // We don't enforce the equality of field names
76                    let data_type_equal = l_key_field.data_type() == r_key_field.data_type()
77                        && l_value_field.data_type() == r_value_field.data_type();
78                    let nullability_equal = l_key_field.is_nullable() == r_key_field.is_nullable()
79                        && l_value_field.is_nullable() == r_value_field.is_nullable();
80                    let metadata_equal = l_key_field.metadata() == r_key_field.metadata()
81                        && l_value_field.metadata() == r_value_field.metadata();
82                    data_type_equal && nullability_equal && metadata_equal
83                }
84                _ => panic!("Map type should have 2 fields Struct in its field"),
85            };
86            field_equal && l_sorted == r_sorted
87        }
88        (l_data_type, r_data_type) => l_data_type == r_data_type,
89    };
90    equal_type && lhs.len() == rhs.len()
91}
92
93// whether the two memory regions are equal
94#[inline]
95pub(super) fn equal_len(
96    lhs_values: &[u8],
97    rhs_values: &[u8],
98    lhs_start: usize,
99    rhs_start: usize,
100    len: usize,
101) -> bool {
102    lhs_values[lhs_start..(lhs_start + len)] == rhs_values[rhs_start..(rhs_start + len)]
103}