diff --git a/cmake_modules/arrow.diff b/cmake_modules/arrow.diff index ae7752097..c11da3429 100644 --- a/cmake_modules/arrow.diff +++ b/cmake_modules/arrow.diff @@ -431,9 +431,305 @@ diff --git a/cpp/cmake_modules/BuildUtils.cmake b/cpp/cmake_modules/BuildUtils.c 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 -@@ -211,7 +211,7 @@ +@@ -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& column_indices, ++ FileColumnIteratorFactory iterator_factory, ++ std::unique_ptr* 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 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() override { return field_; } + + private: +@@ -532,6 +567,22 @@ + return storage_reader_->LoadBatch(number_of_records); + } + ++ std::vector 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* out) override { + std::shared_ptr storage; +@@ -576,6 +627,22 @@ + return item_reader_->LoadBatch(number_of_records); + } + ++ std::vector 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> AssembleArray( + std::shared_ptr data) { + if (field_->type()->id() == ::arrow::Type::MAP) { +@@ -709,6 +776,39 @@ + } + return Status::OK(); + } ++ ++ std::vector LeafColumnIndices() const override { ++ std::vector indices; ++ for (const std::unique_ptr& reader : children_) { ++ std::vector 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& 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& 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& reader : children_) { ++ read += reader->ReadRecords(col_idx, num_records); ++ } ++ return read; ++ } ++ + Status BuildArray(int64_t length_upper_bound, + std::shared_ptr* out) override; + Status GetDefLevels(const int16_t** data, int64_t* length) override; +@@ -1228,6 +1328,23 @@ + std::unique_ptr 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& column_indices, ++ FileColumnIteratorFactory iterator_factory, ++ std::unique_ptr* out) { ++ RETURN_NOT_OK(BoundsCheckColumn(i)); ++ auto ctx = std::make_shared(); ++ 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 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 + #include ++#include + #include + + #include "parquet/file_reader.h" +@@ -48,9 +49,13 @@ + + class ColumnChunkReader; + class ColumnReader; ++class FileColumnIterator; + struct SchemaManifest; + class RowGroupReader; + ++using FileColumnIteratorFactory = ++ std::function; ++ + /// \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* 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& column_indices, ++ FileColumnIteratorFactory iterator_factory, ++ std::unique_ptr* 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 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 + #include + ++#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 row_groups_; + }; + +-using FileColumnIteratorFactory = +- std::function; +- + Status TransferColumnData(::parquet::internal::RecordReader* reader, + const std::shared_ptr<::arrow::Field>& value_field, + const ColumnDescriptor* descr, ::arrow::MemoryPool* pool, diff --git a/src/paimon/format/parquet/file_reader_wrapper.cpp b/src/paimon/format/parquet/file_reader_wrapper.cpp index 21cfa930b..54c70798b 100644 --- a/src/paimon/format/parquet/file_reader_wrapper.cpp +++ b/src/paimon/format/parquet/file_reader_wrapper.cpp @@ -30,6 +30,7 @@ #include "paimon/format/parquet/parquet_format_defs.h" #include "paimon/macros.h" #include "parquet/arrow/reader.h" +#include "parquet/arrow/schema.h" #include "parquet/file_reader.h" #include "parquet/metadata.h" #include "parquet/page_index.h" @@ -230,15 +231,14 @@ Result> FileReaderWrapper::NextPageFiltered( if (!current_page_filtered_reader_) { const auto& target_rg = target_row_groups_[current_row_group_idx_]; auto page_ranges = PageFilteredRowGroupReader::ComputePageRanges( - file_reader_->parquet_reader(), target_rg, target_column_indices_); + target_rg, target_column_indices_, file_reader_->parquet_reader()); bool pre_buffered = !prebuffered_ranges_.empty(); int64_t max_chunksize = batch_size_ > 0 ? batch_size_ : std::numeric_limits::max(); PAIMON_ASSIGN_OR_RAISE( current_page_filtered_reader_, PageFilteredRowGroupReader::ReadFilteredRowGroup( - file_reader_->parquet_reader(), target_rg, target_column_indices_, - page_filtered_read_schema_, file_reader_->properties().cache_options(), - pre_buffered, page_ranges, max_chunksize, pool_)); + target_rg, target_column_indices_, file_reader_->properties().cache_options(), + pre_buffered, page_ranges, max_chunksize, pool_, file_reader_.get())); current_filtered_row_ranges_ = target_rg.GetRowRanges(); current_filtered_rg_start_ = all_row_group_ranges_[rg_id].first; filtered_global_offset_ = 0; @@ -345,29 +345,6 @@ Status FileReaderWrapper::PrepareForReadingLazy( return Status::OK(); } -Status FileReaderWrapper::BuildPageFilteredSchema(const std::vector& column_indices) { - if (page_filtered_read_schema_) { - return Status::OK(); - } - std::shared_ptr schema; - PAIMON_RETURN_NOT_OK_FROM_ARROW(file_reader_->GetSchema(&schema)); - auto parquet_schema = file_reader_->parquet_reader()->metadata()->schema(); - std::vector> fields; - for (int32_t col_idx : column_indices) { - const std::string& col_name = parquet_schema->Column(col_idx)->name(); - auto field = schema->GetFieldByName(col_name); - if (!field) { - return Status::Invalid(fmt::format( - "PrepareForReading: Parquet column {} ('{}') has no matching Arrow field in " - "file schema", - col_idx, col_name)); - } - fields.push_back(field); - } - page_filtered_read_schema_ = arrow::schema(fields); - return Status::OK(); -} - std::vector<::arrow::io::ReadRange> FileReaderWrapper::CollectPreBufferRanges( const std::vector& column_indices) { std::vector<::arrow::io::ReadRange> ranges; @@ -379,7 +356,7 @@ std::vector<::arrow::io::ReadRange> FileReaderWrapper::CollectPreBufferRanges( if (trg.IsPartiallyMatched()) { // Page-filtered RGs: only matching page byte ranges. auto page_ranges = PageFilteredRowGroupReader::ComputePageRanges( - file_reader_->parquet_reader(), trg, column_indices); + trg, column_indices, file_reader_->parquet_reader()); ranges.insert(ranges.end(), std::make_move_iterator(page_ranges.begin()), std::make_move_iterator(page_ranges.end())); } else { @@ -416,7 +393,6 @@ Status FileReaderWrapper::PrepareForReading(const std::vector& t try { target_row_groups_ = target_row_groups; target_column_indices_ = column_indices; - page_filtered_read_schema_.reset(); // Partition into fully-matched and page-filtered row groups, skipping excluded ones. std::vector fully_matched_row_groups; @@ -432,9 +408,6 @@ Status FileReaderWrapper::PrepareForReading(const std::vector& t } bool has_partially_matched = fully_matched_row_groups.size() != active_count; - if (has_partially_matched) { - PAIMON_RETURN_NOT_OK(BuildPageFilteredSchema(column_indices)); - } WaitForPendingPreBuffer(); diff --git a/src/paimon/format/parquet/file_reader_wrapper.h b/src/paimon/format/parquet/file_reader_wrapper.h index 3221dbf12..5160ac19b 100644 --- a/src/paimon/format/parquet/file_reader_wrapper.h +++ b/src/paimon/format/parquet/file_reader_wrapper.h @@ -158,9 +158,6 @@ class FileReaderWrapper { /// Read next batch from the fully-matched batch_reader_. Returns nullptr when exhausted. Result> NextFullyMatched(); - /// Build page_filtered_read_schema_ from the given column indices. No-op if already built. - Status BuildPageFilteredSchema(const std::vector& column_indices); - /// Collect all byte ranges that need pre-buffering (page-filtered + fully-matched). std::vector<::arrow::io::ReadRange> CollectPreBufferRanges( const std::vector& column_indices); @@ -194,11 +191,6 @@ class FileReaderWrapper { // Target row groups with row ranges for none page-level filtering and page-level filtering std::vector target_row_groups_; - // Arrow schema covering target_column_indices_, used when constructing the per-RG - // page-filtered reader. Cached in PrepareForReading because it's identical across - // all page-filtered RGs in a session. - std::shared_ptr page_filtered_read_schema_; - // Track pre-buffered ranges so we can wait on destruction std::vector<::arrow::io::ReadRange> prebuffered_ranges_; }; diff --git a/src/paimon/format/parquet/page_filtered_row_group_reader.cpp b/src/paimon/format/parquet/page_filtered_row_group_reader.cpp index 5c6fb8653..07b057b23 100644 --- a/src/paimon/format/parquet/page_filtered_row_group_reader.cpp +++ b/src/paimon/format/parquet/page_filtered_row_group_reader.cpp @@ -28,7 +28,9 @@ #include "fmt/format.h" #include "paimon/common/utils/arrow/arrow_utils.h" #include "paimon/common/utils/arrow/status_utils.h" +#include "parquet/arrow/reader.h" #include "parquet/arrow/reader_internal.h" +#include "parquet/arrow/schema.h" #include "parquet/metadata.h" #include "parquet/schema.h" @@ -74,6 +76,29 @@ class TableRecordBatchReader : public arrow::RecordBatchReader { std::shared_ptr pool_; }; +/// A FileColumnIterator that installs a data_page_filter on every PageReader it +/// produces, enabling I/O-level page skipping. The base class handles row group +/// iteration; this subclass only decorates the PageReader returned by NextChunk(). +class PageFilteringColumnIterator : public ::parquet::arrow::FileColumnIterator { + public: + PageFilteringColumnIterator( + int column_index, ::parquet::ParquetFileReader* reader, std::vector row_groups, + std::function data_page_filter) + : FileColumnIterator(column_index, reader, std::move(row_groups)), + data_page_filter_(std::move(data_page_filter)) {} + + std::unique_ptr<::parquet::PageReader> NextChunk() override { + std::unique_ptr<::parquet::PageReader> page_reader = FileColumnIterator::NextChunk(); + if (page_reader && data_page_filter_) { + page_reader->set_data_page_filter(data_page_filter_); + } + return page_reader; + } + + private: + std::function data_page_filter_; +}; + } // namespace std::pair PageFilteredRowGroupReader::GetPageRowRange( @@ -139,92 +164,35 @@ std::pair PageFilteredRowGroupReader::ComputeCompressedRowRa } Status PageFilteredRowGroupReader::ExecuteSkipReadPattern( - const std::shared_ptr<::parquet::internal::RecordReader>& record_reader, - const RowRanges& ranges, int64_t total_row_count, int32_t row_group_index, - int32_t column_index) { - int64_t current_row = 0; + int col_idx, const RowRanges& ranges, int64_t total, + ::parquet::arrow::ColumnReader* column_reader) { + PAIMON_RETURN_NOT_OK_FROM_ARROW(column_reader->ResetLeaf(col_idx, total)); + int64_t current = 0; for (const auto& range : ranges.GetRanges()) { - if (range.from > current_row) { - int64_t to_skip = range.from - current_row; - int64_t skipped = record_reader->SkipRecords(to_skip); - if (skipped != to_skip) { - return Status::Invalid(fmt::format( - "PageFilteredRowGroupReader: expected to skip {} records but skipped {} " - "(row_group={}, column={})", - to_skip, skipped, row_group_index, column_index)); - } - current_row = range.from; + int64_t skip = range.from > current ? range.from - current : 0; + int64_t skipped = column_reader->SkipRecords(col_idx, skip); + if (skipped != skip) { + return Status::Invalid(fmt::format( + "PageFilteredRowGroupReader: leaf {} expected to skip {} records but skipped {}", + col_idx, skip, skipped)); } int64_t to_read = range.Count(); - int64_t read = record_reader->ReadRecords(to_read); + int64_t read = column_reader->ReadRecords(col_idx, to_read); if (read != to_read) { - return Status::Invalid( - fmt::format("PageFilteredRowGroupReader: expected to read {} records but read {} " - "(row_group={}, column={}, range=[{},{}])", - to_read, read, row_group_index, column_index, range.from, range.to)); + return Status::Invalid(fmt::format( + "PageFilteredRowGroupReader: leaf {} expected to read {} records but read {}", + col_idx, to_read, read)); } - current_row += to_read; - } - if (current_row < total_row_count) { - record_reader->SkipRecords(total_row_count - current_row); + current = range.to + 1; } return Status::OK(); } -Result> PageFilteredRowGroupReader::ReadFilteredColumn( - const std::shared_ptr<::parquet::RowGroupReader>& row_group_reader, - ::parquet::ParquetFileReader* parquet_reader, - const std::shared_ptr<::parquet::RowGroupPageIndexReader>& rg_page_index_reader, - int32_t row_group_index, int32_t column_index, const RowRanges& row_ranges, - const std::shared_ptr& field, int64_t row_group_row_count, - std::shared_ptr<::arrow::MemoryPool> pool) { - auto file_metadata = parquet_reader->metadata(); - const auto* col_descriptor = file_metadata->schema()->Column(column_index); - - // Try to get OffsetIndex for I/O-level page skipping - RowRanges effective_ranges = row_ranges; - int64_t effective_row_count = row_group_row_count; - - std::shared_ptr<::parquet::OffsetIndex> offset_index; - if (rg_page_index_reader) { - offset_index = rg_page_index_reader->GetOffsetIndex(column_index); - } - - auto page_reader = row_group_reader->GetColumnPageReader(column_index); - - if (offset_index) { - // Set data_page_filter for I/O-level page skipping - page_reader->set_data_page_filter( - MakePageFilter(row_ranges, offset_index, row_group_row_count)); - // Compute compressed RowRanges for the decode-level skip/read pattern - auto [compressed_ranges, compressed_total] = - ComputeCompressedRowRanges(row_ranges, offset_index, row_group_row_count); - effective_ranges = std::move(compressed_ranges); - effective_row_count = compressed_total; - } - - // Create RecordReader - ::parquet::internal::LevelInfo leaf_info = - ::parquet::internal::LevelInfo::ComputeLevelInfo(col_descriptor); - auto record_reader = - ::parquet::internal::RecordReader::Make(col_descriptor, leaf_info, pool.get()); - record_reader->SetPageReader(std::move(page_reader)); - - PAIMON_RETURN_NOT_OK(ExecuteSkipReadPattern( - record_reader, effective_ranges, effective_row_count, row_group_index, column_index)); - - std::shared_ptr chunked_array; - PAIMON_RETURN_NOT_OK_FROM_ARROW(::parquet::arrow::TransferColumnData( - record_reader.get(), field, col_descriptor, pool.get(), &chunked_array)); - - return chunked_array; -} - Status PageFilteredRowGroupReader::WaitForPreBuffer( - ::parquet::ParquetFileReader* parquet_reader, int32_t row_group_index, - const std::vector& column_indices, const ::arrow::io::CacheOptions& cache_options, - bool pre_buffered, const std::vector<::arrow::io::ReadRange>& page_ranges, - std::shared_ptr<::arrow::MemoryPool> pool) { + int32_t row_group_index, const std::vector& column_indices, + const ::arrow::io::CacheOptions& cache_options, bool pre_buffered, + const std::vector<::arrow::io::ReadRange>& page_ranges, + std::shared_ptr<::arrow::MemoryPool> pool, ::parquet::ParquetFileReader* parquet_reader) { std::vector rg_vec = {row_group_index}; std::vector col_vec(column_indices.begin(), column_indices.end()); if (!pre_buffered) { @@ -244,28 +212,84 @@ Status PageFilteredRowGroupReader::WaitForPreBuffer( return Status::OK(); } +Result> PageFilteredRowGroupReader::ReadFilteredField( + const std::shared_ptr<::parquet::RowGroupPageIndexReader>& rg_page_index_reader, + int32_t row_group_index, int32_t field_index, const std::vector& column_indices, + const RowRanges& row_ranges, int64_t row_group_row_count, + ::parquet::arrow::FileReader* arrow_file_reader) { + // Factory: set data_page_filter on every leaf (per-leaf OffsetIndex). + // data_page_filter enables I/O-level page skipping for all leaves. + auto factory = + [row_group_index, &rg_page_index_reader, &row_ranges, row_group_row_count]( + int col_idx, + ::parquet::ParquetFileReader* reader) -> ::parquet::arrow::FileColumnIterator* { + std::function data_page_filter; + if (rg_page_index_reader) { + auto offset_index = rg_page_index_reader->GetOffsetIndex(col_idx); + if (offset_index) { + data_page_filter = MakePageFilter(row_ranges, offset_index, row_group_row_count); + } + } + return new PageFilteringColumnIterator(col_idx, reader, std::vector{row_group_index}, + std::move(data_page_filter)); + }; + + // Build reader tree with leaf column filtering + std::unique_ptr<::parquet::arrow::ColumnReader> column_reader; + PAIMON_RETURN_NOT_OK_FROM_ARROW( + arrow_file_reader->GetColumn(field_index, column_indices, factory, &column_reader)); + + if (!column_reader) { + return Status::Invalid( + fmt::format("PageFilteredRowGroupReader: field {} has no matching leaf columns " + "(row_group={})", + field_index, row_group_index)); + } + + // Since leaf columns may have misaligned pages, we compute compressed row ranges and drive each + // leaf column independently + + for (int col_idx : column_reader->LeafColumnIndices()) { + RowRanges effective_ranges = row_ranges; + int64_t effective_total = row_group_row_count; + if (rg_page_index_reader) { + auto offset_index = rg_page_index_reader->GetOffsetIndex(col_idx); + if (offset_index) { + auto [compressed, total] = + ComputeCompressedRowRanges(row_ranges, offset_index, row_group_row_count); + effective_ranges = std::move(compressed); + effective_total = total; + } + } + + PAIMON_RETURN_NOT_OK(ExecuteSkipReadPattern(col_idx, effective_ranges, effective_total, + column_reader.get())); + } + + // Build the Arrow array (TransferColumnData for leaves + assemble for nested) + std::shared_ptr chunked_array; + PAIMON_RETURN_NOT_OK_FROM_ARROW( + column_reader->BuildArray(row_ranges.RowCount(), &chunked_array)); + + return chunked_array; +} + Result> PageFilteredRowGroupReader::ReadFilteredRowGroup( - ::parquet::ParquetFileReader* parquet_reader, const TargetRowGroup& target_row_group, - const std::vector& column_indices, const std::shared_ptr& arrow_schema, + const TargetRowGroup& target_row_group, const std::vector& column_indices, const ::arrow::io::CacheOptions& cache_options, bool pre_buffered, const std::vector<::arrow::io::ReadRange>& page_ranges, int64_t max_chunksize, - std::shared_ptr<::arrow::MemoryPool> pool) { + std::shared_ptr<::arrow::MemoryPool> pool, ::parquet::arrow::FileReader* arrow_file_reader) { + auto parquet_reader = arrow_file_reader->parquet_reader(); const auto& row_ranges = target_row_group.GetRowRanges(); int32_t row_group_index = target_row_group.GetRowGroupIndex(); - if (row_ranges.IsEmpty()) { - PAIMON_ASSIGN_OR_RAISE_FROM_ARROW(std::shared_ptr empty_table, - arrow::Table::MakeEmpty(arrow_schema, pool.get())); - return std::make_unique(std::move(empty_table), max_chunksize, - pool); - } - int64_t expected_rows = row_ranges.RowCount(); - PAIMON_RETURN_NOT_OK(WaitForPreBuffer(parquet_reader, row_group_index, column_indices, - cache_options, pre_buffered, page_ranges, pool)); + if (!row_ranges.IsEmpty()) { + PAIMON_RETURN_NOT_OK(WaitForPreBuffer(row_group_index, column_indices, cache_options, + pre_buffered, page_ranges, pool, parquet_reader)); + } - auto row_group_reader = parquet_reader->RowGroup(row_group_index); auto rg_metadata = parquet_reader->metadata()->RowGroup(row_group_index); int64_t row_group_row_count = rg_metadata->num_rows(); @@ -277,36 +301,45 @@ Result> PageFilteredRowGroupReader::Re rg_page_index_reader = page_index_reader->RowGroup(row_group_index); } - // Read each column with page filtering - std::vector> columns; - columns.reserve(column_indices.size()); + const auto& manifest = arrow_file_reader->manifest(); + PAIMON_ASSIGN_OR_RAISE_FROM_ARROW( + std::vector field_indices, + manifest.GetFieldIndices(std::vector(column_indices.begin(), column_indices.end()))); - for (size_t i = 0; i < column_indices.size(); ++i) { + std::vector> result_arrays; + result_arrays.reserve(field_indices.size()); + + for (int field_idx : field_indices) { PAIMON_ASSIGN_OR_RAISE( std::shared_ptr chunked_array, - ReadFilteredColumn(row_group_reader, parquet_reader, rg_page_index_reader, - row_group_index, column_indices[i], row_ranges, - arrow_schema->field(static_cast(i)), row_group_row_count, - pool)); + ReadFilteredField(rg_page_index_reader, row_group_index, field_idx, column_indices, + row_ranges, row_group_row_count, arrow_file_reader)); if (chunked_array->length() != expected_rows) { - return Status::Invalid(fmt::format( - "PageFilteredRowGroupReader: column {} produced {} rows but expected {} " - "(row_group={})", - column_indices[i], chunked_array->length(), expected_rows, row_group_index)); + return Status::Invalid( + fmt::format("PageFilteredRowGroupReader: field {} produced {} rows but expected {} " + "(row_group={})", + field_idx, chunked_array->length(), expected_rows, row_group_index)); } - columns.push_back(std::move(chunked_array)); + result_arrays.push_back(std::move(chunked_array)); + } + + std::vector> result_fields; + for (size_t i = 0; i < result_arrays.size(); ++i) { + const auto& field = manifest.schema_fields[field_indices[i]].field; + result_fields.push_back(arrow::field(field->name(), result_arrays[i]->type(), + field->nullable(), field->metadata())); } + auto result_schema = arrow::schema(result_fields); - auto table = arrow::Table::Make(arrow_schema, std::move(columns), expected_rows); - return std::make_unique(std::move(table), max_chunksize, - std::move(pool)); + auto table = arrow::Table::Make(result_schema, std::move(result_arrays), expected_rows); + return std::make_unique(std::move(table), max_chunksize, pool); } std::vector<::arrow::io::ReadRange> PageFilteredRowGroupReader::ComputePageRanges( - ::parquet::ParquetFileReader* parquet_reader, const TargetRowGroup& target_row_group, - const std::vector& column_indices) { + const TargetRowGroup& target_row_group, const std::vector& column_indices, + ::parquet::ParquetFileReader* parquet_reader) { int32_t row_group_index = target_row_group.GetRowGroupIndex(); const auto& row_ranges = target_row_group.GetRowRanges(); diff --git a/src/paimon/format/parquet/page_filtered_row_group_reader.h b/src/paimon/format/parquet/page_filtered_row_group_reader.h index c7376512f..1b907cf7d 100644 --- a/src/paimon/format/parquet/page_filtered_row_group_reader.h +++ b/src/paimon/format/parquet/page_filtered_row_group_reader.h @@ -29,6 +29,7 @@ #include "paimon/format/parquet/row_ranges.h" #include "paimon/format/parquet/target_row_group.h" #include "paimon/result.h" +#include "parquet/arrow/reader.h" #include "parquet/column_reader.h" #include "parquet/file_reader.h" #include "parquet/page_index.h" @@ -38,39 +39,35 @@ namespace paimon::parquet { /// Reads a single row group using page-level filtering. /// Non-matching rows are skipped at the decoding level via RecordReader::SkipRecords, /// using RowRanges computed from the page index (ColumnIndex + OffsetIndex). -/// MakePageFilter is available for future I/O-level page skipping optimization. class PageFilteredRowGroupReader { public: PageFilteredRowGroupReader() = delete; ~PageFilteredRowGroupReader() = delete; /// Read a row group with page-level filtering. - /// @param parquet_reader The underlying ParquetFileReader /// @param target_row_group Target row group with index and row ranges /// @param column_indices Leaf column indices to read - /// @param arrow_schema The target Arrow schema for output columns /// @param pool Memory pool /// @param cache_options Cache options for PreBuffer /// @param pre_buffered If true, assumes PreBuffer was already called externally /// and only waits via WhenBuffered (no redundant PreBuffer). /// @param page_ranges If non-empty, wait via WhenBufferedRanges instead of WhenBuffered /// @param max_chunksize Per-batch row cap for the returned reader. + /// @param arrow_file_reader The Arrow FileReader for ColumnReader tree creation /// @return A RecordBatchReader streaming the filtered rows. static Result> ReadFilteredRowGroup( - ::parquet::ParquetFileReader* parquet_reader, const TargetRowGroup& target_row_group, - const std::vector& column_indices, - const std::shared_ptr& arrow_schema, + const TargetRowGroup& target_row_group, const std::vector& column_indices, const ::arrow::io::CacheOptions& cache_options, bool pre_buffered, const std::vector<::arrow::io::ReadRange>& page_ranges, int64_t max_chunksize, - std::shared_ptr<::arrow::MemoryPool> pool); + std::shared_ptr<::arrow::MemoryPool> pool, ::parquet::arrow::FileReader* arrow_file_reader); /// Compute the byte ranges of pages that overlap with the given RowRanges. /// Uses OffsetIndex to determine per-page file offsets and sizes. /// Includes dictionary pages unconditionally. /// Falls back to entire column chunk range if OffsetIndex is unavailable. static std::vector<::arrow::io::ReadRange> ComputePageRanges( - ::parquet::ParquetFileReader* parquet_reader, const TargetRowGroup& target_row_group, - const std::vector& column_indices); + const TargetRowGroup& target_row_group, const std::vector& column_indices, + ::parquet::ParquetFileReader* parquet_reader); private: /// Get the [first_row, last_row] range of a page given page locations. @@ -79,38 +76,38 @@ class PageFilteredRowGroupReader { int64_t row_group_row_count); /// Wait for pre-buffered data to become available before reading. - static Status WaitForPreBuffer(::parquet::ParquetFileReader* parquet_reader, - int32_t row_group_index, + static Status WaitForPreBuffer(int32_t row_group_index, const std::vector& column_indices, const ::arrow::io::CacheOptions& cache_options, bool pre_buffered, const std::vector<::arrow::io::ReadRange>& page_ranges, - std::shared_ptr<::arrow::MemoryPool> pool); - - /// Execute the skip/read pattern on a RecordReader based on RowRanges. - static Status ExecuteSkipReadPattern( - const std::shared_ptr<::parquet::internal::RecordReader>& record_reader, - const RowRanges& ranges, int64_t total_row_count, int32_t row_group_index, - int32_t column_index); + std::shared_ptr<::arrow::MemoryPool> pool, + ::parquet::ParquetFileReader* parquet_reader); /// Create a data_page_filter callback for a column based on RowRanges + OffsetIndex. static std::function MakePageFilter( const RowRanges& row_ranges, const std::shared_ptr<::parquet::OffsetIndex>& offset_index, int64_t row_group_row_count); - /// Read a single column using skip/read pattern driven by RowRanges. - static Result> ReadFilteredColumn( - const std::shared_ptr<::parquet::RowGroupReader>& row_group_reader, - ::parquet::ParquetFileReader* parquet_reader, - const std::shared_ptr<::parquet::RowGroupPageIndexReader>& rg_page_index_reader, - int32_t row_group_index, int32_t column_index, const RowRanges& row_ranges, - const std::shared_ptr& field, int64_t row_group_row_count, - std::shared_ptr<::arrow::MemoryPool> pool); - /// Compute compressed RowRanges after data_page_filter skips non-matching pages. static std::pair ComputeCompressedRowRanges( const RowRanges& original_ranges, const std::shared_ptr<::parquet::OffsetIndex>& offset_index, int64_t row_group_row_count); + + /// Reset the given leaf and replay the skip/read pattern derived from `ranges` + /// directly against the ColumnReader (ResetLeaf + SkipRecords/ReadRecords). + static Status ExecuteSkipReadPattern(int col_idx, const RowRanges& ranges, int64_t total, + ::parquet::arrow::ColumnReader* column_reader); + + /// Read a field (flat or nested) using ColumnReader tree. + /// Sets data_page_filter on all leaves via factory, then drives each leaf + /// independently via ResetLeaf/SkipRecords/ReadRecords using its own + /// compressed_ranges. + static Result> ReadFilteredField( + const std::shared_ptr<::parquet::RowGroupPageIndexReader>& rg_page_index_reader, + int32_t row_group_index, int32_t field_index, const std::vector& column_indices, + const RowRanges& row_ranges, int64_t row_group_row_count, + ::parquet::arrow::FileReader* arrow_file_reader); }; } // namespace paimon::parquet diff --git a/src/paimon/format/parquet/page_filtered_row_group_reader_test.cpp b/src/paimon/format/parquet/page_filtered_row_group_reader_test.cpp index 361c8b199..e1c90efd8 100644 --- a/src/paimon/format/parquet/page_filtered_row_group_reader_test.cpp +++ b/src/paimon/format/parquet/page_filtered_row_group_reader_test.cpp @@ -17,6 +17,7 @@ #include "paimon/format/parquet/page_filtered_row_group_reader.h" #include +#include #include #include #include @@ -30,6 +31,8 @@ #include "arrow/c/bridge.h" #include "arrow/ipc/json_simple.h" #include "gtest/gtest.h" +#include "paimon/common/data/variant/generic_variant.h" +#include "paimon/common/data/variant/variant_type_utils.h" #include "paimon/common/utils/arrow/arrow_input_stream_adapter.h" #include "paimon/common/utils/arrow/mem_utils.h" #include "paimon/defs.h" @@ -44,6 +47,7 @@ #include "paimon/status.h" #include "paimon/testing/utils/read_result_collector.h" #include "paimon/testing/utils/testharness.h" +#include "paimon/testing/utils/variant_test_data.h" #include "paimon/utils/roaring_bitmap32.h" #include "parquet/arrow/reader.h" #include "parquet/file_reader.h" @@ -76,7 +80,7 @@ class PageFilteredRowGroupReaderTest : public ::testing::Test { void WriteTestFile(const std::string& file_name, const std::shared_ptr& struct_array, int32_t write_batch_size, int64_t max_row_group_length, - bool enable_dictionary = false) { + bool enable_dictionary = false, int64_t data_page_size = 1) { auto data_type = struct_array->struct_type(); auto data_schema = arrow::schema(data_type->fields()); auto data_arrow_array = std::make_unique(); @@ -92,10 +96,12 @@ class PageFilteredRowGroupReaderTest : public ::testing::Test { builder.disable_dictionary(); // Ensure page index min/max are meaningful } builder.enable_write_page_index(); // Enable page index for page-level filtering - // Set data page size to 1 byte to force a new page after every write_batch_size rows. - // The writer flushes a page when accumulated data exceeds data_pagesize, so setting - // it to 1 ensures each batch of write_batch_size rows becomes exactly one page. - builder.data_pagesize(1); + // Data page size controls when a page is flushed. The default of 1 byte forces a new + // page after every write_batch_size rows (each batch becomes one page), giving pages + // aligned across columns. A larger byte-based value combined with write_batch_size=1 + // instead lets columns of different physical widths flush pages at different row + // counts, producing intentionally misaligned pages across leaves. + builder.data_pagesize(data_page_size); auto writer_properties = builder.build(); ASSERT_OK_AND_ASSIGN( auto format_writer, @@ -551,9 +557,8 @@ TEST_F(PageFilteredRowGroupReaderTest, ComputePageRangesPartialMatch) { row_ranges.Add(RowRanges::Range(50, 59)); auto ranges = PageFilteredRowGroupReader::ComputePageRanges( - parquet_reader.get(), TargetRowGroup(/*rg_index=*/0, /*is_partially_matched=*/true, /*ranges=*/row_ranges), - /*column_indices=*/{0}); + /*column_indices=*/{0}, parquet_reader.get()); // Should have exactly 1 range (page 5 of column 0, no dictionary since disabled) ASSERT_EQ(1, ranges.size()); @@ -577,8 +582,8 @@ TEST_F(PageFilteredRowGroupReaderTest, ComputePageRangesAllMatch) { row_ranges.Add(RowRanges::Range(0, 99)); auto ranges = PageFilteredRowGroupReader::ComputePageRanges( - parquet_reader.get(), - TargetRowGroup(/*rg_index=*/0, /*is_partially_matched=*/true, /*ranges=*/row_ranges), {0}); + TargetRowGroup(/*rg_index=*/0, /*is_partially_matched=*/true, /*ranges=*/row_ranges), {0}, + parquet_reader.get()); // 10 pages, all matching ASSERT_EQ(10, ranges.size()); @@ -602,8 +607,8 @@ TEST_F(PageFilteredRowGroupReaderTest, ComputePageRangesNoMatch) { RowRanges row_ranges; // empty auto ranges = PageFilteredRowGroupReader::ComputePageRanges( - parquet_reader.get(), - TargetRowGroup(/*rg_index=*/0, /*is_partially_matched=*/true, /*ranges=*/row_ranges), {0}); + TargetRowGroup(/*rg_index=*/0, /*is_partially_matched=*/true, /*ranges=*/row_ranges), {0}, + parquet_reader.get()); ASSERT_EQ(0, ranges.size()); } @@ -624,9 +629,8 @@ TEST_F(PageFilteredRowGroupReaderTest, ComputePageRangesMultiColumn) { row_ranges.Add(RowRanges::Range(50, 59)); auto ranges = PageFilteredRowGroupReader::ComputePageRanges( - parquet_reader.get(), TargetRowGroup(/*rg_index=*/0, /*is_partially_matched=*/true, /*ranges=*/row_ranges), - {0, 1}); + {0, 1}, parquet_reader.get()); // 1 matching page per column = 2 ranges total ASSERT_EQ(2, ranges.size()); @@ -652,8 +656,8 @@ TEST_F(PageFilteredRowGroupReaderTest, ComputePageRangesMultiplePages) { row_ranges.Add(RowRanges::Range(70, 79)); auto ranges = PageFilteredRowGroupReader::ComputePageRanges( - parquet_reader.get(), - TargetRowGroup(/*rg_index=*/0, /*is_partially_matched=*/true, /*ranges=*/row_ranges), {0}); + TargetRowGroup(/*rg_index=*/0, /*is_partially_matched=*/true, /*ranges=*/row_ranges), {0}, + parquet_reader.get()); // 2 matching pages for 1 column ASSERT_EQ(2, ranges.size()); @@ -837,9 +841,8 @@ TEST_F(PageFilteredRowGroupReaderTest, ComputePageRangesWithDictionaryEncoding) row_ranges.Add(RowRanges::Range(0, 99)); auto ranges = PageFilteredRowGroupReader::ComputePageRanges( - parquet_reader.get(), TargetRowGroup(/*rg_index=*/0, /*is_partially_matched=*/false, /*ranges=*/row_ranges), - /*column_indices=*/{0}); + /*column_indices=*/{0}, parquet_reader.get()); ASSERT_FALSE(ranges.empty()); @@ -914,56 +917,56 @@ TEST_F(PageFilteredRowGroupReaderTest, ComputePageRangesWithDictionaryEncoding) auto partial_concat = arrow::Concatenate(result_partial->chunks()).ValueOrDie(); ASSERT_TRUE(partial_concat->Equals(expected_struct)); } -/// Helper: build a StructArray with a top-level int32 "id" column and a nested struct column -/// "info" containing two int32 fields: "x" and "y". -/// id[i] = i, info.x[i] = i * 100, info.y[i] = i * 100 + 1, for i in [0, N). -/// -/// Arrow schema: { id: int32, info: struct } -/// Parquet leaf columns: [id (index 0), info.x (index 1), info.y (index 2)] -static std::shared_ptr MakeNestedStructData(int32_t num_rows) { - arrow::Int32Builder id_builder, x_builder, y_builder; +/// Helper: build an Int32Array with sequential values 0..N-1. +static std::shared_ptr MakeIdColumn(int32_t num_rows) { + arrow::Int32Builder id_builder; EXPECT_TRUE(id_builder.Reserve(num_rows).ok()); + for (int32_t i = 0; i < num_rows; ++i) { + id_builder.UnsafeAppend(i); + } + return id_builder.Finish().ValueOrDie(); +} + +/// Helper: build a struct array (without id column). +/// x[i] = i * 100, y[i] = i * 100 + 1, for i in [0, N). +static std::shared_ptr MakeNestedStructData(int32_t num_rows) { + arrow::Int32Builder x_builder, y_builder; EXPECT_TRUE(x_builder.Reserve(num_rows).ok()); EXPECT_TRUE(y_builder.Reserve(num_rows).ok()); for (int32_t i = 0; i < num_rows; ++i) { - id_builder.UnsafeAppend(i); x_builder.UnsafeAppend(i * 100); y_builder.UnsafeAppend(i * 100 + 1); } - auto id_array = id_builder.Finish().ValueOrDie(); auto x_array = x_builder.Finish().ValueOrDie(); auto y_array = y_builder.Finish().ValueOrDie(); auto field_x = arrow::field("x", arrow::int32()); auto field_y = arrow::field("y", arrow::int32()); - auto inner_struct = - arrow::StructArray::Make({x_array, y_array}, {field_x, field_y}).ValueOrDie(); - - auto field_id = arrow::field("id", arrow::int32()); - auto field_info = arrow::field("info", arrow::struct_({field_x, field_y})); - return arrow::StructArray::Make({id_array, inner_struct}, {field_id, field_info}).ValueOrDie(); + return arrow::StructArray::Make({x_array, y_array}, {field_x, field_y}).ValueOrDie(); } /// Test: rowgroup-level filtering on a file with nested struct columns. /// -/// This test exposes the bug where BuildPageFilteredSchema fails to correctly map -/// Parquet leaf column indices to Arrow fields for nested types, and -/// ReadFilteredRowGroup cannot correctly assemble nested column results. -/// /// Schema: { id: int32, info: struct } /// Parquet leaf columns: [id=0, info.x=1, info.y=2] /// 100 rows, 10 per page, 2 row groups. -/// Predicate: id >= 70 → row groups 0 skipped, row groups 1 read → 50 rows expected. +/// Predicate: id >= 70 → page 0-7 skipped, paged 8-9 read → 30 rows expected. /// The read schema requests both "id" and "info" columns. -TEST_F(PageFilteredRowGroupReaderTest, NestedStructColumnRowGroupFilter) { +TEST_F(PageFilteredRowGroupReaderTest, NestedStructColumnPageFilter) { std::string file_name = dir_->Str() + "/nested_struct_filter.parquet"; - auto data = MakeNestedStructData(100); - WriteTestFile(file_name, data, /*write_batch_size=*/10, /*max_row_group_length=*/50); auto field_x = arrow::field("x", arrow::int32()); auto field_y = arrow::field("y", arrow::int32()); - auto read_schema = arrow::schema({arrow::field("id", arrow::int32()), - arrow::field("info", arrow::struct_({field_x, field_y}))}); + auto field_id = arrow::field("id", arrow::int32()); + auto field_info = arrow::field("info", arrow::struct_({field_x, field_y})); + + auto id_array = MakeIdColumn(100); + auto info_array = MakeNestedStructData(100); + auto data = + arrow::StructArray::Make({id_array, info_array}, {field_id, field_info}).ValueOrDie(); + WriteTestFile(file_name, data, /*write_batch_size=*/10, /*max_row_group_length=*/50); + + auto read_schema = arrow::schema({field_id, field_info}); auto predicate = PredicateBuilder::GreaterOrEqual( /*field_index=*/0, /*field_name=*/"id", FieldType::INT, Literal(70)); @@ -971,32 +974,36 @@ TEST_F(PageFilteredRowGroupReaderTest, NestedStructColumnRowGroupFilter) { std::shared_ptr result; ReadWithPredicateImpl(file_name, read_schema, predicate, &result); - // Should get rows 50-99 = 50 rows + // Should get rows 70-99 = 30 rows ASSERT_TRUE(result); - ASSERT_EQ(50, result->length()); + ASSERT_EQ(30, result->length()); // Build expected result: rows 50-99 from the original data - auto expected = data->Slice(50, 50); + auto expected = data->Slice(70, 30); ASSERT_TRUE(expected->Equals(result->chunk(0))); } /// Test: Page-level filtering reading only the predicate column (no nested column in read schema). /// -/// This verifies that when reading only the "id" column (without the nested struct), -/// page-level filtering works correctly since the read schema contains no nested types. +/// This verifies that when reading only the "id" column (without the nested struct) /// /// Schema: { id: int32, info: struct } /// Read schema: { id: int32 } /// Predicate on "id": id >= 70. Page-level filtering active → rows 70-99 (30 rows). TEST_F(PageFilteredRowGroupReaderTest, NestedStructColumnOnlyReadIdField) { std::string file_name = dir_->Str() + "/nested_struct_only_nested.parquet"; - auto data = MakeNestedStructData(100); - WriteTestFile(file_name, data, /*write_batch_size=*/10, /*max_row_group_length=*/50); - auto field_id = arrow::field("id", arrow::int32()); auto field_x = arrow::field("x", arrow::int32()); auto field_y = arrow::field("y", arrow::int32()); + auto field_id = arrow::field("id", arrow::int32()); auto field_info = arrow::field("info", arrow::struct_({field_x, field_y})); + + auto id_array = MakeIdColumn(100); + auto info_array = MakeNestedStructData(100); + auto data = + arrow::StructArray::Make({id_array, info_array}, {field_id, field_info}).ValueOrDie(); + WriteTestFile(file_name, data, /*write_batch_size=*/10, /*max_row_group_length=*/50); + // Read "id" column only auto read_schema = arrow::schema({field_id}); @@ -1016,19 +1023,9 @@ TEST_F(PageFilteredRowGroupReaderTest, NestedStructColumnOnlyReadIdField) { ASSERT_TRUE(data->field(0)->Slice(70, 30)->Equals(result_struct->field(0))); } -/// Helper: build a StructArray with an int32 "id" column and a list "tags" column. -/// id[i] = i, tags[i] = [i*10, i*10+1], for i in [0, N). -/// -/// Arrow schema: { id: int32, tags: list } -/// Parquet leaf columns: [id (index 0), tags.item (index 1)] -static std::shared_ptr MakeListColumnData(int32_t num_rows) { - arrow::Int32Builder id_builder; - EXPECT_TRUE(id_builder.Reserve(num_rows).ok()); - for (int32_t i = 0; i < num_rows; ++i) { - id_builder.UnsafeAppend(i); - } - auto id_array = id_builder.Finish().ValueOrDie(); - +/// Helper: build a list array (without id column). +/// tags[i] = [i*10, i*10+1], for i in [0, N). +static std::shared_ptr MakeListColumnData(int32_t num_rows) { auto value_builder = std::make_shared(); arrow::ListBuilder list_builder(arrow::default_memory_pool(), value_builder); for (int32_t i = 0; i < num_rows; ++i) { @@ -1036,26 +1033,12 @@ static std::shared_ptr MakeListColumnData(int32_t num_rows) EXPECT_TRUE(value_builder->Append(i * 10).ok()); EXPECT_TRUE(value_builder->Append(i * 10 + 1).ok()); } - auto list_array = list_builder.Finish().ValueOrDie(); - - auto field_id = arrow::field("id", arrow::int32()); - auto field_tags = arrow::field("tags", arrow::list(arrow::field("item", arrow::int32()))); - return arrow::StructArray::Make({id_array, list_array}, {field_id, field_tags}).ValueOrDie(); + return list_builder.Finish().ValueOrDie(); } -/// Helper: build a StructArray with an int32 "id" column and a map "props" column. -/// id[i] = i, props[i] = {"k_i": i * 100}, for i in [0, N). -/// -/// Arrow schema: { id: int32, props: map } -/// Parquet leaf columns: [id (index 0), props.key (index 1), props.value (index 2)] -static std::shared_ptr MakeMapColumnData(int32_t num_rows) { - arrow::Int32Builder id_builder; - EXPECT_TRUE(id_builder.Reserve(num_rows).ok()); - for (int32_t i = 0; i < num_rows; ++i) { - id_builder.UnsafeAppend(i); - } - auto id_array = id_builder.Finish().ValueOrDie(); - +/// Helper: build a map array (without id column). +/// props[i] = {"k_i": i * 100}, for i in [0, N). +static std::shared_ptr MakeMapColumnData(int32_t num_rows) { auto key_builder = std::make_shared(); auto value_builder = std::make_shared(); arrow::MapBuilder map_builder(arrow::default_memory_pool(), key_builder, value_builder); @@ -1065,26 +1048,27 @@ static std::shared_ptr MakeMapColumnData(int32_t num_rows) { EXPECT_TRUE(key_builder->Append(key).ok()); EXPECT_TRUE(value_builder->Append(i * 100).ok()); } - auto map_array = map_builder.Finish().ValueOrDie(); - - auto field_id = arrow::field("id", arrow::int32()); - auto field_props = arrow::field("props", arrow::map(arrow::utf8(), arrow::int32())); - return arrow::StructArray::Make({id_array, map_array}, {field_id, field_props}).ValueOrDie(); + return map_builder.Finish().ValueOrDie(); } /// Test: rowgroup-level filtering on a file with a list column. /// /// Schema: { id: int32, tags: list } /// 100 rows, 10 per page, 2 row groups. -/// Predicate: id >= 70 → row groups 0 skipped, row groups 1 read → 50 rows expected. -TEST_F(PageFilteredRowGroupReaderTest, NestedListColumnRowGroupFilter) { +/// Predicate: id >= 70 → page 0-7 skipped, page 8-9 read → 30 rows expected. +TEST_F(PageFilteredRowGroupReaderTest, NestedListColumnPageFilter) { std::string file_name = dir_->Str() + "/nested_list_filter.parquet"; - auto data = MakeListColumnData(100); + + auto field_id = arrow::field("id", arrow::int32()); + auto field_tags = arrow::field("tags", arrow::list(arrow::field("item", arrow::int32()))); + + auto id_array = MakeIdColumn(100); + auto tags_array = MakeListColumnData(100); + auto data = + arrow::StructArray::Make({id_array, tags_array}, {field_id, field_tags}).ValueOrDie(); WriteTestFile(file_name, data, /*write_batch_size=*/10, /*max_row_group_length=*/50); - auto read_schema = - arrow::schema({arrow::field("id", arrow::int32()), - arrow::field("tags", arrow::list(arrow::field("item", arrow::int32())))}); + auto read_schema = arrow::schema({field_id, field_tags}); auto predicate = PredicateBuilder::GreaterOrEqual( /*field_index=*/0, /*field_name=*/"id", FieldType::INT, Literal(70)); @@ -1093,10 +1077,10 @@ TEST_F(PageFilteredRowGroupReaderTest, NestedListColumnRowGroupFilter) { ReadWithPredicateImpl(file_name, read_schema, predicate, &result); ASSERT_TRUE(result); - ASSERT_EQ(50, result->length()); + ASSERT_EQ(30, result->length()); - // Build expected result: rows 50-99 from the original data - auto expected = data->Slice(50, 50); + // Build expected result: rows 70-99 from the original data + auto expected = data->Slice(70, 30); ASSERT_TRUE(expected->Equals(result->chunk(0))); } @@ -1104,116 +1088,32 @@ TEST_F(PageFilteredRowGroupReaderTest, NestedListColumnRowGroupFilter) { /// /// Schema: { id: int32, props: map } /// 100 rows, 10 per page, 2 row groups. -/// Predicate: id >= 70 → row groups 0 skipped, row groups 1 read → 50 rows expected. -TEST_F(PageFilteredRowGroupReaderTest, NestedMapColumnRowGroupFilter) { +/// Predicate: id >= 70 → page 0-7 skipped, page 8-9 read → 30 rows expected. +TEST_F(PageFilteredRowGroupReaderTest, NestedMapColumnPageFilter) { std::string file_name = dir_->Str() + "/nested_map_filter.parquet"; - auto data = MakeMapColumnData(100); - WriteTestFile(file_name, data, /*write_batch_size=*/10, /*max_row_group_length=*/50); - - auto read_schema = - arrow::schema({arrow::field("id", arrow::int32()), - arrow::field("props", arrow::map(arrow::utf8(), arrow::int32()))}); - - auto predicate = PredicateBuilder::GreaterOrEqual( - /*field_index=*/0, /*field_name=*/"id", FieldType::INT, Literal(70)); - - std::shared_ptr result; - ReadWithPredicateImpl(file_name, read_schema, predicate, &result); - - ASSERT_TRUE(result); - ASSERT_EQ(50, result->length()); - - // Build expected result: rows 50-99 from the original data - auto expected = data->Slice(50, 50); - ASSERT_TRUE(expected->Equals(result->chunk(0))); -} - -/// Test: nested map projection falls back to row-group-level filtering when page index filter is -/// unavailable for nested read schemas. -/// -/// Schema: { id: int32, props: map } -/// 100 rows, 10 per page, 2 row group. -/// Bitmap: {70..99} hits the second row group (50..99). -/// Because nested schema disables page-level filtering, the entire row group 1 (50..99) is read, -/// so rows [50, 99] should all be returned. -TEST_F(PageFilteredRowGroupReaderTest, NestedMapBitmapFallback) { - std::string file_name = dir_->Str() + "/nested_map_projection_fallback.parquet"; - auto data = MakeMapColumnData(100); - WriteTestFile(file_name, data, /*write_batch_size=*/10, /*max_row_group_length=*/50); + auto field_id = arrow::field("id", arrow::int32()); auto field_props = arrow::field("props", arrow::map(arrow::utf8(), arrow::int32())); - auto read_schema = arrow::schema({arrow::field("id", arrow::int32()), field_props}); - RoaringBitmap32 bitmap; - bitmap.AddRange(70, 100); - - std::shared_ptr result; - ReadWithPredicateAndBitmapImpl(file_name, read_schema, nullptr, bitmap, &result); - - ASSERT_TRUE(result); - // Because page-level filtering is skipped for nested schemas, we read full row groups. - ASSERT_EQ(50, result->length()); - - auto expected = data->Slice(50, 50); - ASSERT_TRUE(expected->Equals(result->chunk(0))); -} - -/// Test: nested list projection falls back to row-group-level filtering when page index filter is -/// unavailable for nested read schemas. -/// -/// Schema: { id: int32, tags: list } -/// 100 rows, 10 per page, 2 row group. -/// Bitmap: {70..99} hits the second row group (50..99). -/// Because nested schema disables page-level filtering, the entire row group 1 (50..99) is read, -/// so rows [50, 99] should all be returned. -TEST_F(PageFilteredRowGroupReaderTest, NestedListBitmapFallback) { - std::string file_name = dir_->Str() + "/nested_list_projection_fallback.parquet"; - auto data = MakeListColumnData(100); + auto id_array = MakeIdColumn(100); + auto props_array = MakeMapColumnData(100); + auto data = + arrow::StructArray::Make({id_array, props_array}, {field_id, field_props}).ValueOrDie(); WriteTestFile(file_name, data, /*write_batch_size=*/10, /*max_row_group_length=*/50); - auto field_tags = arrow::field("tags", arrow::list(arrow::field("item", arrow::int32()))); - auto read_schema = arrow::schema({arrow::field("id", arrow::int32()), field_tags}); - - RoaringBitmap32 bitmap; - bitmap.AddRange(70, 100); + auto read_schema = arrow::schema({field_id, field_props}); - std::shared_ptr result; - ReadWithPredicateAndBitmapImpl(file_name, read_schema, nullptr, bitmap, &result); - - ASSERT_TRUE(result); - ASSERT_EQ(50, result->length()); - - auto expected = data->Slice(50, 50); - ASSERT_TRUE(expected->Equals(result->chunk(0))); -} - -/// Test: nested struct projection falls back to row-group-level filtering when page index filter is -/// unavailable for nested read schemas. -/// -/// Schema: { id: int32, info: struct } -/// Bitmap: {70..99} hits the second row group (50..99). -/// Because nested schema disables page-level filtering, the entire second row group (50..99) is -/// read. -TEST_F(PageFilteredRowGroupReaderTest, NestedStructBitmapFallback) { - std::string file_name = dir_->Str() + "/nested_struct_projection_fallback.parquet"; - auto data = MakeNestedStructData(100); - WriteTestFile(file_name, data, /*write_batch_size=*/10, /*max_row_group_length=*/50); - - auto field_x = arrow::field("x", arrow::int32()); - auto field_y = arrow::field("y", arrow::int32()); - auto field_info = arrow::field("info", arrow::struct_({field_x, field_y})); - auto read_schema = arrow::schema({arrow::field("id", arrow::int32()), field_info}); - - RoaringBitmap32 bitmap; - bitmap.AddRange(70, 100); + auto predicate = PredicateBuilder::GreaterOrEqual( + /*field_index=*/0, /*field_name=*/"id", FieldType::INT, Literal(70)); std::shared_ptr result; - ReadWithPredicateAndBitmapImpl(file_name, read_schema, nullptr, bitmap, &result); + ReadWithPredicateImpl(file_name, read_schema, predicate, &result); ASSERT_TRUE(result); - ASSERT_EQ(50, result->length()); + ASSERT_EQ(30, result->length()); - auto expected = data->Slice(50, 50); + // Build expected result: rows 70-99 from the original data + auto expected = data->Slice(70, 30); ASSERT_TRUE(expected->Equals(result->chunk(0))); } @@ -1221,39 +1121,20 @@ TEST_F(PageFilteredRowGroupReaderTest, NestedStructBitmapFallback) { /// /// Schema: { id: int32, info: struct, tags: list } /// This tests the boundary handling when two nested fields are adjacent in the schema. -/// Predicate: id >= 70 → row groups 0 skipped, row groups 1 read → 50 rows expected. +/// Predicate: id >= 70 → page 0-7 skipped, page 8-9 read → 30 rows expected. TEST_F(PageFilteredRowGroupReaderTest, MultipleAdjacentNestedColumns) { std::string file_name = dir_->Str() + "/multi_nested.parquet"; - // Build data with id, info (struct), tags (list) - arrow::Int32Builder id_builder, x_builder, y_builder; - ASSERT_TRUE(id_builder.Reserve(100).ok()); - ASSERT_TRUE(x_builder.Reserve(100).ok()); - ASSERT_TRUE(y_builder.Reserve(100).ok()); - auto value_builder = std::make_shared(); - arrow::ListBuilder list_builder(arrow::default_memory_pool(), value_builder); - - for (int32_t i = 0; i < 100; ++i) { - id_builder.UnsafeAppend(i); - x_builder.UnsafeAppend(i * 100); - y_builder.UnsafeAppend(i * 100 + 1); - ASSERT_TRUE(list_builder.Append().ok()); - ASSERT_TRUE(value_builder->Append(i * 10).ok()); - } - auto id_array = id_builder.Finish().ValueOrDie(); - auto x_array = x_builder.Finish().ValueOrDie(); - auto y_array = y_builder.Finish().ValueOrDie(); - auto list_array = list_builder.Finish().ValueOrDie(); - auto field_x = arrow::field("x", arrow::int32()); auto field_y = arrow::field("y", arrow::int32()); - auto inner_struct = - arrow::StructArray::Make({x_array, y_array}, {field_x, field_y}).ValueOrDie(); - auto field_id = arrow::field("id", arrow::int32()); auto field_info = arrow::field("info", arrow::struct_({field_x, field_y})); auto field_tags = arrow::field("tags", arrow::list(arrow::field("item", arrow::int32()))); - auto data = arrow::StructArray::Make({id_array, inner_struct, list_array}, + + auto id_array = MakeIdColumn(100); + auto info_array = MakeNestedStructData(100); + auto tags_array = MakeListColumnData(100); + auto data = arrow::StructArray::Make({id_array, info_array, tags_array}, {field_id, field_info, field_tags}) .ValueOrDie(); @@ -1267,10 +1148,10 @@ TEST_F(PageFilteredRowGroupReaderTest, MultipleAdjacentNestedColumns) { ReadWithPredicateImpl(file_name, read_schema, predicate, &result); ASSERT_TRUE(result); - ASSERT_EQ(50, result->length()); + ASSERT_EQ(30, result->length()); - // Build expected result: rows 50-99 from the original data - auto expected = data->Slice(50, 50); + // Build expected result: rows 70-99 from the original data + auto expected = data->Slice(70, 30); ASSERT_TRUE(expected->Equals(result->chunk(0))); } /// Test: bitmap hits all pages of a subset of row groups (no predicate). @@ -1719,4 +1600,412 @@ TEST_F(PageFilteredRowGroupReaderTest, BitmapTrimMultiColumnTest) { } } +/// Test: predicate pushdown with all nested column types (struct, list, map). +/// +/// Schema: { id: int32, info: struct, +/// tags: list, props: map } +/// 100 rows, 10 rows per page, 50 rows per row group → 2 row groups. +/// Predicate: id in [15, 29] or id in [80, 99] (Between is inclusive). +/// Read schema: full schema (all columns). +/// Page-level filtering (10 rows/page): +/// Between(15, 29) → pages 1-2 (rows 10-29) +/// Between(80, 99) → pages 8-9 (rows 80-99) +/// Total: 40 rows. +TEST_F(PageFilteredRowGroupReaderTest, MultipleNestedColumns) { + std::string file_name = dir_->Str() + "/multi_nested_columns.parquet"; + + auto field_x = arrow::field("x", arrow::int32()); + auto field_y = arrow::field("y", arrow::int32()); + auto field_id = arrow::field("id", arrow::int32()); + auto field_info = arrow::field("info", arrow::struct_({field_x, field_y})); + auto field_tags = arrow::field("tags", arrow::list(arrow::field("item", arrow::int32()))); + auto field_props = arrow::field("props", arrow::map(arrow::utf8(), arrow::int32())); + + // Build data with all nested column types using shared helpers + auto id_array = MakeIdColumn(100); + auto info_array = MakeNestedStructData(100); + auto tags_array = MakeListColumnData(100); + auto props_array = MakeMapColumnData(100); + auto data = arrow::StructArray::Make({id_array, info_array, tags_array, props_array}, + {field_id, field_info, field_tags, field_props}) + .ValueOrDie(); + + // Write: 10 rows per page, 50 rows per row group → 2 row groups + WriteTestFile(file_name, data, /*write_batch_size=*/10, /*max_row_group_length=*/50); + + // Read full schema + auto read_schema = arrow::schema({field_id, field_info, field_tags, field_props}); + + // predicate: id in [15, 29] or id in [80, 99] + ASSERT_OK_AND_ASSIGN( + auto predicate, PredicateBuilder::Or( + {PredicateBuilder::Between(/*field_index=*/0, /*field_name=*/"id", + FieldType::INT, Literal(15), Literal(29)), + PredicateBuilder::Between(/*field_index=*/0, /*field_name=*/"id", + FieldType::INT, Literal(80), Literal(99))})); + + std::shared_ptr result; + ReadWithPredicateImpl(file_name, read_schema, predicate, &result, + /*batch_size=*/1024); + + // Page-level filtering (10 rows/page): + // Between(15, 29) → pages 1-2 (rows 10-29) + // Between(80, 99) → pages 8-9 (rows 80-99) + // Total: 40 rows + ASSERT_TRUE(result); + ASSERT_EQ(40, result->length()); + + auto expected = + arrow::ChunkedArray::Make({data->Slice(10, 20), data->Slice(80, 20)}).ValueOrDie(); + ASSERT_TRUE(result->Equals(expected)); +} + +/// Test: sub-column projection of a struct type with page-level filtering. +/// +/// Schema: { id: int32, info: struct } +/// Read schema: { info: struct } — project only x, not y. +/// Predicate: id >= 70 → 30 rows expected. +/// Verifies that reading a sub-column of a nested struct works correctly +/// with page-level filtering and the ColumnReader tree (GetColumn + filter_leaves). +TEST_F(PageFilteredRowGroupReaderTest, NestedStructSubColumnProjection) { + std::string file_name = dir_->Str() + "/nested_struct_subcol.parquet"; + + auto field_x = arrow::field("x", arrow::int32()); + auto field_y = arrow::field("y", arrow::int32()); + auto field_id = arrow::field("id", arrow::int32()); + auto field_info = arrow::field("info", arrow::struct_({field_x, field_y})); + + auto id_array = MakeIdColumn(100); + auto info_array = MakeNestedStructData(100); + auto data = + arrow::StructArray::Make({id_array, info_array}, {field_id, field_info}).ValueOrDie(); + WriteTestFile(file_name, data, /*write_batch_size=*/10, /*max_row_group_length=*/50); + + // Read only info.x (sub-column projection: only x, not y) + auto read_schema = arrow::schema({arrow::field("info", arrow::struct_({field_x}))}); + + auto predicate = PredicateBuilder::GreaterOrEqual( + /*field_index=*/0, /*field_name=*/"id", FieldType::INT, Literal(70)); + + std::shared_ptr result; + ReadWithPredicateImpl(file_name, read_schema, predicate, &result); + + ASSERT_TRUE(result); + ASSERT_EQ(30, result->length()); + + // Result is struct> + auto result_struct = std::dynamic_pointer_cast(result->chunk(0)); + ASSERT_TRUE(result_struct); + ASSERT_EQ(1, result_struct->num_fields()); + + auto info_result = std::dynamic_pointer_cast(result_struct->field(0)); + ASSERT_TRUE(info_result); + ASSERT_EQ(1, info_result->num_fields()); + + auto x_arr = std::dynamic_pointer_cast(info_result->field(0)); + ASSERT_TRUE(x_arr); + for (int32_t i = 0; i < 30; ++i) { + ASSERT_EQ((70 + i) * 100, x_arr->Value(i)) << "Mismatch at index " << i; + } +} + +/// Helper: build a struct with a flat key plus several nested columns of different +/// physical widths / repetition, so that with a byte-based data page size their +/// leaves flush pages at different row counts (misaligned pages): +/// key: int64 (fixed 8B -> ~5 rows/page) +/// s: struct (x ~10 rows/page, y ~5 rows/page) +/// tags: list (2 values/row -> ~5 rows/page) +/// props: map (variable-width utf8 key -> irregular rows/page) +/// key/s.x/s.y encode the row index (= i); tags/props reuse the shared list/map +/// helpers. Correctness is verified per row by deep-comparing against this array. +static std::shared_ptr MakeMisalignedNestedData(int32_t num_rows) { + arrow::Int64Builder key_builder; + arrow::Int32Builder x_builder; + arrow::Int64Builder y_builder; + EXPECT_TRUE(key_builder.Reserve(num_rows).ok()); + EXPECT_TRUE(x_builder.Reserve(num_rows).ok()); + EXPECT_TRUE(y_builder.Reserve(num_rows).ok()); + for (int32_t i = 0; i < num_rows; ++i) { + key_builder.UnsafeAppend(i); + x_builder.UnsafeAppend(i); + y_builder.UnsafeAppend(i); + } + auto key_array = key_builder.Finish().ValueOrDie(); + auto x_array = x_builder.Finish().ValueOrDie(); + auto y_array = y_builder.Finish().ValueOrDie(); + + auto field_x = arrow::field("x", arrow::int32()); + auto field_y = arrow::field("y", arrow::int64()); + auto s_array = arrow::StructArray::Make({x_array, y_array}, {field_x, field_y}).ValueOrDie(); + + // Repeated nested leaves (list and map) paginate on value + // bytes, so they misalign with the struct's fixed-width leaves too. + auto tags_array = MakeListColumnData(num_rows); + auto props_array = MakeMapColumnData(num_rows); + + auto field_key = arrow::field("key", arrow::int64()); + auto field_s = arrow::field("s", arrow::struct_({field_x, field_y})); + auto field_tags = arrow::field("tags", arrow::list(arrow::field("item", arrow::int32()))); + auto field_props = arrow::field("props", arrow::map(arrow::utf8(), arrow::int32())); + return arrow::StructArray::Make({key_array, s_array, tags_array, props_array}, + {field_key, field_s, field_tags, field_props}) + .ValueOrDie(); +} + +/// Test: page-level filtering across multiple nested columns whose leaf pages are +/// MISALIGNED, within a SINGLE row group. The file mixes a flat key, a +/// struct, a list and a map; with write_batch_size=1 +/// and a byte-based data page size every leaf flushes pages at a different (and, for +/// the utf8 map key, irregular) row count. +/// Bitmap: [0,15), [77, 87) (to avoid bitmap hole filling) +/// Expected: 25 rows +TEST_F(PageFilteredRowGroupReaderTest, NestedColumnsMisalignedPagesSingleRowGroup) { + std::string file_name = dir_->Str() + "/nested_misaligned_single_rg.parquet"; + constexpr int32_t kNumRows = 100; + auto data = MakeMisalignedNestedData(kNumRows); + + // write_batch_size=1 + a byte-based data page size makes the int32 and int64 leaves + // flush pages at different row counts, i.e. deliberately misaligned pages. + // max_row_group_length=kNumRows keeps all rows in a single row group. + WriteTestFile(file_name, data, /*write_batch_size=*/1, /*max_row_group_length=*/kNumRows, + /*enable_dictionary=*/false, /*data_page_size=*/40); + + auto read_schema = + arrow::schema({arrow::field("key", arrow::int64()), + arrow::field("s", arrow::struct_({arrow::field("x", arrow::int32()), + arrow::field("y", arrow::int64())})), + arrow::field("tags", arrow::list(arrow::field("item", arrow::int32()))), + arrow::field("props", arrow::map(arrow::utf8(), arrow::int32()))}); + + // bitmap: [0,15), [77, 87) + RoaringBitmap32 bitmap; + bitmap.AddRange(0, 15); + bitmap.AddRange(77, 87); + + std::shared_ptr result; + ReadWithPredicateAndBitmapImpl(file_name, read_schema, nullptr, bitmap, &result); + ASSERT_TRUE(result); + + int64_t total = 0; + auto top = std::dynamic_pointer_cast(result->chunk(0)); + ASSERT_TRUE(top); + auto key_arr = std::dynamic_pointer_cast(top->field(0)); + for (int64_t i = 0; i < key_arr->length(); ++i) { + int64_t k = key_arr->Value(i); + // Full-row deep compare across ALL columns (struct + list + map): the returned + // row must equal the original row identified by key k. This is what catches a + // desync in the repeated list/map leaves, whose reassembly also relies on the + // per-leaf skip/read staying row-consistent. + ASSERT_TRUE(data->Slice(k, 1)->Equals(*top->Slice(i, 1))) + << "row content mismatch at key " << k; + ++total; + } + + for (int64_t i = 0; i < 15; ++i) { + ASSERT_EQ(i, key_arr->Value(i)); + } + for (int64_t i = 0; i < 10; ++i) { + ASSERT_EQ(77 + i, key_arr->Value(15 + i)); + } + + ASSERT_EQ(total, 25); +} + +/// Test: same misaligned nested layout as the single-row-group case above, but split +/// across MULTIPLE row groups (max_row_group_length=40 -> row groups of 40/40/20). +/// The selection bitmap spans row-group boundaries so the reader must keep the +/// per-leaf skip/read row-consistent both across misaligned pages and across row +/// groups. +/// Bitmap: [0,15), [77, 87) (to avoid bitmap hole filling) +/// Expected: 25 rows -> keys 0..14 and 77..86 +TEST_F(PageFilteredRowGroupReaderTest, NestedColumnsMisalignedPagesMultiRowGroup) { + std::string file_name = dir_->Str() + "/nested_misaligned_multi_rg.parquet"; + constexpr int32_t kNumRows = 100; + auto data = MakeMisalignedNestedData(kNumRows); + + // write_batch_size=1 + byte-based data page size -> misaligned leaf pages. + // max_row_group_length=40 -> 3 row groups (40, 40, 20). + WriteTestFile(file_name, data, /*write_batch_size=*/1, /*max_row_group_length=*/40, + /*enable_dictionary=*/false, /*data_page_size=*/40); + + auto read_schema = + arrow::schema({arrow::field("key", arrow::int64()), + arrow::field("s", arrow::struct_({arrow::field("x", arrow::int32()), + arrow::field("y", arrow::int64())})), + arrow::field("tags", arrow::list(arrow::field("item", arrow::int32()))), + arrow::field("props", arrow::map(arrow::utf8(), arrow::int32()))}); + + // bitmap: [0,15), [77, 87) -> spans all three row groups. + RoaringBitmap32 bitmap; + bitmap.AddRange(0, 15); + bitmap.AddRange(77, 87); + + std::shared_ptr result; + ReadWithPredicateAndBitmapImpl(file_name, read_schema, nullptr, bitmap, &result); + ASSERT_TRUE(result); + + std::vector keys; + + auto concated = arrow::Concatenate(result->chunks()).ValueOrDie(); + auto top = std::dynamic_pointer_cast(concated); + ASSERT_TRUE(top); + auto key_arr = std::dynamic_pointer_cast(top->field(0)); + ASSERT_TRUE(key_arr); + for (int64_t i = 0; i < key_arr->length(); ++i) { + int64_t k = key_arr->Value(i); + ASSERT_TRUE(data->Slice(k, 1)->Equals(*top->Slice(i, 1))) + << "row content mismatch at key " << k; + keys.push_back(k); + } + + std::vector expected; + for (int64_t i = 0; i < 15; ++i) { + expected.push_back(i); + } + for (int64_t i = 77; i < 87; ++i) { + expected.push_back(i); + } + ASSERT_EQ(keys, expected); +} + +/// Helper: like MakeMisalignedNestedData but sprinkles nulls into the nested columns +/// so that definition levels vary per row (which also perturbs page boundaries): +/// key: int64, always non-null (= i, used to identify the row) +/// s: struct; whole struct null when i%11==0, otherwise +/// x null when i%5==0 and y null when i%7==0 +/// tags: list; null when i%6==0, otherwise [i*10, i*10+1] +/// props: map; null when i%8==0, otherwise {"k_i": i*100} +/// Correctness is verified per row by deep-comparing against this array. +static std::shared_ptr MakeMisalignedNestedDataWithNulls(int32_t num_rows) { + arrow::Int64Builder key_builder; + EXPECT_TRUE(key_builder.Reserve(num_rows).ok()); + for (int32_t i = 0; i < num_rows; ++i) { + key_builder.UnsafeAppend(i); + } + auto key_array = key_builder.Finish().ValueOrDie(); + + // s: struct with nulls at both leaf and struct level. + auto field_x = arrow::field("x", arrow::int32()); + auto field_y = arrow::field("y", arrow::int64()); + auto x_builder = std::make_shared(); + auto y_builder = std::make_shared(); + arrow::StructBuilder s_builder(arrow::struct_({field_x, field_y}), arrow::default_memory_pool(), + {x_builder, y_builder}); + for (int32_t i = 0; i < num_rows; ++i) { + if (i % 11 == 0) { + // AppendNull() also appends nulls to the child builders, keeping lengths in sync. + EXPECT_TRUE(s_builder.AppendNull().ok()); + continue; + } + EXPECT_TRUE(s_builder.Append().ok()); + if (i % 5 == 0) { + EXPECT_TRUE(x_builder->AppendNull().ok()); + } else { + EXPECT_TRUE(x_builder->Append(i).ok()); + } + if (i % 7 == 0) { + EXPECT_TRUE(y_builder->AppendNull().ok()); + } else { + EXPECT_TRUE(y_builder->Append(i).ok()); + } + } + auto s_array = s_builder.Finish().ValueOrDie(); + + // tags: list with null lists. + auto item_builder = std::make_shared(); + arrow::ListBuilder tags_builder(arrow::default_memory_pool(), item_builder); + for (int32_t i = 0; i < num_rows; ++i) { + if (i % 6 == 0) { + EXPECT_TRUE(tags_builder.AppendNull().ok()); + } else { + EXPECT_TRUE(tags_builder.Append().ok()); + EXPECT_TRUE(item_builder->Append(i * 10).ok()); + EXPECT_TRUE(item_builder->Append(i * 10 + 1).ok()); + } + } + auto tags_array = tags_builder.Finish().ValueOrDie(); + + // props: map with null maps. + auto map_key_builder = std::make_shared(); + auto map_val_builder = std::make_shared(); + arrow::MapBuilder props_builder(arrow::default_memory_pool(), map_key_builder, map_val_builder); + for (int32_t i = 0; i < num_rows; ++i) { + if (i % 8 == 0) { + EXPECT_TRUE(props_builder.AppendNull().ok()); + } else { + EXPECT_TRUE(props_builder.Append().ok()); + EXPECT_TRUE(map_key_builder->Append("k_" + std::to_string(i)).ok()); + EXPECT_TRUE(map_val_builder->Append(i * 100).ok()); + } + } + auto props_array = props_builder.Finish().ValueOrDie(); + + auto field_key = arrow::field("key", arrow::int64()); + auto field_s = arrow::field("s", arrow::struct_({field_x, field_y})); + auto field_tags = arrow::field("tags", arrow::list(arrow::field("item", arrow::int32()))); + auto field_props = arrow::field("props", arrow::map(arrow::utf8(), arrow::int32())); + return arrow::StructArray::Make({key_array, s_array, tags_array, props_array}, + {field_key, field_s, field_tags, field_props}) + .ValueOrDie(); +} + +/// Test: nested columns containing NULLs, with MISALIGNED leaf pages, split across +/// MULTIPLE row groups. This combines the three stress dimensions: null-driven +/// definition levels, byte-based misaligned pages, and row-group boundaries that the +/// selection bitmap crosses. Every returned row is deep-compared (nulls included) +/// against the original data identified by its non-null key. +/// Bitmap: [0,15), [77, 87) (to avoid bitmap hole filling) +/// Expected: 25 rows -> keys 0..14 and 77..86 +TEST_F(PageFilteredRowGroupReaderTest, NestedColumnsWithNullsMisalignedPagesMultiRowGroup) { + std::string file_name = dir_->Str() + "/nested_nulls_misaligned_multi_rg.parquet"; + constexpr int32_t kNumRows = 100; + auto data = MakeMisalignedNestedDataWithNulls(kNumRows); + + // write_batch_size=1 + byte-based data page size -> misaligned leaf pages. + // max_row_group_length=40 -> 3 row groups (40, 40, 20). + WriteTestFile(file_name, data, /*write_batch_size=*/1, /*max_row_group_length=*/40, + /*enable_dictionary=*/false, /*data_page_size=*/40); + + auto read_schema = + arrow::schema({arrow::field("key", arrow::int64()), + arrow::field("s", arrow::struct_({arrow::field("x", arrow::int32()), + arrow::field("y", arrow::int64())})), + arrow::field("tags", arrow::list(arrow::field("item", arrow::int32()))), + arrow::field("props", arrow::map(arrow::utf8(), arrow::int32()))}); + + // bitmap: [0,15), [77, 87) -> spans all three row groups. + RoaringBitmap32 bitmap; + bitmap.AddRange(0, 15); + bitmap.AddRange(77, 87); + + std::shared_ptr result; + ReadWithPredicateAndBitmapImpl(file_name, read_schema, nullptr, bitmap, &result); + ASSERT_TRUE(result); + + std::vector keys; + + auto concated = arrow::Concatenate(result->chunks()).ValueOrDie(); + auto top = std::dynamic_pointer_cast(concated); + ASSERT_TRUE(top); + auto key_arr = std::dynamic_pointer_cast(top->field(0)); + ASSERT_TRUE(key_arr); + for (int64_t i = 0; i < key_arr->length(); ++i) { + ASSERT_FALSE(key_arr->IsNull(i)) << "key column must stay non-null"; + int64_t k = key_arr->Value(i); + // Deep compare including nulls across struct/list/map leaves. + ASSERT_TRUE(data->Slice(k, 1)->Equals(*top->Slice(i, 1))) + << "row content mismatch at key " << k; + keys.push_back(k); + } + + std::vector expected; + for (int64_t i = 0; i < 15; ++i) { + expected.push_back(i); + } + for (int64_t i = 77; i < 87; ++i) { + expected.push_back(i); + } + ASSERT_EQ(keys, expected); +} + } // namespace paimon::parquet::test diff --git a/src/paimon/format/parquet/parquet_file_batch_reader.cpp b/src/paimon/format/parquet/parquet_file_batch_reader.cpp index f77d41ce9..2863d6b78 100644 --- a/src/paimon/format/parquet/parquet_file_batch_reader.cpp +++ b/src/paimon/format/parquet/parquet_file_batch_reader.cpp @@ -139,13 +139,6 @@ Status ParquetFileBatchReader::SetReadSchema( arrow::ImportSchema(schema)); PAIMON_ASSIGN_OR_RAISE(std::shared_ptr file_schema, reader_->GetSchema()); - bool has_nested_field = false; - for (const auto& field : read_schema->fields()) { - if (ArrowSchemaValidator::IsNestedType(field->type())) { - has_nested_field = true; - break; - } - } // Recursively match read_schema against file_schema by field names. // STRUCT supports sub-field projection; LIST/MAP require exact type match. @@ -178,9 +171,7 @@ Status ParquetFileBatchReader::SetReadSchema( PAIMON_ASSIGN_OR_RAISE( target_row_groups, FilterRowGroupsByBitmap(selection_bitmap.value(), target_row_groups)); - // workaround: page index filter does not support nested fields for now, skip page index - // bitmap pushdown if there is any nested field in the schema - if (!has_nested_field && enable_page_index_filter) { + if (enable_page_index_filter) { // To decide which strategy to use, "trim" or "coalesce". "Coalesce" By default. PAIMON_ASSIGN_OR_RAISE( std::string strategy, @@ -208,9 +199,7 @@ Status ParquetFileBatchReader::SetReadSchema( // pages for row groups that the bitmap already excluded. // If no predicate is provided, skip page-level filtering if (predicate && !target_row_groups.empty()) { - // workaround: page index filter does not support nested fields for now, skip page index - // filter if there is any nested field in the schema - if (enable_page_index_filter && !has_nested_field) { + if (enable_page_index_filter) { // Build column name to index map for page-level filtering. // For leaf columns, indices[0] is the correct leaf column index in Parquet. // For nested types (struct/list/map), FlattenSchema produces multiple leaf indices,