GH-34626: [C++] Add ordered/segmented aggregation Substrait extension

cpp/proto/substrait/extension_rels.proto

@@ -58,3 +58,14 @@ message NamedTapRel {
   // If empty, field names will be automatically generated.
   repeated string columns = 3;
 }
+
+message SegmentedAggregateRel {
+  // Grouping keys of the aggregation
+  repeated substrait.Expression.FieldReference grouping_keys = 1;
+
+  // Segment keys of the aggregation
+  repeated substrait.Expression.FieldReference segment_keys = 2;
+
+  // A list of one or more aggregate expressions along with an optional filter.
+  repeated substrait.AggregateRel.Measure measures = 3;
+}

cpp/src/arrow/compute/exec/source_node.cc

@@ -28,8 +28,10 @@
 #include "arrow/compute/expression.h"
 #include "arrow/datum.h"
 #include "arrow/io/util_internal.h"
+#include "arrow/ipc/util.h"
 #include "arrow/result.h"
 #include "arrow/table.h"
+#include "arrow/util/align_util.h"
 #include "arrow/util/async_generator.h"
 #include "arrow/util/async_util.h"
 #include "arrow/util/checked_cast.h"
@@ -102,6 +104,13 @@ struct SourceNode : ExecNode, public TracedNode {
               batch_size = morsel_length;
             }
             ExecBatch batch = morsel.Slice(offset, batch_size);
+            for (auto& value : batch.values) {
+              if (value.is_array()) {
+                ARROW_ASSIGN_OR_RAISE(
+                    value, util::EnsureAlignment(value.make_array(), ipc::kArrowAlignment,
+                                                 default_memory_pool()));
+              }
+            }
             if (has_ordering) {
               batch.index = batch_index;
             }

cpp/src/arrow/engine/substrait/expression_internal.cc

@@ -138,6 +138,105 @@ std::string EnumToString(int value, const google::protobuf::EnumDescriptor* desc
   return value_desc->name();
 }
 
+Result<compute::Expression> FromProto(const substrait::Expression::ReferenceSegment* ref,
+                                      const ExtensionSet& ext_set,
+                                      const ConversionOptions& conversion_options,
+                                      std::optional<compute::Expression> in_expr) {
+  auto in_ref = ref;
+  auto& current = in_expr;
+  while (ref != nullptr) {
+    switch (ref->reference_type_case()) {
+      case substrait::Expression::ReferenceSegment::kStructField: {
+        auto index = ref->struct_field().field();
+        if (!current) {
+          // Root StructField (column selection)
+          current = compute::field_ref(FieldRef(index));
+        } else if (auto current_ref = current->field_ref()) {
+          // Nested StructFields on the root (selection of struct-typed column
+          // combined with selecting struct fields)
+          current = compute::field_ref(FieldRef(*current_ref, index));
+        } else if (current->call() && current->call()->function_name == "struct_field") {
+          // Nested StructFields on top of an arbitrary expression
+          auto* field_options =
+              checked_cast<compute::StructFieldOptions*>(current->call()->options.get());
+          field_options->field_ref = FieldRef(std::move(field_options->field_ref), index);
+        } else {
+          // First StructField on top of an arbitrary expression
+          current = compute::call("struct_field", {std::move(*current)},
+                                  arrow::compute::StructFieldOptions({index}));
+        }
+
+        // Segment handled, continue with child segment (if any)
+        if (ref->struct_field().has_child()) {
+          ref = &ref->struct_field().child();
+        } else {
+          ref = nullptr;
+        }
+        break;
+      }
+      case substrait::Expression::ReferenceSegment::kListElement: {
+        if (!current) {
+          // Root ListField (illegal)
+          return Status::Invalid(
+              "substrait::ListElement cannot take a Relation as an argument");
+        }
+
+        // ListField on top of an arbitrary expression
+        current = compute::call(
+            "list_element",
+            {std::move(*current), compute::literal(ref->list_element().offset())});
+
+        // Segment handled, continue with child segment (if any)
+        if (ref->list_element().has_child()) {
+          ref = &ref->list_element().child();
+        } else {
+          ref = nullptr;
+        }
+        break;
+      }
+      default:
+        // Unimplemented construct, break current of loop
+        current.reset();
+        ref = nullptr;
+    }
+  }
+  if (current) {
+    return *std::move(current);
+  }
+
+  return Status::NotImplemented(
+      "conversion to arrow::compute::Expression from Substrait reference segment: ",
+      in_ref ? in_ref->DebugString() : "null");
+}
+
+Result<compute::Expression> FromProto(const substrait::Expression::FieldReference* fref,
+                                      const ExtensionSet& ext_set,
+                                      const ConversionOptions& conversion_options,
+                                      std::optional<compute::Expression> in_expr) {
+  if (fref->reference_type_case() !=
+          substrait::Expression::FieldReference::kDirectReference ||
+      fref->root_type_case() != substrait::Expression::FieldReference::kRootReference) {
+    return Status::NotImplemented("substrait::FieldReference not direct root reference");
+  }
+  auto& dref = fref->direct_reference();
+  return FromProto(&dref, ext_set, conversion_options, std::move(in_expr));
+}
+
+Result<FieldRef> DirectReferenceFromProto(
+    const substrait::Expression::FieldReference* fref, const ExtensionSet& ext_set,
+    const ConversionOptions& conversion_options) {
+  ARROW_ASSIGN_OR_RAISE(compute::Expression expr,
+                        FromProto(fref, ext_set, conversion_options, {}));
+  const FieldRef* field_ref = expr.field_ref();
+  if (field_ref) {
+    return *field_ref;
+  } else {
+    return Status::Invalid(
+        "A direct reference was expected but a more complex expression was given "
+        "instead");
+  }
+}
+
 Result<SubstraitCall> FromProto(const substrait::AggregateFunction& func, bool is_hash,
                                 const ExtensionSet& ext_set,
                                 const ConversionOptions& conversion_options) {
@@ -193,67 +292,7 @@ Result<compute::Expression> FromProto(const substrait::Expression& expr,
       }
 
       const auto* ref = &expr.selection().direct_reference();
-      while (ref != nullptr) {
-        switch (ref->reference_type_case()) {
-          case substrait::Expression::ReferenceSegment::kStructField: {
-            auto index = ref->struct_field().field();
-            if (!out) {
-              // Root StructField (column selection)
-              out = compute::field_ref(FieldRef(index));
-            } else if (auto out_ref = out->field_ref()) {
-              // Nested StructFields on the root (selection of struct-typed column
-              // combined with selecting struct fields)
-              out = compute::field_ref(FieldRef(*out_ref, index));
-            } else if (out->call() && out->call()->function_name == "struct_field") {
-              // Nested StructFields on top of an arbitrary expression
-              auto* field_options =
-                  checked_cast<compute::StructFieldOptions*>(out->call()->options.get());
-              field_options->field_ref =
-                  FieldRef(std::move(field_options->field_ref), index);
-            } else {
-              // First StructField on top of an arbitrary expression
-              out = compute::call("struct_field", {std::move(*out)},
-                                  arrow::compute::StructFieldOptions({index}));
-            }
-
-            // Segment handled, continue with child segment (if any)
-            if (ref->struct_field().has_child()) {
-              ref = &ref->struct_field().child();
-            } else {
-              ref = nullptr;
-            }
-            break;
-          }
-          case substrait::Expression::ReferenceSegment::kListElement: {
-            if (!out) {
-              // Root ListField (illegal)
-              return Status::Invalid(
-                  "substrait::ListElement cannot take a Relation as an argument");
-            }
-
-            // ListField on top of an arbitrary expression
-            out = compute::call(
-                "list_element",
-                {std::move(*out), compute::literal(ref->list_element().offset())});
-
-            // Segment handled, continue with child segment (if any)
-            if (ref->list_element().has_child()) {
-              ref = &ref->list_element().child();
-            } else {
-              ref = nullptr;
-            }
-            break;
-          }
-          default:
-            // Unimplemented construct, break out of loop
-            out.reset();
-            ref = nullptr;
-        }
-      }
-      if (out) {
-        return *std::move(out);
-      }
-      break;
+      return FromProto(ref, ext_set, conversion_options, std::move(out));
     }
 
     case substrait::Expression::kIfThen: {

cpp/src/arrow/engine/substrait/expression_internal.h

@@ -20,6 +20,7 @@
 #pragma once
 
 #include <memory>
+#include <optional>
 
 #include "arrow/compute/type_fwd.h"
 #include "arrow/datum.h"
@@ -34,6 +35,15 @@ namespace engine {
 
 class SubstraitCall;
 
+ARROW_ENGINE_EXPORT
+Result<FieldRef> DirectReferenceFromProto(const substrait::Expression::FieldReference*,
+                                          const ExtensionSet&, const ConversionOptions&);
+
+ARROW_ENGINE_EXPORT
+Result<compute::Expression> FromProto(const substrait::Expression::FieldReference*,
+                                      const ExtensionSet&, const ConversionOptions&,
+                                      std::optional<compute::Expression>);
+
 ARROW_ENGINE_EXPORT
 Result<compute::Expression> FromProto(const substrait::Expression&, const ExtensionSet&,
                                       const ConversionOptions&);

cpp/src/arrow/engine/substrait/options.cc

@@ -75,6 +75,11 @@ class DefaultExtensionProvider : public BaseExtensionProvider {
       rel.UnpackTo(&named_tap_rel);
       return MakeNamedTapRel(conv_opts, inputs, named_tap_rel, ext_set);
     }
+    if (rel.Is<substrait_ext::SegmentedAggregateRel>()) {
+      substrait_ext::SegmentedAggregateRel seg_agg_rel;
+      rel.UnpackTo(&seg_agg_rel);
+      return MakeSegmentedAggregateRel(conv_opts, inputs, seg_agg_rel, ext_set);
+    }
     return Status::NotImplemented("Unrecognized extension in Susbstrait plan: ",
                                   rel.DebugString());
   }
@@ -165,6 +170,74 @@ class DefaultExtensionProvider : public BaseExtensionProvider {
                                                 named_tap_rel.name(), renamed_schema));
     return RelationInfo{{std::move(decl), std::move(renamed_schema)}, std::nullopt};
   }
+
+  Result<RelationInfo> MakeSegmentedAggregateRel(
+      const ConversionOptions& conv_opts, const std::vector<DeclarationInfo>& inputs,
+      const substrait_ext::SegmentedAggregateRel& seg_agg_rel,
+      const ExtensionSet& ext_set) {
+    if (inputs.size() != 1) {
+      return Status::Invalid(
+          "substrait_ext::SegmentedAggregateRel requires a single input but got: ",
+          inputs.size());
+    }
+    if (seg_agg_rel.segment_keys_size() == 0) {
+      return Status::Invalid(
+          "substrait_ext::SegmentedAggregateRel requires at least one segment key");
+    }
+
+    auto input_schema = inputs[0].output_schema;
+
+    // store key fields to be used when output schema is created
+    std::vector<int> key_field_ids;
+    std::vector<FieldRef> keys;
+    for (auto& ref : seg_agg_rel.grouping_keys()) {
+      ARROW_ASSIGN_OR_RAISE(auto field_ref,
+                            DirectReferenceFromProto(&ref, ext_set, conv_opts));
+      ARROW_ASSIGN_OR_RAISE(auto match, field_ref.FindOne(*input_schema));
+      key_field_ids.emplace_back(std::move(match[0]));
+      keys.emplace_back(std::move(field_ref));
+    }
+
+    // store segment key fields to be used when output schema is created
+    std::vector<int> segment_key_field_ids;
+    std::vector<FieldRef> segment_keys;
+    for (auto& ref : seg_agg_rel.segment_keys()) {
+      ARROW_ASSIGN_OR_RAISE(auto field_ref,
+                            DirectReferenceFromProto(&ref, ext_set, conv_opts));
+      ARROW_ASSIGN_OR_RAISE(auto match, field_ref.FindOne(*input_schema));
+      segment_key_field_ids.emplace_back(std::move(match[0]));
+      segment_keys.emplace_back(std::move(field_ref));
+    }
+
+    std::vector<compute::Aggregate> aggregates;
+    aggregates.reserve(seg_agg_rel.measures_size());
+    std::vector<std::vector<int>> agg_src_fieldsets;
+    agg_src_fieldsets.reserve(seg_agg_rel.measures_size());
+    for (auto agg_measure : seg_agg_rel.measures()) {
+      ARROW_ASSIGN_OR_RAISE(
+          auto parsed_measure,
+          internal::ParseAggregateMeasure(agg_measure, ext_set, conv_opts,
+                                          /*is_hash=*/!keys.empty(), input_schema));
+      aggregates.push_back(std::move(parsed_measure.aggregate));
+      agg_src_fieldsets.push_back(std::move(parsed_measure.fieldset));
+    }
+
+    ARROW_ASSIGN_OR_RAISE(auto decl_info,
+                          internal::MakeAggregateDeclaration(
+                              std::move(inputs[0].declaration), std::move(input_schema),
+                              seg_agg_rel.measures_size(), std::move(aggregates),
+                              std::move(agg_src_fieldsets), std::move(keys),
+                              std::move(key_field_ids), std::move(segment_keys),
+                              std::move(segment_key_field_ids), ext_set, conv_opts));
+
+    const auto& output_schema = decl_info.output_schema;
+    size_t out_size = output_schema->num_fields();
+    std::vector<int> field_output_indices(out_size);
+    for (int i = 0; i < static_cast<int>(out_size); i++) {
+      field_output_indices[i] = i;
+    }
+    return RelationInfo{decl_info, std::move(field_output_indices)};
+  }
 };
 
 namespace {