Method Merge

cpp/src/arrow/acero/groupby_aggregate_node.cc:255–298

}

Status GroupByNode::Merge() {
  arrow::util::tracing::Span span;
  START_COMPUTE_SPAN(span, "Merge",
                     {{"group_by", ToStringExtra(0)}, {"node.label", label()}});
  ThreadLocalState* state0 = &local_states_[0];
  for (size_t i = 1; i < local_states_.size(); ++i) {
    ThreadLocalState* state = &local_states_[i];
    if (!state->grouper) {
      continue;
    }

    ARROW_ASSIGN_OR_RAISE(ExecBatch other_keys, state->grouper->GetUniques());
    ARROW_ASSIGN_OR_RAISE(Datum transposition,
                          state0->grouper->Consume(ExecSpan(other_keys)));
    state->grouper.reset();

    for (size_t span_i = 0; span_i < agg_kernels_.size(); ++span_i) {
      arrow::util::tracing::Span span_item;
      START_COMPUTE_SPAN(
          span_item, aggs_[span_i].function,
          {{"function.name", aggs_[span_i].function},
           {"function.options",
            aggs_[span_i].options ? aggs_[span_i].options->ToString() : "<NULLPTR>"},
           {"function.kind", std::string(kind_name()) + "::Merge"}});

      auto ctx = plan_->query_context()->exec_context();
      KernelContext batch_ctx{ctx};
      DCHECK(state0->agg_states[span_i]);
      batch_ctx.SetState(state0->agg_states[span_i].get());

      // XXX this resizes each KernelState (state0->agg_states[span_i]) multiple times.
      // An alternative would be a two-pass algorithm:
      // 1. Compute all transpositions (one per local state) and the final number of
      // groups.
      // 2. Process all agg kernels, resizing each KernelState only once.
      RETURN_NOT_OK(
          agg_kernels_[span_i]->resize(&batch_ctx, state0->grouper->num_groups()));
      RETURN_NOT_OK(agg_kernels_[span_i]->merge(
          &batch_ctx, std::move(*state->agg_states[span_i]), *transposition.array()));
      state->agg_states[span_i].reset();
    }
  }
  return Status::OK();
}

Result<ExecBatch> GroupByNode::Finalize() {
  arrow::util::tracing::Span span;