Check for duplicate pipe sinks

Author: Rafał Hibner

cpp/src/arrow/acero/pipe_node.cc

@@ -186,7 +186,10 @@ class PipeTeeNode : public PipeSource, public PipeSinkNode {
   PipeTeeNode(ExecPlan* plan, std::vector<ExecNode*> inputs, std::string pipe_name)
       : PipeSinkNode(plan, inputs, pipe_name) {
     output_schema_ = inputs[0]->output_schema();
-    PipeSinkNode::pipe_->addSource(this);
+    auto st = PipeSinkNode::pipe_->addSource(this);
+    if (ARROW_PREDICT_FALSE(!st.ok())) {  // this should never happen
+      internal::DieWithMessage(std::string("PipeTee unexpected error: ") + st.ToString());
+    }
   }
 
   static Result<ExecNode*> Make(ExecPlan* plan, std::vector<ExecNode*> inputs,
@@ -236,7 +239,11 @@ const char PipeTeeNode::kKindName[] = "PipeTeeNode";
 }  // namespace
 
 PipeSource::PipeSource() {}
-void PipeSource::Initialize(Pipe* pipe) { pipe_ = pipe; }
+Status PipeSource::Initialize(Pipe* pipe) {
+  if (pipe_) return Status::Invalid("Pipe:" + pipe->PipeName() + " has multiple sinks");
+  pipe_ = pipe;
+  return Status::OK();
+}
 
 void PipeSource::Pause(int32_t counter) { pipe_->Pause(this, counter); }
 void PipeSource::Resume(int32_t counter) { pipe_->Resume(this, counter); }
@@ -307,15 +314,16 @@ Status Pipe::InputFinished(int total_batches) {
   return Status::OK();
 }
 
-void Pipe::addSource(PipeSource* source) {
-  source->Initialize(this);
+Status Pipe::addSource(PipeSource* source) {
+  ARROW_RETURN_NOT_OK(source->Initialize(this));
   // First added source is handled in receiving task. All additional sources are delivered
   // in their own sutmit tasks
   if (!last_source_node_)
     last_source_node_ = source;
   else {
     source_nodes_.push_back(source);
   }
+  return Status::OK();
 }
 
 Status Pipe::Init(const std::shared_ptr<Schema> schema) {
@@ -326,7 +334,7 @@ Status Pipe::Init(const std::shared_ptr<Schema> schema) {
         if (!schema->Equals(node->output_schema())) {
           return Status::Invalid("Pipe schema does not match for " + pipe_name_);
         }
-        addSource(pipe_source);
+        ARROW_RETURN_NOT_OK(addSource(pipe_source));
       }
     }
   }

cpp/src/arrow/acero/pipe_node.h

@@ -40,7 +40,7 @@ class PipeSource {
 
  private:
   friend class Pipe;
-  void Initialize(Pipe* pipe);
+  Status Initialize(Pipe* pipe);
 
   virtual Status HandleInputReceived(ExecBatch batch) = 0;
   virtual Status HandleInputFinished(int total_batches) = 0;
@@ -66,13 +66,15 @@ class ARROW_ACERO_EXPORT Pipe {
   // Called from pipe_sink
   Status InputFinished(int total_batches);
 
-  void addSource(PipeSource* source);
+  Status addSource(PipeSource* source);
 
   // Called from pipe_sink Init
   Status Init(const std::shared_ptr<Schema> schema);
 
   bool HasSources() const;
 
+  std::string PipeName() const { return pipe_name_; }
+
  private:
   // pipe
   ExecPlan* plan_;

cpp/src/arrow/acero/pipe_node_test.cc

@@ -131,6 +131,32 @@ TEST(ExecPlanExecution, PipeErrorSink) {
              HasSubstr("Pipe 'named_pipe_missing_sink' error: Pipe does not have sink")));
 }
 
+TEST(ExecPlanExecution, PipeErrorDuplicateSink) {
+  auto basic_data = MakeBasicBatches();
+
+  AsyncGenerator<std::optional<ExecBatch>> main_sink_gen;
+  AsyncGenerator<std::optional<ExecBatch>> dup_sink_gen;
+  Declaration decl = Declaration::Sequence(
+      {{"source", SourceNodeOptions{basic_data.schema, basic_data.gen(/*parallel=*/false,
+                                                                      /*slow=*/false)}},
+       {"pipe_tee", PipeSinkNodeOptions{"named_pipe_1"}},
+       {"pipe_tee", PipeSinkNodeOptions{"named_pipe_1"}},
+       {"sink", SinkNodeOptions{&main_sink_gen}}});
+
+  Declaration dup = Declaration::Sequence(
+      {{"pipe_source", PipeSourceNodeOptions{"named_pipe_1", basic_data.schema}},
+       {"sink", SinkNodeOptions{&dup_sink_gen}}});
+
+  // fail on planning
+  ASSERT_OK_AND_ASSIGN(auto plan, ExecPlan::Make());
+  ASSERT_OK(decl.AddToPlan(plan.get()));
+  ASSERT_OK(dup.AddToPlan(plan.get()));
+  plan->StartProducing();
+  ASSERT_THAT(
+      plan->finished().result().status(),
+      Raises(StatusCode::Invalid, HasSubstr("Pipe:named_pipe_1 has multiple sinks")));
+}
+
 TEST(ExecPlanExecution, PipeFilterSink) {
   auto basic_data = MakeBasicBatches();
   AsyncGenerator<std::optional<ExecBatch>> main_sink_gen;