wgtmac commented on code in PR #36779:
URL: https://github.com/apache/arrow/pull/36779#discussion_r1299514336
##########
cpp/src/parquet/arrow/reader.h:
##########
@@ -249,6 +250,63 @@ class PARQUET_EXPORT FileReader {
virtual ::arrow::Status ReadRowGroups(const std::vector<int>& row_groups,
std::shared_ptr<::arrow::Table>* out)
= 0;
+ using AsyncBatchGenerator =
Review Comment:
cpp/src/arrow/util/async_generator_fwd.h has defined AsyncGenerator below,
should we reuse it?
```cpp
template <typename T>
using AsyncGenerator = std::function<Future<T>()>;
```
##########
cpp/src/parquet/arrow/reader.cc:
##########
@@ -1233,6 +1285,127 @@ Status FileReaderImpl::ReadRowGroups(const
std::vector<int>& row_groups,
return Status::OK();
}
+struct AsyncBatchGeneratorState {
+ ::arrow::internal::Executor* io_executor;
+ ::arrow::internal::Executor* cpu_executor;
+ std::vector<std::shared_ptr<ColumnReaderImpl>> column_readers;
+ std::vector<std::shared_ptr<ChunkedArray>> current_cols;
+ std::queue<std::shared_ptr<RecordBatch>> overflow;
+ std::shared_ptr<::arrow::Schema> schema;
+ int64_t batch_size;
+ int64_t rows_remaining;
+ bool use_threads;
+ bool allow_sliced_batches;
+};
+
+class AsyncBatchGeneratorImpl {
+ public:
+ explicit AsyncBatchGeneratorImpl(std::shared_ptr<AsyncBatchGeneratorState>
state)
+ : state_(std::move(state)) {}
+ Future<std::shared_ptr<RecordBatch>> operator()() {
+ if (!state_->overflow.empty()) {
+ std::shared_ptr<RecordBatch> next = std::move(state_->overflow.front());
+ state_->overflow.pop();
+ return next;
+ }
+
+ if (state_->rows_remaining == 0) {
+ // Exhausted
+ return Future<std::shared_ptr<RecordBatch>>::MakeFinished(
+ ::arrow::IterationEnd<std::shared_ptr<RecordBatch>>());
+ }
+
+ int64_t rows_in_batch = std::min(state_->rows_remaining,
state_->batch_size);
+ state_->rows_remaining -= rows_in_batch;
+
+ // We read the columns in parallel. Each reader returns a chunked array.
This is
+ // because we might need to chunk a column if that column is too large. We
+ // do provide a batch size but even for a small batch size it is possible
that a
+ // column has extremely large strings which don't fit in a single batch.
+ std::vector<Future<>> column_futs;
+ column_futs.reserve(state_->column_readers.size());
+ for (std::size_t column_index = 0; column_index <
state_->column_readers.size();
+ column_index++) {
+ const auto& column_reader = state_->column_readers[column_index];
+ column_futs.push_back(
+ column_reader
+ ->NextBatchAsync(rows_in_batch, state_->io_executor,
state_->cpu_executor)
+ .Then([state = state_,
+ column_index](const std::shared_ptr<ChunkedArray>&
chunked_array) {
+ state->current_cols[column_index] = chunked_array;
+ }));
+ }
+ Future<> all_columns_finished_fut =
::arrow::AllFinished(std::move(column_futs));
+
+ // Grab the first batch of data and return it. If there is more than one
batch then
+ // throw the reamining batches into overflow and they will be fetched on
the next call
Review Comment:
```suggestion
// throw the remaining batches into overflow and they will be fetched on
the next call
```
##########
cpp/src/parquet/arrow/arrow_reader_writer_test.cc:
##########
@@ -2514,6 +2516,137 @@ TEST(TestArrowReadWrite, GetRecordBatchGenerator) {
}
}
+TEST(TestArrowReadWrite, ReadRowGroupsAsync) {
+ constexpr int kNumRows = 1024;
+ constexpr int kRowGroupSize = 512;
+ constexpr int kNumColumns = 2;
+
+ std::shared_ptr<Table> table;
+ ASSERT_NO_FATAL_FAILURE(MakeDoubleTable(kNumColumns, kNumRows, 1, &table));
+
+ std::shared_ptr<Buffer> buffer;
+ ASSERT_NO_FATAL_FAILURE(WriteTableToBuffer(table, kRowGroupSize,
+
default_arrow_writer_properties(), &buffer));
+
+ for (std::vector<int> row_groups : std::vector<std::vector<int>>{{}, {0},
{0, 1}}) {
+ ARROW_SCOPED_TRACE("# row_groups = ", row_groups.size());
+ int32_t expected_total_rows = static_cast<int32_t>(row_groups.size()) *
kRowGroupSize;
+
+ for (std::vector<int> columns : std::vector<std::vector<int>>{{}, {0}, {0,
1}}) {
+ ARROW_SCOPED_TRACE("# columns = ", columns.size());
+
+ for (int row_group_size : {128, 512, 1024, 2048}) {
Review Comment:
nit: row_group_size -> batch_size
##########
cpp/src/parquet/arrow/reader.h:
##########
@@ -249,6 +250,63 @@ class PARQUET_EXPORT FileReader {
virtual ::arrow::Status ReadRowGroups(const std::vector<int>& row_groups,
std::shared_ptr<::arrow::Table>* out)
= 0;
+ using AsyncBatchGenerator =
+ std::function<::arrow::Future<std::shared_ptr<::arrow::RecordBatch>>()>;
+
+ /// \brief Read row groups from the file
+ ///
+ /// \see ReadRowGroupsAsync for operation details
+ ///
+ /// \param row_groups indices of the row groups to read
+ /// \param cpu_executor an executor to use to run CPU tasks
+ /// \param allow_sliced_batches if false, an error is raised if a batch has
too much
+ /// data for the given batch size. If true,
smaller
+ /// batches will be returned instead.
+ virtual AsyncBatchGenerator ReadRowGroupsAsync(
+ const std::vector<int>& row_groups, ::arrow::internal::Executor*
cpu_executor,
+ bool allow_sliced_batches = true) = 0;
+
+ /// \brief Read some columns from the given rows groups from the file
+ ///
+ /// If pre-buffering is enabled then all of the data will be read using the
pre-buffer
+ /// cache. See ParquetFileReader::PreBuffer for details on how this affects
memory and
+ /// performance.
+ ///
+ /// This operation is not perfectly async. The read from disk will be done
on an I/O
+ /// thread, which is correct. However, decompression is also done on the
I/O thread
+ /// which may not be ideal.
+ ///
+ /// The returned generator will respect the batch size set in the
ArrowReaderProperties.
+ /// Batches will not be larger than the given batch size. However, batches
may be
+ /// smaller. This can happen, for example, when there is not enough data or
when a
+ /// string column is too large to fit into a single batch. The parameter
+ /// `allow_sliced_batches` can be set to false to disallow this later case.
This can be
+ /// useful when you need to know exactly how many batches you will get from
the
+ /// operation before you start.
+ ///
+ /// The returned generator is not async-reentrant. You must not fetch
another future
+ /// from it until the previously fetched future has completed.
+ ///
+ /// Note: When reading multiple row groups there is no guarantee you will
get one
+ /// record batch per row group. Data from multiple row groups could get
combined into
+ /// a single batch.
+ ///
+ /// Note: If a row group has 0 rows it will effectively be ignored. If you
are only
+ /// reading empty row groups then the returned generator will immediately
finish. This
+ /// method should never generate an empty record batch.
+ ///
+ /// The I/O executor is obtained from the I/O context in the reader
properties.
+ ///
+ /// \param row_groups indices of the row groups to read
+ /// \param column_indices indices of the columns to read
+ /// \param cpu_executor an executor to use to run CPU tasks
+ /// \param allow_sliced_batches if false, an error is raised if a batch has
too much
+ /// data for the given batch size. If false,
smaller
+ /// batches will be returned instead.
+ virtual AsyncBatchGenerator ReadRowGroupsAsync(
Review Comment:
It does not return arrow::Result, so the docstring would be good to include
expected return value, especially when `allow_sliced_batches` is false.
##########
cpp/src/parquet/arrow/reader.h:
##########
@@ -316,6 +374,14 @@ class PARQUET_EXPORT ColumnReader {
// the data available in the file.
virtual ::arrow::Status NextBatch(int64_t batch_size,
std::shared_ptr<::arrow::ChunkedArray>*
out) = 0;
+
+ // Overload of NextBatch that returns a Result
+ virtual ::arrow::Result<std::shared_ptr<::arrow::ChunkedArray>> NextBatch(
Review Comment:
Is it required to expose this?
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
