GH-34315: [C++] Correct is_null kernel for Union and RunEndEncoded logical nulls

[jorisvandenbossche]

Rationale for this change

Currently the is_null kernel always returns all-False for union and run-end-encoded types, since those don't have a top-level validity buffer. Update the kernel to take the logical nulls into account.

Are these changes tested?

Yes, both in Python and C++

Are there any user-facing changes?

Yes, this changes (corrects) the behaviour of the is_null kernel.

• Closes: [Python] pa.compute.is_null() returns incorrect answer for dense union arrays and segfaults for dense union scalars #34315

[jorisvandenbossche]

cc @felipecrv (this is very much a draft (still needs more tests, doc comments, applying the same changes to is_valid, etc), but already would like to get your thoughts on the implementation approach)

[jorisvandenbossche]

I followed the same pattern as we had for SetNanBits (used just below, i.e. updating the previously created bitmap), not that this is necessarily the best approach though

[jorisvandenbossche]

This test is passing, except for the variant that runs it with scalars, because ArraySpan::FillFromScalar is not implemented for RunEndEncoded cc @felipecrv

[felipecrv]

I thought I had implemented that 🤔

[felipecrv]

Thanks for adding it! I'm gonna use it for something I'm doing as well.

[felipecrv]

this comment can now be removed as the if checks wrapping the null_count assignments document the context of each assignment

[felipecrv]

auto& to elide one refcount increment

[felipecrv]

This whole block can be moved into the else above to make it more clear that this only gets called after a complete zeroing of out_bitmap.

Now we have to know about a logical implication that is not explicit in the code: t \in {UNION, REE} \implies arr.GetNullCount() == 0. That is what in turn guarantees out_bitmap is zeroed before the Set*LogicalNullBits. Too much logic. 😅

[felipecrv]

The // Input has no nulls => output is entirely false. also needs to be more nuanced after that move.

[felipecrv]

nitpick: To be consistent with REE code elsewhere: RunEndType.

[felipecrv]

run_end_type (run_type is ambiguous -- could be the type of the value in the run).

[felipecrv]

length typo

[felipecrv]

I think the comment can be removed after the other changes that clarify the code are made.

[felipecrv]

FillRunEndsArrayFromScalar? s/type/run_end_type

[felipecrv]

std::move(type) and std::move(data_buffer)

[felipecrv]

I think it's better if the caller passes this->child_data[0] to this instead of the span. Making it obvious at the callsite how both child_data are filled.

[felipecrv]

You could derive the scratch_space parameter from that as well: making you use the child_data[0].scratch_space instead of the parent's scratch_space.

[felipecrv]

You're getting that stack trace you posted above from this new implementation?

I think this will be less risky if you use the this->child_data[0].scratch_space instead of this->scratch_space for the run-ends array buffer. I think that's the contract: an ArraySpan is free to use the scratch_space it directly owns.

[jorisvandenbossche]

You're getting that stack trace you posted above from this new implementation?

No, that was with the previous implementation. This last commit was to fix that

[jorisvandenbossche]

But thanks for the comments! Will clean-up

[westonpace]

I don't know if it's related or not but there is some danger with the way we handle scratch space (#35581 )

[felipecrv]

Hmmmm... ArraySpan is non-owning, so it won't keep this data available. Sorry for not noticing this before.

[felipecrv]

Feel free to set the members directly. One by one here. That skips allocating and deallocating an ArrayData which is the point of using ArraySpans in the first-place.

[felipecrv]

The run_end_type should be a pointer coming directly from the parent REE type, so don't pass it as a shared_ptr copy.

[jorisvandenbossche]

The same is done for Dictionary though:

      // Populate dictionary data
      const auto& dict_scalar = checked_cast<const DictionaryScalar&>(value);
      this->child_data.resize(1);
      this->child_data[0].SetMembers(*dict_scalar.value.dictionary->data());

https://github.com/apache/arrow/blob/fbe5f641d327ee81db00ce5f056940a69f4d8603/cpp/src/arrow/array/data.cc#L325C1-L328

Or is that fine because that is data owned by the scalar?

[felipecrv]

Or is that fine because that is data owned by the scalar?

Yes.

[pitrou]

AFAIK, it is not forbidden to have nested REE values. So this should be turned into an error return (perhaps Status::NotImplemented) at a higher level.

However, instead of refusing to implement this, you could use another strategy:

• call IsNull on the values array
• REE-decode the REE array comprised of (run ends, run values IsNull)

[felipecrv]

I've been meaning to send a PR forbidding nested REEs because it's pointless to run-end encode a values array since all the runs already have length 1 and the run-end arrays contains a strictly increasing sequence of indexes.

[pitrou]

@felipecrv I think you're talking about something else? "Nested" here means any type with child fields (e.g. list, struct...).

[jorisvandenbossche]

you could use another strategy:

That would indeed be a good alternative and would be more robust for whathever type is used for the REE values. I did a quick benchmark comparing both strategies in python:

In [2]: run_lengths = np.random.randint(1, 10, 100_000)

In [3]: run_values = [1, 2, 3, 4, None] * 20000

In [4]: arr = pa.RunEndEncodedArray.from_arrays(run_lengths.cumsum(), run_values)

In [5]: res1 = pc.is_null(arr)

In [6]: res2 = pc.run_end_decode(pa.RunEndEncodedArray.from_arrays(np.asarray(arr.run_ends), pc.is_null(arr.values)))

In [7]: res1.equals(res2)
Out[7]: True

In [8]: %timeit pc.is_null(arr)
309 µs ± 843 ns per loop (mean ± std. dev. of 7 runs, 1,000 loops each)

In [9]: %timeit pc.run_end_decode(pa.RunEndEncodedArray.from_arrays(np.asarray(arr.run_ends), pc.is_null(arr.values)))
1.07 ms ± 17.7 µs per loop (mean ± std. dev. of 7 runs, 1,000 loops each)

This is running with this branch (in release mode), so pc.is_null is using this PR's implementation, and the other is the python equivalent of what you propose (IIUC).
The alternative seems significantly slower, although I don't know how much of that is due to overhead of going through python several times.

[pitrou]

No idea, but the proposal here is to make things correct first ;-) It's also possible that run-end-decoding could be optimized...

[jorisvandenbossche]

Based on a quick profile, a significant part of the time is spent in the actual RunEndDecode impl, indicating it's not just from the python overhead. But it's also true that this implementation is not necessary optimized ;)

Now maybe the DCHECK(!is_nested(values.type->id())); isn't actually correct. What it is protecting from is that I am getting the first buffer of the data and assuming this is a validity bitmap. But that's only not true for REE/union, and not for "nested" types in general.

Although for supporting nan_is_null, that will still be easier to do through recursively calling the is_null kernel. Will take a look at doing it that way.

[felipecrv]

@pitrou I misunderstood "nested REE" as a REE inside another REE.

[pitrou]

Can you add a test with floating-point run end values and clearing/setting the nan_is_null option?

[pitrou]

Can you add floating-point child union values and test with/without the nan_is_null option set?

[pitrou]

Is it necessary to add these tests on the Python side? You have similar tests in C++ already.

[pitrou]

Are these changes related for this PR?

[pitrou]

If so, can you perhaps add a test for FillFromScalar on a run-end-encoded scalar? For example in arrow/array/array_run_end_test.cc.

[jorisvandenbossche]

Yes, the tests for scalar kernels automatically run the kernel also on actual scalars, and then this ends up creating an array from the scalar. So those changes to FillFromScalar are needed to be able to run any kernel on a REE scalar.

At the moment we don't actually have any test for FillFromScalar directly (also not for other types), only through testing the kernels on scalars.

[pitrou]

Suggested change

	} else if (is_union(type_id) || type_id == Type::RUN_END_ENCODED) {
	} else if (!HasValidityBitmap(type_id)) {

[pitrou]

We do, but they have the same semantics as dictionary arrays: the top-level validity refers to the indices, regardless of the underlying dictionary values.

[pitrou]

Nit: style

Suggested change

	auto& run_end_type = scalar.run_end_type();
	const auto& run_end_type = scalar.run_end_type();

[github-actions]

Thank you for your contribution. Unfortunately, this pull request has been marked as stale because it has had no activity in the past 365 days. Please remove the stale label or comment below, or this PR will be closed in 14 days. Feel free to re-open this if it has been closed in error. If you do not have repository permissions to reopen the PR, please tag a maintainer.