Function CreateAgg

cpp/src/arrow/engine/substrait/test_plan_builder.cc:125–165

}

Result<std::unique_ptr<substrait::AggregateRel>> CreateAgg(Id function_id,
                                                           const std::vector<int>& keys,
                                                           std::vector<int> arg_idxs,
                                                           const DataType& output_type,
                                                           ExtensionSet* ext_set) {
  auto agg = std::make_unique<substrait::AggregateRel>();

  if (!keys.empty()) {
    substrait::AggregateRel::Grouping* grouping = agg->add_groupings();
    for (int key : keys) {
      substrait::Expression* key_expr = grouping->add_grouping_expressions();
      CreateDirectReference(key, key_expr);
    }
  }

  substrait::AggregateRel::Measure* measure_wrapper = agg->add_measures();
  auto agg_func = std::make_unique<substrait::AggregateFunction>();
  ARROW_ASSIGN_OR_RAISE(uint32_t function_anchor, ext_set->EncodeFunction(function_id));

  agg_func->set_function_reference(function_anchor);

  for (int arg_idx : arg_idxs) {
    substrait::FunctionArgument* arg = agg_func->add_arguments();
    auto arg_expr = std::make_unique<substrait::Expression>();
    CreateDirectReference(arg_idx, arg_expr.get());
    arg->set_allocated_value(arg_expr.release());
  }

  agg_func->set_phase(substrait::AggregationPhase::AGGREGATION_PHASE_INITIAL_TO_RESULT);
  agg_func->set_invocation(
      substrait::AggregateFunction::AggregationInvocation::
          AggregateFunction_AggregationInvocation_AGGREGATION_INVOCATION_ALL);

  ARROW_ASSIGN_OR_RAISE(
      std::unique_ptr<substrait::Type> output_type_substrait,
      ToProto(output_type, /*nullable=*/true, ext_set, kPlanBuilderConversionOptions));
  agg_func->set_allocated_output_type(output_type_substrait.release());
  measure_wrapper->set_allocated_measure(agg_func.release());

  return agg;
}

std::unique_ptr<substrait::Version> CreateTestVersion() {
  auto version = std::make_unique<substrait::Version>();