You are viewing a plain text version of this content. The canonical link for it is here.
Posted to github@arrow.apache.org by GitBox <gi...@apache.org> on 2021/08/24 13:26:12 UTC

[GitHub] [arrow] lidavidm commented on a change in pull request #10985: ARROW-13681: [C++] Fix list_parent_indices behaviour on chunked array

lidavidm commented on a change in pull request #10985:
URL: https://github.com/apache/arrow/pull/10985#discussion_r694844184



##########
File path: cpp/src/arrow/compute/kernels/vector_nested.cc
##########
@@ -77,24 +146,53 @@ const FunctionDoc list_parent_indices_doc(
      "is emitted."),
     {"lists"});
 
+class ListParentIndicesFunction : public MetaFunction {
+ public:
+  ListParentIndicesFunction()
+      : MetaFunction("list_parent_indices", Arity::Unary(), &list_parent_indices_doc) {}
+
+  Result<Datum> ExecuteImpl(const std::vector<Datum>& args,
+                            const FunctionOptions* options,
+                            ExecContext* ctx) const override {
+    KernelContext kernel_ctx(ctx);
+    switch (args[0].kind()) {
+      case Datum::ARRAY:
+        return ListParentIndicesArray::Exec(&kernel_ctx, args[0].array());
+      case Datum::CHUNKED_ARRAY: {
+        const auto& input = args[0].chunked_array();
+        ARROW_ASSIGN_OR_RAISE(auto out_ty, ListParentIndicesType(*input->type()));
+
+        int64_t base_output_offset = 0;
+        ArrayVector out_chunks;
+        for (const auto& chunk : input->chunks()) {
+          ARROW_ASSIGN_OR_RAISE(auto out_chunk,
+                                ListParentIndicesArray::Exec(&kernel_ctx, chunk->data(),
+                                                             base_output_offset));
+          out_chunks.push_back(MakeArray(std::move(out_chunk)));
+          base_output_offset += chunk->length();
+        }
+        return std::make_shared<ChunkedArray>(std::move(out_chunks), std::move(out_ty));
+      }
+      default:
+        return Status::NotImplemented(
+            "Unsupported input type for function 'list_parent_indices': ",
+            args[0].ToString());
+    }
+  }
+};
+
 }  // namespace
 
 void RegisterVectorNested(FunctionRegistry* registry) {
   auto flatten =
       std::make_shared<VectorFunction>("list_flatten", Arity::Unary(), &list_flatten_doc);
-  DCHECK_OK(flatten->AddKernel({InputType::Array(Type::LIST)}, OutputType(ValuesType),
+  DCHECK_OK(flatten->AddKernel({InputType::Array(Type::LIST)}, OutputType(ListValuesType),
                                ListFlatten<ListType>));
   DCHECK_OK(flatten->AddKernel({InputType::Array(Type::LARGE_LIST)},
-                               OutputType(ValuesType), ListFlatten<LargeListType>));
+                               OutputType(ListValuesType), ListFlatten<LargeListType>));
   DCHECK_OK(registry->AddFunction(std::move(flatten)));
 
-  auto list_parent_indices = std::make_shared<VectorFunction>(
-      "list_parent_indices", Arity::Unary(), &list_parent_indices_doc);
-  DCHECK_OK(list_parent_indices->AddKernel({InputType::Array(Type::LIST)}, int32(),
-                                           ListParentIndices<ListType>));
-  DCHECK_OK(list_parent_indices->AddKernel({InputType::Array(Type::LARGE_LIST)}, int64(),
-                                           ListParentIndices<LargeListType>));

Review comment:
       The old ListParentIndices implementation can be removed, no?




-- 
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-unsubscribe@arrow.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org