[GitHub] [arrow] lidavidm commented on a change in pull request #10358: ARROW-2665: [C++][Python] Add index() kernel

[GitBox <gi...@apache.org>]

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



##########
File path: cpp/src/arrow/compute/kernels/aggregate_basic.cc
##########
@@ -218,6 +218,65 @@ Result<std::unique_ptr<KernelState>> AllInit(KernelContext*, const KernelInitArg
   return ::arrow::internal::make_unique<BooleanAllImpl>();
 }
 
+// ----------------------------------------------------------------------
+// Index implementation
+
+struct IndexImpl : public ScalarAggregator {
+  explicit IndexImpl(IndexOptions options, int64_t seen, int64_t index)
+      : options(std::move(options)), seen{seen}, index{index} {}
+
+  Status Consume(KernelContext* ctx, const ExecBatch& batch) override {
+    // short-circuit
+    if (index >= 0 || !options.value->is_valid) {
+      return Status::OK();
+    }
+
+    const auto& data = *batch[0].array();
+    seen = data.length;
+    ARROW_ASSIGN_OR_RAISE(
+        auto result, CallFunction("equal", {data, options.value}, ctx->exec_context()));

Review comment:
       Yup, that is fair. I figured for now I didn't want to create a lot of generated code for this use case. It also at least helps that it'll short-circuit. 




-- 
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.

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