ARROW-13344: [R] Initial bindings for ExecPlan/ExecNode

cpp/src/arrow/compute/exec.cc

@@ -115,7 +115,7 @@ ExecBatch ExecBatch::Slice(int64_t offset, int64_t length) const {
     if (value.is_scalar()) continue;
     value = value.array()->Slice(offset, length);
   }
-  out.length = length;
+  out.length = std::min(length, this->length - offset);
   return out;
 }
 

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

@@ -719,11 +719,13 @@ struct ScalarAggregateNode : ExecNode {
   ScalarAggregateNode(ExecNode* input, std::string label,
                       std::shared_ptr<Schema> output_schema,
                       std::vector<const ScalarAggregateKernel*> kernels,
+                      std::vector<int> argument_indices,
                       std::vector<std::vector<std::unique_ptr<KernelState>>> states)
       : ExecNode(input->plan(), std::move(label), {input}, {"target"},
                  /*output_schema=*/std::move(output_schema),
                  /*num_outputs=*/1),
         kernels_(std::move(kernels)),
+        argument_indices_(std::move(argument_indices)),
         states_(std::move(states)) {}
 
   const char* kind_name() override { return "ScalarAggregateNode"; }
@@ -733,7 +735,7 @@ struct ScalarAggregateNode : ExecNode {
       KernelContext batch_ctx{plan()->exec_context()};
       batch_ctx.SetState(states_[i][thread_index].get());
 
-      ExecBatch single_column_batch{{batch.values[i]}, batch.length};
+      ExecBatch single_column_batch{{batch[argument_indices_[i]]}, batch.length};
       RETURN_NOT_OK(kernels_[i]->consume(&batch_ctx, single_column_batch));
     }
     return Status::OK();
@@ -807,7 +809,8 @@ struct ScalarAggregateNode : ExecNode {
   }
 
   Future<> finished_ = Future<>::MakeFinished();
-  std::vector<const ScalarAggregateKernel*> kernels_;
+  const std::vector<const ScalarAggregateKernel*> kernels_;
+  const std::vector<int> argument_indices_;
 
   std::vector<std::vector<std::unique_ptr<KernelState>>> states_;
 
@@ -816,20 +819,34 @@ struct ScalarAggregateNode : ExecNode {
 };
 
 Result<ExecNode*> MakeScalarAggregateNode(ExecNode* input, std::string label,
-                                          std::vector<internal::Aggregate> aggregates) {
-  if (input->output_schema()->num_fields() != static_cast<int>(aggregates.size())) {
-    return Status::Invalid("Provided ", aggregates.size(),
-                           " aggregates, expected one for each field of ",
-                           input->output_schema()->ToString());
+                                          std::vector<internal::Aggregate> aggregates,
+                                          std::vector<FieldRef> arguments,
+                                          std::vector<std::string> out_field_names) {
+  if (aggregates.size() != arguments.size()) {
+    return Status::Invalid("Provided ", aggregates.size(), " aggregates but ",
+                           arguments.size(), " arguments.");
+  }
+
+  if (aggregates.size() != out_field_names.size()) {
+    return Status::Invalid("Provided ", aggregates.size(), " aggregates but ",
+                           out_field_names.size(), " field names for the output.");
   }
 
   auto exec_ctx = input->plan()->exec_context();
 
   std::vector<const ScalarAggregateKernel*> kernels(aggregates.size());
   std::vector<std::vector<std::unique_ptr<KernelState>>> states(kernels.size());
   FieldVector fields(kernels.size());
+  std::vector<int> argument_indices(kernels.size());
 
   for (size_t i = 0; i < kernels.size(); ++i) {
+    if (!arguments[i].IsName()) {
+      return Status::NotImplemented("Non name field refs");
+    }
+    ARROW_ASSIGN_OR_RAISE(auto match,
+                          arguments[i].FindOneOrNone(*input->output_schema()));
+    argument_indices[i] = match[0];
+
     ARROW_ASSIGN_OR_RAISE(auto function,
                           exec_ctx->func_registry()->GetFunction(aggregates[i].function));
 
@@ -862,12 +879,12 @@ Result<ExecNode*> MakeScalarAggregateNode(ExecNode* input, std::string label,
     ARROW_ASSIGN_OR_RAISE(
         auto descr, kernels[i]->signature->out_type().Resolve(&kernel_ctx, {in_type}));
 
-    fields[i] = field(aggregates[i].function, std::move(descr.type));
+    fields[i] = field(std::move(out_field_names[i]), std::move(descr.type));
   }
 
   return input->plan()->EmplaceNode<ScalarAggregateNode>(
       input, std::move(label), schema(std::move(fields)), std::move(kernels),
-      std::move(states));
+      std::move(argument_indices), std::move(states));
 }
 
 namespace internal {

cpp/src/arrow/compute/exec/exec_plan.h

@@ -285,7 +285,9 @@ Result<ExecNode*> MakeProjectNode(ExecNode* input, std::string label,
 
 ARROW_EXPORT
 Result<ExecNode*> MakeScalarAggregateNode(ExecNode* input, std::string label,
-                                          std::vector<internal::Aggregate> aggregates);
+                                          std::vector<internal::Aggregate> aggregates,
+                                          std::vector<FieldRef> arguments,
+                                          std::vector<std::string> out_field_names);
 
 /// \brief Make a node which groups input rows based on key fields and computes
 /// aggregates for each group

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

@@ -531,9 +531,11 @@ TEST(ExecPlanExecution, SourceScalarAggSink) {
                        MakeTestSourceNode(plan.get(), "source", basic_data,
                                           /*parallel=*/false, /*slow=*/false));
 
-  ASSERT_OK_AND_ASSIGN(auto scalar_agg,
-                       MakeScalarAggregateNode(source, "scalar_agg",
-                                               {{"sum", nullptr}, {"any", nullptr}}));
+  ASSERT_OK_AND_ASSIGN(
+      auto scalar_agg,
+      MakeScalarAggregateNode(source, "scalar_agg", {{"sum", nullptr}, {"any", nullptr}},
+                              /*targets=*/{"i32", "bool"},
+                              /*out_field_names=*/{"sum(i32)", "any(bool)"}));
 
   auto sink_gen = MakeSinkNode(scalar_agg, "sink");
 
@@ -565,7 +567,8 @@ TEST(ExecPlanExecution, ScalarSourceScalarAggSink) {
   ASSERT_OK_AND_ASSIGN(
       auto scalar_agg,
       MakeScalarAggregateNode(source, "scalar_agg",
-                              {{"count", nullptr}, {"sum", nullptr}, {"mean", nullptr}}));
+                              {{"count", nullptr}, {"sum", nullptr}, {"mean", nullptr}},
+                              {"a", "b", "c"}, {"sum a", "sum b", "sum c"}));
 
   auto sink_gen = MakeSinkNode(scalar_agg, "sink");
 

cpp/src/arrow/dataset/scanner.cc

@@ -816,14 +816,15 @@ Result<int64_t> AsyncScanner::CountRows() {
   ARROW_ASSIGN_OR_RAISE(auto scan,
                         MakeScanNode(plan.get(), std::move(fragment_gen), options));
 
-  ARROW_ASSIGN_OR_RAISE(
-      auto get_selection,
-      compute::MakeProjectNode(scan, "get_selection", {options->filter}));
+  ARROW_ASSIGN_OR_RAISE(auto get_selection,
+                        compute::MakeProjectNode(scan, "get_selection", {options->filter},
+                                                 {"selection_mask"}));
 
   ARROW_ASSIGN_OR_RAISE(
       auto sum_selection,
       compute::MakeScalarAggregateNode(get_selection, "sum_selection",
-                                       {compute::internal::Aggregate{"sum", nullptr}}));
+                                       {compute::internal::Aggregate{"sum", nullptr}},
+                                       {"selection_mask"}, {"sum"}));
 
   AsyncGenerator<util::optional<compute::ExecBatch>> sink_gen =
       compute::MakeSinkNode(sum_selection, "sink");

cpp/src/arrow/dataset/scanner_test.cc

@@ -1471,14 +1471,16 @@ TEST(ScanNode, MinimalScalarAggEndToEnd) {
   ASSERT_OK_AND_ASSIGN(
       compute::ExecNode * sum,
       compute::MakeScalarAggregateNode(project, "scalar_agg",
-                                       {compute::internal::Aggregate{"sum", nullptr}}));
+                                       {compute::internal::Aggregate{"sum", nullptr}},
+                                       {a_times_2.ToString()}, {"a*2 sum"}));
 
   // finally, pipe the project node into a sink node
   auto sink_gen = compute::MakeSinkNode(sum, "sink");
 
   // translate sink_gen (async) to sink_reader (sync)
-  std::shared_ptr<RecordBatchReader> sink_reader = compute::MakeGeneratorReader(
-      schema({field("sum", int64())}), std::move(sink_gen), exec_context.memory_pool());
+  std::shared_ptr<RecordBatchReader> sink_reader =
+      compute::MakeGeneratorReader(schema({field("a*2 sum", int64())}),
+                                   std::move(sink_gen), exec_context.memory_pool());
 
   // start the ExecPlan
   ASSERT_OK(plan->StartProducing());
@@ -1489,9 +1491,9 @@ TEST(ScanNode, MinimalScalarAggEndToEnd) {
   // wait 1s for completion
   ASSERT_TRUE(plan->finished().Wait(/*seconds=*/1)) << "ExecPlan didn't finish within 1s";
 
-  auto expected = TableFromJSON(schema({field("sum", int64())}), {
-                                                                     R"([
-                                               {"sum": 4}
+  auto expected = TableFromJSON(schema({field("a*2 sum", int64())}), {
+                                                                         R"([
+                                               {"a*2 sum": 4}
                                           ])"});
   AssertTablesEqual(*expected, *collected, /*same_chunk_layout=*/false);
 }

r/DESCRIPTION

@@ -109,6 +109,7 @@ Collate:
     'metadata.R'
     'parquet.R'
     'python.R'
+    'query-engine.R'
     'record-batch-reader.R'
     'record-batch-writer.R'
     'reexports-bit64.R'