diff --git a/cpp/src/arrow/dataset/file_parquet_test.cc b/cpp/src/arrow/dataset/file_parquet_test.cc index 0d86f5eaa0d3..9a982db59133 100644 --- a/cpp/src/arrow/dataset/file_parquet_test.cc +++ b/cpp/src/arrow/dataset/file_parquet_test.cc @@ -458,6 +458,57 @@ TEST_F(TestParquetFileFormat, MultithreadedScan) { ASSERT_EQ(batches.size(), kNumRowGroups); } +// Verify that a threaded, pre-buffered Dataset scan makes progress with only one I/O +// worker while multiple fragments and batches are scheduled ahead. +TEST_F(TestParquetFileFormat, MultithreadedScanWithSingleIOThread) { + // Constrain the global I/O pool for this test and restore its original capacity on + // every exit path. + struct PoolResetGuard { + int original_capacity = io::GetIOThreadPoolCapacity(); + ~PoolResetGuard() { DCHECK_OK(io::SetIOThreadPoolCapacity(original_capacity)); } + } guard; + ASSERT_OK(io::SetIOThreadPoolCapacity(1)); + + // Create several fragments with multiple row groups so fragment and row-group work + // can overlap and exercise the asynchronous scheduling path. + constexpr int64_t kNumFiles = 3; + constexpr int64_t kNumRowGroups = 4; + FragmentVector fragments; + for (int64_t i = 0; i < kNumFiles; ++i) { + auto reader = ArithmeticDatasetFixture::GetRecordBatchReader(kNumRowGroups); + auto source = GetFileSource(reader.get()); + fragments.push_back(MakeFragment(*source)); + } + + FragmentDataset dataset(ArithmeticDatasetFixture::schema(), std::move(fragments)); + ScannerBuilder builder({&dataset, [](...) {}}); + auto fragment_scan_options = std::make_shared(); + + // Select the incremental pre-buffered reader and apply enough batch and fragment + // readahead to keep the single I/O worker under scheduling pressure. + fragment_scan_options->arrow_reader_properties->set_pre_buffer(true); + + ASSERT_OK(builder.UseThreads(true)); + ASSERT_OK(builder.BatchSize(2)); + ASSERT_OK(builder.BatchReadahead(4)); + ASSERT_OK(builder.FragmentReadahead(kNumFiles)); + ASSERT_OK(builder.FragmentScanOptions(std::move(fragment_scan_options))); + ASSERT_OK_AND_ASSIGN(auto scanner, builder.Finish()); + ASSERT_OK_AND_ASSIGN(auto generator, scanner->ScanBatchesUnorderedAsync()); + + // Collecting the complete generator proves that all scheduled work makes progress; + // an executor deadlock would prevent this future from finishing. + ASSERT_FINISHES_OK_AND_ASSIGN(auto batches, + CollectAsyncGenerator(std::move(generator))); + + // Confirm that constrained scheduling did not drop or duplicate any rows. + int64_t num_rows = 0; + for (const auto& batch : batches) { + num_rows += batch.record_batch.value->num_rows(); + } + ASSERT_EQ(num_rows, kNumFiles * kNumRowGroups * (kNumRowGroups + 1) / 2); +} + TEST_F(TestParquetFileFormat, SingleThreadExecutor) { // Reset capacity for io executor struct PoolResetGuard { diff --git a/cpp/src/parquet/arrow/arrow_reader_writer_test.cc b/cpp/src/parquet/arrow/arrow_reader_writer_test.cc index 2bdbc38b3647..c17c9fb6ad91 100644 --- a/cpp/src/parquet/arrow/arrow_reader_writer_test.cc +++ b/cpp/src/parquet/arrow/arrow_reader_writer_test.cc @@ -38,20 +38,24 @@ #include "arrow/array/builder_dict.h" #include "arrow/array/builder_nested.h" #include "arrow/array/builder_primitive.h" +#include "arrow/array/util.h" #include "arrow/chunked_array.h" #include "arrow/compute/api.h" #include "arrow/extension/json.h" #include "arrow/io/api.h" +#include "arrow/memory_pool.h" #include "arrow/record_batch.h" #include "arrow/scalar.h" #include "arrow/table.h" #include "arrow/testing/builder.h" #include "arrow/testing/extension_type.h" +#include "arrow/testing/future_util.h" #include "arrow/testing/gtest_util.h" #include "arrow/testing/random.h" #include "arrow/testing/util.h" #include "arrow/type_fwd.h" #include "arrow/type_traits.h" +#include "arrow/util/async_generator.h" #include "arrow/util/checked_cast.h" #include "arrow/util/config.h" // for ARROW_CSV and PARQUET_REQUIRE_ENCRYPTION #include "arrow/util/decimal.h" @@ -59,6 +63,7 @@ #include "arrow/util/key_value_metadata.h" #include "arrow/util/logging_internal.h" #include "arrow/util/range.h" +#include "arrow/util/thread_pool.h" #ifdef ARROW_CSV # include "arrow/csv/api.h" @@ -2726,11 +2731,18 @@ TEST(TestArrowReadWrite, GetRecordBatchReaderNoColumns) { ASSERT_EQ(actual_batch->num_rows(), num_rows); } +// Verify incremental decoding across row groups, including each final partial batch. TEST(TestArrowReadWrite, GetRecordBatchGenerator) { ArrowReaderProperties properties = default_arrow_reader_properties(); - const int num_rows = 1024; - const int row_group_size = 512; + const int num_rows = 1000; + const int row_group_size = 500; + const int batch_size = 128; const int num_columns = 2; + const std::vector expected_batch_sizes = {128, 128, 128, 116, + 128, 128, 128, 116}; + properties.set_batch_size(batch_size); + properties.set_pre_buffer(true); + properties.set_use_threads(true); std::shared_ptr table; ASSERT_NO_FATAL_FAILURE(MakeDoubleTable(num_columns, num_rows, 1, &table)); @@ -2755,37 +2767,170 @@ TEST(TestArrowReadWrite, GetRecordBatchGenerator) { ASSERT_EQ(batch->num_rows(), num_rows); }; { + // Check batch boundaries and reconstruct the original table from the results. ASSERT_OK_AND_ASSIGN(auto batch_generator, reader->GetRecordBatchGenerator(reader, {0, 1}, {0, 1})); - auto fut1 = batch_generator(); - auto fut2 = batch_generator(); - auto fut3 = batch_generator(); - ASSERT_OK_AND_ASSIGN(auto batch1, fut1.result()); - ASSERT_OK_AND_ASSIGN(auto batch2, fut2.result()); - ASSERT_OK_AND_ASSIGN(auto batch3, fut3.result()); - ASSERT_EQ(batch3, nullptr); - check_batches(batch1, num_columns, row_group_size); - check_batches(batch2, num_columns, row_group_size); - ASSERT_OK_AND_ASSIGN(auto actual, ::arrow::Table::FromRecordBatches( - batch1->schema(), {batch1, batch2})); + ASSERT_FINISHES_OK_AND_ASSIGN(auto batches, + ::arrow::CollectAsyncGenerator(batch_generator)); + ASSERT_EQ(batches.size(), expected_batch_sizes.size()); + for (size_t index = 0; index < batches.size(); ++index) { + check_batches(batches[index], num_columns, expected_batch_sizes[index]); + } + ASSERT_OK_AND_ASSIGN( + auto actual, + ::arrow::Table::FromRecordBatches(batches.front()->schema(), batches)); AssertTablesEqual(*table, *actual, /*same_chunk_layout=*/false); } { - // No columns case + // Preserve the generator's existing support for pulls issued before earlier pulls + // complete. + ASSERT_OK_AND_ASSIGN(auto batch_generator, + reader->GetRecordBatchGenerator(reader, {0, 1}, {0, 1})); + auto future1 = batch_generator(); + auto future2 = batch_generator(); + auto future3 = batch_generator(); + ASSERT_OK_AND_ASSIGN(auto batch1, future1.result()); + ASSERT_OK_AND_ASSIGN(auto batch2, future2.result()); + ASSERT_OK_AND_ASSIGN(auto batch3, future3.result()); + check_batches(batch1, num_columns, batch_size); + check_batches(batch2, num_columns, batch_size); + check_batches(batch3, num_columns, batch_size); + } + { + // Guard against nested-executor deadlock: Next() could occupy the only worker, + // submit column work to the same pool, and wait while no worker is free to run it. + ASSERT_OK_AND_ASSIGN(auto shared_executor, ::arrow::internal::ThreadPool::Make(1)); + auto shared_executor_properties = properties; + shared_executor_properties.set_io_context( + ::arrow::io::IOContext(shared_executor.get())); + std::unique_ptr unique_reader; + FileReaderBuilder builder; + ASSERT_OK(builder.Open(std::make_shared(buffer))); + ASSERT_OK(builder.properties(shared_executor_properties)->Build(&unique_reader)); + std::shared_ptr shared_executor_reader = std::move(unique_reader); + ASSERT_OK_AND_ASSIGN( + auto batch_generator, + shared_executor_reader->GetRecordBatchGenerator( + shared_executor_reader, {0, 1}, {0, 1}, shared_executor.get(), num_rows)); + ASSERT_FINISHES_OK_AND_ASSIGN(auto batches, + ::arrow::CollectAsyncGenerator(batch_generator)); + ASSERT_EQ(batches.size(), expected_batch_sizes.size()); + for (size_t index = 0; index < batches.size(); ++index) { + check_batches(batches[index], num_columns, expected_batch_sizes[index]); + } + } + { + // Exercise the same deadlock guard when a null executor resolves to the default + // CPU pool, which is also serving as the single-worker I/O executor. + struct CpuThreadPoolCapacityGuard { + int capacity = ::arrow::GetCpuThreadPoolCapacity(); + ~CpuThreadPoolCapacityGuard() { + ARROW_CHECK_OK(::arrow::SetCpuThreadPoolCapacity(capacity)); + } + } capacity_guard; + ASSERT_OK(::arrow::SetCpuThreadPoolCapacity(1)); + auto* cpu_executor = ::arrow::internal::GetCpuThreadPool(); + auto default_executor_properties = properties; + default_executor_properties.set_io_context(::arrow::io::IOContext(cpu_executor)); + std::unique_ptr unique_reader; + FileReaderBuilder builder; + ASSERT_OK(builder.Open(std::make_shared(buffer))); + ASSERT_OK(builder.properties(default_executor_properties)->Build(&unique_reader)); + std::shared_ptr default_executor_reader = std::move(unique_reader); + ASSERT_OK_AND_ASSIGN( + auto batch_generator, + default_executor_reader->GetRecordBatchGenerator(default_executor_reader, + {0, 1}, {0, 1})); + auto batches_future = ::arrow::DeferNotOk(cpu_executor->Submit( + [batch_generator = std::move(batch_generator)]() mutable { + return ::arrow::CollectAsyncGenerator(batch_generator); + })); + ASSERT_FINISHES_OK_AND_ASSIGN(auto batches, batches_future); + ASSERT_EQ(batches.size(), expected_batch_sizes.size()); + for (size_t index = 0; index < batches.size(); ++index) { + check_batches(batches[index], num_columns, expected_batch_sizes[index]); + } + } + { + // Preserve row counts and batch boundaries when no columns are selected. ASSERT_OK_AND_ASSIGN(auto batch_generator, reader->GetRecordBatchGenerator(reader, {0, 1}, {})); - auto fut1 = batch_generator(); - auto fut2 = batch_generator(); - auto fut3 = batch_generator(); - ASSERT_OK_AND_ASSIGN(auto batch1, fut1.result()); - ASSERT_OK_AND_ASSIGN(auto batch2, fut2.result()); - ASSERT_OK_AND_ASSIGN(auto batch3, fut3.result()); - ASSERT_EQ(batch3, nullptr); - check_batches(batch1, 0, row_group_size); - check_batches(batch2, 0, row_group_size); + ASSERT_FINISHES_OK_AND_ASSIGN(auto batches, + ::arrow::CollectAsyncGenerator(batch_generator)); + ASSERT_EQ(batches.size(), expected_batch_sizes.size()); + for (size_t index = 0; index < batches.size(); ++index) { + check_batches(batches[index], 0, expected_batch_sizes[index]); + } } } +// Verify that requesting one small batch does not decode the complete wide row group. +// The incremental path should fit under the cap; the old eager path would fail before +// returning its first batch because it materialized every row and column at once. +TEST(TestArrowReadWrite, GetRecordBatchGeneratorDecodesFirstBatchWithinMemoryLimit) { + constexpr int num_rows = 100000; + constexpr int batch_size = 1000; + constexpr int num_columns = 32; + constexpr int64_t memory_limit = 8 * 1024 * 1024; + + // The cap can hold one decoded batch, but not all values in the row group. + ASSERT_LT(static_cast(batch_size) * num_columns * sizeof(double), + memory_limit); + ASSERT_GT(static_cast(num_rows) * num_columns * sizeof(double), + memory_limit); + + // Constant values keep pre-buffered Parquet input small enough to isolate decoded + // Arrow allocations. + ASSERT_OK_AND_ASSIGN( + auto values, + ::arrow::MakeArrayFromScalar(::arrow::DoubleScalar(1.0), num_rows)); + std::vector> columns; + std::vector> fields; + for (int column_index = 0; column_index < num_columns; ++column_index) { + columns.push_back(std::make_shared<::arrow::ChunkedArray>(values)); + fields.push_back(::arrow::field("col" + std::to_string(column_index), + ::arrow::float64())); + } + auto table = Table::Make(::arrow::schema(std::move(fields)), std::move(columns)); + + // Serialize all rows into one wide row group. The source table and Parquet buffer are + // created outside the capped pool so only reader-side allocations are constrained. + std::shared_ptr buffer; + ASSERT_NO_FATAL_FAILURE( + WriteTableToBuffer(table, num_rows, default_arrow_writer_properties(), &buffer)); + + // Wrap the default pool so any reader allocation beyond the test budget returns OOM. + ::arrow::ProxyMemoryPool proxy_pool(::arrow::default_memory_pool()); + ::arrow::CappedMemoryPool capped_pool(&proxy_pool, memory_limit); + ASSERT_OK_AND_ASSIGN(auto cpu_executor, ::arrow::internal::ThreadPool::Make(1)); + + // Use the pre-buffered incremental path and serial decoding for deterministic memory + // accounting. A successful Next() should allocate only the configured batch. + ArrowReaderProperties properties = default_arrow_reader_properties(); + properties.set_batch_size(batch_size); + properties.set_pre_buffer(true); + properties.set_use_threads(false); + + // Build an in-memory Parquet reader whose decode allocations use the capped pool. + std::unique_ptr unique_reader; + FileReaderBuilder builder; + ASSERT_OK(builder.Open(std::make_shared(buffer))); + ASSERT_OK(builder.memory_pool(&capped_pool)->properties(properties)->Build( + &unique_reader)); + std::shared_ptr reader = std::move(unique_reader); + + // Select every column in the only row group, but request just the first output batch. + ASSERT_OK_AND_ASSIGN( + auto batch_generator, + reader->GetRecordBatchGenerator(reader, {0}, Iota(num_columns), + cpu_executor.get())); + // An eager full-row-group decode would exceed the cap before returning this batch. + ASSERT_FINISHES_OK_AND_ASSIGN(auto first_batch, batch_generator()); + ASSERT_NE(first_batch, nullptr); + ASSERT_EQ(first_batch->num_columns(), num_columns); + ASSERT_EQ(first_batch->num_rows(), batch_size); +} + TEST(TestArrowReadWrite, ScanContents) { const int num_columns = 20; const int num_rows = 1000; diff --git a/cpp/src/parquet/arrow/reader.cc b/cpp/src/parquet/arrow/reader.cc index bb8bacc7751a..3227eb20c42c 100644 --- a/cpp/src/parquet/arrow/reader.cc +++ b/cpp/src/parquet/arrow/reader.cc @@ -335,6 +335,11 @@ class FileReaderImpl : public FileReader { const std::vector& row_group_indices, const std::vector& column_indices) override; + Result> GetRecordBatchReaderInternal( + const std::vector& row_group_indices, + const std::vector& column_indices, + ::arrow::internal::Executor* cpu_executor); + Result> GetRecordBatchReader( const std::vector& row_group_indices) override { return GetRecordBatchReader(row_group_indices, @@ -1111,6 +1116,14 @@ Result> FileReaderImpl::GetRecordBatchReader( END_PARQUET_CATCH_EXCEPTIONS } + return GetRecordBatchReaderInternal(row_groups, column_indices, + /*cpu_executor=*/nullptr); +} + +Result> +FileReaderImpl::GetRecordBatchReaderInternal( + const std::vector& row_groups, const std::vector& column_indices, + ::arrow::internal::Executor* cpu_executor) { std::vector> readers; std::shared_ptr<::arrow::Schema> batch_schema; RETURN_NOT_OK(GetFieldReaders(column_indices, row_groups, &readers, &batch_schema)); @@ -1142,6 +1155,12 @@ Result> FileReaderImpl::GetRecordBatchReader( for (int row_group : row_groups) { num_rows += parquet_reader()->metadata()->RowGroup(row_group)->num_rows(); } + // The incremental pre-buffered path bypasses DecodeRowGroups(), so preserve + // per-column parallelism here. + const bool use_threads = + reader_properties_.use_threads() && + (cpu_executor == nullptr || + cpu_executor != reader_properties_.io_context().executor()); using ::arrow::RecordBatchIterator; @@ -1150,7 +1169,7 @@ Result> FileReaderImpl::GetRecordBatchReader( // `this` is a non-owning pointer so we are relying on the parent FileReader outliving // this RecordBatchReader. ::arrow::Iterator batches = ::arrow::MakeFunctionIterator( - [readers, batch_schema, num_rows, + [readers, batch_schema, num_rows, cpu_executor, use_threads, this]() mutable -> ::arrow::Result { ::arrow::ChunkedArrayVector columns(readers.size()); @@ -1159,8 +1178,9 @@ Result> FileReaderImpl::GetRecordBatchReader( num_rows -= batch_size; RETURN_NOT_OK(::arrow::internal::OptionalParallelFor( - reader_properties_.use_threads(), static_cast(readers.size()), - [&](int i) { return readers[i]->NextBatch(batch_size, &columns[i]); })); + use_threads, static_cast(readers.size()), + [&](int i) { return readers[i]->NextBatch(batch_size, &columns[i]); }, + cpu_executor ? cpu_executor : ::arrow::internal::GetCpuThreadPool())); for (const auto& column : columns) { if (column == nullptr || column->length() == 0) { @@ -1259,7 +1279,8 @@ class RowGroupGenerator { ready.Then([cpu_executor = cpu_executor_, reader, row_group, column_indices = std::move( column_indices)]() -> ::arrow::Future { - return ReadOneRowGroup(cpu_executor, reader, row_group, column_indices); + return ReadOneRowGroupInBatches(cpu_executor, reader, row_group, + column_indices); }); } in_flight_reads_.push({std::move(row_group_read), num_rows}); @@ -1296,6 +1317,34 @@ class RowGroupGenerator { }); } + static ::arrow::Future ReadOneRowGroupInBatches( + ::arrow::internal::Executor* cpu_executor, std::shared_ptr self, + const int row_group, const std::vector& column_indices) { + if (!cpu_executor) { + cpu_executor = ::arrow::internal::GetCpuThreadPool(); + } + ARROW_ASSIGN_OR_RAISE(auto batch_reader, + self->GetRecordBatchReaderInternal({row_group}, column_indices, + cpu_executor)); + // Threaded Next() waits for column tasks on the CPU executor, so coordinate it on + // the I/O executor. Serial Next() performs CPU work directly on the CPU executor. + auto next_executor = self->properties().use_threads() + ? self->properties().io_context().executor() + : cpu_executor; + auto shared_reader = std::shared_ptr(std::move(batch_reader)); + RecordBatchGenerator generator = + [self = std::move(self), shared_reader, next_executor, + cpu_executor]() -> ::arrow::Future> { + auto next = ::arrow::DeferNotOk(next_executor->Submit( + [self, shared_reader] { return shared_reader->Next(); })); + if (next_executor == cpu_executor) { + return next; + } + return cpu_executor->TransferAlways(std::move(next)); + }; + return generator; + } + std::shared_ptr arrow_reader_; ::arrow::internal::Executor* cpu_executor_; std::vector row_groups_; diff --git a/cpp/src/parquet/arrow/reader_writer_benchmark.cc b/cpp/src/parquet/arrow/reader_writer_benchmark.cc index 2f288fd2eb0f..55ae1e004b91 100644 --- a/cpp/src/parquet/arrow/reader_writer_benchmark.cc +++ b/cpp/src/parquet/arrow/reader_writer_benchmark.cc @@ -19,6 +19,7 @@ #include #include +#include #include #include @@ -36,6 +37,7 @@ #include "arrow/array/data.h" #include "arrow/compute/cast.h" #include "arrow/io/memory.h" +#include "arrow/memory_pool.h" #include "arrow/table.h" #include "arrow/testing/gtest_util.h" #include "arrow/testing/random.h" @@ -43,6 +45,7 @@ #include "arrow/util/bit_util.h" #include "arrow/util/bitmap_ops.h" #include "arrow/util/logging.h" +#include "arrow/util/thread_pool.h" using arrow::Array; using arrow::ArrayVector; @@ -812,5 +815,77 @@ static void BM_ReadMultipleRowGroupsGenerator(::benchmark::State& state) { BENCHMARK(BM_ReadMultipleRowGroupsGenerator); +static void BM_ReadWideRowGroupGenerator(::benchmark::State& state) { + constexpr int64_t num_rows = 256 * 1024; + constexpr int num_columns = 128; + constexpr int64_t batch_size = 32 * 1024; + + std::vector values(num_rows, 128); + auto single_column = TableFromVector(values, /*nullable=*/false); + std::vector> columns; + FieldVector fields; + columns.reserve(num_columns); + fields.reserve(num_columns); + for (int column_index = 0; column_index < num_columns; ++column_index) { + columns.push_back(single_column->column(0)); + fields.push_back(::arrow::field("column" + std::to_string(column_index), + ::arrow::int64(), /*nullable=*/false)); + } + auto table = Table::Make(::arrow::schema(std::move(fields)), std::move(columns)); + + auto output = CreateOutputStream(); + EXIT_NOT_OK(WriteTable(*table, ::arrow::default_memory_pool(), output, num_rows)); + PARQUET_ASSIGN_OR_THROW(auto buffer, output->Finish()); + + ::parquet::ArrowReaderProperties properties; + properties.set_batch_size(batch_size); + properties.set_pre_buffer(true); + properties.set_use_threads(state.range(0) != 0); + + std::vector column_indices(num_columns); + std::iota(column_indices.begin(), column_indices.end(), 0); + constexpr int64_t num_batches = (num_rows + batch_size - 1) / batch_size; + int64_t peak_memory = 0; + + for (auto _ : state) { + ::arrow::ProxyMemoryPool reader_pool(::arrow::default_memory_pool()); + auto parquet_reader = + ParquetFileReader::Open(std::make_shared<::arrow::io::BufferReader>(buffer)); + auto arrow_reader_result = + FileReader::Make(&reader_pool, std::move(parquet_reader), properties); + EXIT_NOT_OK(arrow_reader_result.status()); + std::shared_ptr arrow_reader = std::move(*arrow_reader_result); + + ASSIGN_OR_ABORT(auto generator, + arrow_reader->GetRecordBatchGenerator( + arrow_reader, {0}, column_indices, + ::arrow::internal::GetCpuThreadPool(), /*rows_to_readahead=*/0)); + int64_t rows_read = 0; + for (int64_t batch_index = 0; batch_index < num_batches; ++batch_index) { + auto batch_future = generator(); + ASSIGN_OR_ABORT(auto batch, batch_future.result()); + rows_read += batch->num_rows(); + ::benchmark::DoNotOptimize(batch); + } + auto end_future = generator(); + ASSIGN_OR_ABORT(auto end, end_future.result()); + ARROW_CHECK(::arrow::IsIterationEnd(end)); + ARROW_CHECK_EQ(rows_read, num_rows); + peak_memory = std::max(peak_memory, reader_pool.max_memory()); + } + + state.counters["rows_per_second"] = ::benchmark::Counter( + static_cast(num_rows * state.iterations()), + ::benchmark::Counter::kIsRate); + state.counters["peak_memory_bytes"] = + ::benchmark::Counter(static_cast(peak_memory)); +} + +BENCHMARK(BM_ReadWideRowGroupGenerator) + ->ArgNames({"use_threads"}) + ->Arg(0) + ->Arg(1) + ->UseRealTime(); + } // namespace benchmarks } // namespace parquet