ARROW-14008: [R][Compute] Running an ExecPlan should yield Reader ins…

…tead of Table This only modifies the binding (c++) so that returning a RecordBatchReader is possible; the R API continues to collect the stream of batches into a Table. Closes apache#11213 from bkietz/14008-ExecPlan-run-should-retur Authored-by: Benjamin Kietzman <[email protected]> Signed-off-by: Benjamin Kietzman <[email protected]>
davidhao3300 · Sep 23, 2021 · eec7b55 · eec7b55
1 parent dd6c325
commit eec7b55
Show file tree

Hide file tree

Showing 3 changed files with 18 additions and 8 deletions.
diff --git a/r/R/query-engine.R b/r/R/query-engine.R
@@ -18,7 +18,7 @@
 do_exec_plan <- function(.data) {
   plan <- ExecPlan$create()
   final_node <- plan$Build(.data)
-  tab <- plan$Run(final_node)
+  tab <- plan$Run(final_node)$read_table()
 
   # If arrange() created $temp_columns, make sure to omit them from the result
   # We can't currently handle this in the ExecPlan itself because sorting

diff --git a/r/src/arrowExports.cpp b/r/src/arrowExports.cpp
diff --git a/r/src/compute-exec.cpp b/r/src/compute-exec.cpp
@@ -55,7 +55,7 @@ std::shared_ptr<compute::ExecNode> MakeExecNodeOrStop(
 }
 
 // [[arrow::export]]
-std::shared_ptr<arrow::Table> ExecPlan_run(
+std::shared_ptr<arrow::RecordBatchReader> ExecPlan_run(
     const std::shared_ptr<compute::ExecPlan>& plan,
     const std::shared_ptr<compute::ExecNode>& final_node, cpp11::list sort_options) {
   // For now, don't require R to construct SinkNodes.
@@ -77,11 +77,21 @@ std::shared_ptr<arrow::Table> ExecPlan_run(
   StopIfNotOk(plan->Validate());
   StopIfNotOk(plan->StartProducing());
 
-  std::shared_ptr<arrow::RecordBatchReader> sink_reader = compute::MakeGeneratorReader(
-      final_node->output_schema(), std::move(sink_gen), gc_memory_pool());
-
-  plan->finished().Wait();
-  return ValueOrStop(arrow::Table::FromRecordBatchReader(sink_reader.get()));
+  // If the generator is destroyed before being completely drained, inform plan
+  std::shared_ptr<void> stop_producing{nullptr, [plan](...) {
+                                         bool not_finished_yet =
+                                             plan->finished().TryAddCallback([&plan] {
+                                               return [plan](const arrow::Status&) {};
+                                             });
+
+                                         if (not_finished_yet) {
+                                           plan->StopProducing();
+                                         }
+                                       }};
+
+  return compute::MakeGeneratorReader(
+      final_node->output_schema(),
+      [stop_producing, plan, sink_gen] { return sink_gen(); }, gc_memory_pool());
 }
 
 #if defined(ARROW_R_WITH_DATASET)