| diff --git a/cpp/src/parquet/arrow/schema.cc b/cpp/src/parquet/arrow/schema.cc |
| index ec3890a41f..943f69bb6c 100644 |
| --- a/cpp/src/parquet/arrow/schema.cc |
| +++ b/cpp/src/parquet/arrow/schema.cc |
| @@ -178,7 +178,7 @@ static Status GetTimestampMetadata(const ::arrow::TimestampType& type, |
| |
| // The user is explicitly asking for Impala int96 encoding, there is no |
| // logical type. |
| - if (arrow_properties.support_deprecated_int96_timestamps()) { |
| + if (arrow_properties.support_deprecated_int96_timestamps() && target_unit == ::arrow::TimeUnit::NANO) { |
| *physical_type = ParquetType::INT96; |
| return Status::OK(); |
| } |
| |
| diff --git a/cpp/src/parquet/arrow/reader.cc b/cpp/src/parquet/arrow/reader.cc |
| index 285e2a5973..aa6f92f077 100644 |
| --- a/cpp/src/parquet/arrow/reader.cc |
| +++ b/cpp/src/parquet/arrow/reader.cc |
| @@ -1013,25 +1013,32 @@ Status FileReaderImpl::GetRecordBatchReader(const std::vector<int>& row_groups, |
| return Status::OK(); |
| } |
| |
| - int64_t num_rows = 0; |
| + std::vector<int64_t> num_rows; |
| for (int row_group : row_groups) { |
| - num_rows += parquet_reader()->metadata()->RowGroup(row_group)->num_rows(); |
| + num_rows.push_back(parquet_reader()->metadata()->RowGroup(row_group)->num_rows()); |
| } |
| |
| using ::arrow::RecordBatchIterator; |
| + int row_group_idx = 0; |
| |
| // NB: This lambda will be invoked outside the scope of this call to |
| // `GetRecordBatchReader()`, so it must capture `readers` and `batch_schema` by value. |
| // `this` is a non-owning pointer so we are relying on the parent FileReader outliving |
| // this RecordBatchReader. |
| ::arrow::Iterator<RecordBatchIterator> batches = ::arrow::MakeFunctionIterator( |
| - [readers, batch_schema, num_rows, |
| + [readers, batch_schema, num_rows, row_group_idx, |
| this]() mutable -> ::arrow::Result<RecordBatchIterator> { |
| ::arrow::ChunkedArrayVector columns(readers.size()); |
| |
| - // don't reserve more rows than necessary |
| - int64_t batch_size = std::min(properties().batch_size(), num_rows); |
| - num_rows -= batch_size; |
| + int64_t batch_size = 0; |
| + if (!num_rows.empty()) { |
| + // don't reserve more rows than necessary |
| + batch_size = std::min(properties().batch_size(), num_rows[row_group_idx]); |
| + num_rows[row_group_idx] -= batch_size; |
| + if (num_rows[row_group_idx] == 0 && (num_rows.size() - 1) != row_group_idx) { |
| + row_group_idx++; |
| + } |
| + } |
| |
| RETURN_NOT_OK(::arrow::internal::OptionalParallelFor( |
| reader_properties_.use_threads(), static_cast<int>(readers.size()), |
| diff --git a/cpp/src/parquet/arrow/writer.cc b/cpp/src/parquet/arrow/writer.cc |
| index 4fd7ef1b47..87326a54f1 100644 |
| --- a/cpp/src/parquet/arrow/writer.cc |
| +++ b/cpp/src/parquet/arrow/writer.cc |
| @@ -314,6 +314,14 @@ class FileWriterImpl : public FileWriter { |
| return Status::OK(); |
| } |
| |
| + int64_t GetBufferedSize() override { |
| + if (row_group_writer_ == nullptr) { |
| + return 0; |
| + } |
| + return row_group_writer_->total_compressed_bytes() + |
| + row_group_writer_->total_compressed_bytes_written(); |
| + } |
| + |
| Status Close() override { |
| if (!closed_) { |
| // Make idempotent |
| @@ -418,10 +426,13 @@ class FileWriterImpl : public FileWriter { |
| |
| // Max number of rows allowed in a row group. |
| const int64_t max_row_group_length = this->properties().max_row_group_length(); |
| + const int64_t max_row_group_size = this->properties().max_row_group_size(); |
| |
| // Initialize a new buffered row group writer if necessary. |
| if (row_group_writer_ == nullptr || !row_group_writer_->buffered() || |
| - row_group_writer_->num_rows() >= max_row_group_length) { |
| + row_group_writer_->num_rows() >= max_row_group_length || |
| + (row_group_writer_->total_compressed_bytes_written() + |
| + row_group_writer_->total_compressed_bytes() >= max_row_group_size)) { |
| RETURN_NOT_OK(NewBufferedRowGroup()); |
| } |
| |
| diff --git a/cpp/src/parquet/arrow/writer.h b/cpp/src/parquet/arrow/writer.h |
| index 4a1a033a7b..0f13d05e44 100644 |
| --- a/cpp/src/parquet/arrow/writer.h |
| +++ b/cpp/src/parquet/arrow/writer.h |
| @@ -138,6 +138,9 @@ class PARQUET_EXPORT FileWriter { |
| /// option in this case. |
| virtual ::arrow::Status WriteRecordBatch(const ::arrow::RecordBatch& batch) = 0; |
| |
| + /// \brief Return the buffered size in bytes. |
| + virtual int64_t GetBufferedSize() = 0; |
| + |
| /// \brief Write the footer and close the file. |
| virtual ::arrow::Status Close() = 0; |
| virtual ~FileWriter(); |
| diff --git a/cpp/src/parquet/properties.h b/cpp/src/parquet/properties.h |
| index 4d3acb491e..3906ff3c59 100644 |
| --- a/cpp/src/parquet/properties.h |
| +++ b/cpp/src/parquet/properties.h |
| @@ -139,6 +139,7 @@ static constexpr bool DEFAULT_IS_DICTIONARY_ENABLED = true; |
| static constexpr int64_t DEFAULT_DICTIONARY_PAGE_SIZE_LIMIT = kDefaultDataPageSize; |
| static constexpr int64_t DEFAULT_WRITE_BATCH_SIZE = 1024; |
| static constexpr int64_t DEFAULT_MAX_ROW_GROUP_LENGTH = 1024 * 1024; |
| +static constexpr int64_t DEFAULT_MAX_ROW_GROUP_SIZE = 128 * 1024 * 1024; |
| static constexpr bool DEFAULT_ARE_STATISTICS_ENABLED = true; |
| static constexpr int64_t DEFAULT_MAX_STATISTICS_SIZE = 4096; |
| static constexpr Encoding::type DEFAULT_ENCODING = Encoding::UNKNOWN; |
| @@ -232,6 +233,7 @@ class PARQUET_EXPORT WriterProperties { |
| dictionary_pagesize_limit_(DEFAULT_DICTIONARY_PAGE_SIZE_LIMIT), |
| write_batch_size_(DEFAULT_WRITE_BATCH_SIZE), |
| max_row_group_length_(DEFAULT_MAX_ROW_GROUP_LENGTH), |
| + max_row_group_size_(DEFAULT_MAX_ROW_GROUP_SIZE), |
| pagesize_(kDefaultDataPageSize), |
| version_(ParquetVersion::PARQUET_2_6), |
| data_page_version_(ParquetDataPageVersion::V1), |
| @@ -244,6 +246,7 @@ class PARQUET_EXPORT WriterProperties { |
| dictionary_pagesize_limit_(properties.dictionary_pagesize_limit()), |
| write_batch_size_(properties.write_batch_size()), |
| max_row_group_length_(properties.max_row_group_length()), |
| + max_row_group_size_(properties.max_row_group_size()), |
| pagesize_(properties.data_pagesize()), |
| version_(properties.version()), |
| data_page_version_(properties.data_page_version()), |
| @@ -321,6 +324,13 @@ class PARQUET_EXPORT WriterProperties { |
| return this; |
| } |
| |
| + /// Specify the max bytes size to put in a single row group. |
| + /// Default 128 M. |
| + Builder* max_row_group_size(int64_t max_row_group_size) { |
| + max_row_group_size_ = max_row_group_size; |
| + return this; |
| + } |
| + |
| /// Specify the data page size. |
| /// Default 1MB. |
| Builder* data_pagesize(int64_t pg_size) { |
| @@ -664,7 +674,7 @@ class PARQUET_EXPORT WriterProperties { |
| |
| return std::shared_ptr<WriterProperties>(new WriterProperties( |
| pool_, dictionary_pagesize_limit_, write_batch_size_, max_row_group_length_, |
| - pagesize_, version_, created_by_, page_checksum_enabled_, |
| + max_row_group_size_, pagesize_, version_, created_by_, page_checksum_enabled_, |
| std::move(file_encryption_properties_), default_column_properties_, |
| column_properties, data_page_version_, store_decimal_as_integer_, |
| std::move(sorting_columns_))); |
| @@ -675,6 +685,7 @@ class PARQUET_EXPORT WriterProperties { |
| int64_t dictionary_pagesize_limit_; |
| int64_t write_batch_size_; |
| int64_t max_row_group_length_; |
| + int64_t max_row_group_size_; |
| int64_t pagesize_; |
| ParquetVersion::type version_; |
| ParquetDataPageVersion data_page_version_; |
| @@ -705,6 +716,8 @@ class PARQUET_EXPORT WriterProperties { |
| |
| inline int64_t max_row_group_length() const { return max_row_group_length_; } |
| |
| + inline int64_t max_row_group_size() const { return max_row_group_size_; } |
| + |
| inline int64_t data_pagesize() const { return pagesize_; } |
| |
| inline ParquetDataPageVersion data_page_version() const { |
| @@ -810,7 +823,7 @@ class PARQUET_EXPORT WriterProperties { |
| private: |
| explicit WriterProperties( |
| MemoryPool* pool, int64_t dictionary_pagesize_limit, int64_t write_batch_size, |
| - int64_t max_row_group_length, int64_t pagesize, ParquetVersion::type version, |
| + int64_t max_row_group_length, int64_t max_row_group_size, int64_t pagesize, ParquetVersion::type version, |
| const std::string& created_by, bool page_write_checksum_enabled, |
| std::shared_ptr<FileEncryptionProperties> file_encryption_properties, |
| const ColumnProperties& default_column_properties, |
| @@ -821,6 +834,7 @@ class PARQUET_EXPORT WriterProperties { |
| dictionary_pagesize_limit_(dictionary_pagesize_limit), |
| write_batch_size_(write_batch_size), |
| max_row_group_length_(max_row_group_length), |
| + max_row_group_size_(max_row_group_size), |
| pagesize_(pagesize), |
| parquet_data_page_version_(data_page_version), |
| parquet_version_(version), |
| @@ -836,6 +850,7 @@ class PARQUET_EXPORT WriterProperties { |
| int64_t dictionary_pagesize_limit_; |
| int64_t write_batch_size_; |
| int64_t max_row_group_length_; |
| + int64_t max_row_group_size_; |
| int64_t pagesize_; |
| ParquetDataPageVersion parquet_data_page_version_; |
| ParquetVersion::type parquet_version_; |
| |
| --- a/cpp/src/parquet/file_reader.h |
| +++ b/cpp/src/parquet/file_reader.h |
| @@ -210,6 +210,17 @@ |
| ::arrow::Future<> WhenBuffered(const std::vector<int>& row_groups, |
| const std::vector<int>& column_indices) const; |
| |
| + /// Pre-buffer arbitrary byte ranges (e.g., page-level ranges from OffsetIndex). |
| + /// Unlike PreBuffer(), this does NOT set the column bitmap, so |
| + /// GetColumnPageReader will use CachedInputStream (page-level cache path). |
| + void PreBufferRanges(const std::vector<::arrow::io::ReadRange>& ranges, |
| + const ::arrow::io::IOContext& ctx, |
| + const ::arrow::io::CacheOptions& options); |
| + |
| + /// Wait for arbitrary byte ranges to be pre-buffered. |
| + ::arrow::Future<> WhenBufferedRanges( |
| + const std::vector<::arrow::io::ReadRange>& ranges) const; |
| + |
| private: |
| // Holds a pointer to an instance of Contents implementation |
| std::unique_ptr<Contents> contents_; |
| |
| --- a/cpp/src/parquet/file_reader.cc |
| +++ b/cpp/src/parquet/file_reader.cc |
| @@ -207,6 +207,117 @@ |
| return {col_start, col_length}; |
| } |
| |
| +// CachedInputStream: InputStream adapter that reads through ReadRangeCache with |
| +// zero-cost skip for non-cached pages. Used for page-level caching where only |
| +// specific pages are pre-buffered. |
| +// |
| +// Key behavior: |
| +// - Read(): On cache hit, returns cached data. On cache miss, returns zero-filled |
| +// buffer (zero I/O). This makes InputStream::Advance() (which calls Read() and |
| +// discards) effectively free for skipped pages. |
| +// - Peek(): Always falls back to source on cache miss, because PageReader uses |
| +// Peek() to read Thrift page headers (~30 bytes) which must have real data. |
| +class CachedInputStream : public ::arrow::io::InputStream { |
| + public: |
| + CachedInputStream( |
| + std::shared_ptr<::arrow::io::internal::ReadRangeCache> cache, |
| + std::shared_ptr<ArrowInputFile> source, |
| + int64_t offset, int64_t length) |
| + : cache_(std::move(cache)), |
| + source_(std::move(source)), |
| + base_offset_(offset), |
| + length_(length) {} |
| + |
| + ::arrow::Status Close() override { |
| + closed_ = true; |
| + return ::arrow::Status::OK(); |
| + } |
| + |
| + bool closed() const override { return closed_; } |
| + |
| + ::arrow::Result<int64_t> Tell() const override { return position_; } |
| + |
| + ::arrow::Result<std::string_view> Peek(int64_t nbytes) override { |
| + int64_t to_read = std::min(nbytes, length_ - position_); |
| + if (to_read <= 0) { |
| + return std::string_view(); |
| + } |
| + ::arrow::io::ReadRange range{base_offset_ + position_, to_read}; |
| + auto result = cache_->Read(range); |
| + if (result.ok()) { |
| + peek_buffer_ = *result; |
| + } else { |
| + // Peek is used for Thrift page headers (~30 bytes) — must read real data |
| + ARROW_ASSIGN_OR_RAISE(peek_buffer_, |
| + source_->ReadAt(range.offset, range.length)); |
| + } |
| + return std::string_view( |
| + reinterpret_cast<const char*>(peek_buffer_->data()), |
| + static_cast<size_t>(peek_buffer_->size())); |
| + } |
| + |
| + ::arrow::Result<int64_t> Read(int64_t nbytes, void* out) override { |
| + int64_t to_read = std::min(nbytes, length_ - position_); |
| + if (to_read <= 0) return 0; |
| + ::arrow::io::ReadRange range{base_offset_ + position_, to_read}; |
| + auto result = cache_->Read(range); |
| + if (result.ok()) { |
| + auto& buf = *result; |
| + memcpy(out, buf->data(), static_cast<size_t>(buf->size())); |
| + position_ += buf->size(); |
| + return buf->size(); |
| + } |
| + // Cache miss: fall back to real I/O from source |
| + ARROW_ASSIGN_OR_RAISE(auto buf, source_->ReadAt(range.offset, range.length)); |
| + memcpy(out, buf->data(), static_cast<size_t>(buf->size())); |
| + position_ += buf->size(); |
| + return buf->size(); |
| + } |
| + |
| + ::arrow::Result<std::shared_ptr<::arrow::Buffer>> Read(int64_t nbytes) override { |
| + int64_t to_read = std::min(nbytes, length_ - position_); |
| + if (to_read <= 0) { |
| + return std::make_shared<::arrow::Buffer>(nullptr, 0); |
| + } |
| + ::arrow::io::ReadRange range{base_offset_ + position_, to_read}; |
| + auto result = cache_->Read(range); |
| + if (result.ok()) { |
| + position_ += (*result)->size(); |
| + return *result; |
| + } |
| + // Cache miss: fall back to real I/O from source |
| + ARROW_ASSIGN_OR_RAISE(auto buf, source_->ReadAt(range.offset, range.length)); |
| + position_ += buf->size(); |
| + return std::shared_ptr<::arrow::Buffer>(std::move(buf)); |
| + } |
| + |
| + // Override Advance to avoid real I/O for skipped pages. |
| + // The default InputStream::Advance() calls Read() and discards the result, |
| + // which would trigger source_->ReadAt() on cache miss — defeating page-level |
| + // I/O skipping via data_page_filter. Since Advance() is only used to skip |
| + // over data that will not be consumed, we can safely just move the position. |
| + ::arrow::Status Advance(int64_t nbytes) override { |
| + if (nbytes <= 0) { |
| + return ::arrow::Status::OK(); |
| + } |
| + int64_t remaining = length_ - position_; |
| + if (remaining <= 0) { |
| + return ::arrow::Status::OK(); |
| + } |
| + position_ += std::min(nbytes, remaining); |
| + return ::arrow::Status::OK(); |
| + } |
| + |
| + private: |
| + std::shared_ptr<::arrow::io::internal::ReadRangeCache> cache_; |
| + std::shared_ptr<ArrowInputFile> source_; |
| + int64_t base_offset_; |
| + int64_t length_; |
| + int64_t position_ = 0; |
| + bool closed_ = false; |
| + std::shared_ptr<::arrow::Buffer> peek_buffer_; |
| +}; |
| + |
| // RowGroupReader::Contents implementation for the Parquet file specification |
| class SerializedRowGroup : public RowGroupReader::Contents { |
| public: |
| @@ -242,6 +343,11 @@ |
| // segments. |
| PARQUET_ASSIGN_OR_THROW(auto buffer, cached_source_->Read(col_range)); |
| stream = std::make_shared<::arrow::io::BufferReader>(buffer); |
| + } else if (cached_source_) { |
| + // Page-level caching: read through cache with fallback to source. |
| + // Advance() is zero-cost for skipped pages via data_page_filter. |
| + stream = std::make_shared<CachedInputStream>( |
| + cached_source_, source_, col_range.offset, col_range.length); |
| } else { |
| stream = properties_.GetStream(source_, col_range.offset, col_range.length); |
| } |
| @@ -417,6 +523,26 @@ |
| return cached_source_->WaitFor(ranges); |
| } |
| |
| + void PreBufferRanges(const std::vector<::arrow::io::ReadRange>& ranges, |
| + const ::arrow::io::IOContext& ctx, |
| + const ::arrow::io::CacheOptions& options) { |
| + cached_source_ = |
| + std::make_shared<::arrow::io::internal::ReadRangeCache>(source_, ctx, options); |
| + // Do NOT set prebuffered_column_chunks_ bitmap — GetColumnPageReader will |
| + // use CachedInputStream path instead of full-chunk BufferReader path. |
| + prebuffered_column_chunks_.clear(); |
| + PARQUET_THROW_NOT_OK(cached_source_->Cache(ranges)); |
| + } |
| + |
| + ::arrow::Future<> WhenBufferedRanges( |
| + const std::vector<::arrow::io::ReadRange>& ranges) const { |
| + if (!cached_source_) { |
| + return ::arrow::Status::Invalid( |
| + "Must call PreBufferRanges before WhenBufferedRanges"); |
| + } |
| + return cached_source_->WaitFor(ranges); |
| + } |
| + |
| // Metadata/footer parsing. Divided up to separate sync/async paths, and to use |
| // exceptions for error handling (with the async path converting to Future/Status). |
| |
| @@ -911,6 +1037,22 @@ |
| return file->WhenBuffered(row_groups, column_indices); |
| } |
| |
| +void ParquetFileReader::PreBufferRanges( |
| + const std::vector<::arrow::io::ReadRange>& ranges, |
| + const ::arrow::io::IOContext& ctx, |
| + const ::arrow::io::CacheOptions& options) { |
| + SerializedFile* file = |
| + ::arrow::internal::checked_cast<SerializedFile*>(contents_.get()); |
| + file->PreBufferRanges(ranges, ctx, options); |
| +} |
| + |
| +::arrow::Future<> ParquetFileReader::WhenBufferedRanges( |
| + const std::vector<::arrow::io::ReadRange>& ranges) const { |
| + SerializedFile* file = |
| + ::arrow::internal::checked_cast<SerializedFile*>(contents_.get()); |
| + return file->WhenBufferedRanges(ranges); |
| +} |
| + |
| // ---------------------------------------------------------------------- |
| // File metadata helpers |
| |
| diff --git a/cpp/cmake_modules/ThirdpartyToolchain.cmake b/cpp/cmake_modules/ThirdpartyToolchain.cmake |
| --- a/cpp/cmake_modules/ThirdpartyToolchain.cmake |
| +++ b/cpp/cmake_modules/ThirdpartyToolchain.cmake |
| @@ -981,6 +981,11 @@ if(CMAKE_TOOLCHAIN_FILE) |
| list(APPEND EP_COMMON_CMAKE_ARGS -DCMAKE_TOOLCHAIN_FILE=${CMAKE_TOOLCHAIN_FILE}) |
| endif() |
| |
| +# Compatibility with bundled dependencies that require old CMake versions. |
| +if(CMAKE_VERSION VERSION_GREATER_EQUAL "3.30") |
| + list(APPEND EP_COMMON_CMAKE_ARGS -DCMAKE_POLICY_VERSION_MINIMUM=3.5) |
| +endif() |
| + |
| # and crosscompiling emulator (for try_run() ) |
| if(CMAKE_CROSSCOMPILING_EMULATOR) |
| string(REPLACE ";" ${EP_LIST_SEPARATOR} EP_CMAKE_CROSSCOMPILING_EMULATOR |
| @@ -1720,6 +1725,7 @@ macro(build_thrift) |
| -DWITH_JAVASCRIPT=OFF |
| -DWITH_LIBEVENT=OFF |
| -DWITH_NODEJS=OFF |
| + -DWITH_OPENSSL=OFF |
| -DWITH_PYTHON=OFF |
| -DWITH_QT5=OFF |
| -DWITH_ZLIB=OFF) |
| diff --git a/cpp/cmake_modules/BuildUtils.cmake b/cpp/cmake_modules/BuildUtils.cmake |
| --- a/cpp/cmake_modules/BuildUtils.cmake |
| +++ b/cpp/cmake_modules/BuildUtils.cmake |
| @@ -112,7 +112,7 @@ function(arrow_create_merged_static_lib output_target) |
| execute_process(COMMAND ${LIBTOOL_MACOS} -V |
| OUTPUT_VARIABLE LIBTOOL_V_OUTPUT |
| OUTPUT_STRIP_TRAILING_WHITESPACE) |
| - if(NOT "${LIBTOOL_V_OUTPUT}" MATCHES ".*cctools-([0-9.]+).*") |
| + if(NOT "${LIBTOOL_V_OUTPUT}" MATCHES ".*cctools(_ld)?-([0-9.]+).*") |
| message(FATAL_ERROR "libtool found appears to be the incompatible GNU libtool: ${LIBTOOL_MACOS}" |
| ) |
| endif() |
| |
| diff --git a/cpp/src/arrow/io/interfaces.h b/cpp/src/arrow/io/interfaces.h |
| --- a/cpp/src/arrow/io/interfaces.h |
| +++ b/cpp/src/arrow/io/interfaces.h |
| @@ -210,7 +210,7 @@ |
| /// \brief Advance or skip stream indicated number of bytes |
| /// \param[in] nbytes the number to move forward |
| /// \return Status |
| - Status Advance(int64_t nbytes); |
| + virtual Status Advance(int64_t nbytes); |
| |
| /// \brief Return zero-copy string_view to upcoming bytes. |
| /// |
| --- a/cpp/src/parquet/arrow/reader.cc |
| +++ b/cpp/src/parquet/arrow/reader.cc |
| @@ -254,6 +254,11 @@ |
| return GetColumn(i, AllRowGroupsFactory(), out); |
| } |
| |
| + ::arrow::Status GetColumn( |
| + int i, const std::vector<int>& column_indices, |
| + FileColumnIteratorFactory iterator_factory, |
| + std::unique_ptr<ColumnReader>* out) override; |
| + |
| Status GetSchema(std::shared_ptr<::arrow::Schema>* out) override { |
| return FromParquetSchema(reader_->metadata()->schema(), reader_properties_, |
| reader_->metadata()->key_value_metadata(), out); |
| @@ -493,10 +498,40 @@ |
| |
| ::arrow::Status BuildArray(int64_t length_upper_bound, |
| std::shared_ptr<::arrow::ChunkedArray>* out) final { |
| + if (!out_) { |
| + BEGIN_PARQUET_CATCH_EXCEPTIONS |
| + RETURN_NOT_OK( |
| + TransferColumnData(record_reader_.get(), field_, descr_, ctx_->pool, &out_)); |
| + END_PARQUET_CATCH_EXCEPTIONS |
| + } |
| *out = out_; |
| return Status::OK(); |
| } |
| |
| + std::vector<int> LeafColumnIndices() const final { |
| + return {input_->column_index()}; |
| + } |
| + |
| + ::arrow::Status ResetLeaf(int col_idx, int64_t reserve) final { |
| + if (col_idx != input_->column_index()) return Status::OK(); |
| + BEGIN_PARQUET_CATCH_EXCEPTIONS |
| + out_ = nullptr; |
| + record_reader_->Reset(); |
| + record_reader_->Reserve(reserve); |
| + return Status::OK(); |
| + END_PARQUET_CATCH_EXCEPTIONS |
| + } |
| + |
| + int64_t SkipRecords(int col_idx, int64_t num_records) final { |
| + if (col_idx != input_->column_index() || num_records <= 0) return 0; |
| + return record_reader_->SkipRecords(num_records); |
| + } |
| + |
| + int64_t ReadRecords(int col_idx, int64_t num_records) final { |
| + if (col_idx != input_->column_index() || num_records <= 0) return 0; |
| + return record_reader_->ReadRecords(num_records); |
| + } |
| + |
| const std::shared_ptr<Field> field() override { return field_; } |
| |
| private: |
| @@ -532,6 +567,22 @@ |
| return storage_reader_->LoadBatch(number_of_records); |
| } |
| |
| + std::vector<int> LeafColumnIndices() const final { |
| + return storage_reader_->LeafColumnIndices(); |
| + } |
| + |
| + ::arrow::Status ResetLeaf(int col_idx, int64_t reserve) final { |
| + return storage_reader_->ResetLeaf(col_idx, reserve); |
| + } |
| + |
| + int64_t SkipRecords(int col_idx, int64_t num_records) final { |
| + return storage_reader_->SkipRecords(col_idx, num_records); |
| + } |
| + |
| + int64_t ReadRecords(int col_idx, int64_t num_records) final { |
| + return storage_reader_->ReadRecords(col_idx, num_records); |
| + } |
| + |
| Status BuildArray(int64_t length_upper_bound, |
| std::shared_ptr<ChunkedArray>* out) override { |
| std::shared_ptr<ChunkedArray> storage; |
| @@ -576,6 +627,22 @@ |
| return item_reader_->LoadBatch(number_of_records); |
| } |
| |
| + std::vector<int> LeafColumnIndices() const final { |
| + return item_reader_->LeafColumnIndices(); |
| + } |
| + |
| + ::arrow::Status ResetLeaf(int col_idx, int64_t reserve) final { |
| + return item_reader_->ResetLeaf(col_idx, reserve); |
| + } |
| + |
| + int64_t SkipRecords(int col_idx, int64_t num_records) final { |
| + return item_reader_->SkipRecords(col_idx, num_records); |
| + } |
| + |
| + int64_t ReadRecords(int col_idx, int64_t num_records) final { |
| + return item_reader_->ReadRecords(col_idx, num_records); |
| + } |
| + |
| virtual ::arrow::Result<std::shared_ptr<ChunkedArray>> AssembleArray( |
| std::shared_ptr<ArrayData> data) { |
| if (field_->type()->id() == ::arrow::Type::MAP) { |
| @@ -709,6 +776,39 @@ |
| } |
| return Status::OK(); |
| } |
| + |
| + std::vector<int> LeafColumnIndices() const override { |
| + std::vector<int> indices; |
| + for (const std::unique_ptr<ColumnReaderImpl>& reader : children_) { |
| + std::vector<int> child_indices = reader->LeafColumnIndices(); |
| + indices.insert(indices.end(), child_indices.begin(), child_indices.end()); |
| + } |
| + return indices; |
| + } |
| + |
| + ::arrow::Status ResetLeaf(int col_idx, int64_t reserve) override { |
| + for (const std::unique_ptr<ColumnReaderImpl>& reader : children_) { |
| + RETURN_NOT_OK(reader->ResetLeaf(col_idx, reserve)); |
| + } |
| + return Status::OK(); |
| + } |
| + |
| + int64_t SkipRecords(int col_idx, int64_t num_records) override { |
| + int64_t skipped = 0; |
| + for (const std::unique_ptr<ColumnReaderImpl>& reader : children_) { |
| + skipped += reader->SkipRecords(col_idx, num_records); |
| + } |
| + return skipped; |
| + } |
| + |
| + int64_t ReadRecords(int col_idx, int64_t num_records) override { |
| + int64_t read = 0; |
| + for (const std::unique_ptr<ColumnReaderImpl>& reader : children_) { |
| + read += reader->ReadRecords(col_idx, num_records); |
| + } |
| + return read; |
| + } |
| + |
| Status BuildArray(int64_t length_upper_bound, |
| std::shared_ptr<ChunkedArray>* out) override; |
| Status GetDefLevels(const int16_t** data, int64_t* length) override; |
| @@ -1228,6 +1328,23 @@ |
| std::unique_ptr<ColumnReaderImpl> result; |
| RETURN_NOT_OK(GetReader(manifest_.schema_fields[i], ctx, &result)); |
| *out = std::move(result); |
| + return Status::OK(); |
| +} |
| + |
| +::arrow::Status FileReaderImpl::GetColumn( |
| + int i, const std::vector<int>& column_indices, |
| + FileColumnIteratorFactory iterator_factory, |
| + std::unique_ptr<ColumnReader>* out) { |
| + RETURN_NOT_OK(BoundsCheckColumn(i)); |
| + auto ctx = std::make_shared<ReaderContext>(); |
| + ctx->reader = reader_.get(); |
| + ctx->pool = pool_; |
| + ctx->iterator_factory = iterator_factory; |
| + ctx->filter_leaves = true; |
| + ctx->included_leaves = VectorToSharedSet(column_indices); |
| + std::unique_ptr<ColumnReaderImpl> result; |
| + RETURN_NOT_OK(GetReader(manifest_.schema_fields[i], ctx, &result)); |
| + *out = std::move(result); |
| return Status::OK(); |
| } |
| |
| --- a/cpp/src/parquet/arrow/reader.h |
| +++ b/cpp/src/parquet/arrow/reader.h |
| @@ -21,6 +21,7 @@ |
| // N.B. we don't include async_generator.h as it's relatively heavy |
| #include <functional> |
| #include <memory> |
| +#include <utility> |
| #include <vector> |
| |
| #include "parquet/file_reader.h" |
| @@ -48,9 +49,13 @@ |
| |
| class ColumnChunkReader; |
| class ColumnReader; |
| +class FileColumnIterator; |
| struct SchemaManifest; |
| class RowGroupReader; |
| |
| +using FileColumnIteratorFactory = |
| + std::function<FileColumnIterator*(int, ParquetFileReader*)>; |
| + |
| /// \brief Arrow read adapter class for deserializing Parquet files as Arrow row batches. |
| /// |
| /// This interfaces caters for different use cases and thus provides different |
| @@ -136,6 +141,27 @@ |
| // The indicated column index is relative to the schema |
| virtual ::arrow::Status GetColumn(int i, std::unique_ptr<ColumnReader>* out) = 0; |
| |
| + /// \brief Return a ColumnReader with a custom FileColumnIteratorFactory |
| + /// and leaf column filtering. |
| + /// |
| + /// This allows callers to customize page reading behavior (e.g., setting |
| + /// data_page_filter for page-level skipping) and to select only specific |
| + /// leaf columns within a nested field. The factory is called once per leaf |
| + /// column included in column_indices. |
| + /// |
| + /// \param i top-level field index (same as GetColumn(int i, ...)) |
| + /// \param column_indices leaf column indices to include (enables sub-column |
| + /// projection within nested types) |
| + /// \param iterator_factory factory to create FileColumnIterator per leaf |
| + /// \param[out] out the ColumnReader (may be nullptr if all leaves are pruned) |
| + virtual ::arrow::Status GetColumn( |
| + int i, const std::vector<int>& column_indices, |
| + FileColumnIteratorFactory iterator_factory, |
| + std::unique_ptr<ColumnReader>* out) { |
| + return ::arrow::Status::NotImplemented( |
| + "GetColumn with factory not implemented"); |
| + } |
| + |
| /// \brief Return arrow schema for all the columns. |
| virtual ::arrow::Status GetSchema(std::shared_ptr<::arrow::Schema>* out) = 0; |
| |
| @@ -316,6 +342,43 @@ |
| // the data available in the file. |
| virtual ::arrow::Status NextBatch(int64_t batch_size, |
| std::shared_ptr<::arrow::ChunkedArray>* out) = 0; |
| + |
| + /// \brief Leaf column indices covered by this (sub)tree, in leaf order. |
| + /// |
| + /// Used to drive per-leaf row filtering: after page-level skipping each leaf |
| + /// lives in its own compressed coordinate space, so callers must reset and |
| + /// skip/read each leaf independently rather than in lockstep. |
| + virtual std::vector<int> LeafColumnIndices() const { return {}; } |
| + |
| + /// \brief Reset the leaf identified by col_idx and reserve space for |
| + /// `reserve` records (in that leaf's post-page-filter compressed space). |
| + /// Must be called before SkipRecords()/ReadRecords() for that leaf, and |
| + /// followed by BuildArray() to get the result. |
| + virtual ::arrow::Status ResetLeaf(int col_idx, int64_t reserve) { |
| + return ::arrow::Status::NotImplemented("ResetLeaf not implemented"); |
| + } |
| + |
| + /// \brief Skip num_records on the leaf identified by col_idx and return the |
| + /// number of records actually skipped. Returns 0 when num_records <= 0 or |
| + /// col_idx does not belong to this (sub)tree. May throw ParquetException on a |
| + /// decode error; callers convert it to Status at the public boundary. |
| + virtual int64_t SkipRecords(int col_idx, int64_t num_records) { return 0; } |
| + |
| + /// \brief Read num_records on the leaf identified by col_idx and return the |
| + /// number of records actually read. Values accumulate across successive calls |
| + /// until BuildArray() is called. Returns 0 when num_records <= 0 or col_idx |
| + /// does not belong to this (sub)tree. May throw ParquetException on a decode |
| + /// error; callers convert it to Status at the public boundary. |
| + virtual int64_t ReadRecords(int col_idx, int64_t num_records) { return 0; } |
| + |
| + /// \brief Build the Arrow array from previously loaded data. |
| + /// For leaf readers, calls TransferColumnData if not already done. |
| + /// For nested readers, assembles the nested array from child arrays. |
| + virtual ::arrow::Status BuildArray( |
| + int64_t length_upper_bound, |
| + std::shared_ptr<::arrow::ChunkedArray>* out) { |
| + return ::arrow::Status::NotImplemented("BuildArray not implemented"); |
| + } |
| }; |
| |
| /// \brief Experimental helper class for bindings (like Python) that struggle |
| --- a/cpp/src/parquet/arrow/reader_internal.h |
| +++ b/cpp/src/parquet/arrow/reader_internal.h |
| @@ -26,6 +26,7 @@ |
| #include <utility> |
| #include <vector> |
| |
| +#include "parquet/arrow/reader.h" |
| #include "parquet/arrow/schema.h" |
| #include "parquet/column_reader.h" |
| #include "parquet/file_reader.h" |
| @@ -70,7 +71,10 @@ |
| |
| virtual ~FileColumnIterator() {} |
| |
| - std::unique_ptr<::parquet::PageReader> NextChunk() { |
| + /// \brief Fetch the PageReader for the next row group in this iterator's |
| + /// range. Virtual so subclasses can decorate the returned PageReader, e.g. |
| + /// to install a data_page_filter for I/O-level page skipping. |
| + virtual std::unique_ptr<::parquet::PageReader> NextChunk() { |
| if (row_groups_.empty()) { |
| return nullptr; |
| } |
| @@ -95,9 +99,6 @@ |
| std::deque<int> row_groups_; |
| }; |
| |
| -using FileColumnIteratorFactory = |
| - std::function<FileColumnIterator*(int, ParquetFileReader*)>; |
| - |
| Status TransferColumnData(::parquet::internal::RecordReader* reader, |
| const std::shared_ptr<::arrow::Field>& value_field, |
| const ColumnDescriptor* descr, ::arrow::MemoryPool* pool, |
| diff --git a/cpp/src/parquet/column_reader.h b/cpp/src/parquet/column_reader.h |
| --- a/cpp/src/parquet/column_reader.h |
| +++ b/cpp/src/parquet/column_reader.h |
| @@ -76,6 +76,18 @@ struct PARQUET_EXPORT DataPageStats { |
| std::optional<int32_t> num_rows; |
| }; |
| |
| +/// \brief Identifies a data page that PageReader should read directly. |
| +/// |
| +/// The offset is relative to the beginning of the column chunk stream passed to |
| +/// PageReader::Open. The compressed size includes both the serialized page header |
| +/// and the compressed page body. The ordinal is the original data page ordinal in |
| +/// the column chunk and does not include the dictionary page. |
| +struct PARQUET_EXPORT DataPageReadPlanEntry { |
| + int32_t page_ordinal; |
| + int64_t offset; |
| + int32_t compressed_page_size; |
| +}; |
| + |
| class PARQUET_EXPORT LevelDecoder { |
| public: |
| LevelDecoder(); |
| @@ -147,9 +159,21 @@ class PARQUET_EXPORT PageReader { |
| // ApplicationVersion::HasCorrectStatistics(). |
| // \note API EXPERIMENTAL |
| void set_data_page_filter(DataPageFilter data_page_filter) { |
| + if (data_page_read_plan_enabled_) { |
| + throw ParquetException( |
| + "data_page_filter and data_page_read_plan cannot be enabled together"); |
| + } |
| data_page_filter_ = std::move(data_page_filter); |
| } |
| |
| + /// Configure PageReader to jump directly to selected data pages before reading |
| + /// their headers. `first_data_page_offset` and each entry offset are relative to |
| + /// the beginning of the column chunk stream. Dictionary pages before |
| + /// `first_data_page_offset` are still read normally. |
| + // \note API EXPERIMENTAL |
| + void set_data_page_read_plan(int64_t first_data_page_offset, |
| + std::vector<DataPageReadPlanEntry> data_pages); |
| + |
| // @returns: shared_ptr<Page>(nullptr) on EOS, std::shared_ptr<Page> |
| // containing new Page otherwise |
| // |
| @@ -162,6 +186,11 @@ class PARQUET_EXPORT PageReader { |
| protected: |
| // Callback that decides if we should skip a page or not. |
| DataPageFilter data_page_filter_; |
| + |
| + bool data_page_read_plan_enabled_ = false; |
| + int64_t first_data_page_offset_ = 0; |
| + std::vector<DataPageReadPlanEntry> data_page_read_plan_; |
| + size_t next_data_page_ = 0; |
| }; |
| |
| class PARQUET_EXPORT ColumnReader { |
| diff --git a/cpp/src/parquet/column_reader.cc b/cpp/src/parquet/column_reader.cc |
| --- a/cpp/src/parquet/column_reader.cc |
| +++ b/cpp/src/parquet/column_reader.cc |
| @@ -207,6 +207,39 @@ ReaderProperties default_reader_properties() { |
| return default_reader_properties; |
| } |
| |
| +void PageReader::set_data_page_read_plan( |
| + int64_t first_data_page_offset, |
| + std::vector<DataPageReadPlanEntry> data_pages) { |
| + if (data_page_filter_) { |
| + throw ParquetException( |
| + "data_page_filter and data_page_read_plan cannot be enabled together"); |
| + } |
| + if (first_data_page_offset < 0) { |
| + throw ParquetException("Invalid negative first data page offset"); |
| + } |
| + |
| + int64_t previous_end = first_data_page_offset; |
| + int32_t previous_ordinal = -1; |
| + for (const auto& page : data_pages) { |
| + int64_t page_end; |
| + if (page.page_ordinal < 0 || page.offset < first_data_page_offset || |
| + page.compressed_page_size <= 0 || |
| + AddWithOverflow(page.offset, page.compressed_page_size, &page_end)) { |
| + throw ParquetException("Invalid data page read plan entry"); |
| + } |
| + if (page.offset < previous_end || page.page_ordinal <= previous_ordinal) { |
| + throw ParquetException("Data page read plan entries must be ordered"); |
| + } |
| + previous_end = page_end; |
| + previous_ordinal = page.page_ordinal; |
| + } |
| + |
| + data_page_read_plan_enabled_ = true; |
| + first_data_page_offset_ = first_data_page_offset; |
| + data_page_read_plan_ = std::move(data_pages); |
| + next_data_page_ = 0; |
| +} |
| + |
| namespace { |
| |
| // Extracts encoded statistics from V1 and V2 data page headers |
| @@ -430,9 +463,43 @@ std::shared_ptr<Page> SerializedPageReader::NextPage() { |
| |
| // Loop here because there may be unhandled page types that we skip until |
| // finding a page that we do know what to do with |
| - while (seen_num_values_ < total_num_values_) { |
| + while (data_page_read_plan_enabled_ || seen_num_values_ < total_num_values_) { |
| + const DataPageReadPlanEntry* planned_data_page = nullptr; |
| + uint32_t page_header_limit = max_page_header_size_; |
| + |
| + if (data_page_read_plan_enabled_) { |
| + if (next_data_page_ >= data_page_read_plan_.size()) { |
| + return nullptr; |
| + } |
| + |
| + PARQUET_ASSIGN_OR_THROW(int64_t current_position, stream_->Tell()); |
| + if (current_position < first_data_page_offset_) { |
| + page_header_limit = static_cast<uint32_t>(std::min<int64_t>( |
| + page_header_limit, first_data_page_offset_ - current_position)); |
| + } else { |
| + planned_data_page = &data_page_read_plan_[next_data_page_]; |
| + if (current_position > planned_data_page->offset) { |
| + throw ParquetException("Data page read plan points behind stream position"); |
| + } |
| + PARQUET_THROW_NOT_OK( |
| + stream_->Advance(planned_data_page->offset - current_position)); |
| + PARQUET_ASSIGN_OR_THROW(int64_t target_position, stream_->Tell()); |
| + if (target_position != planned_data_page->offset) { |
| + throw ParquetException("Failed to seek to planned data page"); |
| + } |
| + page_ordinal_ = planned_data_page->page_ordinal; |
| + page_header_limit = static_cast<uint32_t>(std::min<int64_t>( |
| + page_header_limit, planned_data_page->compressed_page_size)); |
| + } |
| + } |
| + |
| + if (page_header_limit == 0) { |
| + throw ParquetException("No bytes available for page header"); |
| + } |
| + |
| uint32_t header_size = 0; |
| - uint32_t allowed_page_size = kDefaultPageHeaderSize; |
| + uint32_t allowed_page_size = |
| + std::min<uint32_t>(kDefaultPageHeaderSize, page_header_limit); |
| |
| // Page headers can be very large because of page statistics |
| // We try to deserialize a larger buffer progressively |
| @@ -458,11 +525,12 @@ std::shared_ptr<Page> SerializedPageReader::NextPage() { |
| // Failed to deserialize. Double the allowed page header size and try again |
| std::stringstream ss; |
| ss << e.what(); |
| - allowed_page_size *= 2; |
| - if (allowed_page_size > max_page_header_size_) { |
| + if (allowed_page_size >= page_header_limit) { |
| ss << "Deserializing page header failed.\n"; |
| throw ParquetException(ss.str()); |
| } |
| + allowed_page_size = |
| + std::min<uint32_t>(allowed_page_size * 2, page_header_limit); |
| } |
| } |
| // Advance the stream offset |
| @@ -474,6 +542,20 @@ std::shared_ptr<Page> SerializedPageReader::NextPage() { |
| throw ParquetException("Invalid page header"); |
| } |
| |
| + const PageType::type page_type = LoadEnumSafe(¤t_page_header_.type); |
| + if (planned_data_page != nullptr) { |
| + if (page_type != PageType::DATA_PAGE && page_type != PageType::DATA_PAGE_V2) { |
| + throw ParquetException("Data page read plan points to a non-data page"); |
| + } |
| + int64_t total_compressed_size; |
| + if (AddWithOverflow(static_cast<int64_t>(header_size), |
| + static_cast<int64_t>(compressed_len), |
| + &total_compressed_size) || |
| + total_compressed_size != planned_data_page->compressed_page_size) { |
| + throw ParquetException("Planned data page size does not match page header"); |
| + } |
| + } |
| + |
| EncodedStatistics data_page_statistics; |
| if (ShouldSkipPage(&data_page_statistics)) { |
| PARQUET_THROW_NOT_OK(stream_->Advance(compressed_len)); |
| @@ -494,8 +576,6 @@ std::shared_ptr<Page> SerializedPageReader::NextPage() { |
| ParquetException::EofException(ss.str()); |
| } |
| |
| - const PageType::type page_type = LoadEnumSafe(¤t_page_header_.type); |
| - |
| if (properties_.page_checksum_verification() && current_page_header_.__isset.crc && |
| PageCanUseChecksum(page_type)) { |
| // verify crc |
| @@ -534,6 +614,9 @@ std::shared_ptr<Page> SerializedPageReader::NextPage() { |
| LoadEnumSafe(&dict_header.encoding), |
| is_sorted); |
| } else if (page_type == PageType::DATA_PAGE) { |
| + if (planned_data_page != nullptr) { |
| + ++next_data_page_; |
| + } |
| ++page_ordinal_; |
| const format::DataPageHeader& header = current_page_header_.data_page_header; |
| page_buffer = |
| @@ -545,6 +628,9 @@ std::shared_ptr<Page> SerializedPageReader::NextPage() { |
| LoadEnumSafe(&header.repetition_level_encoding), uncompressed_len, |
| std::move(data_page_statistics)); |
| } else if (page_type == PageType::DATA_PAGE_V2) { |
| + if (planned_data_page != nullptr) { |
| + ++next_data_page_; |
| + } |
| ++page_ordinal_; |
| const format::DataPageHeaderV2& header = current_page_header_.data_page_header_v2; |