feat(parquet): support page-level (page index) filtering for nested column types (patch arrow)

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<int>& column_indices,
++      FileColumnIteratorFactory iterator_factory,
++      std::unique_ptr<ColumnReader>* out) override;
++
+   Status GetSchema(std::shared_ptr<::arrow::Schema>* out) override {
+     return FromParquetSchema(reader_->metadata()->schema(), reader_properties_,
+                              reader_->metadata()->key_value_metadata(), out);
+@@ -493,10 +498,40 @@
+
+   ::arrow::Status BuildArray(int64_t length_upper_bound,
+                              std::shared_ptr<::arrow::ChunkedArray>* out) final {
++    if (!out_) {
++      BEGIN_PARQUET_CATCH_EXCEPTIONS
++      RETURN_NOT_OK(
++          TransferColumnData(record_reader_.get(), field_, descr_, ctx_->pool, &out_));
++      END_PARQUET_CATCH_EXCEPTIONS
++    }
+     *out = out_;
+     return Status::OK();
+   }
+
++  std::vector<int> LeafColumnIndices() const final {
++    return {input_->column_index()};
++  }
++
++  ::arrow::Status ResetLeaf(int col_idx, int64_t reserve) final {
++    if (col_idx != input_->column_index()) return Status::OK();
++    BEGIN_PARQUET_CATCH_EXCEPTIONS
++    out_ = nullptr;
++    record_reader_->Reset();
++    record_reader_->Reserve(reserve);
++    return Status::OK();
++    END_PARQUET_CATCH_EXCEPTIONS
++  }
++
++  int64_t SkipRecords(int col_idx, int64_t num_records) final {
++    if (col_idx != input_->column_index() || num_records <= 0) return 0;
++    return record_reader_->SkipRecords(num_records);
++  }
++
++  int64_t ReadRecords(int col_idx, int64_t num_records) final {
++    if (col_idx != input_->column_index() || num_records <= 0) return 0;
++    return record_reader_->ReadRecords(num_records);
++  }
++
+   const std::shared_ptr<Field> field() override { return field_; }
+
+  private:
+@@ -532,6 +567,22 @@
+     return storage_reader_->LoadBatch(number_of_records);
+   }
+
++  std::vector<int> LeafColumnIndices() const final {
++    return storage_reader_->LeafColumnIndices();
++  }
++
++  ::arrow::Status ResetLeaf(int col_idx, int64_t reserve) final {
++    return storage_reader_->ResetLeaf(col_idx, reserve);
++  }
++
++  int64_t SkipRecords(int col_idx, int64_t num_records) final {
++    return storage_reader_->SkipRecords(col_idx, num_records);
++  }
++
++  int64_t ReadRecords(int col_idx, int64_t num_records) final {
++    return storage_reader_->ReadRecords(col_idx, num_records);
++  }
++
+   Status BuildArray(int64_t length_upper_bound,
+                     std::shared_ptr<ChunkedArray>* out) override {
+     std::shared_ptr<ChunkedArray> storage;
+@@ -576,6 +627,22 @@
+     return item_reader_->LoadBatch(number_of_records);
+   }
+
++  std::vector<int> LeafColumnIndices() const final {
++    return item_reader_->LeafColumnIndices();
++  }
++
++  ::arrow::Status ResetLeaf(int col_idx, int64_t reserve) final {
++    return item_reader_->ResetLeaf(col_idx, reserve);
++  }
++
++  int64_t SkipRecords(int col_idx, int64_t num_records) final {
++    return item_reader_->SkipRecords(col_idx, num_records);
++  }
++
++  int64_t ReadRecords(int col_idx, int64_t num_records) final {
++    return item_reader_->ReadRecords(col_idx, num_records);
++  }
++
+   virtual ::arrow::Result<std::shared_ptr<ChunkedArray>> AssembleArray(
+       std::shared_ptr<ArrayData> data) {
+     if (field_->type()->id() == ::arrow::Type::MAP) {
+@@ -709,6 +776,39 @@
+     }
+     return Status::OK();
+   }
++
++  std::vector<int> LeafColumnIndices() const override {
++    std::vector<int> indices;
++    for (const std::unique_ptr<ColumnReaderImpl>& reader : children_) {
++      std::vector<int> child_indices = reader->LeafColumnIndices();
++      indices.insert(indices.end(), child_indices.begin(), child_indices.end());
++    }
++    return indices;
++  }
++
++  ::arrow::Status ResetLeaf(int col_idx, int64_t reserve) override {
++    for (const std::unique_ptr<ColumnReaderImpl>& reader : children_) {
++      RETURN_NOT_OK(reader->ResetLeaf(col_idx, reserve));
++    }
++    return Status::OK();
++  }
++
++  int64_t SkipRecords(int col_idx, int64_t num_records) override {
++    int64_t skipped = 0;
++    for (const std::unique_ptr<ColumnReaderImpl>& reader : children_) {
++      skipped += reader->SkipRecords(col_idx, num_records);
++    }
++    return skipped;
++  }
++
++  int64_t ReadRecords(int col_idx, int64_t num_records) override {
++    int64_t read = 0;
++    for (const std::unique_ptr<ColumnReaderImpl>& reader : children_) {
++      read += reader->ReadRecords(col_idx, num_records);
++    }
++    return read;
++  }
++
+   Status BuildArray(int64_t length_upper_bound,
+                     std::shared_ptr<ChunkedArray>* out) override;
+   Status GetDefLevels(const int16_t** data, int64_t* length) override;
+@@ -1228,6 +1328,23 @@
+   std::unique_ptr<ColumnReaderImpl> result;
+   RETURN_NOT_OK(GetReader(manifest_.schema_fields[i], ctx, &result));
+   *out = std::move(result);
++  return Status::OK();
++}
++
++::arrow::Status FileReaderImpl::GetColumn(
++    int i, const std::vector<int>& column_indices,
++    FileColumnIteratorFactory iterator_factory,
++    std::unique_ptr<ColumnReader>* out) {
++  RETURN_NOT_OK(BoundsCheckColumn(i));
++  auto ctx = std::make_shared<ReaderContext>();
++  ctx->reader = reader_.get();
++  ctx->pool = pool_;
++  ctx->iterator_factory = iterator_factory;
++  ctx->filter_leaves = true;
++  ctx->included_leaves = VectorToSharedSet(column_indices);
++  std::unique_ptr<ColumnReaderImpl> result;
++  RETURN_NOT_OK(GetReader(manifest_.schema_fields[i], ctx, &result));
++  *out = std::move(result);
+   return Status::OK();
+ }
+
+--- a/cpp/src/parquet/arrow/reader.h
++++ b/cpp/src/parquet/arrow/reader.h
+@@ -21,6 +21,7 @@
+ // N.B. we don't include async_generator.h as it's relatively heavy
+ #include <functional>
+ #include <memory>
++#include <utility>
+ #include <vector>
+
+ #include "parquet/file_reader.h"
+@@ -48,9 +49,13 @@
+
+ class ColumnChunkReader;
+ class ColumnReader;
++class FileColumnIterator;
+ struct SchemaManifest;
+ class RowGroupReader;
+
++using FileColumnIteratorFactory =
++    std::function<FileColumnIterator*(int, ParquetFileReader*)>;
++
+ /// \brief Arrow read adapter class for deserializing Parquet files as Arrow row batches.
+ ///
+ /// This interfaces caters for different use cases and thus provides different
+@@ -136,6 +141,27 @@
+   // The indicated column index is relative to the schema
+   virtual ::arrow::Status GetColumn(int i, std::unique_ptr<ColumnReader>* out) = 0;
+
++  /// \brief Return a ColumnReader with a custom FileColumnIteratorFactory
++  /// and leaf column filtering.
++  ///
++  /// This allows callers to customize page reading behavior (e.g., setting
++  /// data_page_filter for page-level skipping) and to select only specific
++  /// leaf columns within a nested field. The factory is called once per leaf
++  /// column included in column_indices.
++  ///
++  /// \param i top-level field index (same as GetColumn(int i, ...))
++  /// \param column_indices leaf column indices to include (enables sub-column
++  ///        projection within nested types)
++  /// \param iterator_factory factory to create FileColumnIterator per leaf
++  /// \param[out] out the ColumnReader (may be nullptr if all leaves are pruned)
++  virtual ::arrow::Status GetColumn(
++      int i, const std::vector<int>& column_indices,
++      FileColumnIteratorFactory iterator_factory,
++      std::unique_ptr<ColumnReader>* out) {
++    return ::arrow::Status::NotImplemented(
++        "GetColumn with factory not implemented");
++  }
++
+   /// \brief Return arrow schema for all the columns.
+   virtual ::arrow::Status GetSchema(std::shared_ptr<::arrow::Schema>* out) = 0;
+
+@@ -316,6 +342,43 @@
+   // the data available in the file.
+   virtual ::arrow::Status NextBatch(int64_t batch_size,
+                                     std::shared_ptr<::arrow::ChunkedArray>* out) = 0;
++
++  /// \brief Leaf column indices covered by this (sub)tree, in leaf order.
++  ///
++  /// Used to drive per-leaf row filtering: after page-level skipping each leaf
++  /// lives in its own compressed coordinate space, so callers must reset and
++  /// skip/read each leaf independently rather than in lockstep.
++  virtual std::vector<int> LeafColumnIndices() const { return {}; }
++
++  /// \brief Reset the leaf identified by col_idx and reserve space for
++  /// `reserve` records (in that leaf's post-page-filter compressed space).
++  /// Must be called before SkipRecords()/ReadRecords() for that leaf, and
++  /// followed by BuildArray() to get the result.
++  virtual ::arrow::Status ResetLeaf(int col_idx, int64_t reserve) {
++    return ::arrow::Status::NotImplemented("ResetLeaf not implemented");
++  }
++
++  /// \brief Skip num_records on the leaf identified by col_idx and return the
++  /// number of records actually skipped. Returns 0 when num_records <= 0 or
++  /// col_idx does not belong to this (sub)tree. May throw ParquetException on a
++  /// decode error; callers convert it to Status at the public boundary.
++  virtual int64_t SkipRecords(int col_idx, int64_t num_records) { return 0; }
++
++  /// \brief Read num_records on the leaf identified by col_idx and return the
++  /// number of records actually read. Values accumulate across successive calls
++  /// until BuildArray() is called. Returns 0 when num_records <= 0 or col_idx
++  /// does not belong to this (sub)tree. May throw ParquetException on a decode
++  /// error; callers convert it to Status at the public boundary.
++  virtual int64_t ReadRecords(int col_idx, int64_t num_records) { return 0; }
++
++  /// \brief Build the Arrow array from previously loaded data.
++  /// For leaf readers, calls TransferColumnData if not already done.
++  /// For nested readers, assembles the nested array from child arrays.
++  virtual ::arrow::Status BuildArray(
++      int64_t length_upper_bound,
++      std::shared_ptr<::arrow::ChunkedArray>* out) {
++    return ::arrow::Status::NotImplemented("BuildArray not implemented");
++  }
+ };
+
+ /// \brief Experimental helper class for bindings (like Python) that struggle
+--- a/cpp/src/parquet/arrow/reader_internal.h
++++ b/cpp/src/parquet/arrow/reader_internal.h
+@@ -26,6 +26,7 @@
+ #include <utility>
+ #include <vector>
+
++#include "parquet/arrow/reader.h"
+ #include "parquet/arrow/schema.h"
+ #include "parquet/column_reader.h"
+ #include "parquet/file_reader.h"
+@@ -70,7 +71,10 @@
+
+   virtual ~FileColumnIterator() {}
+
+-  std::unique_ptr<::parquet::PageReader> NextChunk() {
++  /// \brief Fetch the PageReader for the next row group in this iterator's
++  /// range. Virtual so subclasses can decorate the returned PageReader, e.g.
++  /// to install a data_page_filter for I/O-level page skipping.
++  virtual std::unique_ptr<::parquet::PageReader> NextChunk() {
+     if (row_groups_.empty()) {
+       return nullptr;
+     }
+@@ -95,9 +99,6 @@
+   std::deque<int> row_groups_;
+ };
+
+-using FileColumnIteratorFactory =
+-    std::function<FileColumnIterator*(int, ParquetFileReader*)>;
+-
+ Status TransferColumnData(::parquet::internal::RecordReader* reader,
+                           const std::shared_ptr<::arrow::Field>& value_field,
+                           const ColumnDescriptor* descr, ::arrow::MemoryPool* pool,

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<std::shared_ptr<arrow::RecordBatch>> 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<int64_t>::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<int32_t>& column_indices) {
-    if (page_filtered_read_schema_) {
-        return Status::OK();
-    }
-    std::shared_ptr<arrow::Schema> schema;
-    PAIMON_RETURN_NOT_OK_FROM_ARROW(file_reader_->GetSchema(&schema));
-    auto parquet_schema = file_reader_->parquet_reader()->metadata()->schema();
-    std::vector<std::shared_ptr<arrow::Field>> 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<int32_t>& 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<TargetRowGroup>& 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<int32_t> fully_matched_row_groups;
@@ -432,9 +408,6 @@ Status FileReaderWrapper::PrepareForReading(const std::vector<TargetRowGroup>& t
         }
 
         bool has_partially_matched = fully_matched_row_groups.size() != active_count;
-        if (has_partially_matched) {
-            PAIMON_RETURN_NOT_OK(BuildPageFilteredSchema(column_indices));
-        }
 
         WaitForPendingPreBuffer();