felipecrv commented on code in PR #41700:
URL: https://github.com/apache/arrow/pull/41700#discussion_r1723505610
##########
cpp/src/arrow/compute/kernels/vector_selection_take_internal.cc:
##########
@@ -715,12 +853,247 @@ class TakeMetaFunction : public MetaFunction {
return Status::NotImplemented(
"Unsupported types for take operation: "
"values=",
- args[0].ToString(), "indices=", args[1].ToString());
+ args[0].ToString(), ", indices=", args[1].ToString());
}
};
// ----------------------------------------------------------------------
+/// \brief Prepare the output array like ExecuteArrayKernel::PrepareOutput()
+std::shared_ptr<ArrayData> PrepareOutput(const ExecBatch& batch, int64_t
length) {
+ DCHECK_EQ(batch.length, length);
+ auto out = std::make_shared<ArrayData>(batch.values[0].type(), length);
+ out->buffers.resize(batch.values[0].type()->layout().buffers.size());
+ return out;
+}
+
+Result<std::shared_ptr<Array>> ChunkedArrayAsArray(
+ const std::shared_ptr<ChunkedArray>& values, MemoryPool* pool) {
+ switch (values->num_chunks()) {
+ case 0:
+ return MakeArrayOfNull(values->type(), /*length=*/0, pool);
+ case 1:
+ return values->chunk(0);
+ default:
+ return Concatenate(values->chunks(), pool);
+ }
+}
+
+Status CallAAAKernel(ArrayKernelExec take_aaa_exec, KernelContext* ctx,
+ std::shared_ptr<ArrayData> values,
+ std::shared_ptr<ArrayData> indices, Datum* out) {
+ int64_t batch_length = values->length;
+ std::vector<Datum> args = {std::move(values), std::move(indices)};
+ ExecBatch array_array_batch(std::move(args), batch_length);
+ DCHECK_EQ(out->kind(), Datum::ARRAY);
+ ExecSpan exec_span{array_array_batch};
+ ExecResult result;
+ result.value = out->array();
+ RETURN_NOT_OK(take_aaa_exec(ctx, exec_span, &result));
+ DCHECK(result.is_array_data());
+ out->value = result.array_data();
+ return Status::OK();
+}
+
+Status CallCAAKernel(VectorKernel::ChunkedExec take_caa_exec, KernelContext*
ctx,
+ std::shared_ptr<ChunkedArray> values,
+ std::shared_ptr<ArrayData> indices, Datum* out) {
+ int64_t batch_length = values->length();
+ std::vector<Datum> args = {std::move(values), std::move(indices)};
+ ExecBatch chunked_array_array_batch(std::move(args), batch_length);
+ DCHECK_EQ(out->kind(), Datum::ARRAY);
+ return take_caa_exec(ctx, chunked_array_array_batch, out);
+}
+
+Status TakeACCChunkedExec(ArrayKernelExec take_aaa_exec, KernelContext* ctx,
+ const ExecBatch& batch, Datum* out) {
+ auto& values = batch.values[0].array();
+ auto& indices = batch.values[1].chunked_array();
+ auto num_chunks = indices->num_chunks();
+ std::vector<std::shared_ptr<Array>> new_chunks(num_chunks);
+ for (int i = 0; i < num_chunks; i++) {
+ // Take with that indices chunk
+ auto& indices_chunk = indices->chunk(i)->data();
+ Datum result = PrepareOutput(batch, values->length);
+ RETURN_NOT_OK(CallAAAKernel(take_aaa_exec, ctx, values, indices_chunk,
&result));
+ new_chunks[i] = MakeArray(result.array());
+ }
+ out->value = std::make_shared<ChunkedArray>(std::move(new_chunks),
values->type);
+ return Status::OK();
+}
+
+/// \brief Generic (slower) VectorKernel::exec_chunked (`CA->C`, `CC->C`, and
`AC->C`).
+///
+/// This function concatenates the chunks of the values and then calls the
`AA->A` take
+/// kernel to handle the `CA->C` cases. The ArrayData returned by the `AA->A`
kernel is
+/// converted to a ChunkedArray with a single chunk to honor the `CA->C`
contract.
+///
+/// For `CC->C` cases, it concatenates the chunks of the values and calls the
`AA->A` take
+/// kernel for each chunk of the indices, producing a new chunked array with
the same
+/// shape as the indices.
+///
+/// `AC->C` cases are trivially delegated to TakeACCChunkedExec without any
concatenation.
+///
+/// \param take_aaa_exec The `AA->A` take kernel to use.
+Status GenericTakeChunkedExec(ArrayKernelExec take_aaa_exec, KernelContext*
ctx,
+ const ExecBatch& batch, Datum* out) {
+ const auto& args = batch.values;
+ if (args[0].kind() == Datum::CHUNKED_ARRAY) {
+ auto& values_chunked = args[0].chunked_array();
Review Comment:
No problem. Even though `const` is inferred here and wouldn't be on line 940.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
