Skip to content

Commit db56e9f

Browse files
zhf999zhf999
authored
feat(parquet): support page-level (page index) filtering for nested column types (patch arrow) (alibaba#443)
1 parent 3f3e83b commit db56e9f

7 files changed

Lines changed: 976 additions & 410 deletions

‎cmake_modules/arrow.diff‎

Lines changed: 297 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -439,9 +439,305 @@ diff --git a/cpp/cmake_modules/BuildUtils.cmake b/cpp/cmake_modules/BuildUtils.c
439439
diff --git a/cpp/src/arrow/io/interfaces.h b/cpp/src/arrow/io/interfaces.h
440440
--- a/cpp/src/arrow/io/interfaces.h
441441
+++ b/cpp/src/arrow/io/interfaces.h
442-
@@ -211,7 +211,7 @@
442+
@@ -210,7 +210,7 @@
443443
/// \brief Advance or skip stream indicated number of bytes
444444
/// \param[in] nbytes the number to move forward
445445
/// \return Status
446446
- Status Advance(int64_t nbytes);
447447
+ virtual Status Advance(int64_t nbytes);
448+
449+
/// \brief Return zero-copy string_view to upcoming bytes.
450+
///
451+
--- a/cpp/src/parquet/arrow/reader.cc
452+
+++ b/cpp/src/parquet/arrow/reader.cc
453+
@@ -254,6 +254,11 @@
454+
return GetColumn(i, AllRowGroupsFactory(), out);
455+
}
456+
457+
+ ::arrow::Status GetColumn(
458+
+ int i, const std::vector<int>& column_indices,
459+
+ FileColumnIteratorFactory iterator_factory,
460+
+ std::unique_ptr<ColumnReader>* out) override;
461+
+
462+
Status GetSchema(std::shared_ptr<::arrow::Schema>* out) override {
463+
return FromParquetSchema(reader_->metadata()->schema(), reader_properties_,
464+
reader_->metadata()->key_value_metadata(), out);
465+
@@ -493,10 +498,40 @@
466+
467+
::arrow::Status BuildArray(int64_t length_upper_bound,
468+
std::shared_ptr<::arrow::ChunkedArray>* out) final {
469+
+ if (!out_) {
470+
+ BEGIN_PARQUET_CATCH_EXCEPTIONS
471+
+ RETURN_NOT_OK(
472+
+ TransferColumnData(record_reader_.get(), field_, descr_, ctx_->pool, &out_));
473+
+ END_PARQUET_CATCH_EXCEPTIONS
474+
+ }
475+
*out = out_;
476+
return Status::OK();
477+
}
478+
479+
+ std::vector<int> LeafColumnIndices() const final {
480+
+ return {input_->column_index()};
481+
+ }
482+
+
483+
+ ::arrow::Status ResetLeaf(int col_idx, int64_t reserve) final {
484+
+ if (col_idx != input_->column_index()) return Status::OK();
485+
+ BEGIN_PARQUET_CATCH_EXCEPTIONS
486+
+ out_ = nullptr;
487+
+ record_reader_->Reset();
488+
+ record_reader_->Reserve(reserve);
489+
+ return Status::OK();
490+
+ END_PARQUET_CATCH_EXCEPTIONS
491+
+ }
492+
+
493+
+ int64_t SkipRecords(int col_idx, int64_t num_records) final {
494+
+ if (col_idx != input_->column_index() || num_records <= 0) return 0;
495+
+ return record_reader_->SkipRecords(num_records);
496+
+ }
497+
+
498+
+ int64_t ReadRecords(int col_idx, int64_t num_records) final {
499+
+ if (col_idx != input_->column_index() || num_records <= 0) return 0;
500+
+ return record_reader_->ReadRecords(num_records);
501+
+ }
502+
+
503+
const std::shared_ptr<Field> field() override { return field_; }
504+
505+
private:
506+
@@ -532,6 +567,22 @@
507+
return storage_reader_->LoadBatch(number_of_records);
508+
}
509+
510+
+ std::vector<int> LeafColumnIndices() const final {
511+
+ return storage_reader_->LeafColumnIndices();
512+
+ }
513+
+
514+
+ ::arrow::Status ResetLeaf(int col_idx, int64_t reserve) final {
515+
+ return storage_reader_->ResetLeaf(col_idx, reserve);
516+
+ }
517+
+
518+
+ int64_t SkipRecords(int col_idx, int64_t num_records) final {
519+
+ return storage_reader_->SkipRecords(col_idx, num_records);
520+
+ }
521+
+
522+
+ int64_t ReadRecords(int col_idx, int64_t num_records) final {
523+
+ return storage_reader_->ReadRecords(col_idx, num_records);
524+
+ }
525+
+
526+
Status BuildArray(int64_t length_upper_bound,
527+
std::shared_ptr<ChunkedArray>* out) override {
528+
std::shared_ptr<ChunkedArray> storage;
529+
@@ -576,6 +627,22 @@
530+
return item_reader_->LoadBatch(number_of_records);
531+
}
532+
533+
+ std::vector<int> LeafColumnIndices() const final {
534+
+ return item_reader_->LeafColumnIndices();
535+
+ }
536+
+
537+
+ ::arrow::Status ResetLeaf(int col_idx, int64_t reserve) final {
538+
+ return item_reader_->ResetLeaf(col_idx, reserve);
539+
+ }
540+
+
541+
+ int64_t SkipRecords(int col_idx, int64_t num_records) final {
542+
+ return item_reader_->SkipRecords(col_idx, num_records);
543+
+ }
544+
+
545+
+ int64_t ReadRecords(int col_idx, int64_t num_records) final {
546+
+ return item_reader_->ReadRecords(col_idx, num_records);
547+
+ }
548+
+
549+
virtual ::arrow::Result<std::shared_ptr<ChunkedArray>> AssembleArray(
550+
std::shared_ptr<ArrayData> data) {
551+
if (field_->type()->id() == ::arrow::Type::MAP) {
552+
@@ -709,6 +776,39 @@
553+
}
554+
return Status::OK();
555+
}
556+
+
557+
+ std::vector<int> LeafColumnIndices() const override {
558+
+ std::vector<int> indices;
559+
+ for (const std::unique_ptr<ColumnReaderImpl>& reader : children_) {
560+
+ std::vector<int> child_indices = reader->LeafColumnIndices();
561+
+ indices.insert(indices.end(), child_indices.begin(), child_indices.end());
562+
+ }
563+
+ return indices;
564+
+ }
565+
+
566+
+ ::arrow::Status ResetLeaf(int col_idx, int64_t reserve) override {
567+
+ for (const std::unique_ptr<ColumnReaderImpl>& reader : children_) {
568+
+ RETURN_NOT_OK(reader->ResetLeaf(col_idx, reserve));
569+
+ }
570+
+ return Status::OK();
571+
+ }
572+
+
573+
+ int64_t SkipRecords(int col_idx, int64_t num_records) override {
574+
+ int64_t skipped = 0;
575+
+ for (const std::unique_ptr<ColumnReaderImpl>& reader : children_) {
576+
+ skipped += reader->SkipRecords(col_idx, num_records);
577+
+ }
578+
+ return skipped;
579+
+ }
580+
+
581+
+ int64_t ReadRecords(int col_idx, int64_t num_records) override {
582+
+ int64_t read = 0;
583+
+ for (const std::unique_ptr<ColumnReaderImpl>& reader : children_) {
584+
+ read += reader->ReadRecords(col_idx, num_records);
585+
+ }
586+
+ return read;
587+
+ }
588+
+
589+
Status BuildArray(int64_t length_upper_bound,
590+
std::shared_ptr<ChunkedArray>* out) override;
591+
Status GetDefLevels(const int16_t** data, int64_t* length) override;
592+
@@ -1228,6 +1328,23 @@
593+
std::unique_ptr<ColumnReaderImpl> result;
594+
RETURN_NOT_OK(GetReader(manifest_.schema_fields[i], ctx, &result));
595+
*out = std::move(result);
596+
+ return Status::OK();
597+
+}
598+
+
599+
+::arrow::Status FileReaderImpl::GetColumn(
600+
+ int i, const std::vector<int>& column_indices,
601+
+ FileColumnIteratorFactory iterator_factory,
602+
+ std::unique_ptr<ColumnReader>* out) {
603+
+ RETURN_NOT_OK(BoundsCheckColumn(i));
604+
+ auto ctx = std::make_shared<ReaderContext>();
605+
+ ctx->reader = reader_.get();
606+
+ ctx->pool = pool_;
607+
+ ctx->iterator_factory = iterator_factory;
608+
+ ctx->filter_leaves = true;
609+
+ ctx->included_leaves = VectorToSharedSet(column_indices);
610+
+ std::unique_ptr<ColumnReaderImpl> result;
611+
+ RETURN_NOT_OK(GetReader(manifest_.schema_fields[i], ctx, &result));
612+
+ *out = std::move(result);
613+
return Status::OK();
614+
}
615+
616+
--- a/cpp/src/parquet/arrow/reader.h
617+
+++ b/cpp/src/parquet/arrow/reader.h
618+
@@ -21,6 +21,7 @@
619+
// N.B. we don't include async_generator.h as it's relatively heavy
620+
#include <functional>
621+
#include <memory>
622+
+#include <utility>
623+
#include <vector>
624+
625+
#include "parquet/file_reader.h"
626+
@@ -48,9 +49,13 @@
627+
628+
class ColumnChunkReader;
629+
class ColumnReader;
630+
+class FileColumnIterator;
631+
struct SchemaManifest;
632+
class RowGroupReader;
633+
634+
+using FileColumnIteratorFactory =
635+
+ std::function<FileColumnIterator*(int, ParquetFileReader*)>;
636+
+
637+
/// \brief Arrow read adapter class for deserializing Parquet files as Arrow row batches.
638+
///
639+
/// This interfaces caters for different use cases and thus provides different
640+
@@ -136,6 +141,27 @@
641+
// The indicated column index is relative to the schema
642+
virtual ::arrow::Status GetColumn(int i, std::unique_ptr<ColumnReader>* out) = 0;
643+
644+
+ /// \brief Return a ColumnReader with a custom FileColumnIteratorFactory
645+
+ /// and leaf column filtering.
646+
+ ///
647+
+ /// This allows callers to customize page reading behavior (e.g., setting
648+
+ /// data_page_filter for page-level skipping) and to select only specific
649+
+ /// leaf columns within a nested field. The factory is called once per leaf
650+
+ /// column included in column_indices.
651+
+ ///
652+
+ /// \param i top-level field index (same as GetColumn(int i, ...))
653+
+ /// \param column_indices leaf column indices to include (enables sub-column
654+
+ /// projection within nested types)
655+
+ /// \param iterator_factory factory to create FileColumnIterator per leaf
656+
+ /// \param[out] out the ColumnReader (may be nullptr if all leaves are pruned)
657+
+ virtual ::arrow::Status GetColumn(
658+
+ int i, const std::vector<int>& column_indices,
659+
+ FileColumnIteratorFactory iterator_factory,
660+
+ std::unique_ptr<ColumnReader>* out) {
661+
+ return ::arrow::Status::NotImplemented(
662+
+ "GetColumn with factory not implemented");
663+
+ }
664+
+
665+
/// \brief Return arrow schema for all the columns.
666+
virtual ::arrow::Status GetSchema(std::shared_ptr<::arrow::Schema>* out) = 0;
667+
668+
@@ -316,6 +342,43 @@
669+
// the data available in the file.
670+
virtual ::arrow::Status NextBatch(int64_t batch_size,
671+
std::shared_ptr<::arrow::ChunkedArray>* out) = 0;
672+
+
673+
+ /// \brief Leaf column indices covered by this (sub)tree, in leaf order.
674+
+ ///
675+
+ /// Used to drive per-leaf row filtering: after page-level skipping each leaf
676+
+ /// lives in its own compressed coordinate space, so callers must reset and
677+
+ /// skip/read each leaf independently rather than in lockstep.
678+
+ virtual std::vector<int> LeafColumnIndices() const { return {}; }
679+
+
680+
+ /// \brief Reset the leaf identified by col_idx and reserve space for
681+
+ /// `reserve` records (in that leaf's post-page-filter compressed space).
682+
+ /// Must be called before SkipRecords()/ReadRecords() for that leaf, and
683+
+ /// followed by BuildArray() to get the result.
684+
+ virtual ::arrow::Status ResetLeaf(int col_idx, int64_t reserve) {
685+
+ return ::arrow::Status::NotImplemented("ResetLeaf not implemented");
686+
+ }
687+
+
688+
+ /// \brief Skip num_records on the leaf identified by col_idx and return the
689+
+ /// number of records actually skipped. Returns 0 when num_records <= 0 or
690+
+ /// col_idx does not belong to this (sub)tree. May throw ParquetException on a
691+
+ /// decode error; callers convert it to Status at the public boundary.
692+
+ virtual int64_t SkipRecords(int col_idx, int64_t num_records) { return 0; }
693+
+
694+
+ /// \brief Read num_records on the leaf identified by col_idx and return the
695+
+ /// number of records actually read. Values accumulate across successive calls
696+
+ /// until BuildArray() is called. Returns 0 when num_records <= 0 or col_idx
697+
+ /// does not belong to this (sub)tree. May throw ParquetException on a decode
698+
+ /// error; callers convert it to Status at the public boundary.
699+
+ virtual int64_t ReadRecords(int col_idx, int64_t num_records) { return 0; }
700+
+
701+
+ /// \brief Build the Arrow array from previously loaded data.
702+
+ /// For leaf readers, calls TransferColumnData if not already done.
703+
+ /// For nested readers, assembles the nested array from child arrays.
704+
+ virtual ::arrow::Status BuildArray(
705+
+ int64_t length_upper_bound,
706+
+ std::shared_ptr<::arrow::ChunkedArray>* out) {
707+
+ return ::arrow::Status::NotImplemented("BuildArray not implemented");
708+
+ }
709+
};
710+
711+
/// \brief Experimental helper class for bindings (like Python) that struggle
712+
--- a/cpp/src/parquet/arrow/reader_internal.h
713+
+++ b/cpp/src/parquet/arrow/reader_internal.h
714+
@@ -26,6 +26,7 @@
715+
#include <utility>
716+
#include <vector>
717+
718+
+#include "parquet/arrow/reader.h"
719+
#include "parquet/arrow/schema.h"
720+
#include "parquet/column_reader.h"
721+
#include "parquet/file_reader.h"
722+
@@ -70,7 +71,10 @@
723+
724+
virtual ~FileColumnIterator() {}
725+
726+
- std::unique_ptr<::parquet::PageReader> NextChunk() {
727+
+ /// \brief Fetch the PageReader for the next row group in this iterator's
728+
+ /// range. Virtual so subclasses can decorate the returned PageReader, e.g.
729+
+ /// to install a data_page_filter for I/O-level page skipping.
730+
+ virtual std::unique_ptr<::parquet::PageReader> NextChunk() {
731+
if (row_groups_.empty()) {
732+
return nullptr;
733+
}
734+
@@ -95,9 +99,6 @@
735+
std::deque<int> row_groups_;
736+
};
737+
738+
-using FileColumnIteratorFactory =
739+
- std::function<FileColumnIterator*(int, ParquetFileReader*)>;
740+
-
741+
Status TransferColumnData(::parquet::internal::RecordReader* reader,
742+
const std::shared_ptr<::arrow::Field>& value_field,
743+
const ColumnDescriptor* descr, ::arrow::MemoryPool* pool,

‎src/paimon/format/parquet/file_reader_wrapper.cpp‎

Lines changed: 5 additions & 32 deletions
Original file line numberDiff line numberDiff line change
@@ -30,6 +30,7 @@
3030
#include "paimon/format/parquet/parquet_format_defs.h"
3131
#include "paimon/macros.h"
3232
#include "parquet/arrow/reader.h"
33+
#include "parquet/arrow/schema.h"
3334
#include "parquet/file_reader.h"
3435
#include "parquet/metadata.h"
3536
#include "parquet/page_index.h"
@@ -230,15 +231,14 @@ Result<std::shared_ptr<arrow::RecordBatch>> FileReaderWrapper::NextPageFiltered(
230231
if (!current_page_filtered_reader_) {
231232
const auto& target_rg = target_row_groups_[current_row_group_idx_];
232233
auto page_ranges = PageFilteredRowGroupReader::ComputePageRanges(
233-
file_reader_->parquet_reader(), target_rg, target_column_indices_);
234+
target_rg, target_column_indices_, file_reader_->parquet_reader());
234235
bool pre_buffered = !prebuffered_ranges_.empty();
235236
int64_t max_chunksize = batch_size_ > 0 ? batch_size_ : std::numeric_limits<int64_t>::max();
236237
PAIMON_ASSIGN_OR_RAISE(
237238
current_page_filtered_reader_,
238239
PageFilteredRowGroupReader::ReadFilteredRowGroup(
239-
file_reader_->parquet_reader(), target_rg, target_column_indices_,
240-
page_filtered_read_schema_, file_reader_->properties().cache_options(),
241-
pre_buffered, page_ranges, max_chunksize, pool_));
240+
target_rg, target_column_indices_, file_reader_->properties().cache_options(),
241+
pre_buffered, page_ranges, max_chunksize, pool_, file_reader_.get()));
242242
current_filtered_row_ranges_ = target_rg.GetRowRanges();
243243
current_filtered_rg_start_ = all_row_group_ranges_[rg_id].first;
244244
filtered_global_offset_ = 0;
@@ -345,29 +345,6 @@ Status FileReaderWrapper::PrepareForReadingLazy(
345345
return Status::OK();
346346
}
347347

348-
Status FileReaderWrapper::BuildPageFilteredSchema(const std::vector<int32_t>& column_indices) {
349-
if (page_filtered_read_schema_) {
350-
return Status::OK();
351-
}
352-
std::shared_ptr<arrow::Schema> schema;
353-
PAIMON_RETURN_NOT_OK_FROM_ARROW(file_reader_->GetSchema(&schema));
354-
auto parquet_schema = file_reader_->parquet_reader()->metadata()->schema();
355-
std::vector<std::shared_ptr<arrow::Field>> fields;
356-
for (int32_t col_idx : column_indices) {
357-
const std::string& col_name = parquet_schema->Column(col_idx)->name();
358-
auto field = schema->GetFieldByName(col_name);
359-
if (!field) {
360-
return Status::Invalid(fmt::format(
361-
"PrepareForReading: Parquet column {} ('{}') has no matching Arrow field in "
362-
"file schema",
363-
col_idx, col_name));
364-
}
365-
fields.push_back(field);
366-
}
367-
page_filtered_read_schema_ = arrow::schema(fields);
368-
return Status::OK();
369-
}
370-
371348
std::vector<::arrow::io::ReadRange> FileReaderWrapper::CollectPreBufferRanges(
372349
const std::vector<int32_t>& column_indices) {
373350
std::vector<::arrow::io::ReadRange> ranges;
@@ -379,7 +356,7 @@ std::vector<::arrow::io::ReadRange> FileReaderWrapper::CollectPreBufferRanges(
379356
if (trg.IsPartiallyMatched()) {
380357
// Page-filtered RGs: only matching page byte ranges.
381358
auto page_ranges = PageFilteredRowGroupReader::ComputePageRanges(
382-
file_reader_->parquet_reader(), trg, column_indices);
359+
trg, column_indices, file_reader_->parquet_reader());
383360
ranges.insert(ranges.end(), std::make_move_iterator(page_ranges.begin()),
384361
std::make_move_iterator(page_ranges.end()));
385362
} else {
@@ -416,7 +393,6 @@ Status FileReaderWrapper::PrepareForReading(const std::vector<TargetRowGroup>& t
416393
try {
417394
target_row_groups_ = target_row_groups;
418395
target_column_indices_ = column_indices;
419-
page_filtered_read_schema_.reset();
420396

421397
// Partition into fully-matched and page-filtered row groups, skipping excluded ones.
422398
std::vector<int32_t> fully_matched_row_groups;
@@ -432,9 +408,6 @@ Status FileReaderWrapper::PrepareForReading(const std::vector<TargetRowGroup>& t
432408
}
433409

434410
bool has_partially_matched = fully_matched_row_groups.size() != active_count;
435-
if (has_partially_matched) {
436-
PAIMON_RETURN_NOT_OK(BuildPageFilteredSchema(column_indices));
437-
}
438411

439412
WaitForPendingPreBuffer();
440413

0 commit comments

Comments
 (0)