Skip to main content

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