| diff --git a/cpp/cmake_modules/BuildUtils.cmake b/cpp/cmake_modules/BuildUtils.cmake | |
| index e7523add27..e079a1ad41 100644 | |
| --- 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/cmake_modules/ThirdpartyToolchain.cmake b/cpp/cmake_modules/ThirdpartyToolchain.cmake | |
| index 8cb3ec83f5..9aad83ac95 100644 | |
| --- a/cpp/cmake_modules/ThirdpartyToolchain.cmake | |
| +++ b/cpp/cmake_modules/ThirdpartyToolchain.cmake | |
| @@ -814,6 +814,7 @@ if(DEFINED ENV{ARROW_THRIFT_URL}) | |
| set(THRIFT_SOURCE_URL "$ENV{ARROW_THRIFT_URL}") | |
| else() | |
| set_urls(THRIFT_SOURCE_URL | |
| + "https://archive.apache.org/dist/thrift/${ARROW_THRIFT_BUILD_VERSION}/thrift-${ARROW_THRIFT_BUILD_VERSION}.tar.gz" | |
| "https://www.apache.org/dyn/closer.cgi?action=download&filename=/thrift/${ARROW_THRIFT_BUILD_VERSION}/thrift-${ARROW_THRIFT_BUILD_VERSION}.tar.gz" | |
| "https://downloads.apache.org/thrift/${ARROW_THRIFT_BUILD_VERSION}/thrift-${ARROW_THRIFT_BUILD_VERSION}.tar.gz" | |
| "https://apache.claz.org/thrift/${ARROW_THRIFT_BUILD_VERSION}/thrift-${ARROW_THRIFT_BUILD_VERSION}.tar.gz" | |
| @@ -983,6 +984,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 | |
| @@ -1716,6 +1722,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/src/arrow/io/interfaces.h b/cpp/src/arrow/io/interfaces.h | |
| index b36c38c6d4..f974a33073 100644 | |
| --- a/cpp/src/arrow/io/interfaces.h | |
| +++ b/cpp/src/arrow/io/interfaces.h | |
| @@ -210,7 +210,7 @@ class ARROW_EXPORT InputStream : virtual public FileInterface, virtual public Re | |
| /// \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. | |
| /// | |
| diff --git a/cpp/src/arrow/util/bit_run_reader.h b/cpp/src/arrow/util/bit_run_reader.h | |
| index a436a50b86..27d483978c 100644 | |
| --- a/cpp/src/arrow/util/bit_run_reader.h | |
| +++ b/cpp/src/arrow/util/bit_run_reader.h | |
| @@ -168,6 +168,26 @@ class ARROW_EXPORT BitRunReader { | |
| using BitRunReader = BitRunReaderLinear; | |
| #endif | |
| +template <typename Visit> | |
| +inline Status VisitBitRuns(const uint8_t* bitmap, int64_t offset, int64_t length, | |
| + Visit&& visit) { | |
| + if (bitmap == NULLPTR) { | |
| + // Assuming all set (as in a null bitmap) | |
| + return visit(static_cast<int64_t>(0), length, true); | |
| + } | |
| + BitRunReader reader(bitmap, offset, length); | |
| + int64_t position = 0; | |
| + while (true) { | |
| + const auto run = reader.NextRun(); | |
| + if (run.length == 0) { | |
| + break; | |
| + } | |
| + ARROW_RETURN_NOT_OK(visit(position, run.length, run.set)); | |
| + position += run.length; | |
| + } | |
| + return Status::OK(); | |
| +} | |
| + | |
| struct SetBitRun { | |
| int64_t position; | |
| int64_t length; | |
| diff --git a/cpp/src/parquet/arrow/reader.cc b/cpp/src/parquet/arrow/reader.cc | |
| index 285e2a5973..40bddef195 100644 | |
| --- a/cpp/src/parquet/arrow/reader.cc | |
| +++ b/cpp/src/parquet/arrow/reader.cc | |
| @@ -19,12 +19,14 @@ | |
| #include <algorithm> | |
| #include <cstring> | |
| +#include <iterator> | |
| #include <memory> | |
| #include <unordered_set> | |
| #include <utility> | |
| #include <vector> | |
| #include "arrow/array.h" | |
| +#include "arrow/array/concatenate.h" | |
| #include "arrow/buffer.h" | |
| #include "arrow/extension_type.h" | |
| #include "arrow/io/memory.h" | |
| @@ -32,12 +34,14 @@ | |
| #include "arrow/table.h" | |
| #include "arrow/type.h" | |
| #include "arrow/util/async_generator.h" | |
| +#include "arrow/util/bit_run_reader.h" | |
| #include "arrow/util/bit_util.h" | |
| #include "arrow/util/future.h" | |
| #include "arrow/util/iterator.h" | |
| #include "arrow/util/logging.h" | |
| #include "arrow/util/parallel.h" | |
| #include "arrow/util/range.h" | |
| +#include "arrow/util/span.h" | |
| #include "arrow/util/tracing_internal.h" | |
| #include "parquet/arrow/reader_internal.h" | |
| #include "parquet/column_reader.h" | |
| @@ -254,6 +258,11 @@ class FileReaderImpl : public FileReader { | |
| return GetColumn(i, AllRowGroupsFactory(), out); | |
| } | |
| + ::arrow::Status GetColumn( | |
| + int i, const std::shared_ptr<std::unordered_set<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 +502,43 @@ class LeafReader : public ColumnReaderImpl { | |
| ::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_records, | |
| + int64_t reserve_values, | |
| + int64_t reserve_value_bytes) final { | |
| + if (col_idx != input_->column_index()) return Status::OK(); | |
| + BEGIN_PARQUET_CATCH_EXCEPTIONS | |
| + out_ = nullptr; | |
| + record_reader_->Reset(); | |
| + record_reader_->Reserve(reserve_records); | |
| + record_reader_->ReserveValueBytes(reserve_values, reserve_value_bytes); | |
| + 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 +574,25 @@ class ExtensionReader : public ColumnReaderImpl { | |
| 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_records, | |
| + int64_t reserve_values, | |
| + int64_t reserve_value_bytes) final { | |
| + return storage_reader_->ResetLeaf(col_idx, reserve_records, reserve_values, | |
| + reserve_value_bytes); | |
| + } | |
| + | |
| + 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 +637,25 @@ class ListReader : public ColumnReaderImpl { | |
| 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_records, | |
| + int64_t reserve_values, | |
| + int64_t reserve_value_bytes) final { | |
| + return item_reader_->ResetLeaf(col_idx, reserve_records, reserve_values, | |
| + reserve_value_bytes); | |
| + } | |
| + | |
| + 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) { | |
| @@ -642,8 +722,10 @@ class ListReader : public ColumnReaderImpl { | |
| const std::shared_ptr<Field> field() override { return field_; } | |
| - private: | |
| + protected: | |
| std::shared_ptr<ReaderContext> ctx_; | |
| + | |
| + private: | |
| std::shared_ptr<Field> field_; | |
| ::parquet::internal::LevelInfo level_info_; | |
| std::unique_ptr<ColumnReaderImpl> item_reader_; | |
| @@ -662,12 +744,62 @@ class PARQUET_NO_EXPORT FixedSizeListReader : public ListReader<int32_t> { | |
| DCHECK_EQ(field()->type()->id(), ::arrow::Type::FIXED_SIZE_LIST); | |
| const auto& type = checked_cast<::arrow::FixedSizeListType&>(*field()->type()); | |
| const int32_t* offsets = reinterpret_cast<const int32_t*>(data->buffers[1]->data()); | |
| - for (int x = 1; x <= data->length; x++) { | |
| - int32_t size = offsets[x] - offsets[x - 1]; | |
| - if (size != type.list_size()) { | |
| - return Status::Invalid("Expected all lists to be of size=", type.list_size(), | |
| - " but index ", x, " had size=", size); | |
| + const int32_t list_size = type.list_size(); | |
| + auto validate_offsets = [&](int64_t start, int64_t length, | |
| + bool has_elements) -> Status { | |
| + const int32_t expected_size = has_elements ? list_size : 0; | |
| + ::arrow::util::span<const int32_t> run_offsets( | |
| + offsets + start, static_cast<size_t>(length + 1)); | |
| + const auto first_invalid_offset = std::adjacent_find( | |
| + run_offsets.begin(), run_offsets.end(), | |
| + [&](int32_t left, int32_t right) { return right - left != expected_size; }); | |
| + if (first_invalid_offset != run_offsets.end()) { | |
| + const int64_t x = | |
| + start + std::distance(run_offsets.begin(), first_invalid_offset); | |
| + const int32_t size = offsets[x + 1] - offsets[x]; | |
| + if (has_elements) { | |
| + return Status::Invalid("Expected all lists to be of size=", list_size, | |
| + " but index ", x + 1, " had size=", size); | |
| + } | |
| + return Status::Invalid("Expected null fixed-size list at index ", x + 1, | |
| + " to have no child values but had size=", size); | |
| } | |
| + return Status::OK(); | |
| + }; | |
| + if (data->GetNullCount() != 0) { | |
| + // Rebuild the child array run-by-run so null fixed-size list slots still | |
| + // contribute list_size child values in the final layout. | |
| + ::arrow::ArrayVector child_arrays; | |
| + | |
| + auto visit_run = [&](int64_t start, int64_t length, bool has_elements) -> Status { | |
| + RETURN_NOT_OK(validate_offsets(start, length, has_elements)); | |
| + | |
| + const int64_t child_length = length * list_size; | |
| + // Valid runs reuse the decoded child slice; null runs materialize null | |
| + // children to preserve the fixed-size list shape. | |
| + if (!has_elements) { | |
| + ARROW_ASSIGN_OR_RAISE( | |
| + auto null_array, | |
| + ::arrow::MakeArrayOfNull(type.value_type(), child_length, ctx_->pool)); | |
| + child_arrays.push_back(std::move(null_array)); | |
| + return Status::OK(); | |
| + } | |
| + child_arrays.push_back( | |
| + ::arrow::MakeArray(data->child_data[0]->Slice(offsets[start], child_length))); | |
| + return Status::OK(); | |
| + }; | |
| + | |
| + DCHECK_NE(data->buffers[0], nullptr); | |
| + RETURN_NOT_OK(::arrow::internal::VisitBitRuns( | |
| + data->buffers[0]->data(), data->offset, data->length, visit_run)); | |
| + | |
| + // TODO(GH-50271): Build one padded child array directly instead of creating | |
| + // one temporary Array/ArrayData per validity run and concatenating them. | |
| + ARROW_ASSIGN_OR_RAISE(auto child_array_with_padding, | |
| + ::arrow::Concatenate(child_arrays, ctx_->pool)); | |
| + data->child_data[0] = child_array_with_padding->data(); | |
| + } else { | |
| + RETURN_NOT_OK(validate_offsets(/*start=*/0, data->length, /*valid=*/true)); | |
| } | |
| data->buffers.resize(1); | |
| std::shared_ptr<Array> result = ::arrow::MakeArray(data); | |
| @@ -709,6 +841,42 @@ class PARQUET_NO_EXPORT StructReader : public ColumnReaderImpl { | |
| } | |
| 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_records, | |
| + int64_t reserve_values, | |
| + int64_t reserve_value_bytes) override { | |
| + for (const std::unique_ptr<ColumnReaderImpl>& reader : children_) { | |
| + RETURN_NOT_OK(reader->ResetLeaf(col_idx, reserve_records, reserve_values, | |
| + reserve_value_bytes)); | |
| + } | |
| + 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; | |
| @@ -1013,25 +1181,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()), | |
| @@ -1224,6 +1399,23 @@ Status FileReaderImpl::GetColumn(int i, FileColumnIteratorFactory iterator_facto | |
| return Status::OK(); | |
| } | |
| +::arrow::Status FileReaderImpl::GetColumn( | |
| + int i, const std::shared_ptr<std::unordered_set<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 = column_indices; | |
| + std::unique_ptr<ColumnReaderImpl> result; | |
| + RETURN_NOT_OK(GetReader(manifest_.schema_fields[i], ctx, &result)); | |
| + *out = std::move(result); | |
| + return Status::OK(); | |
| +} | |
| + | |
| Status FileReaderImpl::ReadRowGroups(const std::vector<int>& row_groups, | |
| const std::vector<int>& column_indices, | |
| std::shared_ptr<Table>* out) { | |
| diff --git a/cpp/src/parquet/arrow/reader.h b/cpp/src/parquet/arrow/reader.h | |
| index 6e46ca43f7..e202c6404e 100644 | |
| --- 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 @@ namespace arrow { | |
| 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 @@ class PARQUET_EXPORT FileReader { | |
| // 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::shared_ptr<std::unordered_set<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,51 @@ 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; | |
| + | |
| + /// \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 pre-allocate for the read | |
| + /// that follows. Must be called before SkipRecords()/ReadRecords() for that | |
| + /// leaf, and followed by BuildArray() to get the result. | |
| + /// | |
| + /// `reserve_records` is in that leaf's post-page-filter compressed space, | |
| + /// because SkipRecords walks the levels it bounds. `reserve_values` and | |
| + /// `reserve_value_bytes` describe what will actually be APPENDED, which on a | |
| + /// selective read is far less, and only the variable-width readers use the | |
| + /// byte count. Zero for the latter two means "no estimate" and reserves | |
| + /// nothing beyond `reserve_records`. | |
| + virtual ::arrow::Status ResetLeaf(int col_idx, int64_t reserve_records, | |
| + int64_t reserve_values, | |
| + int64_t reserve_value_bytes) { | |
| + 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 | |
| diff --git a/cpp/src/parquet/arrow/reader_internal.h b/cpp/src/parquet/arrow/reader_internal.h | |
| index cf9dbb8657..9216f18289 100644 | |
| --- 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 @@ class FileColumnIterator { | |
| 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 @@ class FileColumnIterator { | |
| 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/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/writer.cc b/cpp/src/parquet/arrow/writer.cc | |
| index 4fd7ef1b47..feff99c99b 100644 | |
| --- a/cpp/src/parquet/arrow/writer.cc | |
| +++ b/cpp/src/parquet/arrow/writer.cc | |
| @@ -26,6 +26,7 @@ | |
| #include <vector> | |
| #include "arrow/array.h" | |
| +#include "arrow/array/concatenate.h" | |
| #include "arrow/extension_type.h" | |
| #include "arrow/ipc/writer.h" | |
| #include "arrow/record_batch.h" | |
| @@ -142,13 +143,24 @@ class ArrowColumnWriterV2 { | |
| leaf_idx, ctx, [&](const MultipathLevelBuilderResult& result) { | |
| size_t visited_component_size = result.post_list_visited_elements.size(); | |
| DCHECK_GT(visited_component_size, 0); | |
| - if (visited_component_size != 1) { | |
| - return Status::NotImplemented( | |
| - "Lists with non-zero length null components are not supported"); | |
| + std::shared_ptr<Array> values_array; | |
| + if (visited_component_size == 1) { | |
| + const ElementRange& range = result.post_list_visited_elements[0]; | |
| + values_array = result.leaf_array->Slice(range.start, range.Size()); | |
| + } else { | |
| + // Multiple leaf ranges can be produced when child values are | |
| + // skipped, such as null fixed-size-list slots, or when | |
| + // list-view ranges are non-contiguous. Concatenate the slices | |
| + // in logical write order. | |
| + ::arrow::ArrayVector arrays; | |
| + arrays.reserve(visited_component_size); | |
| + for (const auto& range : result.post_list_visited_elements) { | |
| + DCHECK(!range.Empty()); | |
| + arrays.push_back(result.leaf_array->Slice(range.start, range.Size())); | |
| + } | |
| + ARROW_ASSIGN_OR_RAISE(values_array, | |
| + ::arrow::Concatenate(arrays, ctx->memory_pool)); | |
| } | |
| - const ElementRange& range = result.post_list_visited_elements[0]; | |
| - std::shared_ptr<Array> values_array = | |
| - result.leaf_array->Slice(range.start, range.Size()); | |
| return column_writer->WriteArrow(result.def_levels, result.rep_levels, | |
| result.def_rep_level_count, *values_array, | |
| @@ -314,6 +326,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 +438,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/column_reader.cc b/cpp/src/parquet/column_reader.cc | |
| index ebf9515f27..00b96797d9 100644 | |
| --- a/cpp/src/parquet/column_reader.cc | |
| +++ b/cpp/src/parquet/column_reader.cc | |
| @@ -208,6 +208,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; | |
| @@ -2156,6 +2242,22 @@ class ByteArrayChunkedRecordReader final : public TypedRecordReader<ByteArrayTyp | |
| return result; | |
| } | |
| + void ReserveValueBytes(int64_t num_values, int64_t num_bytes) override { | |
| + if (num_values > 0) { | |
| + PARQUET_THROW_NOT_OK(accumulator_.builder->Reserve(num_values)); | |
| + } | |
| + if (num_bytes > 0) { | |
| + // Never reserve past the current chunk: ArrowBinaryHelper pushes a chunk | |
| + // when the builder's data would exceed kBinaryMemoryLimit, and | |
| + // ReserveData() would fail the overflow check on a request past it. The | |
| + // clamp mirrors the one in ArrowBinaryHelper::Prepare. | |
| + const int64_t room = | |
| + ::arrow::kBinaryMemoryLimit - accumulator_.builder->value_data_length(); | |
| + PARQUET_THROW_NOT_OK( | |
| + accumulator_.builder->ReserveData(std::min<int64_t>(num_bytes, room))); | |
| + } | |
| + } | |
| + | |
| void ReadValuesDense(int64_t values_to_read) override { | |
| int64_t num_decoded = this->current_decoder_->DecodeArrowNonNull( | |
| static_cast<int>(values_to_read), &accumulator_); | |
| diff --git a/cpp/src/parquet/column_reader.h b/cpp/src/parquet/column_reader.h | |
| index 29e1b2a25e..a7e68b7dd6 100644 | |
| --- 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 { | |
| @@ -341,6 +370,19 @@ class PARQUET_EXPORT RecordReader { | |
| /// \brief Pre-allocate space for data. Results in better flat read performance | |
| virtual void Reserve(int64_t num_values) = 0; | |
| + /// \brief Pre-allocate the variable-width value buffers for `num_values` | |
| + /// values totalling about `num_bytes` bytes. | |
| + /// | |
| + /// Fixed-width readers size their values buffer from Reserve() alone, so the | |
| + /// default is a no-op. The BYTE_ARRAY readers build into an | |
| + /// ::arrow::BinaryBuilder that Reserve() deliberately leaves alone | |
| + /// (uses_values_ is false for them), which is the only reason this second | |
| + /// hook exists: without it the builder's offsets and data buffers grow by | |
| + /// doubling once per decoded batch, and every doubling copies everything the | |
| + /// row group has accumulated so far. Both counts are additive, like | |
| + /// Reserve(); zero means "no estimate" and reserves nothing. | |
| + virtual void ReserveValueBytes(int64_t num_values, int64_t num_bytes) {} | |
| + | |
| /// \brief Clear consumed values and repetition/definition levels as the | |
| /// result of calling ReadRecords | |
| /// For FLBA and ByteArray types, call GetBuilderChunks() to reset them. | |
| diff --git a/cpp/src/parquet/file_reader.cc b/cpp/src/parquet/file_reader.cc | |
| index 3e9eeea6c6..671ebe4644 100644 | |
| --- a/cpp/src/parquet/file_reader.cc | |
| +++ b/cpp/src/parquet/file_reader.cc | |
| @@ -207,6 +207,117 @@ const RowGroupMetaData* RowGroupReader::metadata() const { return contents_->met | |
| 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 +353,11 @@ class SerializedRowGroup : public RowGroupReader::Contents { | |
| // 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 +533,26 @@ class SerializedFile : public ParquetFileReader::Contents { | |
| 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 +1047,22 @@ void ParquetFileReader::PreBuffer(const std::vector<int>& row_groups, | |
| 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/src/parquet/file_reader.h b/cpp/src/parquet/file_reader.h | |
| index b59b59f95c..657a438a3a 100644 | |
| --- a/cpp/src/parquet/file_reader.h | |
| +++ b/cpp/src/parquet/file_reader.h | |
| @@ -210,6 +210,17 @@ class PARQUET_EXPORT ParquetFileReader { | |
| ::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_; | |
| 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_; |