Merge remote-tracking branch 'gitmodimo/dataset_max_rows_queued' into combined3

Author: Rafał Hibner

cpp/src/arrow/dataset/dataset_writer.h

@@ -29,9 +29,6 @@ namespace arrow {
 namespace dataset {
 namespace internal {
 
-// This lines up with our other defaults in the scanner and execution plan
-constexpr uint64_t kDefaultDatasetWriterMaxRowsQueued = 8 * 1024 * 1024;
-
 /// \brief Utility class that manages a set of writers to different paths
 ///
 /// Writers may be closed and reopened (and a new file created) based on the dataset

cpp/src/arrow/dataset/file_base.cc

@@ -452,9 +452,11 @@ Status ValidateAndPrepareSchema(const WriteNodeOptions& write_node_options,
 class DatasetWritingSinkNodeConsumer : public acero::SinkNodeConsumer {
  public:
   DatasetWritingSinkNodeConsumer(std::shared_ptr<Schema> custom_schema,
-                                 FileSystemDatasetWriteOptions write_options)
+                                 FileSystemDatasetWriteOptions write_options,
+                                 uint64_t max_rows_queued)
       : custom_schema_(std::move(custom_schema)),
-        write_options_(std::move(write_options)) {}
+        write_options_(std::move(write_options)),
+        max_rows_queued_(max_rows_queued) {}
 
   Status Init(const std::shared_ptr<Schema>& schema,
               acero::BackpressureControl* backpressure_control,
@@ -464,12 +466,12 @@ class DatasetWritingSinkNodeConsumer : public acero::SinkNodeConsumer {
     } else {
       schema_ = schema;
     }
-    ARROW_ASSIGN_OR_RAISE(
-        dataset_writer_,
-        internal::DatasetWriter::Make(
-            write_options_, plan->query_context()->async_scheduler(),
-            [backpressure_control] { backpressure_control->Pause(); },
-            [backpressure_control] { backpressure_control->Resume(); }, [] {}));
+    ARROW_ASSIGN_OR_RAISE(dataset_writer_,
+                          internal::DatasetWriter::Make(
+                              write_options_, plan->query_context()->async_scheduler(),
+                              [backpressure_control] { backpressure_control->Pause(); },
+                              [backpressure_control] { backpressure_control->Resume(); },
+                              [] {}, max_rows_queued_));
     return Status::OK();
   }
 
@@ -501,6 +503,7 @@ class DatasetWritingSinkNodeConsumer : public acero::SinkNodeConsumer {
   std::shared_ptr<Schema> custom_schema_;
   std::unique_ptr<internal::DatasetWriter> dataset_writer_;
   FileSystemDatasetWriteOptions write_options_;
+  uint64_t max_rows_queued_;
   Future<> finished_ = Future<>::Make();
   std::shared_ptr<Schema> schema_ = nullptr;
 };
@@ -556,8 +559,9 @@ Result<acero::ExecNode*> MakeWriteNode(acero::ExecPlan* plan,
       ValidateAndPrepareSchema(write_node_options, input_schema, custom_schema));
 
   std::shared_ptr<DatasetWritingSinkNodeConsumer> consumer =
-      std::make_shared<DatasetWritingSinkNodeConsumer>(custom_schema,
-                                                       write_node_options.write_options);
+      std::make_shared<DatasetWritingSinkNodeConsumer>(
+          custom_schema, write_node_options.write_options,
+          write_node_options.max_rows_queued);
   ARROW_ASSIGN_OR_RAISE(
       auto node,
       // to preserve order explicitly, sequence the exec batches
@@ -577,20 +581,21 @@ class TeeNode : public acero::MapNode,
  public:
   TeeNode(acero::ExecPlan* plan, std::vector<acero::ExecNode*> inputs,
           std::shared_ptr<Schema> output_schema,
-          FileSystemDatasetWriteOptions write_options)
+          FileSystemDatasetWriteOptions write_options, uint64_t max_rows_queued)
       : MapNode(plan, std::move(inputs), std::move(output_schema)),
-        write_options_(std::move(write_options)) {
+        write_options_(std::move(write_options)),
+        max_rows_queued_(max_rows_queued) {
     if (write_options.preserve_order) {
       sequencer_ = acero::util::SerialSequencingQueue::Make(this);
     }
   }
 
   Status StartProducing() override {
-    ARROW_ASSIGN_OR_RAISE(
-        dataset_writer_,
-        internal::DatasetWriter::Make(
-            write_options_, plan_->query_context()->async_scheduler(),
-            [this] { Pause(); }, [this] { Resume(); }, [this] { MapNode::Finish(); }));
+    ARROW_ASSIGN_OR_RAISE(dataset_writer_,
+                          internal::DatasetWriter::Make(
+                              write_options_, plan_->query_context()->async_scheduler(),
+                              [this] { Pause(); }, [this] { Resume(); },
+                              [this] { MapNode::Finish(); }, max_rows_queued_));
     return MapNode::StartProducing();
   }
 
@@ -608,7 +613,8 @@ class TeeNode : public acero::MapNode,
         ValidateAndPrepareSchema(write_node_options, input_schema, custom_schema));
 
     return plan->EmplaceNode<TeeNode>(plan, std::move(inputs), std::move(custom_schema),
-                                      std::move(write_node_options.write_options));
+                                      std::move(write_node_options.write_options),
+                                      write_node_options.max_rows_queued);
   }
 
   const char* kind_name() const override { return "TeeNode"; }
@@ -657,6 +663,7 @@ class TeeNode : public acero::MapNode,
  private:
   std::unique_ptr<internal::DatasetWriter> dataset_writer_;
   FileSystemDatasetWriteOptions write_options_;
+  uint64_t max_rows_queued_;
   std::atomic<int32_t> backpressure_counter_ = 0;
   std::unique_ptr<acero::util::SerialSequencingQueue> sequencer_{nullptr};
 };

cpp/src/arrow/dataset/file_base.h

@@ -385,6 +385,9 @@ class ARROW_DS_EXPORT FileWriter {
   std::optional<int64_t> bytes_written_;
 };
 
+// This lines up with our other defaults in the scanner and execution plan
+constexpr uint64_t kDefaultDatasetWriterMaxRowsQueued = 8 * 1024 * 1024;
+
 /// \brief Options for writing a dataset.
 struct ARROW_DS_EXPORT FileSystemDatasetWriteOptions {
   /// Options for individual fragment writing.
@@ -487,6 +490,9 @@ class ARROW_DS_EXPORT WriteNodeOptions : public acero::ExecNodeOptions {
   std::shared_ptr<Schema> custom_schema;
   /// \brief Optional metadata to attach to written batches
   std::shared_ptr<const KeyValueMetadata> custom_metadata;
+
+  /// Maximum rows queued before issuing pasue to upstream node
+  uint64_t max_rows_queued = kDefaultDatasetWriterMaxRowsQueued;
 };
 
 /// @}