refactor(manifest): reuse late materialization for bucket filtering

Author: gripleaf

src/paimon/common/reader/late_materializing_file_batch_reader.cpp

@@ -20,6 +20,7 @@
 #include "paimon/common/reader/late_materializing_file_batch_reader.h"
 
 #include <cassert>
+#include <limits>
 #include <map>
 #include <set>
 #include <string>
@@ -45,28 +46,50 @@
 namespace paimon {
 
 Result<std::unique_ptr<LateMaterializingFileBatchReader>> LateMaterializingFileBatchReader::Create(
-    std::unique_ptr<FileBatchReader> inner, const std::shared_ptr<arrow::MemoryPool>& arrow_pool) {
+    std::unique_ptr<FileBatchReader> inner, const std::shared_ptr<arrow::MemoryPool>& arrow_pool,
+    std::optional<ProbeValidation> validation) {
     if (arrow_pool == nullptr) {
         return Status::Invalid("arrow pool could not be nullptr.");
     }
     if (inner == nullptr) {
         return Status::Invalid("inner could not be nullptr.");
     }
+    if (validation && (!validation->validate || validation->fields.empty())) {
+        return Status::Invalid("probe validation requires fields and a callback");
+    }
     auto* prefetch_inner = dynamic_cast<PrefetchFileBatchReader*>(inner.get());
     auto reader = std::unique_ptr<LateMaterializingFileBatchReader>(
         new LateMaterializingFileBatchReader(std::move(inner), prefetch_inner, arrow_pool));
+    reader->validation_ = std::move(validation);
     return reader;
 }
 
 Result<FileBatchReader::ReadBatch> LateMaterializingFileBatchReader::NextBatch() {
+    if (state_ == kInit && validation_) {
+        return Status::Invalid("SetReadSchema is required before probe validation");
+    }
     if (state_ == kInit) {
         // SetReadSchema has not been called: read with the file schema, matching the
         // FileBatchReader contract for schema-less reads.
         state_ = kNoLatMat;
     }
     if (state_ == kProbing) {
         PAIMON_RETURN_NOT_OK(ReadAndFilterProbeData());
-        if (matched_bitmap_.IsEmpty()) {
+        if (!row_ids_fit_) {
+            // Rebind to full batches and apply the residual predicate without a file bitmap.
+            if (predicate_) {
+                std::map<std::string, int32_t> name_to_idx;
+                for (int32_t i = 0; i < full_schema_->num_fields(); ++i) {
+                    name_to_idx.emplace(full_schema_->field(i)->name(), i);
+                }
+                PAIMON_ASSIGN_OR_RAISE(
+                    std::shared_ptr<Predicate> bound,
+                    PredicateUtils::CreatePickedFieldFilter(predicate_, name_to_idx));
+                probe_filter_ = std::dynamic_pointer_cast<PredicateFilter>(bound);
+            }
+            PAIMON_RETURN_NOT_OK(SetInnerReadSchema(full_schema_, nullptr, std::nullopt));
+            state_ = kFiltering;
+        } else if (matched_bitmap_.IsEmpty()) {
             state_ = kEOF;
         } else {
             // payload pass reads only the matched rows (matched_bitmap_ is non-empty here).
@@ -78,6 +101,8 @@ Result<FileBatchReader::ReadBatch> LateMaterializingFileBatchReader::NextBatch()
 
     if (state_ == kNoLatMat) {
         return inner_->NextBatch();
+    } else if (state_ == kFiltering) {
+        return ReadFilteredBatch();
     } else if (state_ == kRunning) {
         return ReadPayloadBatch();
     } else if (state_ == kEOF) {
@@ -90,8 +115,10 @@ Result<RoaringBitmap32> LateMaterializingFileBatchReader::FilterProbeBatch(
     const std::shared_ptr<arrow::Array>& array,
     const std::shared_ptr<PredicateFilter>& bound_filter) {
     // TODO(zhouhonfeng.zhf): use arrow::compute::Filter instead of PredicateFilter
-    PAIMON_ASSIGN_OR_RAISE(std::vector<char> results,
-                           bound_filter->Test(*array, arrow_pool_.get()));
+    std::vector<char> results(array->length(), true);
+    if (bound_filter) {
+        PAIMON_ASSIGN_OR_RAISE(results, bound_filter->Test(*array, arrow_pool_.get()));
+    }
     if (results.size() != static_cast<size_t>(array->length())) {
         return Status::Invalid(
             fmt::format("predicate result size {} does not match probe batch length {}",
@@ -106,7 +133,12 @@ Result<RoaringBitmap32> LateMaterializingFileBatchReader::FilterProbeBatch(
         // map batch offset to file row id
         PAIMON_ASSIGN_OR_RAISE(uint64_t file_row,
                                inner_->GetPreviousBatchFileRowId(static_cast<uint64_t>(i)));
-        if (selection_ && !selection_->Contains(static_cast<int32_t>(file_row))) {
+        if (selection_ && (file_row > std::numeric_limits<uint32_t>::max() ||
+                           !selection_->Contains(static_cast<uint32_t>(file_row)))) {
+            continue;
+        }
+        if (file_row > std::numeric_limits<uint32_t>::max()) {
+            row_ids_fit_ = false;
             continue;
         }
         batch_matched.Add(static_cast<uint32_t>(i));
@@ -127,8 +159,20 @@ Status LateMaterializingFileBatchReader::ReadAndFilterProbeData() {
         auto& [c_array, c_schema] = batch;
         PAIMON_ASSIGN_OR_RAISE_FROM_ARROW(std::shared_ptr<arrow::Array> array,
                                           arrow::ImportArray(c_array.get(), c_schema.get()));
+        if (validation_) {
+            if (!array || array->type_id() != arrow::Type::STRUCT) {
+                return Status::Invalid("probe validation requires a struct array");
+            }
+            PAIMON_RETURN_NOT_OK(
+                validation_->validate(std::static_pointer_cast<arrow::StructArray>(array)));
+        }
         PAIMON_ASSIGN_OR_RAISE(RoaringBitmap32 batch_matched,
                                FilterProbeBatch(array, probe_filter_));
+        if (!row_ids_fit_) {
+            probe_arrays.clear();
+            matched_bitmap_ = RoaringBitmap32();
+            continue;
+        }
         // Compact each probe batch down to its matched rows so probe_data_ aligns row-for-row
         // (ascending file order) with matched_bitmap_ and the later payload output.
         if (!batch_matched.IsEmpty()) {
@@ -151,6 +195,53 @@ Status LateMaterializingFileBatchReader::ReadAndFilterProbeData() {
     return Status::OK();
 }
 
+Result<FileBatchReader::ReadBatch> LateMaterializingFileBatchReader::ReadFilteredBatch() {
+    while (true) {
+        PAIMON_ASSIGN_OR_RAISE(FileBatchReader::ReadBatch batch, inner_->NextBatch());
+        if (BatchReader::IsEofBatch(batch)) {
+            state_ = kEOF;
+            return batch;
+        }
+        PAIMON_ASSIGN_OR_RAISE_FROM_ARROW(
+            std::shared_ptr<arrow::Array> array,
+            arrow::ImportArray(batch.first.get(), batch.second.get()));
+        std::vector<char> results(array->length(), true);
+        if (probe_filter_) {
+            PAIMON_ASSIGN_OR_RAISE(results, probe_filter_->Test(*array, arrow_pool_.get()));
+        }
+        if (results.size() != static_cast<size_t>(array->length())) {
+            return Status::Invalid("predicate result size does not match full batch length");
+        }
+        RoaringBitmap32 selected;
+        row_mapping_.clear();
+        for (int64_t i = 0; i < array->length(); ++i) {
+            if (!results[i]) {
+                continue;
+            }
+            PAIMON_ASSIGN_OR_RAISE(uint64_t file_row, inner_->GetPreviousBatchFileRowId(i));
+            if (selection_ && (file_row > std::numeric_limits<uint32_t>::max() ||
+                               !selection_->Contains(static_cast<uint32_t>(file_row)))) {
+                continue;
+            }
+            selected.Add(static_cast<uint32_t>(i));
+            row_mapping_.push_back(file_row);
+        }
+        if (selected.IsEmpty()) {
+            continue;
+        }
+        PAIMON_ASSIGN_OR_RAISE(arrow::ArrayVector slices,
+                               ReaderUtils::GenerateFilteredArrayVector(array, selected));
+        PAIMON_ASSIGN_OR_RAISE_FROM_ARROW(std::shared_ptr<arrow::Array> filtered,
+                                          arrow::Concatenate(slices, arrow_pool_.get()));
+        auto c_array = std::make_unique<ArrowArray>();
+        auto c_schema = std::make_unique<ArrowSchema>();
+        PAIMON_RETURN_NOT_OK_FROM_ARROW(
+            arrow::ExportArray(*filtered, c_array.get(), c_schema.get()));
+        PAIMON_RETURN_NOT_OK(AddArrowArrayLifetime(c_array.get(), c_schema.get(), arrow_pool_));
+        return std::make_pair(std::move(c_array), std::move(c_schema));
+    }
+}
+
 Result<FileBatchReader::ReadBatch> LateMaterializingFileBatchReader::ReadPayloadBatch() {
     while (true) {
         PAIMON_ASSIGN_OR_RAISE(FileBatchReader::ReadBatchWithBitmap batch_with_bitmap,
@@ -258,9 +349,19 @@ Status LateMaterializingFileBatchReader::SetReadSchema(
     PAIMON_ASSIGN_OR_RAISE_FROM_ARROW(full_schema_, arrow::ImportSchema(read_schema));
     predicate_ = predicate;
     selection_ = selection_bitmap;
-    if (predicate_ != nullptr) {
+    if (predicate_ != nullptr || validation_) {
         std::set<std::string> probe_names;
-        PAIMON_RETURN_NOT_OK(PredicateUtils::GetAllNames(predicate_, &probe_names));
+        if (predicate_) {
+            PAIMON_RETURN_NOT_OK(PredicateUtils::GetAllNames(predicate_, &probe_names));
+        }
+        if (validation_) {
+            for (const auto& name : validation_->fields) {
+                if (!full_schema_->GetFieldByName(name)) {
+                    return Status::Invalid("probe validation field missing: ", name);
+                }
+                probe_names.insert(name);
+            }
+        }
         arrow::FieldVector probe_fields;
         arrow::FieldVector payload_fields;
         for (const auto& field : full_schema_->fields()) {
@@ -271,30 +372,36 @@ Status LateMaterializingFileBatchReader::SetReadSchema(
             }
         }
         // probing only pays off when the predicate fields are a strict subset of the read schema
-        if (!probe_fields.empty() && !payload_fields.empty()) {
+        if (!probe_fields.empty() && (!payload_fields.empty() || validation_)) {
+            if (payload_fields.empty()) {
+                payload_fields = full_schema_->fields();
+            }
             probe_schema_ = arrow::schema(probe_fields, full_schema_->metadata());
             payload_schema_ = arrow::schema(payload_fields, full_schema_->metadata());
-            PAIMON_RETURN_NOT_OK(PredicateValidator::ValidatePredicateWithSchema(
-                *probe_schema_, predicate_, /*validate_field_idx=*/false));
-            std::map<std::string, int32_t> name_to_idx;
-            for (int32_t i = 0; i < probe_schema_->num_fields(); ++i) {
-                name_to_idx.emplace(probe_schema_->field(i)->name(), i);
-            }
-            PAIMON_ASSIGN_OR_RAISE(
-                std::shared_ptr<Predicate> bound_predicate,
-                PredicateUtils::CreatePickedFieldFilter(predicate_, name_to_idx));
-            probe_filter_ = std::dynamic_pointer_cast<PredicateFilter>(bound_predicate);
-            if (!probe_filter_) {
-                return Status::Invalid("failed to bind predicate to probe schema");
+            if (predicate_) {
+                PAIMON_RETURN_NOT_OK(PredicateValidator::ValidatePredicateWithSchema(
+                    *probe_schema_, predicate_, /*validate_field_idx=*/false));
+                std::map<std::string, int32_t> name_to_idx;
+                for (int32_t i = 0; i < probe_schema_->num_fields(); ++i) {
+                    name_to_idx.emplace(probe_schema_->field(i)->name(), i);
+                }
+                PAIMON_ASSIGN_OR_RAISE(
+                    std::shared_ptr<Predicate> bound_predicate,
+                    PredicateUtils::CreatePickedFieldFilter(predicate_, name_to_idx));
+                probe_filter_ = std::dynamic_pointer_cast<PredicateFilter>(bound_predicate);
+                if (!probe_filter_) {
+                    return Status::Invalid("failed to bind predicate to probe schema");
+                }
             }
         }
     }
 
-    if (predicate_ == nullptr || probe_schema_ == nullptr) {
+    if (probe_schema_ == nullptr) {
         PAIMON_RETURN_NOT_OK(SetInnerReadSchema(full_schema_, predicate_, selection_));
         state_ = kNoLatMat;
     } else {
-        PAIMON_RETURN_NOT_OK(SetInnerReadSchema(probe_schema_, predicate_, selection_));
+        PAIMON_RETURN_NOT_OK(SetInnerReadSchema(probe_schema_, validation_ ? nullptr : predicate_,
+                                                validation_ ? std::nullopt : selection_));
         state_ = kProbing;
     }
     return Status::OK();
@@ -319,6 +426,10 @@ Status LateMaterializingFileBatchReader::SeekToRow(uint64_t row_number) {
     PAIMON_ASSIGN_OR_RAISE(PrefetchFileBatchReader * prefetch_reader,
                            GetPrefetchReaderOrRaise("SeekToRow"));
     PAIMON_RETURN_NOT_OK(prefetch_reader->SeekToRow(row_number));
+    if (!row_ids_fit_) {
+        state_ = kFiltering;
+        return Status::OK();
+    }
     if (state_ == kRunning || state_ == kEOF) {
         if (matched_bitmap_.IsEmpty()) {
             state_ = kEOF;
@@ -350,6 +461,7 @@ Status LateMaterializingFileBatchReader::SetReadRanges(
 
 void LateMaterializingFileBatchReader::Reset() {
     state_ = kInit;
+    row_ids_fit_ = true;
     matched_bitmap_ = RoaringBitmap32();
     probe_data_.reset();
     probe_cursor_ = 0;

src/paimon/common/reader/late_materializing_file_batch_reader.h

@@ -23,7 +23,10 @@
 #include <arrow/c/abi.h>
 
 #include <cstdint>
+#include <functional>
 #include <memory>
+#include <optional>
+#include <string>
 #include <string_view>
 #include <utility>
 #include <vector>
@@ -38,12 +41,20 @@ class PredicateFilter;
 // For convenience, we abbreviate `Later Materializing` as `LatMat`.
 // This reader is installed below the prefetch layer (see
 // AbstractSplitRead::CreateFileBatchReader) and performs probe/payload two-phase reads when a
-// predicate is pushed down through SetReadSchema; without a predicate it is a plain passthrough.
+// predicate is passed through SetReadSchema. Without a predicate or validation it is a passthrough.
 class LateMaterializingFileBatchReader : public PrefetchFileBatchReader {
  public:
+    // Validate every file row before filtering. These fields join the probe projection;
+    // neither predicates nor selections are pushed into the inner probe reader.
+    struct ProbeValidation {
+        std::vector<std::string> fields;
+        std::function<Status(const std::shared_ptr<arrow::StructArray>&)> validate;
+    };
+
     static Result<std::unique_ptr<LateMaterializingFileBatchReader>> Create(
         std::unique_ptr<FileBatchReader> inner,
-        const std::shared_ptr<arrow::MemoryPool>& arrow_pool);
+        const std::shared_ptr<arrow::MemoryPool>& arrow_pool,
+        std::optional<ProbeValidation> validation = std::nullopt);
 
     Result<FileBatchReader::ReadBatch> NextBatch() override;
 
@@ -114,9 +125,10 @@ class LateMaterializingFileBatchReader : public PrefetchFileBatchReader {
 
     enum LatMatState {
         kInit,
-        kProbing,   // schema is set, probing is in progress
-        kNoLatMat,  // no need to late materialization
-        kRunning,   // Lat-mat is enabled and the payload reader is reading data
+        kProbing,    // schema is set, probing is in progress
+        kNoLatMat,   // no need to late materialization
+        kFiltering,  // full batches with residual filtering when file row IDs exceed the bitmap
+        kRunning,    // Lat-mat is enabled and the payload reader is reading data
         kEOF
     };
 
@@ -130,6 +142,7 @@ class LateMaterializingFileBatchReader : public PrefetchFileBatchReader {
 
     /// Read one payload batch with bitmap (matched rows only)
     Result<FileBatchReader::ReadBatch> ReadPayloadBatch();
+    Result<FileBatchReader::ReadBatch> ReadFilteredBatch();
 
     /// Combine the compacted payload columns and the selected probe columns into a single struct
     /// array following full_schema_'s field order.
@@ -163,6 +176,7 @@ class LateMaterializingFileBatchReader : public PrefetchFileBatchReader {
     /// inner_ is never reassigned, so the cast is resolved once in Create().
     PrefetchFileBatchReader* prefetch_inner_ = nullptr;
     std::shared_ptr<arrow::MemoryPool> arrow_pool_;
+    std::optional<ProbeValidation> validation_;
     LatMatState state_ = kInit;
     std::shared_ptr<arrow::Schema> full_schema_;
     // projection holding only the predicate fields; nullptr when probing is not applicable
@@ -176,6 +190,7 @@ class LateMaterializingFileBatchReader : public PrefetchFileBatchReader {
     // the probe_data_ is sliced and compacted with the matched_bitmap_
     std::shared_ptr<arrow::StructArray> probe_data_;
     RoaringBitmap32 matched_bitmap_;
+    bool row_ids_fit_ = true;
     // read cursor into probe_data_ for the payload phase
     int64_t probe_cursor_ = 0;
     // to support GetPreviousBatchFileRowId