ARROW-16525: [C++] Tee node not properly marking node finished

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

@@ -316,7 +316,7 @@ class ARROW_EXPORT MapNode : public ExecNode {
  protected:
   void SubmitTask(std::function<Result<ExecBatch>(ExecBatch)> map_fn, ExecBatch batch);
 
-  void Finish(Status finish_st = Status::OK());
+  virtual void Finish(Status finish_st = Status::OK());
 
  protected:
   // Counter for the number of batches received

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

@@ -363,7 +363,7 @@ class ConsumingSinkNode : public ExecNode, public BackpressureControl {
   }
 
  protected:
-  virtual void Finish(const Status& finish_st) {
+  void Finish(const Status& finish_st) {
     consumer_->Finish().AddCallback([this, finish_st](const Status& st) {
       // Prefer the plan error over the consumer error
       Status final_status = finish_st & st;

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

@@ -165,6 +165,12 @@ ExecBatch ExecBatchFromJSON(const std::vector<ValueDescr>& descrs,
   return batch;
 }
 
+Future<> StartAndFinish(ExecPlan* plan) {
+  RETURN_NOT_OK(plan->Validate());
+  RETURN_NOT_OK(plan->StartProducing());
+  return plan->finished();
+}
+
 Future<std::vector<ExecBatch>> StartAndCollect(
     ExecPlan* plan, AsyncGenerator<util::optional<ExecBatch>> gen) {
   RETURN_NOT_OK(plan->Validate());

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

@@ -82,6 +82,9 @@ struct BatchesWithSchema {
   }
 };
 
+ARROW_TESTING_EXPORT
+Future<> StartAndFinish(ExecPlan* plan);
+
 ARROW_TESTING_EXPORT
 Future<std::vector<ExecBatch>> StartAndCollect(
     ExecPlan* plan, AsyncGenerator<util::optional<ExecBatch>> gen);

cpp/src/arrow/dataset/file_base.cc

@@ -460,6 +460,23 @@ class TeeNode : public compute::MapNode {
 
   const char* kind_name() const override { return "TeeNode"; }
 
+  void Finish(Status finish_st) override {
+    dataset_writer_->Finish().AddCallback([this, finish_st](const Status& dw_status) {
+      // Need to wait for the task group to complete regardless of dw_status
+      task_group_.End().AddCallback(
+          [this, dw_status, finish_st](const Status& tg_status) {
+            // Prefer dw_status then finish_st and then tg_status
+            if (!dw_status.ok()) {
+              finished_.MarkFinished(dw_status);
+            }
+            if (!finish_st.ok()) {
+              finished_.MarkFinished(finish_st);
+            }
+            finished_.MarkFinished(tg_status);
+          });
+    });
+  }
+
   Result<compute::ExecBatch> DoTee(const compute::ExecBatch& batch) {
     ARROW_ASSIGN_OR_RAISE(std::shared_ptr<RecordBatch> record_batch,
                           batch.ToRecordBatch(output_schema()));

cpp/src/arrow/dataset/file_test.cc

@@ -18,14 +18,17 @@
 #include <cstdint>
 #include <memory>
 #include <string>
+#include <tuple>
 #include <vector>
 
 #include <gmock/gmock.h>
 #include <gtest/gtest.h>
 
 #include "arrow/array/array_primitive.h"
+#include "arrow/compute/exec/test_util.h"
 #include "arrow/dataset/api.h"
 #include "arrow/dataset/partition.h"
+#include "arrow/dataset/plan.h"
 #include "arrow/dataset/test_util.h"
 #include "arrow/filesystem/path_util.h"
 #include "arrow/filesystem/test_util.h"
@@ -34,6 +37,8 @@
 #include "arrow/testing/gtest_util.h"
 #include "arrow/util/io_util.h"
 
+namespace cp = arrow::compute;
+
 namespace arrow {
 
 using internal::TemporaryDir;
@@ -342,5 +347,118 @@ TEST_F(TestFileSystemDataset, WriteProjected) {
     }
   }
 }
+
+class FileSystemWriteTest : public testing::TestWithParam<std::tuple<bool, bool>> {
+ protected:
+  bool IsParallel() { return std::get<0>(GetParam()); }
+  bool IsSlow() { return std::get<1>(GetParam()); }
+
+  FileSystemWriteTest() { dataset::internal::Initialize(); }
+
+  void TestDatasetWriteRoundTrip(
+      std::function<Result<std::shared_ptr<cp::ExecPlan>>(
+          const cp::BatchesWithSchema& source_data, const FileSystemDatasetWriteOptions&,
+          std::function<Future<util::optional<cp::ExecBatch>>()>*)>
+          plan_factory,
+      bool has_output) {
+    // Runs in-memory data through the plan and then scans out the written
+    // data to ensure it matches the source data
+    auto format = std::make_shared<IpcFileFormat>();
+    auto fs = std::make_shared<fs::internal::MockFileSystem>(fs::kNoTime);
+    FileSystemDatasetWriteOptions write_options;
+    write_options.file_write_options = format->DefaultWriteOptions();
+    write_options.filesystem = fs;
+    write_options.base_dir = "root";
+    write_options.partitioning = std::make_shared<HivePartitioning>(schema({}));
+    write_options.basename_template = "{i}.feather";
+
+    cp::BatchesWithSchema source_data;
+    source_data.batches = {
+        cp::ExecBatchFromJSON({int32(), boolean()}, "[[null, true], [4, false]]"),
+        cp::ExecBatchFromJSON({int32(), boolean()},
+                              "[[5, null], [6, false], [7, false]]")};
+    source_data.schema = schema({field("i32", int32()), field("bool", boolean())});
+
+    AsyncGenerator<util::optional<cp::ExecBatch>> sink_gen;
+    ASSERT_OK_AND_ASSIGN(auto plan, plan_factory(source_data, write_options, &sink_gen));
+
+    if (has_output) {
+      ASSERT_FINISHES_OK_AND_ASSIGN(auto out_batches,
+                                    cp::StartAndCollect(plan.get(), sink_gen));
+      cp::AssertExecBatchesEqual(source_data.schema, source_data.batches, out_batches);
+    } else {
+      ASSERT_FINISHES_OK(cp::StartAndFinish(plan.get()));
+    }
+
+    // Read written dataset and make sure it matches
+    ASSERT_OK_AND_ASSIGN(auto dataset_factory, FileSystemDatasetFactory::Make(
+                                                   fs, {"root/0.feather"}, format, {}));
+    ASSERT_OK_AND_ASSIGN(auto written_dataset, dataset_factory->Finish(FinishOptions{}));
+    AssertSchemaEqual(*source_data.schema, *written_dataset->schema());
+
+    ASSERT_OK_AND_ASSIGN(plan, cp::ExecPlan::Make());
+    ASSERT_OK_AND_ASSIGN(auto scanner_builder, written_dataset->NewScan());
+    ASSERT_OK_AND_ASSIGN(auto scanner, scanner_builder->Finish());
+    ASSERT_OK(cp::Declaration::Sequence(
+                  {
+                      {"scan", ScanNodeOptions{written_dataset, scanner->options()}},
+                      {"sink", cp::SinkNodeOptions{&sink_gen}},
+                  })
+                  .AddToPlan(plan.get()));
+
+    ASSERT_FINISHES_OK_AND_ASSIGN(auto written_batches,
+                                  cp::StartAndCollect(plan.get(), sink_gen));
+    cp::AssertExecBatchesEqual(source_data.schema, source_data.batches, written_batches);
+  }
+};
+
+TEST_P(FileSystemWriteTest, Write) {
+  auto plan_factory =
+      [this](const cp::BatchesWithSchema& source_data,
+             const FileSystemDatasetWriteOptions& write_options,
+             std::function<Future<util::optional<cp::ExecBatch>>()>* sink_gen)
+      -> Result<std::shared_ptr<cp::ExecPlan>> {
+    ARROW_ASSIGN_OR_RAISE(auto plan, cp::ExecPlan::Make());
+    RETURN_NOT_OK(
+        cp::Declaration::Sequence(
+            {{"source", cp::SourceNodeOptions{source_data.schema,
+                                              source_data.gen(IsParallel(), IsSlow())}},
+             {"write", WriteNodeOptions{write_options}}})
+            .AddToPlan(plan.get()));
+    return plan;
+  };
+  TestDatasetWriteRoundTrip(plan_factory, /*has_output=*/false);
+}
+
+TEST_P(FileSystemWriteTest, TeeWrite) {
+  auto plan_factory =
+      [this](const cp::BatchesWithSchema& source_data,
+             const FileSystemDatasetWriteOptions& write_options,
+             std::function<Future<util::optional<cp::ExecBatch>>()>* sink_gen)
+      -> Result<std::shared_ptr<cp::ExecPlan>> {
+    ARROW_ASSIGN_OR_RAISE(auto plan, cp::ExecPlan::Make());
+    RETURN_NOT_OK(cp::Declaration::Sequence(
+                      {
+                          {"source", cp::SourceNodeOptions{source_data.schema,
+                                                           source_data.gen(IsParallel(),
+                                                                           IsSlow())}},
+                          {"tee", WriteNodeOptions{write_options}},
+                          {"sink", cp::SinkNodeOptions{sink_gen}},
+                      })
+                      .AddToPlan(plan.get()));
+    return plan;
+  };
+  TestDatasetWriteRoundTrip(plan_factory, /*has_output=*/true);
+}
+
+INSTANTIATE_TEST_SUITE_P(
+    FileSystemWrite, FileSystemWriteTest,
+    testing::Combine(testing::Values(false, true), testing::Values(false, true)),
+    [](const testing::TestParamInfo<FileSystemWriteTest::ParamType>& info) {
+      std::string parallel_desc = std::get<0>(info.param) ? "parallel" : "serial";
+      std::string speed_desc = std::get<1>(info.param) ? "slow" : "fast";
+      return parallel_desc + "_" + speed_desc;
+    });
+
 }  // namespace dataset
 }  // namespace arrow