ARROW-13540: [C++] Refactor and clean up sink node

Author: lidavidm

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

@@ -115,8 +115,7 @@ class ARROW_EXPORT SinkNodeOptions : public ExecNodeOptions {
 /// \brief Make a node which sorts rows passed through it
 ///
 /// All batches pushed to this node will be accumulated, then sorted, by the given
-/// fields. Then sorted batches will be pushed to the next node, along a tag
-/// indicating the absolute order of the batches.
+/// fields. Then sorted batches will be forwarded to the generator in sorted order.
 class ARROW_EXPORT OrderBySinkNodeOptions : public SinkNodeOptions {
  public:
   explicit OrderBySinkNodeOptions(

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

@@ -25,6 +25,7 @@
 #include "arrow/compute/exec/expression.h"
 #include "arrow/compute/exec/options.h"
 #include "arrow/compute/exec/test_util.h"
+#include "arrow/compute/exec/util.h"
 #include "arrow/record_batch.h"
 #include "arrow/table.h"
 #include "arrow/testing/future_util.h"
@@ -37,6 +38,7 @@
 #include "arrow/util/vector.h"
 
 using testing::ElementsAre;
+using testing::ElementsAreArray;
 using testing::HasSubstr;
 using testing::Optional;
 using testing::UnorderedElementsAreArray;
@@ -328,7 +330,7 @@ TEST(ExecPlanExecution, SourceOrderBy) {
                     .AddToPlan(plan.get()));
 
       ASSERT_THAT(StartAndCollect(plan.get(), sink_gen),
-                  Finishes(ResultWith(::testing::ElementsAreArray(expected))));
+                  Finishes(ResultWith(ElementsAreArray(expected))));
     }
   }
 }
@@ -414,18 +416,9 @@ TEST(ExecPlanExecution, StressSourceOrderBy) {
       // Check that data is sorted appropriately
       ASSERT_FINISHES_OK_AND_ASSIGN(auto exec_batches,
                                     StartAndCollect(plan.get(), sink_gen));
-      RecordBatchVector batches, original_batches;
-      for (const auto& batch : exec_batches) {
-        ASSERT_OK_AND_ASSIGN(auto rb, batch.ToRecordBatch(input_schema));
-        batches.push_back(std::move(rb));
-      }
-      for (const auto& batch : random_data.batches) {
-        ASSERT_OK_AND_ASSIGN(auto rb, batch.ToRecordBatch(input_schema));
-        original_batches.push_back(std::move(rb));
-      }
-      ASSERT_OK_AND_ASSIGN(auto actual, Table::FromRecordBatches(input_schema, batches));
+      ASSERT_OK_AND_ASSIGN(auto actual, TableFromExecBatches(input_schema, exec_batches));
       ASSERT_OK_AND_ASSIGN(auto original,
-                           Table::FromRecordBatches(input_schema, original_batches));
+                           TableFromExecBatches(input_schema, random_data.batches));
       ASSERT_OK_AND_ASSIGN(auto sort_indices, SortIndices(original, options));
       ASSERT_OK_AND_ASSIGN(auto expected, Take(original, sort_indices));
       AssertTablesEqual(*actual, *expected.table());

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

@@ -187,41 +187,34 @@ struct OrderBySinkNode final : public SinkNode {
   }
 
  protected:
-  Result<std::shared_ptr<Table>> SortData() {
-    std::unique_lock<std::mutex> lock(mutex_);
-    ARROW_ASSIGN_OR_RAISE(
-        auto table,
-        Table::FromRecordBatches(inputs_[0]->output_schema(), std::move(batches_)));
-    ARROW_ASSIGN_OR_RAISE(auto indices,
-                          SortIndices(table, sort_options_, plan()->exec_context()));
-    ARROW_ASSIGN_OR_RAISE(auto sorted, Take(table, indices, TakeOptions::NoBoundsCheck(),
-                                            plan()->exec_context()));
-    return sorted.table();
-  }
-
-  void Finish() override {
-    auto maybe_sorted = SortData();
-    if (ErrorIfNotOk(maybe_sorted.status())) {
-      producer_.Push(maybe_sorted.status());
-      SinkNode::Finish();
-      return;
+  Status DoFinish() {
+    Datum sorted;
+    {
+      std::unique_lock<std::mutex> lock(mutex_);
+      ARROW_ASSIGN_OR_RAISE(
+          auto table,
+          Table::FromRecordBatches(inputs_[0]->output_schema(), std::move(batches_)));
+      ARROW_ASSIGN_OR_RAISE(auto indices,
+                            SortIndices(table, sort_options_, plan()->exec_context()));
+      ARROW_ASSIGN_OR_RAISE(sorted, Take(table, indices, TakeOptions::NoBoundsCheck(),
+                                         plan()->exec_context()));
     }
-    auto sorted = maybe_sorted.MoveValueUnsafe();
-
-    TableBatchReader reader(*sorted);
+    TableBatchReader reader(*sorted.table());
     while (true) {
       std::shared_ptr<RecordBatch> batch;
-      auto status = reader.ReadNext(&batch);
-      if (!status.ok()) {
-        producer_.Push(std::move(status));
-        SinkNode::Finish();
-        return;
-      }
+      RETURN_NOT_OK(reader.ReadNext(&batch));
       if (!batch) break;
       bool did_push = producer_.Push(ExecBatch(*batch));
       if (!did_push) break;  // producer_ was Closed already
     }
+    return Status::OK();
+  }
 
+  void Finish() override {
+    Status st = DoFinish();
+    if (ErrorIfNotOk(st)) {
+      producer_.Push(std::move(st));
+    }
     SinkNode::Finish();
   }
 

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

@@ -18,6 +18,7 @@
 #include "arrow/compute/exec/util.h"
 
 #include "arrow/compute/exec/exec_plan.h"
+#include "arrow/table.h"
 #include "arrow/util/bit_util.h"
 #include "arrow/util/bitmap_ops.h"
 #include "arrow/util/ubsan.h"
@@ -296,5 +297,15 @@ Status ValidateExecNodeInputs(ExecPlan* plan, const std::vector<ExecNode*>& inpu
   return Status::OK();
 }
 
+Result<std::shared_ptr<Table>> TableFromExecBatches(
+    const std::shared_ptr<Schema>& schema, const std::vector<ExecBatch>& exec_batches) {
+  RecordBatchVector batches;
+  for (const auto& batch : exec_batches) {
+    ARROW_ASSIGN_OR_RAISE(auto rb, batch.ToRecordBatch(schema));
+    batches.push_back(std::move(rb));
+  }
+  return Table::FromRecordBatches(schema, batches);
+}
+
 }  // namespace compute
 }  // namespace arrow

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

@@ -188,6 +188,10 @@ ARROW_EXPORT
 Status ValidateExecNodeInputs(ExecPlan* plan, const std::vector<ExecNode*>& inputs,
                               int expected_num_inputs, const char* kind_name);
 
+ARROW_EXPORT
+Result<std::shared_ptr<Table>> TableFromExecBatches(
+    const std::shared_ptr<Schema>& schema, const std::vector<ExecBatch>& exec_batches);
+
 class AtomicCounter {
  public:
   AtomicCounter() = default;