feat(parquet): support page index filter for nested column type (patch arrow)

cmake_modules/arrow.diff

@@ -431,9 +431,259 @@ 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,5 +210,5 @@
    /// \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);
+
+--- 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,42 @@
+
+   ::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();
+   }
+
++  ::arrow::Status LoadBatchWithRowFilter(const LeafRowFilter& get_leaf_filter) final {
++    BEGIN_PARQUET_CATCH_EXCEPTIONS
++    auto [pattern, total] = get_leaf_filter(input_->column_index());
++    out_ = nullptr;
++    record_reader_->Reset();
++    record_reader_->Reserve(total);
++    int64_t current = 0;
++    for (const auto& [skip, read] : pattern) {
++      if (skip > 0) {
++        record_reader_->SkipRecords(skip);
++        current += skip;
++      }
++      if (read > 0) {
++        record_reader_->ReadRecords(read);
++        current += read;
++      }
++    }
++    if (current < total) {
++      record_reader_->SkipRecords(total - current);
++    }
++    RETURN_NOT_OK(
++        TransferColumnData(record_reader_.get(), field_, descr_, ctx_->pool, &out_));
++    return Status::OK();
++    END_PARQUET_CATCH_EXCEPTIONS
++  }
++
+   const std::shared_ptr<Field> field() override { return field_; }
+
+  private:
+@@ -532,6 +569,10 @@
+     return storage_reader_->LoadBatch(number_of_records);
+   }
+
++  ::arrow::Status LoadBatchWithRowFilter(const LeafRowFilter& get_leaf_filter) final {
++    return storage_reader_->LoadBatchWithRowFilter(get_leaf_filter);
++  }
++
+   Status BuildArray(int64_t length_upper_bound,
+                     std::shared_ptr<ChunkedArray>* out) override {
+     std::shared_ptr<ChunkedArray> storage;
+@@ -576,6 +617,10 @@
+     return item_reader_->LoadBatch(number_of_records);
+   }
+
++  ::arrow::Status LoadBatchWithRowFilter(const LeafRowFilter& get_leaf_filter) final {
++    return item_reader_->LoadBatchWithRowFilter(get_leaf_filter);
++  }
++
+   virtual ::arrow::Result<std::shared_ptr<ChunkedArray>> AssembleArray(
+       std::shared_ptr<ArrayData> data) {
+     if (field_->type()->id() == ::arrow::Type::MAP) {
+@@ -709,6 +754,14 @@
+     }
+     return Status::OK();
+   }
++
++  ::arrow::Status LoadBatchWithRowFilter(const LeafRowFilter& get_leaf_filter) override {
++    for (const std::unique_ptr<ColumnReaderImpl>& reader : children_) {
++      RETURN_NOT_OK(reader->LoadBatchWithRowFilter(get_leaf_filter));
++    }
++    return Status::OK();
++  }
++
+   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 +1281,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,28 @@
+   // the data available in the file.
+   virtual ::arrow::Status NextBatch(int64_t batch_size,
+                                     std::shared_ptr<::arrow::ChunkedArray>* out) = 0;
++
++  /// \brief Load batch with per-leaf row filtering.
++  ///
++  /// The callback is called once per leaf column with that leaf's column index,
++  /// returning (skip_read_pattern, total_row_count) for that specific leaf.
++  /// Each pair in skip_read_pattern is (num_records_to_skip, num_records_to_read).
++  /// After processing all pairs, remaining records up to total_row_count are skipped.
++  /// Must be followed by BuildArray() to get the result.
++  using LeafRowFilter = std::function<
++      std::pair<std::vector<std::pair<int64_t, int64_t>>, int64_t>(int)>;
++  virtual ::arrow::Status LoadBatchWithRowFilter(const LeafRowFilter& get_leaf_filter) {
++    return ::arrow::Status::NotImplemented("LoadBatchWithRowFilter not implemented");
++  }
++
++  /// \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,6 +71,13 @@
+
+   virtual ~FileColumnIterator() {}
+
++  /// \brief Set a data_page_filter that will be applied to every PageReader
++  /// created by NextChunk(). This enables I/O-level page skipping.
++  void set_data_page_filter(
++      std::function<bool(const ::parquet::DataPageStats&)> filter) {
++    data_page_filter_ = std::move(filter);
++  }
++
+   std::unique_ptr<::parquet::PageReader> NextChunk() {
+     if (row_groups_.empty()) {
+       return nullptr;
+@@ -77,7 +85,11 @@
+
+     auto row_group_reader = reader_->RowGroup(row_groups_.front());
+     row_groups_.pop_front();
+-    return row_group_reader->GetColumnPageReader(column_index_);
++    auto page_reader = row_group_reader->GetColumnPageReader(column_index_);
++    if (page_reader && data_page_filter_) {
++      page_reader->set_data_page_filter(data_page_filter_);
++    }
++    return page_reader;
+   }
+
+   const SchemaDescriptor* schema() const { return schema_; }
+@@ -93,11 +105,9 @@
+   ParquetFileReader* reader_;
+   const SchemaDescriptor* schema_;
+   std::deque<int> row_groups_;
++  std::function<bool(const ::parquet::DataPageStats&)> data_page_filter_;
+ };
+
+-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

@@ -29,6 +29,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"
@@ -231,12 +232,11 @@ Result<std::shared_ptr<arrow::RecordBatch>> FileReaderWrapper::NextPageFiltered(
             file_reader_->parquet_reader(), target_rg, target_column_indices_);
         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_));
+        PAIMON_ASSIGN_OR_RAISE(current_page_filtered_reader_,
+                               PageFilteredRowGroupReader::ReadFilteredRowGroup(
+                                   file_reader_.get(), target_rg, target_column_indices_,
+                                   file_reader_->properties().cache_options(), pre_buffered,
+                                   page_ranges, max_chunksize, pool_));
         current_filtered_row_ranges_ = target_rg.row_ranges;
         current_filtered_rg_start_ = all_row_group_ranges_[rg_id].first;
         filtered_global_offset_ = 0;
@@ -339,29 +339,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;
@@ -410,7 +387,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;
@@ -426,9 +402,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();
 

src/paimon/format/parquet/file_reader_wrapper.h

@@ -157,9 +157,6 @@ class FileReaderWrapper {
     /// Read next batch from the fully-matched batch_reader_. Returns nullptr when exhausted.
     Result<std::shared_ptr<arrow::RecordBatch>> NextFullyMatched();
 
-    /// Build page_filtered_read_schema_ from the given column indices. No-op if already built.
-    Status BuildPageFilteredSchema(const std::vector<int32_t>& column_indices);
-
     /// Collect all byte ranges that need pre-buffering (page-filtered + fully-matched).
     std::vector<::arrow::io::ReadRange> CollectPreBufferRanges(
         const std::vector<int32_t>& column_indices);
@@ -193,11 +190,6 @@ class FileReaderWrapper {
     // Target row groups with row ranges for none page-level filtering and page-level filtering
     std::vector<TargetRowGroup> 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<arrow::Schema> page_filtered_read_schema_;
-
     // Track pre-buffered ranges so we can wait on destruction
     std::vector<::arrow::io::ReadRange> prebuffered_ranges_;
 };