felipecrv commented on code in PR #41700:
URL: https://github.com/apache/arrow/pull/41700#discussion_r1685742964
##########
cpp/src/arrow/compute/kernels/vector_selection_take_internal.cc:
##########
@@ -715,12 +853,244 @@ 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 take_aaa_exec(ctx, exec_span, &result);
+}
+
+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`).
Review Comment:
Yes, when you have only a hand-written AAA kernel, it uses `Concatenate` to
make the other ones.
--
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: github-unsubscr...@arrow.apache.org
For queries about this service, please contact Infrastructure at:
us...@infra.apache.org
