TeeNode preserve order

Author: Rafał Hibner

cpp/src/arrow/dataset/file_base.cc

@@ -17,6 +17,7 @@
 
 #include "arrow/dataset/file_base.h"
 
+#include "arrow/acero/accumulation_queue.h"
 #include "arrow/acero/exec_plan.h"
 
 #include <algorithm>
@@ -555,24 +556,30 @@ Result<acero::ExecNode*> MakeWriteNode(acero::ExecPlan* plan,
       auto node,
       // to preserve order explicitly sequence the exec batches
       // this requires exec batch index to be set upstream (e.g. by SourceNode)
-      acero::MakeExecNode("consuming_sink", plan, std::move(inputs),
-                          acero::ConsumingSinkNodeOptions{
-                              std::move(consumer),
-                              {},
-                              /*sequence_output=*/write_options.preserve_order}));
+      acero::MakeExecNode(
+          "consuming_sink", plan, std::move(inputs),
+          acero::ConsumingSinkNodeOptions{
+              std::move(consumer),
+              {},
+              /*sequence_output=*/write_node_options.write_options.preserve_order}));
 
   return node;
 }
 
 namespace {
 
-class TeeNode : public acero::MapNode {
+class TeeNode : public acero::MapNode,
+                public arrow::acero::util::SerialSequencingQueue::Processor {
  public:
   TeeNode(acero::ExecPlan* plan, std::vector<acero::ExecNode*> inputs,
           std::shared_ptr<Schema> output_schema,
           FileSystemDatasetWriteOptions write_options)
       : MapNode(plan, std::move(inputs), std::move(output_schema)),
-        write_options_(std::move(write_options)) {}
+        write_options_(std::move(write_options)) {
+    if (write_options.preserve_order) {
+      sequencer_ = acero::util::SerialSequencingQueue::Make(this);
+    }
+  }
 
   Status StartProducing() override {
     ARROW_ASSIGN_OR_RAISE(
@@ -602,6 +609,18 @@ class TeeNode : public acero::MapNode {
 
   const char* kind_name() const override { return "TeeNode"; }
 
+  Status InputReceived(ExecNode* input, ExecBatch batch) override {
+    DCHECK_EQ(input, inputs_[0]);
+    if (sequencer_) {
+      return sequencer_->InsertBatch(std::move(batch));
+    }
+    return Process(std::move(batch));
+  }
+
+  Status Process(ExecBatch batch) override {
+    return acero::MapNode::InputReceived(inputs_[0], batch);
+  }
+
   void Finish() override { dataset_writer_->Finish(); }
 
   Result<compute::ExecBatch> ProcessBatch(compute::ExecBatch batch) override {
@@ -635,6 +654,7 @@ class TeeNode : public acero::MapNode {
   std::unique_ptr<internal::DatasetWriter> dataset_writer_;
   FileSystemDatasetWriteOptions write_options_;
   std::atomic<int32_t> backpressure_counter_ = 0;
+  std::unique_ptr<acero::util::SerialSequencingQueue> sequencer_;
 };
 
 }  // namespace