Method Consume

cpp/src/arrow/acero/groupby_aggregate_node.cc:210–253

}

Status GroupByNode::Consume(ExecSpan batch) {
  size_t thread_index = plan_->query_context()->GetThreadIndex();
  if (thread_index >= local_states_.size()) {
    return Status::IndexError("thread index ", thread_index, " is out of range [0, ",
                              local_states_.size(), ")");
  }

  auto state = &local_states_[thread_index];
  RETURN_NOT_OK(InitLocalStateIfNeeded(state));

  // Create a batch with key columns
  std::vector<ExecValue> keys(key_field_ids_.size());
  for (size_t i = 0; i < key_field_ids_.size(); ++i) {
    keys[i] = batch[key_field_ids_[i]];
  }
  ExecSpan key_batch(std::move(keys), batch.length);

  // Create a batch with group ids
  ARROW_ASSIGN_OR_RAISE(Datum id_batch, state->grouper->Consume(key_batch));

  // Execute aggregate kernels
  for (size_t i = 0; i < agg_kernels_.size(); ++i) {
    arrow::util::tracing::Span span;
    START_COMPUTE_SPAN(span, aggs_[i].function,
                       {{"function.name", aggs_[i].function},
                        {"function.options",
                         aggs_[i].options ? aggs_[i].options->ToString() : "<NULLPTR>"},
                        {"function.kind", std::string(kind_name()) + "::Consume"}});
    auto ctx = plan_->query_context()->exec_context();
    KernelContext kernel_ctx{ctx};
    kernel_ctx.SetState(state->agg_states[i].get());

    std::vector<ExecValue> column_values;
    for (const int field : agg_src_fieldsets_[i]) {
      column_values.push_back(batch[field]);
    }
    column_values.emplace_back(*id_batch.array());
    ExecSpan agg_batch(std::move(column_values), batch.length);
    RETURN_NOT_OK(agg_kernels_[i]->resize(&kernel_ctx, state->grouper->num_groups()));
    RETURN_NOT_OK(agg_kernels_[i]->consume(&kernel_ctx, agg_batch));
  }

  return Status::OK();
}

Status GroupByNode::Merge() {
  arrow::util::tracing::Span span;