Skip to main content

arrow_data/equal/
mod.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
18//! Module containing functionality to compute array equality.
19//! This module uses [ArrayData] and does not
20//! depend on dynamic casting of `Array`.
21
22use crate::data::ArrayData;
23use arrow_buffer::{IntervalMonthDayNano, i256};
24use arrow_schema::{DataType, IntervalUnit};
25use half::f16;
26
27mod boolean;
28mod byte_view;
29mod dictionary;
30mod fixed_binary;
31mod fixed_list;
32mod list;
33mod list_view;
34mod null;
35mod primitive;
36mod run;
37mod structure;
38mod union;
39mod utils;
40mod variable_size;
41
42// these methods assume the same type, len and null count.
43// For this reason, they are not exposed and are instead used
44// to build the generic functions below (`equal_range` and `equal`).
45use self::run::run_equal;
46use crate::equal::list_view::list_view_equal;
47use boolean::boolean_equal;
48use byte_view::byte_view_equal;
49use dictionary::dictionary_equal;
50use fixed_binary::fixed_binary_equal;
51use fixed_list::fixed_list_equal;
52use list::list_equal;
53use null::null_equal;
54use primitive::primitive_equal;
55use structure::struct_equal;
56use union::union_equal;
57use variable_size::variable_sized_equal;
58
59/// Compares the values of two [ArrayData] starting at `lhs_start` and `rhs_start` respectively
60/// for `len` slots.
61///
62/// Assumes the null masks of the two ranges are already known to be equal. To compare
63/// ranges of child data, use [`equal_range`], which also compares the null masks.
64#[inline]
65fn equal_values(
66    lhs: &ArrayData,
67    rhs: &ArrayData,
68    lhs_start: usize,
69    rhs_start: usize,
70    len: usize,
71) -> bool {
72    match lhs.data_type() {
73        DataType::Null => null_equal(lhs, rhs, lhs_start, rhs_start, len),
74        DataType::Boolean => boolean_equal(lhs, rhs, lhs_start, rhs_start, len),
75        DataType::UInt8 => primitive_equal::<u8>(lhs, rhs, lhs_start, rhs_start, len),
76        DataType::UInt16 => primitive_equal::<u16>(lhs, rhs, lhs_start, rhs_start, len),
77        DataType::UInt32 => primitive_equal::<u32>(lhs, rhs, lhs_start, rhs_start, len),
78        DataType::UInt64 => primitive_equal::<u64>(lhs, rhs, lhs_start, rhs_start, len),
79        DataType::Int8 => primitive_equal::<i8>(lhs, rhs, lhs_start, rhs_start, len),
80        DataType::Int16 => primitive_equal::<i16>(lhs, rhs, lhs_start, rhs_start, len),
81        DataType::Int32 => primitive_equal::<i32>(lhs, rhs, lhs_start, rhs_start, len),
82        DataType::Int64 => primitive_equal::<i64>(lhs, rhs, lhs_start, rhs_start, len),
83        DataType::Float32 => primitive_equal::<f32>(lhs, rhs, lhs_start, rhs_start, len),
84        DataType::Float64 => primitive_equal::<f64>(lhs, rhs, lhs_start, rhs_start, len),
85        DataType::Decimal32(_, _) => primitive_equal::<i32>(lhs, rhs, lhs_start, rhs_start, len),
86        DataType::Decimal64(_, _) => primitive_equal::<i64>(lhs, rhs, lhs_start, rhs_start, len),
87        DataType::Decimal128(_, _) => primitive_equal::<i128>(lhs, rhs, lhs_start, rhs_start, len),
88        DataType::Decimal256(_, _) => primitive_equal::<i256>(lhs, rhs, lhs_start, rhs_start, len),
89        DataType::Date32 | DataType::Time32(_) | DataType::Interval(IntervalUnit::YearMonth) => {
90            primitive_equal::<i32>(lhs, rhs, lhs_start, rhs_start, len)
91        }
92        DataType::Date64
93        | DataType::Interval(IntervalUnit::DayTime)
94        | DataType::Time64(_)
95        | DataType::Timestamp(_, _)
96        | DataType::Duration(_) => primitive_equal::<i64>(lhs, rhs, lhs_start, rhs_start, len),
97        DataType::Interval(IntervalUnit::MonthDayNano) => {
98            primitive_equal::<IntervalMonthDayNano>(lhs, rhs, lhs_start, rhs_start, len)
99        }
100        DataType::Utf8 | DataType::Binary => {
101            variable_sized_equal::<i32>(lhs, rhs, lhs_start, rhs_start, len)
102        }
103        DataType::LargeUtf8 | DataType::LargeBinary => {
104            variable_sized_equal::<i64>(lhs, rhs, lhs_start, rhs_start, len)
105        }
106        DataType::FixedSizeBinary(_) => fixed_binary_equal(lhs, rhs, lhs_start, rhs_start, len),
107        DataType::BinaryView | DataType::Utf8View => {
108            byte_view_equal(lhs, rhs, lhs_start, rhs_start, len)
109        }
110        DataType::List(_) => list_equal::<i32>(lhs, rhs, lhs_start, rhs_start, len),
111        DataType::LargeList(_) => list_equal::<i64>(lhs, rhs, lhs_start, rhs_start, len),
112        DataType::ListView(_) => list_view_equal::<i32>(lhs, rhs, lhs_start, rhs_start, len),
113        DataType::LargeListView(_) => list_view_equal::<i64>(lhs, rhs, lhs_start, rhs_start, len),
114        DataType::FixedSizeList(_, _) => fixed_list_equal(lhs, rhs, lhs_start, rhs_start, len),
115        DataType::Struct(_) => struct_equal(lhs, rhs, lhs_start, rhs_start, len),
116        DataType::Union(_, _) => union_equal(lhs, rhs, lhs_start, rhs_start, len),
117        DataType::Dictionary(data_type, _) => match data_type.as_ref() {
118            DataType::Int8 => dictionary_equal::<i8>(lhs, rhs, lhs_start, rhs_start, len),
119            DataType::Int16 => dictionary_equal::<i16>(lhs, rhs, lhs_start, rhs_start, len),
120            DataType::Int32 => dictionary_equal::<i32>(lhs, rhs, lhs_start, rhs_start, len),
121            DataType::Int64 => dictionary_equal::<i64>(lhs, rhs, lhs_start, rhs_start, len),
122            DataType::UInt8 => dictionary_equal::<u8>(lhs, rhs, lhs_start, rhs_start, len),
123            DataType::UInt16 => dictionary_equal::<u16>(lhs, rhs, lhs_start, rhs_start, len),
124            DataType::UInt32 => dictionary_equal::<u32>(lhs, rhs, lhs_start, rhs_start, len),
125            DataType::UInt64 => dictionary_equal::<u64>(lhs, rhs, lhs_start, rhs_start, len),
126            _ => unreachable!(),
127        },
128        DataType::Float16 => primitive_equal::<f16>(lhs, rhs, lhs_start, rhs_start, len),
129        DataType::Map(_, _) => list_equal::<i32>(lhs, rhs, lhs_start, rhs_start, len),
130        DataType::RunEndEncoded(_, _) => run_equal(lhs, rhs, lhs_start, rhs_start, len),
131    }
132}
133
134fn equal_range(
135    lhs: &ArrayData,
136    rhs: &ArrayData,
137    lhs_start: usize,
138    rhs_start: usize,
139    len: usize,
140) -> bool {
141    utils::equal_nulls(lhs, rhs, lhs_start, rhs_start, len)
142        && equal_values(lhs, rhs, lhs_start, rhs_start, len)
143}
144
145/// Logically compares two [ArrayData].
146///
147/// Two arrays are logically equal if and only if:
148/// * their data types are equal
149/// * their lengths are equal
150/// * their null counts are equal
151/// * their null bitmaps are equal
152/// * each of their items are equal
153///
154/// Two items are equal when their in-memory representation is physically equal
155/// (i.e. has the same bit content).
156///
157/// The physical comparison depend on the data type.
158///
159/// # Panics
160///
161/// This function may panic whenever any of the [ArrayData] does not follow the
162/// Arrow specification. (e.g. wrong number of buffers, buffer `len` does not
163/// correspond to the declared `len`)
164pub fn equal(lhs: &ArrayData, rhs: &ArrayData) -> bool {
165    utils::base_equal(lhs, rhs)
166        && lhs.null_count() == rhs.null_count()
167        && utils::equal_nulls(lhs, rhs, 0, 0, lhs.len())
168        && equal_values(lhs, rhs, 0, 0, lhs.len())
169}
170
171// See arrow/tests/array_equal.rs for tests