ARROW-14658: [C++] Add test for project-nested-then-aggregate

apache · Nov 15, 2021 · 7e5f8dd · 7e5f8dd
1 parent cb01f10
commit 7e5f8dd
Showing 1 changed file with 37 additions and 0 deletions.
diff --git a/cpp/src/arrow/compute/exec/plan_test.cc b/cpp/src/arrow/compute/exec/plan_test.cc
@@ -821,6 +821,43 @@ TEST(ExecPlanExecution, NestedSourceFilter) {
   }
 }
 
+TEST(ExecPlanExecution, NestedSourceProjectGroupedSum) {
+  for (bool parallel : {false, true}) {
+    SCOPED_TRACE(parallel ? "parallel/merged" : "serial");
+
+    auto input = MakeNestedBatches();
+    auto expected = ExecBatchFromJSON({int64(), boolean()}, R"([
+      [null, true],
+      [17, false],
+      [5, null]
+])");
+
+    ASSERT_OK_AND_ASSIGN(auto plan, ExecPlan::Make());
+    AsyncGenerator<util::optional<ExecBatch>> sink_gen;
+
+    ASSERT_OK(
+        Declaration::Sequence(
+            {
+                {"source",
+                 SourceNodeOptions{input.schema, input.gen(parallel, /*slow=*/false)}},
+                {"project", ProjectNodeOptions{{
+                                                   field_ref(FieldRef("struct", "i32")),
+                                                   field_ref(FieldRef("struct", "bool")),
+                                               },
+                                               {"i32", "bool"}}},
+                {"aggregate", AggregateNodeOptions{/*aggregates=*/{{"hash_sum", nullptr}},
+                                                   /*targets=*/{"i32"},
+                                                   /*names=*/{"sum(i32)"},
+                                                   /*keys=*/{"bool"}}},
+                {"sink", SinkNodeOptions{&sink_gen}},
+            })
+            .AddToPlan(plan.get()));
+
+    ASSERT_THAT(StartAndCollect(plan.get(), sink_gen),
+                Finishes(ResultWith(UnorderedElementsAreArray({expected}))));
+  }
+}
+
 TEST(ExecPlanExecution, SourceFilterProjectGroupedSumFilter) {
   for (bool parallel : {false, true}) {
     SCOPED_TRACE(parallel ? "parallel/merged" : "serial");