GH-39565: [C++] Do not concatenate chunked values of fixed-width types to run "array_take"

[felipecrv]

Rationale for this change

Concatenating a chunked array into a single array before running the array_take kernels is very inefficient and can lead to out-of-memory crashes. See also #25822.

What changes are included in this PR?

• Implementation of kernels for "array_take" that can receive a ChunkedArray as values and produce an output without concatenating these chunks
• Improvements in the dispatching logic of TakeMetaFunction("take") to make "array_take" able to have a chunked_exec kernel for all types (some specialized and some based on concatenation)

Are these changes tested?

By existing tests. Some tests were added in previous PRs that introduced some of the infrastructure to support this.

• GitHub Issue: Do not concatenate ChunkedArray when running Take kernel #39565

[mapleFU]

May I ask a unrelated question, when would we call assert and when call DCHECK, since I think they would likely to be same?

[felipecrv]

May I ask a unrelated question, when would we call assert and when call DCHECK, since I think they would likely to be same?

We call assert in headers because we don't want to pay the cost of including logging.h everywhere. Think of assert as lighter-weight debug checks. But if you see an assert in a .cc file tell me to change it to DCHECK*.

[felipecrv]

The member variable is called exec_chunked but the type is called ChunkedExec (so confusing). In this PR I ended up sticking to chunked_exec. Once everything is reviewed and merged I could try to unify things to the direction people prefer.

[felipecrv]

@pitrou wouldn't it make sense to keep the responsibility for concatenation to a layer above the kernels? Like a query optimizer? They are in a better position to make memory/time trade-offs than the context-less kernel.

The worst regression (-81%) has the kernel still at 4G items/sec.

TakeChunkedFlatInt64FewRandomIndicesWithNulls
4.173G items/sec 761.551M items/sec   -81.751

I find it very inelegant to put these heuristics at the compute kernel level.

Imagine a pipeline trying to save on memory allocations by keeping the array chunked as much as possible and then a simple filter operation requires allocating enough memory to keep it all in memory.

Another case would be a pipeline where the caller is consolidating a big contiguous array for more operations than just array_take. They should be the ones concatenating.

[pitrou]

@pitrou wouldn't it make sense to keep the responsibility for concatenation to a layer above the kernels? Like a query optimizer? They are in a better position to make memory/time trade-offs than the context-less kernel.

Ideally perhaps. In practice this assumes that 1) there is a query optimizer 2) it has enough information about implementation details to make an informed decision.

In practice, take is probably often called directly in the context of PyArrow's sort methods. So this

arrow/python/pyarrow/table.pxi

Lines 1139 to 1143 in e62fbaa

	indices = _pc().sort_indices(
	self,
	options=_pc().SortOptions(sort_keys=[("", order)], **kwargs)
	)
	return self.take(indices)

The worst regression (-81%) has the kernel still at 4G items/sec.

I might be misreading, but this is the worst regression on the new benchmarks, right (those with a small selection factor)? On the benchmarks with a 1.0 selection factor (such as when sorting), the worst absolute results are around 25 Mitems/sec AFAICT. Or are those results obsolete?

Imagine a pipeline trying to save on memory allocations by keeping the array chunked as much as possible and then a simple filter operation requires allocating enough memory to keep it all in memory.

Well, currently "take" would always concatenate array chunks, so at least there is no regression in that regard.

Still, I understand the concern. We might want to expose an additional option to entirely disable concatenation when possible. But that might be overkill as well.

[felipecrv]

@pitrou what conditional checks should I add here to avoid regressions? I'm giving up on making the non-concatenation versions work well for integer arrays and want to merge this PR sooner rather than later and then start working on the string array implementation which is what will unlock most user value in the first place.

[pitrou]

By building on this (arguably simplified) analysis:

we're trading the concatenation of the chunked values (essentially allocating a new values array) against the resolution of many chunked indices (essentially allocating two new indices arrays). This is only beneficial if the value width is quite large (say a 256-byte FSB) or the number of indices is much smaller than the number of values.

and assuming the following known values:

• n_values: length of the values input
• n_indices: length of the indices input (governing the output length)
• value_width: byte width of the individual values

Then a simple heuristic could be to concatenate iff n_indices * 16 > n_values * value_width. This wouldn't take into account the larger computational cost associated with chunked indexing, but at least it would disable the chunked resolution approach when it doesn't make sense at all.

(btw, a moderate improvement could probably be achieved by using CompressedChunkLocation)