felipecrv commented on code in PR #41700:
URL: https://github.com/apache/arrow/pull/41700#discussion_r1604325538
##########
cpp/src/arrow/compute/kernels/vector_selection_take_internal.cc:
##########
@@ -823,6 +732,129 @@ class TakeMetaFunction : public MetaFunction {
TakeMetaFunction()
: MetaFunction("take", Arity::Binary(), take_doc,
GetDefaultTakeOptions()) {}
+ static Result<Datum> CallArrayTake(const std::vector<Datum>& args,
+ const TakeOptions& options, ExecContext*
ctx) {
+ ARROW_ASSIGN_OR_RAISE(auto array_take_func,
+ ctx->func_registry()->GetFunction("array_take"));
+ return array_take_func->Execute(args, &options, ctx);
+ }
+
+ static Result<std::shared_ptr<ArrayData>> TakeAAA(const std::vector<Datum>&
args,
+ const TakeOptions& options,
+ ExecContext* ctx) {
+ DCHECK_EQ(args[0].kind(), Datum::ARRAY);
+ DCHECK_EQ(args[1].kind(), Datum::ARRAY);
+ ARROW_ASSIGN_OR_RAISE(Datum result, CallArrayTake(args, options, ctx));
+ return result.array();
+ }
+
+ static 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);
+ }
+ }
+
+ static Result<std::shared_ptr<ArrayData>> TakeCAA(
+ const std::shared_ptr<ChunkedArray>& values, const Array& indices,
+ const TakeOptions& options, ExecContext* ctx) {
+ // "array_take" can handle CA->A cases directly
+ // (via their VectorKernel::chunked_exec)
+ ARROW_ASSIGN_OR_RAISE(auto result, CallArrayTake({values, indices},
options, ctx));
+ return result.array();
+ }
+
+ static Result<std::shared_ptr<ChunkedArray>> TakeCAC(
+ const std::shared_ptr<ChunkedArray>& values, const Array& indices,
+ const TakeOptions& options, ExecContext* ctx) {
+ ARROW_ASSIGN_OR_RAISE(auto new_chunk, TakeCAA(values, indices, options,
ctx));
+ return std::make_shared<ChunkedArray>(MakeArray(std::move(new_chunk)));
+ }
+
+ static Result<std::shared_ptr<ChunkedArray>> TakeCCC(
+ const std::shared_ptr<ChunkedArray>& values,
+ const std::shared_ptr<ChunkedArray>& indices, const TakeOptions& options,
+ ExecContext* ctx) {
+ // XXX: for every chunk in indices, values are gathered from all chunks in
values to
+ // form a new chunk in the result. Performing this concatenation is not
ideal, but
+ // greatly simplifies the implementation before something more efficient is
+ // implemented.
+ ARROW_ASSIGN_OR_RAISE(auto values_array,
+ ChunkedArrayAsArray(values, ctx->memory_pool()));
+ std::vector<Datum> args = {std::move(values_array), {}};
+ std::vector<std::shared_ptr<Array>> new_chunks;
+ new_chunks.resize(indices->num_chunks());
+ for (int i = 0; i < indices->num_chunks(); i++) {
+ args[1] = indices->chunk(i);
+ // XXX: this loop can use TakeCAA once it can handle ChunkedArray
+ // without concatenating first
Review Comment:
I'm working on this one.
--
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]
