more tests

apache · Aug 29, 2024 · 135227f · 135227f
1 parent cab4b25
commit 135227f
Showing 1 changed file with 37 additions and 23 deletions.
diff --git a/cpp/src/arrow/acero/aggregate_node_test.cc b/cpp/src/arrow/acero/aggregate_node_test.cc
@@ -215,29 +215,43 @@ TEST(ScalarAggregateNode, AnyAll) {
   // when min_count != 0.
   std::shared_ptr<Schema> in_schema = schema({field("not_used", int32())});
   std::shared_ptr<Schema> out_schema = schema({field("agg_out", boolean())});
-  std::vector<ExecBatch> batches{
-      ExecBatchFromJSON({int32()}, "[[42], [42], [42], [42]]")};
-  for (auto& func_name : {"any", "all"}) {
-    SCOPED_TRACE(func_name);
-    std::vector<Aggregate> aggregates = {
-        Aggregate(func_name,
-                  std::make_shared<compute::ScalarAggregateOptions>(/*skip_nulls=*/false,
-                                                                    /*min_count=*/2),
-                  FieldRef("literal_true"))};
-
-    // And a projection to make the input including a Scalar Boolean
-    Declaration plan = Declaration::Sequence(
-        {{"exec_batch_source", ExecBatchSourceNodeOptions(in_schema, batches)},
-         {"project", ProjectNodeOptions({literal(true)}, {"literal_true"})},
-         {"aggregate", AggregateNodeOptions(aggregates)}});
-
-    ASSERT_OK_AND_ASSIGN(BatchesWithCommonSchema out_batches,
-                         DeclarationToExecBatches(plan));
-
-    ExecBatch expected_batch = ExecBatchFromJSON({boolean()}, "[[true]]");
-
-    AssertExecBatchesEqualIgnoringOrder(out_schema, {expected_batch},
-                                        out_batches.batches);
+  struct AnyAllCase {
+    std::string batches_json;
+    Expression literal;
+    std::string expected_json;
+  };
+  std::vector<AnyAllCase> cases{
+      {"[[42], [42], [42], [42]]", literal(true), "[[true]]"},
+      {"[[42], [42], [42], [42]]", literal(false), "[[false]]"},
+      {"[[42], [42], [42], [42]]", literal(BooleanScalar{}), "[[null]]"},
+      {"[[42]]", literal(true), "[[null]]"},
+      {"[[42], [42], [42]]", literal(true), "[[true]]"},
+  };
+  for (const AnyAllCase& any_all_case : cases) {
+    for (const std::string& func_name : {"any", "all"}) {
+      std::vector<ExecBatch> batches{
+          ExecBatchFromJSON({int32()}, any_all_case.batches_json)};
+      std::vector<Aggregate> aggregates = {Aggregate(
+          "any",
+          std::make_shared<compute::ScalarAggregateOptions>(/*skip_nulls=*/false,
+                                                            /*min_count=*/2),
+          FieldRef("literal"))};
+
+      // And a projection to make the input including a Scalar Boolean
+      Declaration plan = Declaration::Sequence(
+          {{"exec_batch_source", ExecBatchSourceNodeOptions(in_schema, batches)},
+           {"project", ProjectNodeOptions({any_all_case.literal}, {"literal"})},
+           {"aggregate", AggregateNodeOptions(aggregates)}});
+
+      ASSERT_OK_AND_ASSIGN(BatchesWithCommonSchema out_batches,
+                           DeclarationToExecBatches(plan));
+
+      ExecBatch expected_batch =
+          ExecBatchFromJSON({boolean()}, any_all_case.expected_json);
+
+      AssertExecBatchesEqualIgnoringOrder(out_schema, {expected_batch},
+                                          out_batches.batches);
+    }
   }
 }