EIC code displayed by LXR master/​include/​arrow/​dataset/​dataset_writer.h	Source navigation Diff markup Identifier search General search
	Version: master test Revert   Change

File indexing completed on 2025-08-28 08:26:57

 // Licensed to the Apache Software Foundation (ASF) under one
 // or more contributor license agreements.  See the NOTICE file
 // distributed with this work for additional information
 // regarding copyright ownership.  The ASF licenses this file
 // to you under the Apache License, Version 2.0 (the
 // "License"); you may not use this file except in compliance
 // with the License.  You may obtain a copy of the License at
 //
 //   http://www.apache.org/licenses/LICENSE-2.0
 //
 // Unless required by applicable law or agreed to in writing,
 // software distributed under the License is distributed on an
 // "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
 // KIND, either express or implied.  See the License for the
 // specific language governing permissions and limitations
 // under the License.
 
 #pragma once
 
 #include <string>
 
 #include "arrow/dataset/file_base.h"
 #include "arrow/record_batch.h"
 #include "arrow/status.h"
 #include "arrow/util/async_util.h"
 #include "arrow/util/future.h"
 
 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
 /// write options (for example, max_rows_per_file or max_open_files)
 ///
 /// The dataset writer enforces its own back pressure based on the # of rows (as opposed
 /// to # of batches which is how it is typically enforced elsewhere) and # of files.
 class ARROW_DS_EXPORT DatasetWriter {
  public:
   /// \brief Create a dataset writer
   ///
   /// Will fail if basename_template is invalid or if there is existing data and
   /// existing_data_behavior is kError
   ///
   /// \param write_options options to control how the data should be written
   /// \param max_rows_queued max # of rows allowed to be queued before the dataset_writer
   ///                        will ask for backpressure
   static Result<std::unique_ptr<DatasetWriter>> Make(
       FileSystemDatasetWriteOptions write_options, util::AsyncTaskScheduler* scheduler,
       std::function<void()> pause_callback, std::function<void()> resume_callback,
       std::function<void()> finish_callback,
       uint64_t max_rows_queued = kDefaultDatasetWriterMaxRowsQueued);
 
   ~DatasetWriter();
 
   /// \brief Write a batch to the dataset
   /// \param[in] batch The batch to write
   /// \param[in] directory The directory to write to
   ///
   /// Note: The written filename will be {directory}/{filename_factory(i)} where i is a
   /// counter controlled by `max_open_files` and `max_rows_per_file`
   ///
   /// If multiple WriteRecordBatch calls arrive with the same `directory` then the batches
   /// may be written to the same file.
   ///
   /// The returned future will be marked finished when the record batch has been queued
   /// to be written.  If the returned future is unfinished then this indicates the dataset
   /// writer's queue is full and the data provider should pause.
   ///
   /// This method is NOT async reentrant.  The returned future will only be unfinished
   /// if back pressure needs to be applied.  Async reentrancy is not necessary for
   /// concurrent writes to happen.  Calling this method again before the previous future
   /// completes will not just violate max_rows_queued but likely lead to race conditions.
   ///
   /// One thing to note is that the ordering of your data can affect your maximum
   /// potential parallelism.  If this seems odd then consider a dataset where the first
   /// 1000 batches go to the same directory and then the 1001st batch goes to a different
   /// directory.  The only way to get two parallel writes immediately would be to queue
   /// all 1000 pending writes to the first directory.
   void WriteRecordBatch(std::shared_ptr<RecordBatch> batch, const std::string& directory,
                         const std::string& prefix = "");
 
   /// Finish all pending writes and close any open files
   void Finish();
 
  protected:
   DatasetWriter(FileSystemDatasetWriteOptions write_options,
                 util::AsyncTaskScheduler* scheduler, std::function<void()> pause_callback,
                 std::function<void()> resume_callback,
                 std::function<void()> finish_callback,
                 uint64_t max_rows_queued = kDefaultDatasetWriterMaxRowsQueued);
 
   class DatasetWriterImpl;
   std::unique_ptr<DatasetWriterImpl> impl_;
 };
 
 }  // namespace internal
 }  // namespace dataset
 }  // namespace arrow

This page was automatically generated by the 2.3.7 LXR engine. The LXR team