This is an automated email from the ASF dual-hosted git repository.
Jefffrey pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/arrow-rs.git
The following commit(s) were added to refs/heads/main by this push:
new 0b4847a0c5 fix: Sparse UnionArray equality should ignore unselected
children (#11009)
0b4847a0c5 is described below
commit 0b4847a0c52fa1628ee18a78bd0530fdc942359b
Author: Neil Conway <[email protected]>
AuthorDate: Sun Sep 6 21:17:17 2026 -0400
fix: Sparse UnionArray equality should ignore unselected children (#11009)
# Which issue does this PR close?
- Closes #11008.
# Rationale for this change
The Arrow spec allows sparse unions to contain arbitrary values in
"unselected" children:
> Only the slot in the array corresponding to the type index is
considered. All "unselected" values are ignored and could be any
semantically correct array value.
Array equality for sparse unions (`ArrayData::eq`) compares every child
array over the full slot range, regardless of which child each slot's
type id selects. This is inconsistent with the spec, as well as the
Arrow C++ implementation.
# What changes are included in this PR?
* Fix equality semantics for sparse unions
* Add unit tests
# Are these changes tested?
Yes, new tests added.
# Are there any user-facing changes?
Yes, equality semantics for sparse unions have changed. However, it is
pretty unlikely that user code would depend on the previous behavior (if
anything, it would probably result in bugs in user code).
# AI usage
Developed with Claude Code Fable 5.1. I reviewed and revised the
resulting code.
---
arrow-data/src/equal/union.rs | 42 +++++++++++------
arrow/tests/array_equal.rs | 104 ++++++++++++++++++++++++++++++++++++++++--
2 files changed, 128 insertions(+), 18 deletions(-)
diff --git a/arrow-data/src/equal/union.rs b/arrow-data/src/equal/union.rs
index 3e889cd455..0a80f50e9c 100644
--- a/arrow-data/src/equal/union.rs
+++ b/arrow-data/src/equal/union.rs
@@ -59,25 +59,37 @@ fn equal_dense(
})
}
+/// Compares `type_ids.len()` slots of two sparse unions whose type ids over
that
+/// range are already known to be equal.
+///
+/// Only the child selected by each slot's type id contributes to the logical
+/// value of a sparse union; the values held by the other children at that slot
+/// are arbitrary and must not affect equality. Consecutive slots that select
the
+/// same child are compared as a single range of that child.
fn equal_sparse(
lhs: &ArrayData,
rhs: &ArrayData,
+ type_ids: &[i8],
+ fields: &UnionFields,
lhs_start: usize,
rhs_start: usize,
- len: usize,
) -> bool {
- lhs.child_data()
- .iter()
- .zip(rhs.child_data())
- .all(|(lhs_values, rhs_values)| {
- equal_range(
- lhs_values,
- rhs_values,
- lhs_start + lhs.offset(),
- rhs_start + rhs.offset(),
- len,
- )
- })
+ let mut run_start = 0;
+ type_ids.chunk_by(|a, b| a == b).all(|run| {
+ let start = run_start;
+ run_start += run.len();
+
+ // Both sides share the same `UnionFields` (checked when comparing the
+ // data types), so the type id maps to the same child index on each
side.
+ let child_index = fields.iter().position(|(id, _)| id ==
run[0]).unwrap();
+ equal_range(
+ &lhs.child_data()[child_index],
+ &rhs.child_data()[child_index],
+ lhs_start + lhs.offset() + start,
+ rhs_start + rhs.offset() + start,
+ run.len(),
+ )
+ })
}
pub(super) fn union_equal(
@@ -116,9 +128,9 @@ pub(super) fn union_equal(
rhs_fields,
)
}
- (DataType::Union(_, UnionMode::Sparse), DataType::Union(_,
UnionMode::Sparse)) => {
+ (DataType::Union(fields, UnionMode::Sparse), DataType::Union(_,
UnionMode::Sparse)) => {
lhs_type_id_range == rhs_type_id_range
- && equal_sparse(lhs, rhs, lhs_start, rhs_start, len)
+ && equal_sparse(lhs, rhs, lhs_type_id_range, fields,
lhs_start, rhs_start)
}
_ => unimplemented!(
"Logical equality not yet implemented between dense and sparse
union arrays"
diff --git a/arrow/tests/array_equal.rs b/arrow/tests/array_equal.rs
index f7f75a02b9..c108cc9670 100644
--- a/arrow/tests/array_equal.rs
+++ b/arrow/tests/array_equal.rs
@@ -19,7 +19,7 @@ use arrow::array::{
Array, ArrayRef, BooleanArray, Decimal128Array, FixedSizeBinaryArray,
FixedSizeBinaryBuilder,
FixedSizeListBuilder, GenericBinaryArray, GenericStringArray, Int32Array,
Int32Builder,
Int64Builder, ListArray, ListBuilder, NullArray, OffsetSizeTrait,
StringArray,
- StringDictionaryBuilder, StructArray, UnionBuilder, make_array,
+ StringDictionaryBuilder, StructArray, UnionArray, UnionBuilder, make_array,
};
use arrow::datatypes::{Int16Type, Int32Type};
use arrow_array::builder::{
@@ -29,9 +29,9 @@ use arrow_array::cast::AsArray;
use arrow_array::{
DictionaryArray, FixedSizeListArray, GenericListViewArray, PrimitiveArray,
StringViewArray,
};
-use arrow_buffer::{Buffer, ToByteSlice};
+use arrow_buffer::{Buffer, ScalarBuffer, ToByteSlice};
use arrow_data::{ArrayData, ArrayDataBuilder};
-use arrow_schema::{DataType, Field, Fields};
+use arrow_schema::{DataType, Field, Fields, UnionFields};
use arrow_select::take::take;
use std::sync::Arc;
@@ -1273,6 +1273,104 @@ fn test_union_equal_sparse_slice() {
test_equal(&a1.slice(1, 2), &a2, true)
}
+/// Builds a sparse union with `Int32` children `a` (type id 0) and `b` (type
id 1)
+/// from its component parts, so the values in unselected child slots can be
+/// chosen freely. `UnionBuilder` always pads unselected slots with nulls and
so
+/// cannot produce arrays that differ only in unselected values.
+fn sparse_union(type_ids: &[i8], a: Int32Array, b: Int32Array) -> UnionArray {
+ let fields = UnionFields::try_new(
+ [0, 1],
+ [
+ Field::new("a", DataType::Int32, true),
+ Field::new("b", DataType::Int32, true),
+ ],
+ )
+ .unwrap();
+ UnionArray::try_new(
+ fields,
+ ScalarBuffer::from(type_ids.to_vec()),
+ None,
+ vec![Arc::new(a), Arc::new(b)],
+ )
+ .unwrap()
+}
+
+#[test]
+fn test_union_equal_sparse_ignores_unselected_children() {
+ // Slot 0 selects `a` and slot 1 selects `b`. The unselected child at each
+ // slot holds a different value, or a null, on each side.
+ let lhs = sparse_union(
+ &[0, 1],
+ Int32Array::from(vec![Some(1), Some(10)]),
+ Int32Array::from(vec![Some(20), Some(2)]),
+ );
+ let rhs = sparse_union(
+ &[0, 1],
+ Int32Array::from(vec![Some(1), None]),
+ Int32Array::from(vec![None, Some(2)]),
+ );
+ test_equal(&lhs, &rhs, true);
+ test_equal(&lhs.slice(0, 1), &rhs.slice(0, 1), true);
+ test_equal(&lhs.slice(1, 1), &rhs.slice(1, 1), true);
+
+ // A differing selected value is still detected.
+ let rhs = sparse_union(
+ &[0, 1],
+ Int32Array::from(vec![Some(3), Some(10)]),
+ Int32Array::from(vec![Some(20), Some(2)]),
+ );
+ test_equal(&lhs, &rhs, false);
+ test_equal(&lhs.slice(0, 1), &rhs.slice(0, 1), false);
+ test_equal(&lhs.slice(1, 1), &rhs.slice(1, 1), true);
+
+ // A selected null against a selected value is still detected.
+ let rhs = sparse_union(
+ &[0, 1],
+ Int32Array::from(vec![None, Some(10)]),
+ Int32Array::from(vec![Some(20), Some(2)]),
+ );
+ test_equal(&lhs, &rhs, false);
+}
+
+#[test]
+fn test_union_equal_sparse_runs() {
+ // The type ids form the runs [a, a], [b, b, b], [a].
+ let type_ids = [0, 0, 1, 1, 1, 0];
+ let lhs = sparse_union(
+ &type_ids,
+ Int32Array::from(vec![1, 2, 0, 0, 0, 6]),
+ Int32Array::from(vec![0, 0, 3, 4, 5, 0]),
+ );
+ let rhs = sparse_union(
+ &type_ids,
+ Int32Array::from(vec![1, 2, 9, 9, 9, 6]),
+ Int32Array::from(vec![9, 9, 3, 4, 5, 9]),
+ );
+ test_equal(&lhs, &rhs, true);
+ // Slices starting and ending in the middle of a run.
+ test_equal(&lhs.slice(1, 4), &rhs.slice(1, 4), true);
+ test_equal(&lhs.slice(3, 3), &rhs.slice(3, 3), true);
+
+ // A slice compared against an unsliced array with a different offset.
+ let unsliced = sparse_union(
+ &[0, 1, 1, 1],
+ Int32Array::from(vec![2, 7, 7, 7]),
+ Int32Array::from(vec![7, 3, 4, 5]),
+ );
+ test_equal(&lhs.slice(1, 4), &unsliced, true);
+
+ // A differing selected value in the middle of a run is still detected.
+ let rhs = sparse_union(
+ &type_ids,
+ Int32Array::from(vec![1, 2, 9, 9, 9, 6]),
+ Int32Array::from(vec![9, 9, 3, 7, 5, 9]),
+ );
+ test_equal(&lhs, &rhs, false);
+ test_equal(&lhs.slice(1, 4), &rhs.slice(1, 4), false);
+ test_equal(&lhs.slice(0, 3), &rhs.slice(0, 3), true);
+ test_equal(&lhs.slice(4, 2), &rhs.slice(4, 2), true);
+}
+
#[test]
fn test_boolean_slice() {
let array = BooleanArray::from(vec![true; 32]);
