[C++][Python] Sliced unions can report incorrect null counts in `count`

[fenfeng9]

Describe the bug, including details regarding any error messages, version, and platform.

Sliced union arrays can report incorrect logical null counts in count.

Expected behavior:

For sliced sparse and dense unions, count(..., mode="only_valid") and count(..., mode="only_null") should reflect the logical nullness of the slice.

Actual behavior:

For both sliced sparse and dense unions, count(..., mode="only_valid") and count(..., mode="only_null") return incorrect results.

Reproduce

import pyarrow as pa
import pyarrow.compute as pc


def main():
    # sparse_arr == [0.5, None, None, 3.0, True].
    sparse_arr = pa.UnionArray.from_sparse(
        pa.array([0, 1, 0, 0, 1], type=pa.int8()),
        [
            pa.array([0.5, 99.0, None, 3.0, 88.0]),
            pa.array([False, None, True, False, True]),
        ],
    )
    # sparse == [None, None, 3.0, True].
    sparse = sparse_arr.slice(1, 4)

    # dense_arr == [0.5, True, 1.5, None, None, False].
    dense_arr = pa.UnionArray.from_dense(
        pa.array([0, 1, 0, 0, 1, 1], type=pa.int8()),
        pa.array([0, 0, 1, 2, 1, 2], type=pa.int32()),
        [
            pa.array([0.5, 1.5, None]),
            pa.array([True, None, False]),
        ],
    )
    # dense == [True, 1.5, None, None].
    dense = dense_arr.slice(1, 4)

    print(f"pyarrow: {pa.__version__}")
    print()

    # Logical sparse slice: [None, None, 3.0, True].
    print("sparse count only_valid expected: 2")
    print(f"sparse count only_valid actual:   {pc.count(sparse, mode='only_valid').as_py()}")
    print("sparse count only_null expected:  2")
    print(f"sparse count only_null actual:    {pc.count(sparse, mode='only_null').as_py()}")
    print()

    # Logical dense slice: [True, 1.5, None, None].
    print("dense count only_valid expected:  2")
    print(f"dense count only_valid actual:    {pc.count(dense, mode='only_valid').as_py()}")
    print("dense count only_null expected:   2")
    print(f"dense count only_null actual:     {pc.count(dense, mode='only_null').as_py()}")


if __name__ == "__main__":
    main()

Result

pyarrow: 24.0.0

sparse count only_valid expected: 2
sparse count only_valid actual:   4
sparse count only_null expected:  2
sparse count only_null actual:    0

dense count only_valid expected:  2
dense count only_valid actual:    4
dense count only_null expected:   2
dense count only_null actual:     0

Component(s)

C++, Python

[fenfeng9]

I think this bug is in the union logical null count path.

count calls ComputeLogicalNullCount() for array inputs here:

arrow/cpp/src/arrow/compute/kernels/aggregate_basic.cc

Lines 101 to 108 in e980b7e

	Status Consume(KernelContext*, const ExecSpan& batch) override {
	if (options.mode == CountOptions::ALL) {
	this->non_nulls += batch.length;
	} else if (batch[0].is_array()) {
	const ArraySpan& input = batch[0].array;
	const int64_t nulls = input.ComputeLogicalNullCount();
	this->nulls += nulls;
	this->non_nulls += input.length - nulls;

For union arrays, ArraySpan::ComputeLogicalNullCount() dispatches to the union logical null count helpers here:

arrow/cpp/src/arrow/array/data.cc

Lines 624 to 630 in e980b7e

	int64_t ArraySpan::ComputeLogicalNullCount() const {
	const auto t = this->type->id();
	if (t == Type::SPARSE_UNION) {
	return union_util::LogicalSparseUnionNullCount(*this);
	}
	if (t == Type::DENSE_UNION) {
	return union_util::LogicalDenseUnionNullCount(*this);

The sparse and dense union implementations are here:

arrow/cpp/src/arrow/util/union_util.cc

Lines 28 to 38 in e980b7e

	int64_t LogicalSparseUnionNullCount(const ArraySpan& span) {
	const auto* sparse_union_type =
	internal::checked_cast<const SparseUnionType*>(span.type);
	DCHECK_LE(span.child_data.size(), 128);
	const int8_t* types = span.GetValues<int8_t>(1); // NOLINT
	int64_t null_count = 0;
	for (int64_t i = 0; i < span.length; i++) {
	const int8_t child_id = sparse_union_type->child_ids()[types[span.offset + i]];
	null_count += span.child_data[child_id].IsNull(i);

arrow/cpp/src/arrow/util/union_util.cc

Lines 43 to 53 in e980b7e

	int64_t LogicalDenseUnionNullCount(const ArraySpan& span) {
	const auto* dense_union_type = internal::checked_cast<const DenseUnionType*>(span.type);
	DCHECK_LE(span.child_data.size(), 128);
	const int8_t* types = span.GetValues<int8_t>(1); // NOLINT
	const int32_t* offsets = span.GetValues<int32_t>(2); // NOLINT
	int64_t null_count = 0;
	for (int64_t i = 0; i < span.length; i++) {
	const int8_t child_id = dense_union_type->child_ids()[types[span.offset + i]];
	const int32_t offset = offsets[span.offset + i];
	null_count += span.child_data[child_id].IsNull(offset);

LogicalSparseUnionNullCount uses types = span.GetValues<int8_t>(1), which already accounts for span.offset, but it still indexes types[span.offset + i]. It also checks child nullness with IsNull(i) instead of IsNull(span.offset + i).

Likewise, LogicalDenseUnionNullCount uses GetValues(...) pointers that already account for span.offset, but it still indexes types[span.offset + i] and offsets[span.offset + i].

As a result, sliced unions can report incorrect logical null counts, which affects count.

I verified the issue and a local fix, and I plan to open a PR shortly.

[pitrou]

Issue resolved by pull request 50114
#50114