From 31736efd8cb58acc4d8dca8312ace99625f7ac8a Mon Sep 17 00:00:00 2001 From: zhouhongfeng Date: Fri, 21 Aug 2026 15:09:37 +0800 Subject: [PATCH 01/34] feat: support lat-mat (not tested) --- src/paimon/CMakeLists.txt | 1 + .../late_materializing_file_batch_reader.cpp | 353 ++++++++++++++++++ .../late_materializing_file_batch_reader.h | 119 ++++++ 3 files changed, 473 insertions(+) create mode 100644 src/paimon/common/reader/late_materializing_file_batch_reader.cpp create mode 100644 src/paimon/common/reader/late_materializing_file_batch_reader.h diff --git a/src/paimon/CMakeLists.txt b/src/paimon/CMakeLists.txt index 9b0807b64..557605859 100644 --- a/src/paimon/CMakeLists.txt +++ b/src/paimon/CMakeLists.txt @@ -133,6 +133,7 @@ set(PAIMON_COMMON_SRCS common/predicate/starts_with.cpp common/reader/batch_reader.cpp common/reader/concat_batch_reader.cpp + common/reader/late_materializing_file_batch_reader.cpp common/reader/predicate_batch_reader.cpp common/reader/prefetch_file_batch_reader_impl.cpp common/reader/reader_utils.cpp diff --git a/src/paimon/common/reader/late_materializing_file_batch_reader.cpp b/src/paimon/common/reader/late_materializing_file_batch_reader.cpp new file mode 100644 index 000000000..5dd5b36c0 --- /dev/null +++ b/src/paimon/common/reader/late_materializing_file_batch_reader.cpp @@ -0,0 +1,353 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you under the Apache License, Version 2.0 (the + * "License"); you may not use this file except in compliance + * with the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, + * software distributed under the License is distributed on an + * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY + * KIND, either express or implied. See the License for the + * specific language governing permissions and limitations + * under the License. + */ + +#include "paimon/common/reader/late_materializing_file_batch_reader.h" + +#include +#include +#include +#include +#include + +#include "arrow/array/concatenate.h" +#include "arrow/array/util.h" +#include "arrow/c/bridge.h" +#include "arrow/type.h" +#include "arrow/util/checked_cast.h" +#include "fmt/format.h" +#include "paimon/common/predicate/predicate_filter.h" +#include "paimon/common/reader/reader_utils.h" +#include "paimon/common/utils/arrow/status_utils.h" +#include "paimon/predicate/predicate_utils.h" +#include "paimon/status.h" + +namespace paimon { + +Result> LateMaterializingFileBatchReader::Create( + std::unique_ptr inner) { + auto reader = std::unique_ptr( + new LateMaterializingFileBatchReader(std::move(inner))); + return reader; +} + +Result LateMaterializingFileBatchReader::NextBatch() { + return Status::Invalid( + "paimon inner reader PrefetchFileBatchReader should use NextBatchWithBitmap"); +} + +Result +LateMaterializingFileBatchReader::NextBatchWithBitmap() { + if (state_ == kProbing) { + // This calling will update matched_bitmap_ and probe_data_ + PAIMON_RETURN_NOT_OK(ReadAndFilterProbeData()); + if (matched_bitmap_.IsEmpty()) { + state_ = kEOF; + } else { + // inner_->SetReadSchema with matched_bitmap_ + PAIMON_RETURN_NOT_OK(SetInnerPayloadSchema()); + state_ = kRunning; + } + } + + if (state_ == kNoLatMat) { + return inner_->NextBatchWithBitmap(); + } else if (state_ == kRunning) { + return ReadPayloadBatch(); + } else if (state_ == kEOF) { + return MakeEofBatchWithBitmap(); + } + return Status::Invalid("invalid state when calling NextBatchWithBitmap: " + + std::to_string(state_)); +} + +Result> LateMaterializingFileBatchReader::BindProbeFilter() { + // non_partition_filter's field_index_ is relative to the full data field list, so the + // predicate must be rebound by name to the probe schema before evaluating it. + std::map 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 bound_predicate, + PredicateUtils::CreatePickedFieldFilter(predicate_, name_to_idx)); + auto bound_filter = std::dynamic_pointer_cast(bound_predicate); + if (!bound_filter) { + return Status::Invalid("failed to bind predicate to probe schema"); + } + return bound_filter; +} + +Result LateMaterializingFileBatchReader::FilterProbeBatch( + const std::shared_ptr& array, + const std::shared_ptr& bound_filter) { + PAIMON_ASSIGN_OR_RAISE(std::vector results, bound_filter->Test(*array)); + if (results.size() != static_cast(array->length())) { + return Status::Invalid( + fmt::format("predicate result size {} does not match probe batch length {}", + results.size(), array->length())); + } + // batch-local offsets of the rows passing both the predicate and the selection + RoaringBitmap32 batch_matched; + for (int64_t i = 0; i < array->length(); ++i) { + if (!results[static_cast(i)]) { + continue; + } + // The format reader prunes row groups/pages, so a batch offset is not a file row id. + PAIMON_ASSIGN_OR_RAISE(uint64_t file_row, + inner_->GetPreviousBatchFileRowId(static_cast(i))); + // keep matched_bitmap_ ⊆ selection_ (file index ∩ ¬DV), which the probe reader applies + // imprecisely + if (selection_ && !selection_->Contains(file_row)) { + continue; + } + batch_matched.Add(static_cast(i)); + matched_bitmap_.Add(file_row); + } + return batch_matched; +} + +Status LateMaterializingFileBatchReader::ReadAndFilterProbeData() { + PAIMON_ASSIGN_OR_RAISE(std::shared_ptr bound_filter, BindProbeFilter()); + matched_bitmap_ = RoaringBitmap32(); + probe_cursor_ = 0; + arrow::ArrayVector probe_arrays; + while (true) { + PAIMON_ASSIGN_OR_RAISE(FileBatchReader::ReadBatch batch, inner_->NextBatch()); + if (BatchReader::IsEofBatch(batch)) { + break; + } + auto& [c_array, c_schema] = batch; + PAIMON_ASSIGN_OR_RAISE_FROM_ARROW(std::shared_ptr array, + arrow::ImportArray(c_array.get(), c_schema.get())); + PAIMON_ASSIGN_OR_RAISE(RoaringBitmap32 batch_matched, + FilterProbeBatch(array, bound_filter)); + // 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()) { + PAIMON_ASSIGN_OR_RAISE(arrow::ArrayVector matched_slices, + ReaderUtils::GenerateFilteredArrayVector(array, batch_matched)); + probe_arrays.insert(probe_arrays.end(), std::make_move_iterator(matched_slices.begin()), + std::make_move_iterator(matched_slices.end())); + } + } + + std::shared_ptr probe_array; + if (probe_arrays.empty()) { + PAIMON_ASSIGN_OR_RAISE_FROM_ARROW( + probe_array, arrow::MakeEmptyArray(arrow::struct_(probe_schema_->fields()))); + } else { + PAIMON_ASSIGN_OR_RAISE_FROM_ARROW(probe_array, arrow::Concatenate(probe_arrays)); + } + probe_data_ = arrow::internal::checked_pointer_cast(probe_array); + return Status::OK(); +} + +Result LateMaterializingFileBatchReader::ReadPayloadBatch() { + while (true) { + PAIMON_ASSIGN_OR_RAISE(FileBatchReader::ReadBatchWithBitmap batch_with_bitmap, + inner_->NextBatchWithBitmap()); + if (BatchReader::IsEofBatch(batch_with_bitmap)) { + state_ = kEOF; + return MakeEofBatchWithBitmap(); + } + auto& [batch, bitmap] = batch_with_bitmap; + // The payload reader prunes at row group/page granularity, so the returned batch is a + // superset; the bitmap marks the rows that are actually in matched_bitmap_. + if (bitmap.IsEmpty()) { + ReaderUtils::ReleaseReadBatch(std::move(batch)); + continue; + } + auto& [c_array, c_schema] = batch; + PAIMON_ASSIGN_OR_RAISE_FROM_ARROW(std::shared_ptr payload_array, + arrow::ImportArray(c_array.get(), c_schema.get())); + + // Record the file row id of each emitted row so GetPreviousBatchFileRowId stays correct + // after compaction/reassembly. + row_mapping_.clear(); + for (auto it = bitmap.Begin(); it != bitmap.End(); ++it) { + PAIMON_ASSIGN_OR_RAISE(uint64_t file_row, + inner_->GetPreviousBatchFileRowId(static_cast(*it))); + row_mapping_.push_back(file_row); + } + + // Compact the payload superset down to the matched rows (ascending file row order). + PAIMON_ASSIGN_OR_RAISE(arrow::ArrayVector payload_slices, + ReaderUtils::GenerateFilteredArrayVector(payload_array, bitmap)); + PAIMON_ASSIGN_OR_RAISE_FROM_ARROW(std::shared_ptr payload_compacted, + arrow::Concatenate(payload_slices)); + + // probe_data_ holds the matched probe rows in the same ascending file order, so a running + // cursor yields row-for-row alignment with the compacted payload. + int64_t card = static_cast(bitmap.Cardinality()); + if (probe_cursor_ + card > probe_data_->length()) { + return Status::Invalid( + fmt::format("probe cache underflow: cursor {} + {} exceeds probe rows {}", + probe_cursor_, card, probe_data_->length())); + } + std::shared_ptr probe_selected = probe_data_->Slice(probe_cursor_, card); + probe_cursor_ += card; + + PAIMON_ASSIGN_OR_RAISE(FileBatchReader::ReadBatch assembled, + AssembleFullBatch(payload_compacted, probe_selected)); + return ReaderUtils::AddAllValidBitmap(std::move(assembled)); + } +} + +Result LateMaterializingFileBatchReader::AssembleFullBatch( + const std::shared_ptr& payload_array, + const std::shared_ptr& probe_array) { + auto payload_struct = arrow::internal::checked_pointer_cast(payload_array); + auto probe_struct = arrow::internal::checked_pointer_cast(probe_array); + arrow::ArrayVector children; + children.reserve(full_schema_->num_fields()); + for (const auto& field : full_schema_->fields()) { + std::shared_ptr col = payload_struct->GetFieldByName(field->name()); + if (!col) { + col = probe_struct->GetFieldByName(field->name()); + } + if (!col) { + return Status::Invalid( + fmt::format("field {} missing in both payload and probe columns", field->name())); + } + children.push_back(std::move(col)); + } + PAIMON_ASSIGN_OR_RAISE_FROM_ARROW(std::shared_ptr full_struct, + arrow::StructArray::Make(children, full_schema_->fields())); + std::unique_ptr<::ArrowArray> c_array = std::make_unique<::ArrowArray>(); + std::unique_ptr<::ArrowSchema> c_schema = std::make_unique<::ArrowSchema>(); + PAIMON_RETURN_NOT_OK_FROM_ARROW( + arrow::ExportArray(*full_struct, c_array.get(), c_schema.get())); + return std::make_pair(std::move(c_array), std::move(c_schema)); +} + +Status LateMaterializingFileBatchReader::SetInnerProbeSchema() { + ::ArrowSchema c_probe_schema; + PAIMON_RETURN_NOT_OK_FROM_ARROW(arrow::ExportSchema(*probe_schema_, &c_probe_schema)); + return inner_->SetReadSchema(&c_probe_schema, predicate_, selection_); +} + +Status LateMaterializingFileBatchReader::SetInnerPayloadSchema() { + ::ArrowSchema c_payload_schema; + PAIMON_RETURN_NOT_OK_FROM_ARROW(arrow::ExportSchema(*payload_schema_, &c_payload_schema)); + if (matched_bitmap_.IsEmpty()) { + return Status::Invalid("late materialization bitmap is empty. Should return EOF."); + } + return inner_->SetReadSchema(&c_payload_schema, /*predicate=*/nullptr, matched_bitmap_); +} + +Status LateMaterializingFileBatchReader::SetInnerFullSchema() { + ::ArrowSchema c_full_schema; + PAIMON_RETURN_NOT_OK_FROM_ARROW(arrow::ExportSchema(*full_schema_, &c_full_schema)); + return inner_->SetReadSchema(&c_full_schema, predicate_, selection_); +} + +std::shared_ptr LateMaterializingFileBatchReader::GetReaderMetrics() const { + return inner_->GetReaderMetrics(); +} + +void LateMaterializingFileBatchReader::Close() { + inner_->Close(); +} + +Result> LateMaterializingFileBatchReader::GetFileSchema() const { + return inner_->GetFileSchema(); +} + +Status LateMaterializingFileBatchReader::SetReadSchema( + ::ArrowSchema* read_schema, const std::shared_ptr& predicate, + const std::optional& selection_bitmap) { + PAIMON_ASSIGN_OR_RAISE_FROM_ARROW(full_schema_, arrow::ImportSchema(read_schema)); + predicate_ = predicate; + selection_ = selection_bitmap; + matched_bitmap_ = RoaringBitmap32(); + probe_data_.reset(); + probe_cursor_ = 0; + row_mapping_.clear(); + probe_schema_.reset(); + payload_schema_.reset(); + if (predicate_ != nullptr) { + std::set probe_names; + PAIMON_RETURN_NOT_OK(PredicateUtils::GetAllNames(predicate_, &probe_names)); + arrow::FieldVector probe_fields; + arrow::FieldVector payload_fields; + for (const auto& field : full_schema_->fields()) { + if (probe_names.count(field->name()) > 0) { + probe_fields.push_back(field); + } else { + payload_fields.push_back(field); + } + } + // probing only pays off when the predicate fields are a strict subset of the read schema + if (!probe_fields.empty() && !payload_fields.empty()) { + probe_schema_ = arrow::schema(probe_fields, full_schema_->metadata()); + payload_schema_ = arrow::schema(payload_fields, full_schema_->metadata()); + } + } + + if (predicate_ == nullptr || probe_schema_ == nullptr) { + PAIMON_RETURN_NOT_OK(SetInnerFullSchema()); + state_ = kNoLatMat; + } else { + PAIMON_RETURN_NOT_OK(SetInnerProbeSchema()); + state_ = kProbing; + } + return Status::OK(); +} + +Result LateMaterializingFileBatchReader::GetPreviousBatchFileRowId( + uint64_t batch_row_id) const { + if (state_ == kNoLatMat) { + return inner_->GetPreviousBatchFileRowId(batch_row_id); + } + // In kRunning the emitted batch is compacted/reassembled, so row ids come from row_mapping_ + // instead of the inner reader. + if (batch_row_id >= row_mapping_.size()) { + return Status::Invalid( + fmt::format("batch_row_id {} is out of range, last batch row count is {}", batch_row_id, + row_mapping_.size())); + } + return row_mapping_[batch_row_id]; +} + +Result LateMaterializingFileBatchReader::GetNumberOfRows() const { + return inner_->GetNumberOfRows(); +} + +bool LateMaterializingFileBatchReader::SupportPreciseBitmapSelection() const { + return inner_->SupportPreciseBitmapSelection(); +} + +Status LateMaterializingFileBatchReader::SeekToRow(uint64_t row_number) { + return inner_->SeekToRow(row_number); +} + +uint64_t LateMaterializingFileBatchReader::GetNextRowToRead() const { + return inner_->GetNextRowToRead(); +} + +Result>> LateMaterializingFileBatchReader::GenReadRanges( + bool* need_prefetch) const { + return inner_->GenReadRanges(need_prefetch); +} + +Status LateMaterializingFileBatchReader::SetReadRanges( + const std::vector>& read_ranges) { + return inner_->SetReadRanges(read_ranges); +} + +} // namespace paimon diff --git a/src/paimon/common/reader/late_materializing_file_batch_reader.h b/src/paimon/common/reader/late_materializing_file_batch_reader.h new file mode 100644 index 000000000..65df657fb --- /dev/null +++ b/src/paimon/common/reader/late_materializing_file_batch_reader.h @@ -0,0 +1,119 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you under the Apache License, Version 2.0 (the + * "License"); you may not use this file except in compliance + * with the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, + * software distributed under the License is distributed on an + * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY + * KIND, either express or implied. See the License for the + * specific language governing permissions and limitations + * under the License. + */ + +#pragma once + +#include + +#include +#include +#include +#include + +#include "paimon/reader/prefetch_file_batch_reader.h" + +namespace paimon { + +class PredicateFilter; + +// For convenience, we abbreviate `Later Materializing` as `LatMat`. +class LateMaterializingFileBatchReader : public PrefetchFileBatchReader { + public: + static Result> Create( + std::unique_ptr inner); + + Result NextBatch() override; + Result NextBatchWithBitmap() override; + + std::shared_ptr GetReaderMetrics() const override; + void Close() override; + + Result> GetFileSchema() const override; + Status SetReadSchema(::ArrowSchema* read_schema, const std::shared_ptr& predicate, + const std::optional& selection_bitmap) override; + Result GetPreviousBatchFileRowId(uint64_t batch_row_id) const override; + Result GetNumberOfRows() const override; + bool SupportPreciseBitmapSelection() const override; + + Status SeekToRow(uint64_t row_number) override; + uint64_t GetNextRowToRead() const override; + Result>> GenReadRanges( + bool* need_prefetch) const override; + Status SetReadRanges(const std::vector>& read_ranges) override; + + private: + explicit LateMaterializingFileBatchReader(std::unique_ptr inner) + : inner_(std::move(inner)) {} + + enum LatMatState { + kInit, + kProbing, // schema is set + kNoLatMat, // no need to late materialization + kRunning, // Lat-mat is enable an is reading data + kEOF + }; + /// Scan the probe projection once, evaluating the predicate batch by batch: the matched file + /// row ids go into matched_bitmap_ and the matched probe values into probe_data_. + Status ReadAndFilterProbeData(); + + /// Rebind the predicate to probe_schema_'s field indices, so that it can be evaluated over + /// the probe batches. + Result> BindProbeFilter(); + + /// Evaluate bound_filter over a single probe batch, adding the matched file row ids (that also + /// pass selection_) to matched_bitmap_. Returns the batch-local offsets of the matched rows so + /// the caller can compact the probe batch down to those rows. + Result FilterProbeBatch(const std::shared_ptr& array, + const std::shared_ptr& bound_filter); + + /// In kRunning state, read one payload batch, map its matched rows back to the cached probe + /// rows by file row id, and reassemble the full read schema. + Result ReadPayloadBatch(); + + /// Combine the compacted payload columns and the selected probe columns into a single struct + /// array following full_schema_'s field order. + Result AssembleFullBatch( + const std::shared_ptr& payload_array, + const std::shared_ptr& probe_array); + + Status SetInnerProbeSchema(); + Status SetInnerPayloadSchema(); + Status SetInnerFullSchema(); + + LatMatState state_ = kInit; + std::unique_ptr inner_; + std::shared_ptr full_schema_; + // projection holding only the predicate fields; nullptr when probing is not applicable + std::shared_ptr probe_schema_; + // projection holding the payload (non-probe) fields; nullptr when probing is not applicable + std::shared_ptr payload_schema_; + std::shared_ptr predicate_; + std::optional selection_; + // matched probe rows only, concatenated in ascending file row order (aligned 1:1 with + // matched_bitmap_) + std::shared_ptr probe_data_; + // file-level row ids that pass the predicate (and selection_); drives the payload read + RoaringBitmap32 matched_bitmap_; + // read cursor into probe_data_ for the payload phase + int64_t probe_cursor_ = 0; + // file row id of each row in the batch last emitted in kRunning state + std::vector row_mapping_; +}; + +} // namespace paimon From c57e1c34bc1e8956b2dc38018631e0ae890f85b0 Mon Sep 17 00:00:00 2001 From: zhouhongfeng Date: Fri, 21 Aug 2026 16:03:14 +0800 Subject: [PATCH 02/34] cache and forward read ranges --- .../late_materializing_file_batch_reader.cpp | 17 +++++++++++++---- .../late_materializing_file_batch_reader.h | 1 + 2 files changed, 14 insertions(+), 4 deletions(-) diff --git a/src/paimon/common/reader/late_materializing_file_batch_reader.cpp b/src/paimon/common/reader/late_materializing_file_batch_reader.cpp index 5dd5b36c0..7931f27b1 100644 --- a/src/paimon/common/reader/late_materializing_file_batch_reader.cpp +++ b/src/paimon/common/reader/late_materializing_file_batch_reader.cpp @@ -238,7 +238,10 @@ Result LateMaterializingFileBatchReader::AssembleFul Status LateMaterializingFileBatchReader::SetInnerProbeSchema() { ::ArrowSchema c_probe_schema; PAIMON_RETURN_NOT_OK_FROM_ARROW(arrow::ExportSchema(*probe_schema_, &c_probe_schema)); - return inner_->SetReadSchema(&c_probe_schema, predicate_, selection_); + PAIMON_RETURN_NOT_OK(inner_->SetReadSchema(&c_probe_schema, predicate_, selection_)); + // SetReadSchema may refresh the read ranges of the inner reader, so we set it again. + PAIMON_RETURN_NOT_OK(SetReadRanges(read_ranges_)); + return Status::OK(); } Status LateMaterializingFileBatchReader::SetInnerPayloadSchema() { @@ -247,13 +250,17 @@ Status LateMaterializingFileBatchReader::SetInnerPayloadSchema() { if (matched_bitmap_.IsEmpty()) { return Status::Invalid("late materialization bitmap is empty. Should return EOF."); } - return inner_->SetReadSchema(&c_payload_schema, /*predicate=*/nullptr, matched_bitmap_); + PAIMON_RETURN_NOT_OK(inner_->SetReadSchema(&c_payload_schema, /*predicate=*/nullptr, matched_bitmap_)); + PAIMON_RETURN_NOT_OK(SetReadRanges(read_ranges_)); + return Status::OK(); } Status LateMaterializingFileBatchReader::SetInnerFullSchema() { ::ArrowSchema c_full_schema; PAIMON_RETURN_NOT_OK_FROM_ARROW(arrow::ExportSchema(*full_schema_, &c_full_schema)); - return inner_->SetReadSchema(&c_full_schema, predicate_, selection_); + PAIMON_RETURN_NOT_OK(inner_->SetReadSchema(&c_full_schema, predicate_, selection_)); + PAIMON_RETURN_NOT_OK(SetReadRanges(read_ranges_)); + return Status::OK(); } std::shared_ptr LateMaterializingFileBatchReader::GetReaderMetrics() const { @@ -347,7 +354,9 @@ Result>> LateMaterializingFileBatchRea Status LateMaterializingFileBatchReader::SetReadRanges( const std::vector>& read_ranges) { - return inner_->SetReadRanges(read_ranges); + read_ranges_ = read_ranges; + PAIMON_RETURN_NOT_OK(inner_->SetReadRanges(read_ranges_)); + return Status::OK(); } } // namespace paimon diff --git a/src/paimon/common/reader/late_materializing_file_batch_reader.h b/src/paimon/common/reader/late_materializing_file_batch_reader.h index 65df657fb..0ca9fc733 100644 --- a/src/paimon/common/reader/late_materializing_file_batch_reader.h +++ b/src/paimon/common/reader/late_materializing_file_batch_reader.h @@ -98,6 +98,7 @@ class LateMaterializingFileBatchReader : public PrefetchFileBatchReader { LatMatState state_ = kInit; std::unique_ptr inner_; + std::vector> read_ranges_; std::shared_ptr full_schema_; // projection holding only the predicate fields; nullptr when probing is not applicable std::shared_ptr probe_schema_; From 76b83175de1b456a0e2f0a3819e6a49fdd5e5661 Mon Sep 17 00:00:00 2001 From: zhouhongfeng Date: Fri, 21 Aug 2026 16:19:07 +0800 Subject: [PATCH 03/34] fix: not use the bitmap from inner read --- .../late_materializing_file_batch_reader.cpp | 24 ++++++++++++------- 1 file changed, 16 insertions(+), 8 deletions(-) diff --git a/src/paimon/common/reader/late_materializing_file_batch_reader.cpp b/src/paimon/common/reader/late_materializing_file_batch_reader.cpp index 7931f27b1..5c67f9378 100644 --- a/src/paimon/common/reader/late_materializing_file_batch_reader.cpp +++ b/src/paimon/common/reader/late_materializing_file_batch_reader.cpp @@ -166,8 +166,6 @@ Result LateMaterializingFileBatchReader::R return MakeEofBatchWithBitmap(); } auto& [batch, bitmap] = batch_with_bitmap; - // The payload reader prunes at row group/page granularity, so the returned batch is a - // superset; the bitmap marks the rows that are actually in matched_bitmap_. if (bitmap.IsEmpty()) { ReaderUtils::ReleaseReadBatch(std::move(batch)); continue; @@ -176,24 +174,34 @@ Result LateMaterializingFileBatchReader::R PAIMON_ASSIGN_OR_RAISE_FROM_ARROW(std::shared_ptr payload_array, arrow::ImportArray(c_array.get(), c_schema.get())); - // Record the file row id of each emitted row so GetPreviousBatchFileRowId stays correct - // after compaction/reassembly. + // Recompute the precise selection: keep only the candidate rows whose file row id is in + // matched_bitmap_. Also record their file row ids so GetPreviousBatchFileRowId stays + // correct after compaction/reassembly. + RoaringBitmap32 valid; row_mapping_.clear(); for (auto it = bitmap.Begin(); it != bitmap.End(); ++it) { - PAIMON_ASSIGN_OR_RAISE(uint64_t file_row, - inner_->GetPreviousBatchFileRowId(static_cast(*it))); + uint64_t offset = static_cast(*it); + PAIMON_ASSIGN_OR_RAISE(uint64_t file_row, inner_->GetPreviousBatchFileRowId(offset)); + if (!matched_bitmap_.Contains(file_row)) { + continue; + } + valid.Add(static_cast(offset)); row_mapping_.push_back(file_row); } + if (valid.IsEmpty()) { + ReaderUtils::ReleaseReadBatch(std::move(batch)); + continue; + } // Compact the payload superset down to the matched rows (ascending file row order). PAIMON_ASSIGN_OR_RAISE(arrow::ArrayVector payload_slices, - ReaderUtils::GenerateFilteredArrayVector(payload_array, bitmap)); + ReaderUtils::GenerateFilteredArrayVector(payload_array, valid)); PAIMON_ASSIGN_OR_RAISE_FROM_ARROW(std::shared_ptr payload_compacted, arrow::Concatenate(payload_slices)); // probe_data_ holds the matched probe rows in the same ascending file order, so a running // cursor yields row-for-row alignment with the compacted payload. - int64_t card = static_cast(bitmap.Cardinality()); + int64_t card = static_cast(valid.Cardinality()); if (probe_cursor_ + card > probe_data_->length()) { return Status::Invalid( fmt::format("probe cache underflow: cursor {} + {} exceeds probe rows {}", From 4a5fa3bfb1c193d4d7eacf2e038c08287be4ca6f Mon Sep 17 00:00:00 2001 From: zhouhongfeng Date: Fri, 21 Aug 2026 16:30:12 +0800 Subject: [PATCH 04/34] fix: support SeekToRow --- .../late_materializing_file_batch_reader.cpp | 14 +++++++++++++- 1 file changed, 13 insertions(+), 1 deletion(-) diff --git a/src/paimon/common/reader/late_materializing_file_batch_reader.cpp b/src/paimon/common/reader/late_materializing_file_batch_reader.cpp index 5c67f9378..ab36c4adc 100644 --- a/src/paimon/common/reader/late_materializing_file_batch_reader.cpp +++ b/src/paimon/common/reader/late_materializing_file_batch_reader.cpp @@ -348,7 +348,19 @@ bool LateMaterializingFileBatchReader::SupportPreciseBitmapSelection() const { } Status LateMaterializingFileBatchReader::SeekToRow(uint64_t row_number) { - return inner_->SeekToRow(row_number); + PAIMON_RETURN_NOT_OK(inner_->SeekToRow(row_number)); + if (state_ == kRunning || state_ == kEOF) { + int64_t cursor = 0; + for (auto it = matched_bitmap_.Begin(); it != matched_bitmap_.End(); ++it) { + if (static_cast(*it) >= row_number) { + break; + } + ++cursor; + } + probe_cursor_ = cursor; + state_ = kRunning; // a seek after EOF re-activates payload reading + } + return Status::OK(); } uint64_t LateMaterializingFileBatchReader::GetNextRowToRead() const { From b74258def19592061fd5da8727f35184e8c6febd Mon Sep 17 00:00:00 2001 From: zhouhongfeng Date: Fri, 21 Aug 2026 17:42:39 +0800 Subject: [PATCH 05/34] test: add test cases --- src/paimon/CMakeLists.txt | 1 + ...e_materializing_file_batch_reader_test.cpp | 647 ++++++++++++++++++ .../testing/mock/mock_file_batch_reader.h | 82 ++- 3 files changed, 729 insertions(+), 1 deletion(-) create mode 100644 src/paimon/common/reader/late_materializing_file_batch_reader_test.cpp diff --git a/src/paimon/CMakeLists.txt b/src/paimon/CMakeLists.txt index 557605859..b666f20e6 100644 --- a/src/paimon/CMakeLists.txt +++ b/src/paimon/CMakeLists.txt @@ -596,6 +596,7 @@ if(PAIMON_BUILD_TESTS) common/predicate/predicate_utils_test.cpp common/predicate/predicate_validator_test.cpp common/reader/concat_batch_reader_test.cpp + common/reader/late_materializing_file_batch_reader_test.cpp common/reader/predicate_batch_reader_test.cpp common/reader/prefetch_file_batch_reader_impl_test.cpp common/reader/reader_utils_test.cpp diff --git a/src/paimon/common/reader/late_materializing_file_batch_reader_test.cpp b/src/paimon/common/reader/late_materializing_file_batch_reader_test.cpp new file mode 100644 index 000000000..84b52b0c3 --- /dev/null +++ b/src/paimon/common/reader/late_materializing_file_batch_reader_test.cpp @@ -0,0 +1,647 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you under the Apache License, Version 2.0 (the + * "License"); you may not use this file except in compliance + * with the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, + * software distributed under the License is distributed on an + * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY + * KIND, either express or implied. See the License for the + * specific language governing permissions and limitations + * under the License. + */ + +#include "paimon/common/reader/late_materializing_file_batch_reader.h" + +#include +#include +#include +#include +#include +#include +#include +#include + +#include "arrow/api.h" +#include "arrow/array/builder_nested.h" +#include "arrow/c/bridge.h" +#include "gtest/gtest.h" +#include "paimon/common/reader/prefetch_file_batch_reader_impl.h" +#include "paimon/common/reader/reader_utils.h" +#include "paimon/common/utils/arrow/status_utils.h" +#include "paimon/common/utils/checked_cast.h" +#include "paimon/common/utils/read_ahead_cache.h" +#include "paimon/executor.h" +#include "paimon/format/reader_builder.h" +#include "paimon/memory/memory_pool.h" +#include "paimon/predicate/literal.h" +#include "paimon/predicate/predicate.h" +#include "paimon/predicate/predicate_builder.h" +#include "paimon/reader/prefetch_file_batch_reader.h" +#include "paimon/status.h" +#include "paimon/testing/mock/mock_file_batch_reader.h" +#include "paimon/testing/mock/mock_file_system.h" +#include "paimon/testing/utils/testharness.h" +#include "paimon/utils/roaring_bitmap32.h" + +namespace paimon::test { + +// A ReaderBuilder whose Build() wraps a projecting, range-honoring mock in a +// LateMaterializingFileBatchReader, so the late-materialization reader can be exercised as an +// inner reader of PrefetchFileBatchReaderImpl. +class LmReaderBuilder : public ReaderBuilder { + public: + LmReaderBuilder(std::shared_ptr data, std::shared_ptr type, + int32_t batch_size) + : data_(std::move(data)), type_(std::move(type)), batch_size_(batch_size) {} + + ReaderBuilder* WithMemoryPool(const std::shared_ptr& /*pool*/) override { + return this; + } + + Result> Build( + const std::shared_ptr& /*stream*/) const override { + auto mock = std::make_unique(data_, type_, batch_size_); + PAIMON_ASSIGN_OR_RAISE(std::unique_ptr reader, + LateMaterializingFileBatchReader::Create(std::move(mock))); + return std::unique_ptr(std::move(reader)); + } + + private: + std::shared_ptr data_; + std::shared_ptr type_; + int32_t batch_size_ = 0; +}; + +class LateMaterializingFileBatchReaderTest : public ::testing::Test { + public: + void SetUp() override { + k_field_ = arrow::field("k", arrow::int64()); + v_field_ = arrow::field("v", arrow::utf8()); + full_fields_ = {k_field_, v_field_}; + full_type_ = arrow::struct_(full_fields_); + } + + // Build a struct array with column k (int64, values = ks) and column v (utf8, "v_"). + std::shared_ptr BuildData(const std::vector& ks) { + arrow::StructBuilder builder( + full_type_, arrow::default_memory_pool(), + {std::make_shared(), std::make_shared()}); + auto* k_builder = checked_cast(builder.field_builder(0)); + auto* v_builder = checked_cast(builder.field_builder(1)); + for (size_t i = 0; i < ks.size(); ++i) { + EXPECT_TRUE(builder.Append().ok()); + EXPECT_TRUE(k_builder->Append(ks[i]).ok()); + EXPECT_TRUE(v_builder->Append("v_" + std::to_string(i)).ok()); + } + std::shared_ptr array; + EXPECT_TRUE(builder.Finish(&array).ok()); + return array; + } + + struct Row { + int64_t k; + std::string v; + uint64_t file_row; + }; + + // Drive the reader through NextBatchWithBitmap to EOF, decoding the full-schema output rows. + Result> Collect(LateMaterializingFileBatchReader* reader) { + std::vector rows; + while (true) { + PAIMON_ASSIGN_OR_RAISE(BatchReader::ReadBatchWithBitmap batch_with_bitmap, + reader->NextBatchWithBitmap()); + if (BatchReader::IsEofBatch(batch_with_bitmap)) { + break; + } + auto& [batch, bitmap] = batch_with_bitmap; + auto& [c_array, c_schema] = batch; + PAIMON_ASSIGN_OR_RAISE_FROM_ARROW(std::shared_ptr array, + arrow::ImportArray(c_array.get(), c_schema.get())); + auto struct_array = arrow::internal::checked_pointer_cast(array); + EXPECT_EQ(bitmap.Cardinality(), static_cast(struct_array->length())); + auto k_array = arrow::internal::checked_pointer_cast( + struct_array->GetFieldByName("k")); + if (!k_array) { + return Status::Invalid("output batch missing k column"); + } + // v is only present when it belongs to the read schema (payload projection). + auto v_array = arrow::internal::checked_pointer_cast( + struct_array->GetFieldByName("v")); + for (int64_t i = 0; i < struct_array->length(); ++i) { + PAIMON_ASSIGN_OR_RAISE(uint64_t file_row, + reader->GetPreviousBatchFileRowId(static_cast(i))); + rows.push_back(Row{k_array->Value(i), + v_array ? v_array->GetString(i) : std::string(), file_row}); + } + } + return rows; + } + + Status SetReadSchema(LateMaterializingFileBatchReader* reader, + const std::shared_ptr& schema, + const std::shared_ptr& predicate, + const std::optional& selection) { + ::ArrowSchema c_schema; + PAIMON_RETURN_NOT_OK_FROM_ARROW(arrow::ExportSchema(*schema, &c_schema)); + return reader->SetReadSchema(&c_schema, predicate, selection); + } + + // Collect all output rows as a single concatenated struct array (for schema/nested checks). + Result> CollectStruct(FileBatchReader* reader) { + arrow::ArrayVector chunks; + while (true) { + PAIMON_ASSIGN_OR_RAISE(BatchReader::ReadBatchWithBitmap batch_with_bitmap, + reader->NextBatchWithBitmap()); + if (BatchReader::IsEofBatch(batch_with_bitmap)) { + break; + } + PAIMON_ASSIGN_OR_RAISE( + BatchReader::ReadBatch batch, + ReaderUtils::ApplyBitmapToReadBatch(std::move(batch_with_bitmap), + arrow::default_memory_pool())); + auto& [c_array, c_schema] = batch; + PAIMON_ASSIGN_OR_RAISE_FROM_ARROW(std::shared_ptr array, + arrow::ImportArray(c_array.get(), c_schema.get())); + chunks.push_back(array); + } + if (chunks.empty()) { + return std::shared_ptr(); + } + PAIMON_ASSIGN_OR_RAISE_FROM_ARROW(std::shared_ptr combined, + arrow::Concatenate(chunks)); + return arrow::internal::checked_pointer_cast(combined); + } + + // Build a struct with 5 columns [a:int64, b:utf8, c:int64, d:utf8, e:int64], each carrying a + // distinct value pattern so any column reordering is detected. + std::shared_ptr BuildMultiFieldData(int32_t n) { + auto type = arrow::struct_({arrow::field("a", arrow::int64()), + arrow::field("b", arrow::utf8()), + arrow::field("c", arrow::int64()), + arrow::field("d", arrow::utf8()), + arrow::field("e", arrow::int64())}); + arrow::StructBuilder builder( + type, arrow::default_memory_pool(), + {std::make_shared(), std::make_shared(), + std::make_shared(), std::make_shared(), + std::make_shared()}); + auto* a = checked_cast(builder.field_builder(0)); + auto* b = checked_cast(builder.field_builder(1)); + auto* c = checked_cast(builder.field_builder(2)); + auto* d = checked_cast(builder.field_builder(3)); + auto* e = checked_cast(builder.field_builder(4)); + for (int32_t i = 0; i < n; ++i) { + EXPECT_TRUE(builder.Append().ok()); + EXPECT_TRUE(a->Append(i).ok()); + EXPECT_TRUE(b->Append("b_" + std::to_string(i)).ok()); + EXPECT_TRUE(c->Append(static_cast(i) * 100).ok()); + EXPECT_TRUE(d->Append("d_" + std::to_string(i)).ok()); + EXPECT_TRUE(e->Append(static_cast(i) * 10000).ok()); + } + std::shared_ptr array; + EXPECT_TRUE(builder.Finish(&array).ok()); + return array; + } + + // Build a struct with a nested payload column [k:int64, arr:list, tag:utf8]. + std::shared_ptr BuildNestedData(int32_t n) { + auto type = arrow::struct_({arrow::field("k", arrow::int64()), + arrow::field("arr", arrow::list(arrow::int64())), + arrow::field("tag", arrow::utf8())}); + auto arr_value_builder = std::make_shared(); + arrow::StructBuilder builder( + type, arrow::default_memory_pool(), + {std::make_shared(), + std::make_shared(arrow::default_memory_pool(), arr_value_builder), + std::make_shared()}); + auto* k = checked_cast(builder.field_builder(0)); + auto* arr = checked_cast(builder.field_builder(1)); + auto* arr_values = checked_cast(arr->value_builder()); + auto* tag = checked_cast(builder.field_builder(2)); + for (int32_t i = 0; i < n; ++i) { + EXPECT_TRUE(builder.Append().ok()); + EXPECT_TRUE(k->Append(i).ok()); + EXPECT_TRUE(arr->Append().ok()); + EXPECT_TRUE(arr_values->Append(i).ok()); + EXPECT_TRUE(arr_values->Append(i + 1).ok()); + EXPECT_TRUE(tag->Append("t_" + std::to_string(i)).ok()); + } + std::shared_ptr array; + EXPECT_TRUE(builder.Finish(&array).ok()); + return array; + } + + protected: + std::shared_ptr k_field_; + std::shared_ptr v_field_; + arrow::FieldVector full_fields_; + std::shared_ptr full_type_; +}; + +// No predicate: the reader must pass through the inner reader unchanged (all rows, all columns). +TEST_F(LateMaterializingFileBatchReaderTest, PassThroughWhenNoPredicate) { + auto data = BuildData({0, 1, 2, 3, 4}); + auto mock = std::make_unique(data, full_type_, /*batch_size=*/2); + ASSERT_OK_AND_ASSIGN(auto reader, LateMaterializingFileBatchReader::Create(std::move(mock))); + ASSERT_OK(SetReadSchema(reader.get(), arrow::schema(full_fields_), /*predicate=*/nullptr, + std::nullopt)); + + ASSERT_OK_AND_ASSIGN(std::vector rows, Collect(reader.get())); + ASSERT_EQ(rows.size(), 5u); + for (int64_t i = 0; i < 5; ++i) { + EXPECT_EQ(rows[i].k, i); + EXPECT_EQ(rows[i].v, "v_" + std::to_string(i)); + EXPECT_EQ(rows[i].file_row, static_cast(i)); + } +} + +// The predicate references every projected column, so the payload set is empty: no late +// materialization, plain pass-through. +TEST_F(LateMaterializingFileBatchReaderTest, PassThroughWhenPayloadEmpty) { + auto data = BuildData({0, 1, 2, 3, 4}); + auto mock = std::make_unique(data, full_type_, /*batch_size=*/2); + ASSERT_OK_AND_ASSIGN(auto reader, LateMaterializingFileBatchReader::Create(std::move(mock))); + // read schema is just {k}; the predicate on k covers all columns -> payload empty + auto predicate = PredicateBuilder::GreaterOrEqual(/*field_index=*/0, /*field_name=*/"k", + FieldType::BIGINT, Literal(0l)); + ASSERT_OK(SetReadSchema(reader.get(), arrow::schema({k_field_}), predicate, std::nullopt)); + + ASSERT_OK_AND_ASSIGN(std::vector rows, Collect(reader.get())); + ASSERT_EQ(rows.size(), 5u); +} + +// Contiguous matched subset spanning multiple batches. +TEST_F(LateMaterializingFileBatchReaderTest, ContiguousSubsetAcrossBatches) { + auto data = BuildData({0, 1, 2, 3, 4, 5, 6, 7, 8, 9}); + auto mock = std::make_unique(data, full_type_, /*batch_size=*/3); + ASSERT_OK_AND_ASSIGN(auto reader, LateMaterializingFileBatchReader::Create(std::move(mock))); + auto predicate = PredicateBuilder::GreaterThan(/*field_index=*/0, /*field_name=*/"k", + FieldType::BIGINT, Literal(4l)); + ASSERT_OK(SetReadSchema(reader.get(), arrow::schema(full_fields_), predicate, std::nullopt)); + + ASSERT_OK_AND_ASSIGN(std::vector rows, Collect(reader.get())); + ASSERT_EQ(rows.size(), 5u); // k = 5..9 + for (size_t idx = 0; idx < rows.size(); ++idx) { + int64_t expected = 5 + static_cast(idx); + EXPECT_EQ(rows[idx].k, expected); + EXPECT_EQ(rows[idx].v, "v_" + std::to_string(expected)); + EXPECT_EQ(rows[idx].file_row, static_cast(expected)); + } +} + +// Scattered (alternating) matched rows: predicate matches every other row. +TEST_F(LateMaterializingFileBatchReaderTest, ScatteredAlternatingMatch) { + // k = 0,1,0,1,... ; predicate k == 1 matches all odd file rows. + std::vector ks; + for (int i = 0; i < 12; ++i) { + ks.push_back(i % 2); + } + auto data = BuildData(ks); + auto mock = std::make_unique(data, full_type_, /*batch_size=*/3); + ASSERT_OK_AND_ASSIGN(auto reader, LateMaterializingFileBatchReader::Create(std::move(mock))); + auto predicate = PredicateBuilder::Equal(/*field_index=*/0, /*field_name=*/"k", + FieldType::BIGINT, Literal(1l)); + ASSERT_OK(SetReadSchema(reader.get(), arrow::schema(full_fields_), predicate, std::nullopt)); + + ASSERT_OK_AND_ASSIGN(std::vector rows, Collect(reader.get())); + ASSERT_EQ(rows.size(), 6u); // odd rows 1,3,5,7,9,11 + for (size_t idx = 0; idx < rows.size(); ++idx) { + uint64_t expected_row = 2 * idx + 1; + EXPECT_EQ(rows[idx].k, 1); + EXPECT_EQ(rows[idx].v, "v_" + std::to_string(expected_row)); + EXPECT_EQ(rows[idx].file_row, expected_row); + } +} + +// The selection bitmap further restricts the matched rows: matched must be a subset of selection. +TEST_F(LateMaterializingFileBatchReaderTest, MatchedIntersectsSelection) { + std::vector ks; + for (int i = 0; i < 12; ++i) { + ks.push_back(i % 2); + } + auto data = BuildData(ks); + auto mock = std::make_unique(data, full_type_, /*batch_size=*/4); + ASSERT_OK_AND_ASSIGN(auto reader, LateMaterializingFileBatchReader::Create(std::move(mock))); + auto predicate = PredicateBuilder::Equal(/*field_index=*/0, /*field_name=*/"k", + FieldType::BIGINT, Literal(1l)); + // predicate hits {1,3,5,7,9,11}; selection keeps only {1,5,9} + RoaringBitmap32 selection; + selection.Add(1); + selection.Add(5); + selection.Add(9); + ASSERT_OK(SetReadSchema(reader.get(), arrow::schema(full_fields_), predicate, + std::optional(selection))); + + ASSERT_OK_AND_ASSIGN(std::vector rows, Collect(reader.get())); + ASSERT_EQ(rows.size(), 3u); + EXPECT_EQ(rows[0].file_row, 1u); + EXPECT_EQ(rows[1].file_row, 5u); + EXPECT_EQ(rows[2].file_row, 9u); + EXPECT_EQ(rows[0].v, "v_1"); + EXPECT_EQ(rows[1].v, "v_5"); + EXPECT_EQ(rows[2].v, "v_9"); +} + +// No matched rows: the reader returns EOF immediately. +TEST_F(LateMaterializingFileBatchReaderTest, EmptyMatchReturnsEof) { + auto data = BuildData({0, 1, 2, 3, 4}); + auto mock = std::make_unique(data, full_type_, /*batch_size=*/2); + ASSERT_OK_AND_ASSIGN(auto reader, LateMaterializingFileBatchReader::Create(std::move(mock))); + auto predicate = PredicateBuilder::GreaterThan(/*field_index=*/0, /*field_name=*/"k", + FieldType::BIGINT, Literal(100l)); + ASSERT_OK(SetReadSchema(reader.get(), arrow::schema(full_fields_), predicate, std::nullopt)); + + ASSERT_OK_AND_ASSIGN(std::vector rows, Collect(reader.get())); + ASSERT_TRUE(rows.empty()); +} + +// SeekToRow during payload emission must re-align the probe cursor so probe/payload stay matched. +TEST_F(LateMaterializingFileBatchReaderTest, SeekToRowRealignsProbeCursor) { + auto data = BuildData({0, 1, 2, 3, 4, 5, 6, 7, 8, 9}); + auto mock = std::make_unique(data, full_type_, /*batch_size=*/4); + ASSERT_OK_AND_ASSIGN(auto reader, LateMaterializingFileBatchReader::Create(std::move(mock))); + auto predicate = PredicateBuilder::GreaterOrEqual(/*field_index=*/0, /*field_name=*/"k", + FieldType::BIGINT, Literal(5l)); + ASSERT_OK(SetReadSchema(reader.get(), arrow::schema(full_fields_), predicate, std::nullopt)); + + // First payload batch triggers the probe scan; matched rows are 5..9. + ASSERT_OK_AND_ASSIGN(BatchReader::ReadBatchWithBitmap first, reader->NextBatchWithBitmap()); + ASSERT_FALSE(BatchReader::IsEofBatch(first)); + ReaderUtils::ReleaseReadBatch(std::move(first.first)); + + // Seek forward to file row 8: subsequent output must be exactly rows 8 and 9, correctly paired. + ASSERT_OK(reader->SeekToRow(8)); + ASSERT_OK_AND_ASSIGN(std::vector rows, Collect(reader.get())); + ASSERT_EQ(rows.size(), 2u); + EXPECT_EQ(rows[0].k, 8); + EXPECT_EQ(rows[0].v, "v_8"); + EXPECT_EQ(rows[0].file_row, 8u); + EXPECT_EQ(rows[1].k, 9); + EXPECT_EQ(rows[1].v, "v_9"); + EXPECT_EQ(rows[1].file_row, 9u); +} + +// SetReadRanges must be cached and re-forwarded to the inner reader across the probe/payload +// schema switches (SetReadSchema resets the inner reader's ranges). +TEST_F(LateMaterializingFileBatchReaderTest, ReadRangesForwardedAcrossPhases) { + auto data = BuildData({0, 1, 2, 3, 4, 5, 6, 7}); + auto mock = std::make_unique(data, full_type_, /*batch_size=*/4); + auto* mock_ptr = mock.get(); + ASSERT_OK_AND_ASSIGN(auto reader, LateMaterializingFileBatchReader::Create(std::move(mock))); + auto predicate = PredicateBuilder::GreaterOrEqual(/*field_index=*/0, /*field_name=*/"k", + FieldType::BIGINT, Literal(2l)); + ASSERT_OK(SetReadSchema(reader.get(), arrow::schema(full_fields_), predicate, std::nullopt)); + + std::vector> ranges = {{0, 8}}; + ASSERT_OK(reader->SetReadRanges(ranges)); + + // Drive to EOF; this performs the probe pass and the payload schema switch. + ASSERT_OK_AND_ASSIGN(std::vector rows, Collect(reader.get())); + ASSERT_EQ(rows.size(), 6u); // k = 2..7 + + // The inner reader must have received the cached ranges again after the payload switch. + // SetReadSchema (invoked on the payload switch) clears the inner reader's ranges, so the + // cached ranges still being present at the end proves LM re-forwarded them after the switch. + ASSERT_EQ(mock_ptr->GetReadRanges(), ranges); +} + +// SetReadSchema is re-entrant: a second call with a different predicate resets probe state. +TEST_F(LateMaterializingFileBatchReaderTest, ReentrantSetReadSchema) { + auto data = BuildData({0, 1, 2, 3, 4, 5, 6, 7, 8, 9}); + auto mock = std::make_unique(data, full_type_, /*batch_size=*/3); + ASSERT_OK_AND_ASSIGN(auto reader, LateMaterializingFileBatchReader::Create(std::move(mock))); + + auto predicate1 = PredicateBuilder::GreaterThan(/*field_index=*/0, /*field_name=*/"k", + FieldType::BIGINT, Literal(7l)); + ASSERT_OK(SetReadSchema(reader.get(), arrow::schema(full_fields_), predicate1, std::nullopt)); + ASSERT_OK_AND_ASSIGN(std::vector rows1, Collect(reader.get())); + ASSERT_EQ(rows1.size(), 2u); // k = 8,9 + + auto predicate2 = PredicateBuilder::LessThan(/*field_index=*/0, /*field_name=*/"k", + FieldType::BIGINT, Literal(3l)); + ASSERT_OK(SetReadSchema(reader.get(), arrow::schema(full_fields_), predicate2, std::nullopt)); + ASSERT_OK_AND_ASSIGN(std::vector rows2, Collect(reader.get())); + ASSERT_EQ(rows2.size(), 3u); // k = 0,1,2 + for (size_t idx = 0; idx < rows2.size(); ++idx) { + EXPECT_EQ(rows2[idx].k, static_cast(idx)); + EXPECT_EQ(rows2[idx].file_row, static_cast(idx)); + } +} + +// Forwarded metadata accessors should reflect the inner reader. +TEST_F(LateMaterializingFileBatchReaderTest, ForwardsRowCountAndFileSchema) { + auto data = BuildData({0, 1, 2, 3}); + auto mock = std::make_unique(data, full_type_, /*batch_size=*/2); + ASSERT_OK_AND_ASSIGN(auto reader, LateMaterializingFileBatchReader::Create(std::move(mock))); + + ASSERT_OK_AND_ASSIGN(uint64_t num_rows, reader->GetNumberOfRows()); + EXPECT_EQ(num_rows, 4u); + ASSERT_OK_AND_ASSIGN(std::unique_ptr<::ArrowSchema> c_file_schema, reader->GetFileSchema()); + auto import_result = arrow::ImportType(c_file_schema.get()); + ASSERT_TRUE(import_result.ok()); + EXPECT_TRUE(import_result.ValueOrDie()->Equals(full_type_)); +} + +// With many columns and a predicate over two non-adjacent probe columns, the output must keep the +// full read-schema field order (and each probe/payload column's values must not be scrambled). +TEST_F(LateMaterializingFileBatchReaderTest, MultiFieldPreservesColumnOrder) { + auto data = BuildMultiFieldData(10); + auto type = data->type(); + auto mock = std::make_unique(data, type, /*batch_size=*/3); + ASSERT_OK_AND_ASSIGN(auto reader, LateMaterializingFileBatchReader::Create(std::move(mock))); + // probe columns = {a (idx0), c (idx2)}; payload columns = {b, d, e} + auto pred_a = PredicateBuilder::GreaterOrEqual(/*field_index=*/0, "a", FieldType::BIGINT, + Literal(3l)); + auto pred_c = PredicateBuilder::LessThan(/*field_index=*/2, "c", FieldType::BIGINT, + Literal(700l)); + ASSERT_OK_AND_ASSIGN(auto predicate, PredicateBuilder::And({pred_a, pred_c})); + ASSERT_OK(SetReadSchema(reader.get(), arrow::schema(type->fields()), predicate, std::nullopt)); + + ASSERT_OK_AND_ASSIGN(std::shared_ptr result, CollectStruct(reader.get())); + ASSERT_TRUE(result); + // a >= 3 and c(=i*100) < 700 -> i in {3,4,5,6} + ASSERT_EQ(result->length(), 4); + // output field order must equal the requested full schema order + ASSERT_EQ(result->num_fields(), 5); + auto out_type = arrow::internal::checked_pointer_cast(result->type()); + EXPECT_EQ(out_type->field(0)->name(), "a"); + EXPECT_EQ(out_type->field(1)->name(), "b"); + EXPECT_EQ(out_type->field(2)->name(), "c"); + EXPECT_EQ(out_type->field(3)->name(), "d"); + EXPECT_EQ(out_type->field(4)->name(), "e"); + auto a = arrow::internal::checked_pointer_cast(result->GetFieldByName("a")); + auto b = arrow::internal::checked_pointer_cast(result->GetFieldByName("b")); + auto c = arrow::internal::checked_pointer_cast(result->GetFieldByName("c")); + auto d = arrow::internal::checked_pointer_cast(result->GetFieldByName("d")); + auto e = arrow::internal::checked_pointer_cast(result->GetFieldByName("e")); + const std::vector expected = {3, 4, 5, 6}; + for (size_t j = 0; j < expected.size(); ++j) { + int64_t i = expected[j]; + EXPECT_EQ(a->Value(j), i); + EXPECT_EQ(b->GetString(j), "b_" + std::to_string(i)); + EXPECT_EQ(c->Value(j), i * 100); + EXPECT_EQ(d->GetString(j), "d_" + std::to_string(i)); + EXPECT_EQ(e->Value(j), i * 10000); + } +} + +// A nested (list) payload column must round-trip unchanged for the matched rows. +TEST_F(LateMaterializingFileBatchReaderTest, NestedPayloadColumn) { + auto data = BuildNestedData(8); + auto type = data->type(); + auto mock = std::make_unique(data, type, /*batch_size=*/3); + ASSERT_OK_AND_ASSIGN(auto reader, LateMaterializingFileBatchReader::Create(std::move(mock))); + // probe = {k}; payload = {arr (list), tag} + auto predicate = PredicateBuilder::GreaterOrEqual(/*field_index=*/0, "k", FieldType::BIGINT, + Literal(5l)); + ASSERT_OK(SetReadSchema(reader.get(), arrow::schema(type->fields()), predicate, std::nullopt)); + + ASSERT_OK_AND_ASSIGN(std::shared_ptr result, CollectStruct(reader.get())); + ASSERT_TRUE(result); + ASSERT_EQ(result->length(), 3); // k = 5,6,7 + auto k = arrow::internal::checked_pointer_cast(result->GetFieldByName("k")); + auto arr = arrow::internal::checked_pointer_cast(result->GetFieldByName("arr")); + auto tag = + arrow::internal::checked_pointer_cast(result->GetFieldByName("tag")); + ASSERT_TRUE(k && arr && tag); + for (int64_t j = 0; j < result->length(); ++j) { + int64_t i = 5 + j; + EXPECT_EQ(k->Value(j), i); + EXPECT_EQ(tag->GetString(j), "t_" + std::to_string(i)); + auto sub = arrow::internal::checked_pointer_cast(arr->value_slice(j)); + ASSERT_EQ(sub->length(), 2); + EXPECT_EQ(sub->Value(0), i); + EXPECT_EQ(sub->Value(1), i + 1); + } +} + +// The late-materialization reader must work correctly as an inner reader driven by +// PrefetchFileBatchReaderImpl (schema broadcast, range dispatch, seek, row-id tracking). +TEST_F(LateMaterializingFileBatchReaderTest, WorksAsInnerOfPrefetchReader) { + auto data = BuildData({0, 1, 2, 3, 4, 5, 6, 7, 8, 9}); + LmReaderBuilder builder(data, full_type_, /*batch_size=*/3); + auto mock_fs = std::make_shared(); + ASSERT_OK_AND_ASSIGN(std::shared_ptr executor, CreateDefaultExecutor(2)); + ASSERT_OK_AND_ASSIGN( + std::unique_ptr impl, + PrefetchFileBatchReaderImpl::Create( + /*data_file_path=*/"", /*data_file_size=*/0, &builder, mock_fs, + /*prefetch_max_parallel_num=*/1, /*batch_size=*/3, /*prefetch_batch_count=*/2, + /*enable_adaptive_prefetch_strategy=*/false, executor, + /*initialize_read_ranges=*/false, /*read_ahead_cache_enabled=*/false, CacheConfig(), + GetDefaultPool())); + auto predicate = PredicateBuilder::GreaterOrEqual(/*field_index=*/0, "k", FieldType::BIGINT, + Literal(4l)); + ::ArrowSchema c_schema; + ASSERT_TRUE(arrow::ExportSchema(*arrow::schema(full_fields_), &c_schema).ok()); + ASSERT_OK(impl->SetReadSchema(&c_schema, predicate, std::nullopt)); + + ASSERT_OK_AND_ASSIGN(std::shared_ptr result, CollectStruct(impl.get())); + ASSERT_TRUE(result); + ASSERT_EQ(result->length(), 6); // k = 4..9 + auto k = arrow::internal::checked_pointer_cast(result->GetFieldByName("k")); + auto v = arrow::internal::checked_pointer_cast(result->GetFieldByName("v")); + ASSERT_TRUE(k && v); + for (int64_t j = 0; j < result->length(); ++j) { + EXPECT_EQ(k->Value(j), 4 + j); + EXPECT_EQ(v->GetString(j), "v_" + std::to_string(4 + j)); + } + impl->Close(); +} + +// Re-setting the read schema on the prefetch impl (which re-broadcasts to the inner LM readers and +// re-plans ranges) must reset the probe state and produce correct results for the new predicate. +TEST_F(LateMaterializingFileBatchReaderTest, PrefetchInnerReentrantSetReadSchema) { + auto data = BuildData({0, 1, 2, 3, 4, 5, 6, 7, 8, 9}); + LmReaderBuilder builder(data, full_type_, /*batch_size=*/3); + auto mock_fs = std::make_shared(); + ASSERT_OK_AND_ASSIGN(std::shared_ptr executor, CreateDefaultExecutor(2)); + ASSERT_OK_AND_ASSIGN( + std::unique_ptr impl, + PrefetchFileBatchReaderImpl::Create( + /*data_file_path=*/"", /*data_file_size=*/0, &builder, mock_fs, + /*prefetch_max_parallel_num=*/1, /*batch_size=*/3, /*prefetch_batch_count=*/2, + /*enable_adaptive_prefetch_strategy=*/false, executor, + /*initialize_read_ranges=*/false, /*read_ahead_cache_enabled=*/false, CacheConfig(), + GetDefaultPool())); + + auto full_schema = arrow::schema(full_fields_); + auto predicate1 = + PredicateBuilder::GreaterThan(/*field_index=*/0, "k", FieldType::BIGINT, Literal(6l)); + ::ArrowSchema c_schema1; + ASSERT_TRUE(arrow::ExportSchema(*full_schema, &c_schema1).ok()); + ASSERT_OK(impl->SetReadSchema(&c_schema1, predicate1, std::nullopt)); + ASSERT_OK_AND_ASSIGN(std::shared_ptr result1, CollectStruct(impl.get())); + ASSERT_TRUE(result1); + ASSERT_EQ(result1->length(), 3); // k = 7,8,9 + auto k1 = arrow::internal::checked_pointer_cast(result1->GetFieldByName("k")); + for (int64_t j = 0; j < result1->length(); ++j) { + EXPECT_EQ(k1->Value(j), 7 + j); + } + + auto predicate2 = + PredicateBuilder::LessThan(/*field_index=*/0, "k", FieldType::BIGINT, Literal(3l)); + ::ArrowSchema c_schema2; + ASSERT_TRUE(arrow::ExportSchema(*full_schema, &c_schema2).ok()); + ASSERT_OK(impl->SetReadSchema(&c_schema2, predicate2, std::nullopt)); + ASSERT_OK_AND_ASSIGN(std::shared_ptr result2, CollectStruct(impl.get())); + ASSERT_TRUE(result2); + ASSERT_EQ(result2->length(), 3); // k = 0,1,2 + auto k2 = arrow::internal::checked_pointer_cast(result2->GetFieldByName("k")); + auto v2 = arrow::internal::checked_pointer_cast(result2->GetFieldByName("v")); + for (int64_t j = 0; j < result2->length(); ++j) { + EXPECT_EQ(k2->Value(j), j); + EXPECT_EQ(v2->GetString(j), "v_" + std::to_string(j)); + } + impl->Close(); +} + +// With multiple parallel inner readers and per-batch ranges, the prefetch impl dispatches disjoint +// ranges to each LM reader and drives them via EnsureReaderPosition/SeekToRow. The merged output +// must still be exactly the matched rows in ascending file order. +TEST_F(LateMaterializingFileBatchReaderTest, PrefetchInnerParallelReadersWithSeek) { + std::vector ks; + for (int i = 0; i < 20; ++i) { + ks.push_back(i); + } + auto data = BuildData(ks); + // Per-batch ranges (the mock's default) let the impl split work across the parallel readers, + // and each range-honoring reader only reads its assigned slice. + LmReaderBuilder builder(data, full_type_, /*batch_size=*/3); + auto mock_fs = std::make_shared(); + ASSERT_OK_AND_ASSIGN(std::shared_ptr executor, CreateDefaultExecutor(3)); + ASSERT_OK_AND_ASSIGN( + std::unique_ptr impl, + PrefetchFileBatchReaderImpl::Create( + /*data_file_path=*/"", /*data_file_size=*/0, &builder, mock_fs, + /*prefetch_max_parallel_num=*/3, /*batch_size=*/3, /*prefetch_batch_count=*/6, + /*enable_adaptive_prefetch_strategy=*/false, executor, + /*initialize_read_ranges=*/false, /*read_ahead_cache_enabled=*/false, CacheConfig(), + GetDefaultPool())); + auto predicate = PredicateBuilder::GreaterOrEqual(/*field_index=*/0, "k", FieldType::BIGINT, + Literal(5l)); + ::ArrowSchema c_schema; + ASSERT_TRUE(arrow::ExportSchema(*arrow::schema(full_fields_), &c_schema).ok()); + ASSERT_OK(impl->SetReadSchema(&c_schema, predicate, std::nullopt)); + + ASSERT_OK_AND_ASSIGN(std::shared_ptr result, CollectStruct(impl.get())); + ASSERT_TRUE(result); + ASSERT_EQ(result->length(), 15); // k = 5..19 + auto k = arrow::internal::checked_pointer_cast(result->GetFieldByName("k")); + auto v = arrow::internal::checked_pointer_cast(result->GetFieldByName("v")); + ASSERT_TRUE(k && v); + for (int64_t j = 0; j < result->length(); ++j) { + EXPECT_EQ(k->Value(j), 5 + j); + EXPECT_EQ(v->GetString(j), "v_" + std::to_string(5 + j)); + } + impl->Close(); +} + +} // namespace paimon::test diff --git a/src/paimon/testing/mock/mock_file_batch_reader.h b/src/paimon/testing/mock/mock_file_batch_reader.h index f05a2347b..7253b8578 100644 --- a/src/paimon/testing/mock/mock_file_batch_reader.h +++ b/src/paimon/testing/mock/mock_file_batch_reader.h @@ -29,6 +29,7 @@ #include "arrow/api.h" #include "arrow/c/bridge.h" +#include "arrow/util/checked_cast.h" #include "paimon/common/metrics/metrics_impl.h" #include "paimon/common/reader/reader_utils.h" #include "paimon/common/utils/arrow/status_utils.h" @@ -181,7 +182,7 @@ class MockFileBatchReader : public PrefetchFileBatchReader { return false; } - private: + protected: static uint64_t ToReaderRowNumber(int32_t row_number) { if (row_number < 0) { return std::numeric_limits::max(); @@ -203,4 +204,83 @@ class MockFileBatchReader : public PrefetchFileBatchReader { std::mt19937 random_engine_{std::random_device{}()}; // NOLINT(whitespace/braces) }; +// A projecting, range-honoring variant of MockFileBatchReader for exercising readers that split +// the read schema into probe/payload passes (e.g. LateMaterializingFileBatchReader). Compared to +// MockFileBatchReader it: +// * projects each batch down to the columns requested via SetReadSchema (matched by name); +// * honors the FileBatchReader contract that SetReadSchema restarts reading from the first row +// and drops the previously assigned read ranges; +// * restricts NextBatch output to the assigned read ranges (skipping gaps between them), so it +// behaves like a real format reader when PrefetchFileBatchReaderImpl dispatches disjoint +// ranges to parallel readers. An empty range set means "no restriction" (read the whole file). +class ProjectingMockFileBatchReader : public MockFileBatchReader { + public: + using MockFileBatchReader::MockFileBatchReader; + + Status SetReadSchema(::ArrowSchema* read_schema, const std::shared_ptr& /*predicate*/, + const std::optional& /*selection_bitmap*/) override { + PAIMON_ASSIGN_OR_RAISE_FROM_ARROW(proj_schema_, arrow::ImportSchema(read_schema)); + current_pos_ = 0; + previous_batch_first_row_num_ = std::numeric_limits::max(); + read_ranges_.clear(); + return Status::OK(); + } + + Result NextBatchWithBitmap() override { + int64_t begin = current_pos_; + int64_t range_end = read_end_pos_; + if (!read_ranges_.empty()) { + // read_ranges_ are ascending, half-open [first, second); find the first still-unread + // range and skip any gap before it. + const std::pair* selected = nullptr; + for (const auto& range : read_ranges_) { + if (static_cast(range.second) > begin) { + selected = ⦥ + break; + } + } + if (selected == nullptr) { + previous_batch_first_row_num_ = static_cast(begin); + return BatchReader::MakeEofBatchWithBitmap(); + } + begin = std::max(begin, static_cast(selected->first)); + range_end = static_cast(selected->second); + } + if (begin >= read_end_pos_) { + previous_batch_first_row_num_ = static_cast(begin); + return BatchReader::MakeEofBatchWithBitmap(); + } + int64_t end = + std::min({static_cast(read_end_pos_), range_end, begin + batch_size_}); + auto full = arrow::internal::checked_pointer_cast( + data_->Slice(begin, end - begin)); + std::shared_ptr out = full; + if (proj_schema_) { + arrow::ArrayVector children; + arrow::FieldVector fields; + for (const auto& field : proj_schema_->fields()) { + std::shared_ptr col = full->GetFieldByName(field->name()); + if (!col) { + return Status::Invalid("projecting mock: unknown field " + field->name()); + } + children.push_back(col); + fields.push_back(field); + } + PAIMON_ASSIGN_OR_RAISE_FROM_ARROW(out, arrow::StructArray::Make(children, fields)); + } + RoaringBitmap32 bitmap; + bitmap.AddRange(0, static_cast(end - begin)); + previous_batch_first_row_num_ = static_cast(begin); + current_pos_ = static_cast(end); + auto c_array = std::make_unique<::ArrowArray>(); + auto c_schema = std::make_unique<::ArrowSchema>(); + PAIMON_RETURN_NOT_OK_FROM_ARROW(arrow::ExportArray(*out, c_array.get(), c_schema.get())); + return std::make_pair(std::make_pair(std::move(c_array), std::move(c_schema)), + std::move(bitmap)); + } + + private: + std::shared_ptr proj_schema_; +}; + } // namespace paimon::test From 4b0a106bc462de0535ccd5678084c063529743bb Mon Sep 17 00:00:00 2001 From: zhouhongfeng Date: Mon, 24 Aug 2026 10:31:56 +0800 Subject: [PATCH 06/34] test: add projection feature for MockFileBatchReader --- ...e_materializing_file_batch_reader_test.cpp | 26 +-- .../testing/mock/mock_file_batch_reader.h | 162 ++++++++---------- 2 files changed, 87 insertions(+), 101 deletions(-) diff --git a/src/paimon/common/reader/late_materializing_file_batch_reader_test.cpp b/src/paimon/common/reader/late_materializing_file_batch_reader_test.cpp index 84b52b0c3..92c11e9fc 100644 --- a/src/paimon/common/reader/late_materializing_file_batch_reader_test.cpp +++ b/src/paimon/common/reader/late_materializing_file_batch_reader_test.cpp @@ -67,7 +67,7 @@ class LmReaderBuilder : public ReaderBuilder { Result> Build( const std::shared_ptr& /*stream*/) const override { - auto mock = std::make_unique(data_, type_, batch_size_); + auto mock = std::make_unique(data_, type_, batch_size_); PAIMON_ASSIGN_OR_RAISE(std::unique_ptr reader, LateMaterializingFileBatchReader::Create(std::move(mock))); return std::unique_ptr(std::move(reader)); @@ -248,7 +248,7 @@ class LateMaterializingFileBatchReaderTest : public ::testing::Test { // No predicate: the reader must pass through the inner reader unchanged (all rows, all columns). TEST_F(LateMaterializingFileBatchReaderTest, PassThroughWhenNoPredicate) { auto data = BuildData({0, 1, 2, 3, 4}); - auto mock = std::make_unique(data, full_type_, /*batch_size=*/2); + auto mock = std::make_unique(data, full_type_, /*batch_size=*/2); ASSERT_OK_AND_ASSIGN(auto reader, LateMaterializingFileBatchReader::Create(std::move(mock))); ASSERT_OK(SetReadSchema(reader.get(), arrow::schema(full_fields_), /*predicate=*/nullptr, std::nullopt)); @@ -266,7 +266,7 @@ TEST_F(LateMaterializingFileBatchReaderTest, PassThroughWhenNoPredicate) { // materialization, plain pass-through. TEST_F(LateMaterializingFileBatchReaderTest, PassThroughWhenPayloadEmpty) { auto data = BuildData({0, 1, 2, 3, 4}); - auto mock = std::make_unique(data, full_type_, /*batch_size=*/2); + auto mock = std::make_unique(data, full_type_, /*batch_size=*/2); ASSERT_OK_AND_ASSIGN(auto reader, LateMaterializingFileBatchReader::Create(std::move(mock))); // read schema is just {k}; the predicate on k covers all columns -> payload empty auto predicate = PredicateBuilder::GreaterOrEqual(/*field_index=*/0, /*field_name=*/"k", @@ -280,7 +280,7 @@ TEST_F(LateMaterializingFileBatchReaderTest, PassThroughWhenPayloadEmpty) { // Contiguous matched subset spanning multiple batches. TEST_F(LateMaterializingFileBatchReaderTest, ContiguousSubsetAcrossBatches) { auto data = BuildData({0, 1, 2, 3, 4, 5, 6, 7, 8, 9}); - auto mock = std::make_unique(data, full_type_, /*batch_size=*/3); + auto mock = std::make_unique(data, full_type_, /*batch_size=*/3); ASSERT_OK_AND_ASSIGN(auto reader, LateMaterializingFileBatchReader::Create(std::move(mock))); auto predicate = PredicateBuilder::GreaterThan(/*field_index=*/0, /*field_name=*/"k", FieldType::BIGINT, Literal(4l)); @@ -304,7 +304,7 @@ TEST_F(LateMaterializingFileBatchReaderTest, ScatteredAlternatingMatch) { ks.push_back(i % 2); } auto data = BuildData(ks); - auto mock = std::make_unique(data, full_type_, /*batch_size=*/3); + auto mock = std::make_unique(data, full_type_, /*batch_size=*/3); ASSERT_OK_AND_ASSIGN(auto reader, LateMaterializingFileBatchReader::Create(std::move(mock))); auto predicate = PredicateBuilder::Equal(/*field_index=*/0, /*field_name=*/"k", FieldType::BIGINT, Literal(1l)); @@ -327,7 +327,7 @@ TEST_F(LateMaterializingFileBatchReaderTest, MatchedIntersectsSelection) { ks.push_back(i % 2); } auto data = BuildData(ks); - auto mock = std::make_unique(data, full_type_, /*batch_size=*/4); + auto mock = std::make_unique(data, full_type_, /*batch_size=*/4); ASSERT_OK_AND_ASSIGN(auto reader, LateMaterializingFileBatchReader::Create(std::move(mock))); auto predicate = PredicateBuilder::Equal(/*field_index=*/0, /*field_name=*/"k", FieldType::BIGINT, Literal(1l)); @@ -352,7 +352,7 @@ TEST_F(LateMaterializingFileBatchReaderTest, MatchedIntersectsSelection) { // No matched rows: the reader returns EOF immediately. TEST_F(LateMaterializingFileBatchReaderTest, EmptyMatchReturnsEof) { auto data = BuildData({0, 1, 2, 3, 4}); - auto mock = std::make_unique(data, full_type_, /*batch_size=*/2); + auto mock = std::make_unique(data, full_type_, /*batch_size=*/2); ASSERT_OK_AND_ASSIGN(auto reader, LateMaterializingFileBatchReader::Create(std::move(mock))); auto predicate = PredicateBuilder::GreaterThan(/*field_index=*/0, /*field_name=*/"k", FieldType::BIGINT, Literal(100l)); @@ -365,7 +365,7 @@ TEST_F(LateMaterializingFileBatchReaderTest, EmptyMatchReturnsEof) { // SeekToRow during payload emission must re-align the probe cursor so probe/payload stay matched. TEST_F(LateMaterializingFileBatchReaderTest, SeekToRowRealignsProbeCursor) { auto data = BuildData({0, 1, 2, 3, 4, 5, 6, 7, 8, 9}); - auto mock = std::make_unique(data, full_type_, /*batch_size=*/4); + auto mock = std::make_unique(data, full_type_, /*batch_size=*/4); ASSERT_OK_AND_ASSIGN(auto reader, LateMaterializingFileBatchReader::Create(std::move(mock))); auto predicate = PredicateBuilder::GreaterOrEqual(/*field_index=*/0, /*field_name=*/"k", FieldType::BIGINT, Literal(5l)); @@ -392,7 +392,7 @@ TEST_F(LateMaterializingFileBatchReaderTest, SeekToRowRealignsProbeCursor) { // schema switches (SetReadSchema resets the inner reader's ranges). TEST_F(LateMaterializingFileBatchReaderTest, ReadRangesForwardedAcrossPhases) { auto data = BuildData({0, 1, 2, 3, 4, 5, 6, 7}); - auto mock = std::make_unique(data, full_type_, /*batch_size=*/4); + auto mock = std::make_unique(data, full_type_, /*batch_size=*/4); auto* mock_ptr = mock.get(); ASSERT_OK_AND_ASSIGN(auto reader, LateMaterializingFileBatchReader::Create(std::move(mock))); auto predicate = PredicateBuilder::GreaterOrEqual(/*field_index=*/0, /*field_name=*/"k", @@ -415,7 +415,7 @@ TEST_F(LateMaterializingFileBatchReaderTest, ReadRangesForwardedAcrossPhases) { // SetReadSchema is re-entrant: a second call with a different predicate resets probe state. TEST_F(LateMaterializingFileBatchReaderTest, ReentrantSetReadSchema) { auto data = BuildData({0, 1, 2, 3, 4, 5, 6, 7, 8, 9}); - auto mock = std::make_unique(data, full_type_, /*batch_size=*/3); + auto mock = std::make_unique(data, full_type_, /*batch_size=*/3); ASSERT_OK_AND_ASSIGN(auto reader, LateMaterializingFileBatchReader::Create(std::move(mock))); auto predicate1 = PredicateBuilder::GreaterThan(/*field_index=*/0, /*field_name=*/"k", @@ -438,7 +438,7 @@ TEST_F(LateMaterializingFileBatchReaderTest, ReentrantSetReadSchema) { // Forwarded metadata accessors should reflect the inner reader. TEST_F(LateMaterializingFileBatchReaderTest, ForwardsRowCountAndFileSchema) { auto data = BuildData({0, 1, 2, 3}); - auto mock = std::make_unique(data, full_type_, /*batch_size=*/2); + auto mock = std::make_unique(data, full_type_, /*batch_size=*/2); ASSERT_OK_AND_ASSIGN(auto reader, LateMaterializingFileBatchReader::Create(std::move(mock))); ASSERT_OK_AND_ASSIGN(uint64_t num_rows, reader->GetNumberOfRows()); @@ -454,7 +454,7 @@ TEST_F(LateMaterializingFileBatchReaderTest, ForwardsRowCountAndFileSchema) { TEST_F(LateMaterializingFileBatchReaderTest, MultiFieldPreservesColumnOrder) { auto data = BuildMultiFieldData(10); auto type = data->type(); - auto mock = std::make_unique(data, type, /*batch_size=*/3); + auto mock = std::make_unique(data, type, /*batch_size=*/3); ASSERT_OK_AND_ASSIGN(auto reader, LateMaterializingFileBatchReader::Create(std::move(mock))); // probe columns = {a (idx0), c (idx2)}; payload columns = {b, d, e} auto pred_a = PredicateBuilder::GreaterOrEqual(/*field_index=*/0, "a", FieldType::BIGINT, @@ -496,7 +496,7 @@ TEST_F(LateMaterializingFileBatchReaderTest, MultiFieldPreservesColumnOrder) { TEST_F(LateMaterializingFileBatchReaderTest, NestedPayloadColumn) { auto data = BuildNestedData(8); auto type = data->type(); - auto mock = std::make_unique(data, type, /*batch_size=*/3); + auto mock = std::make_unique(data, type, /*batch_size=*/3); ASSERT_OK_AND_ASSIGN(auto reader, LateMaterializingFileBatchReader::Create(std::move(mock))); // probe = {k}; payload = {arr (list), tag} auto predicate = PredicateBuilder::GreaterOrEqual(/*field_index=*/0, "k", FieldType::BIGINT, diff --git a/src/paimon/testing/mock/mock_file_batch_reader.h b/src/paimon/testing/mock/mock_file_batch_reader.h index 7253b8578..fb3cc3ada 100644 --- a/src/paimon/testing/mock/mock_file_batch_reader.h +++ b/src/paimon/testing/mock/mock_file_batch_reader.h @@ -79,11 +79,15 @@ class MockFileBatchReader : public PrefetchFileBatchReader { Status SetReadSchema(::ArrowSchema* read_schema, const std::shared_ptr& predicate, const std::optional& selection_bitmap) override { - // Noted that SetReadSchema only change inner read_schema_, but take no effective on - // NextBatch PAIMON_ASSIGN_OR_RAISE_FROM_ARROW(std::shared_ptr arrow_schema, arrow::ImportSchema(read_schema)); read_schema_ = arrow_schema; + // A real FileBatchReader restarts from the first row and drops its assigned read ranges + // when the read schema is (re)set. Readers that switch schemas mid-file, such as the + // late-materialization reader moving from its probe pass to its payload pass, rely on it. + current_pos_ = 0; + previous_batch_first_row_num_ = std::numeric_limits::max(); + read_ranges_.clear(); return Status::OK(); } @@ -120,8 +124,29 @@ class MockFileBatchReader : public PrefetchFileBatchReader { Result NextBatchWithBitmap() override { while (true) { PAIMON_RETURN_NOT_OK(next_batch_status_); - if (current_pos_ >= read_end_pos_) { - previous_batch_first_row_num_ = current_pos_; + int32_t begin_pos = current_pos_; + int32_t range_end_pos = read_end_pos_; + if (!read_ranges_.empty()) { + // Reading is restricted to the assigned ranges (ascending and half-open), like a + // real format reader, so that a prefetch reader may dispatch disjoint ranges to + // parallel readers. An empty range set means the whole file may be read. + const std::pair* selected = nullptr; + for (const auto& range : read_ranges_) { + if (static_cast(range.second) > begin_pos) { + selected = ⦥ + break; + } + } + if (selected == nullptr) { + previous_batch_first_row_num_ = ToReaderRowNumber(begin_pos); + return BatchReader::MakeEofBatchWithBitmap(); + } + // Skip the gap in front of the first range that has not been read yet. + begin_pos = std::max(begin_pos, static_cast(selected->first)); + range_end_pos = std::min(range_end_pos, static_cast(selected->second)); + } + if (begin_pos >= read_end_pos_) { + previous_batch_first_row_num_ = ToReaderRowNumber(begin_pos); return BatchReader::MakeEofBatchWithBitmap(); } int32_t actual_batch_size = batch_size_; @@ -129,21 +154,23 @@ class MockFileBatchReader : public PrefetchFileBatchReader { std::uniform_int_distribution distribution(1, batch_size_); actual_batch_size = distribution(random_engine_); } - int32_t batch_end_pos = std::min(read_end_pos_, current_pos_ + actual_batch_size); - auto slice = data_->Slice(current_pos_, batch_end_pos - current_pos_); + int32_t batch_end_pos = + std::min({read_end_pos_, range_end_pos, begin_pos + actual_batch_size}); + auto slice = data_->Slice(begin_pos, batch_end_pos - begin_pos); PAIMON_ASSIGN_OR_RAISE_FROM_ARROW( std::shared_ptr concat_slice, arrow::Concatenate({slice}, arrow::default_memory_pool())); RoaringBitmap32 bitmap; - for (auto iter = bitmap_.EqualOrLarger(current_pos_); + for (auto iter = bitmap_.EqualOrLarger(begin_pos); iter != bitmap_.End() && *iter < batch_end_pos; ++iter) { - bitmap.Add(*iter - current_pos_); + bitmap.Add(*iter - begin_pos); } - previous_batch_first_row_num_ = current_pos_; + previous_batch_first_row_num_ = ToReaderRowNumber(begin_pos); current_pos_ = batch_end_pos; if (bitmap.IsEmpty()) { continue; } + PAIMON_ASSIGN_OR_RAISE(concat_slice, ProjectBatch(concat_slice)); std::unique_ptr c_array = std::make_unique(); std::unique_ptr c_schema = std::make_unique(); PAIMON_RETURN_NOT_OK_FROM_ARROW( @@ -190,6 +217,44 @@ class MockFileBatchReader : public PrefetchFileBatchReader { return static_cast(row_number); } + /// Pick the columns requested by `read_schema_` out of `batch`, in the requested order. + /// + /// `batch` is returned as is unless the requested schema is a genuine re-selection of the + /// columns this file has. Requesting a field the file does not have means the read schema is a + /// logical view over some other physical layout, as the shredding and the row tracking readers + /// do, and those map the raw batch themselves. + /// `batch` is expected to have a zero offset, so its validity buffer can be reused as is. + Result> ProjectBatch( + const std::shared_ptr& batch) const { + auto struct_batch = std::dynamic_pointer_cast(batch); + if (struct_batch == nullptr) { + return batch; + } + arrow::ArrayVector children; + arrow::FieldVector fields; + for (const auto& field : read_schema_->fields()) { + std::shared_ptr column = struct_batch->GetFieldByName(field->name()); + if (column == nullptr) { + return batch; + } + children.push_back(column); + fields.push_back(field); + } + const arrow::FieldVector& batch_fields = struct_batch->type()->fields(); + bool keeps_every_column = fields.size() == batch_fields.size(); + for (size_t i = 0; keeps_every_column && i < fields.size(); i++) { + keeps_every_column = fields[i]->name() == batch_fields[i]->name(); + } + if (keeps_every_column) { + return batch; + } + PAIMON_ASSIGN_OR_RAISE_FROM_ARROW( + std::shared_ptr projected, + arrow::StructArray::Make(children, fields, struct_batch->null_bitmap(), + struct_batch->null_count())); + return projected; + } + std::shared_ptr data_; std::shared_ptr file_schema_; std::shared_ptr read_schema_; @@ -204,83 +269,4 @@ class MockFileBatchReader : public PrefetchFileBatchReader { std::mt19937 random_engine_{std::random_device{}()}; // NOLINT(whitespace/braces) }; -// A projecting, range-honoring variant of MockFileBatchReader for exercising readers that split -// the read schema into probe/payload passes (e.g. LateMaterializingFileBatchReader). Compared to -// MockFileBatchReader it: -// * projects each batch down to the columns requested via SetReadSchema (matched by name); -// * honors the FileBatchReader contract that SetReadSchema restarts reading from the first row -// and drops the previously assigned read ranges; -// * restricts NextBatch output to the assigned read ranges (skipping gaps between them), so it -// behaves like a real format reader when PrefetchFileBatchReaderImpl dispatches disjoint -// ranges to parallel readers. An empty range set means "no restriction" (read the whole file). -class ProjectingMockFileBatchReader : public MockFileBatchReader { - public: - using MockFileBatchReader::MockFileBatchReader; - - Status SetReadSchema(::ArrowSchema* read_schema, const std::shared_ptr& /*predicate*/, - const std::optional& /*selection_bitmap*/) override { - PAIMON_ASSIGN_OR_RAISE_FROM_ARROW(proj_schema_, arrow::ImportSchema(read_schema)); - current_pos_ = 0; - previous_batch_first_row_num_ = std::numeric_limits::max(); - read_ranges_.clear(); - return Status::OK(); - } - - Result NextBatchWithBitmap() override { - int64_t begin = current_pos_; - int64_t range_end = read_end_pos_; - if (!read_ranges_.empty()) { - // read_ranges_ are ascending, half-open [first, second); find the first still-unread - // range and skip any gap before it. - const std::pair* selected = nullptr; - for (const auto& range : read_ranges_) { - if (static_cast(range.second) > begin) { - selected = ⦥ - break; - } - } - if (selected == nullptr) { - previous_batch_first_row_num_ = static_cast(begin); - return BatchReader::MakeEofBatchWithBitmap(); - } - begin = std::max(begin, static_cast(selected->first)); - range_end = static_cast(selected->second); - } - if (begin >= read_end_pos_) { - previous_batch_first_row_num_ = static_cast(begin); - return BatchReader::MakeEofBatchWithBitmap(); - } - int64_t end = - std::min({static_cast(read_end_pos_), range_end, begin + batch_size_}); - auto full = arrow::internal::checked_pointer_cast( - data_->Slice(begin, end - begin)); - std::shared_ptr out = full; - if (proj_schema_) { - arrow::ArrayVector children; - arrow::FieldVector fields; - for (const auto& field : proj_schema_->fields()) { - std::shared_ptr col = full->GetFieldByName(field->name()); - if (!col) { - return Status::Invalid("projecting mock: unknown field " + field->name()); - } - children.push_back(col); - fields.push_back(field); - } - PAIMON_ASSIGN_OR_RAISE_FROM_ARROW(out, arrow::StructArray::Make(children, fields)); - } - RoaringBitmap32 bitmap; - bitmap.AddRange(0, static_cast(end - begin)); - previous_batch_first_row_num_ = static_cast(begin); - current_pos_ = static_cast(end); - auto c_array = std::make_unique<::ArrowArray>(); - auto c_schema = std::make_unique<::ArrowSchema>(); - PAIMON_RETURN_NOT_OK_FROM_ARROW(arrow::ExportArray(*out, c_array.get(), c_schema.get())); - return std::make_pair(std::make_pair(std::move(c_array), std::move(c_schema)), - std::move(bitmap)); - } - - private: - std::shared_ptr proj_schema_; -}; - } // namespace paimon::test From f7fa51348690b019e8ee2455db3afb123be19f4b Mon Sep 17 00:00:00 2001 From: zhouhongfeng Date: Mon, 24 Aug 2026 11:12:34 +0800 Subject: [PATCH 07/34] fix: set read ranges only when read range are not empty --- .../late_materializing_file_batch_reader.cpp | 15 +++++++++++---- .../reader/late_materializing_file_batch_reader.h | 1 + 2 files changed, 12 insertions(+), 4 deletions(-) diff --git a/src/paimon/common/reader/late_materializing_file_batch_reader.cpp b/src/paimon/common/reader/late_materializing_file_batch_reader.cpp index ab36c4adc..96a94b1ec 100644 --- a/src/paimon/common/reader/late_materializing_file_batch_reader.cpp +++ b/src/paimon/common/reader/late_materializing_file_batch_reader.cpp @@ -243,12 +243,19 @@ Result LateMaterializingFileBatchReader::AssembleFul return std::make_pair(std::move(c_array), std::move(c_schema)); } +Status LateMaterializingFileBatchReader::ReapplyReadRanges() { + if (read_ranges_.empty()) { + return Status::OK(); + } + return inner_->SetReadRanges(read_ranges_); +} + Status LateMaterializingFileBatchReader::SetInnerProbeSchema() { ::ArrowSchema c_probe_schema; PAIMON_RETURN_NOT_OK_FROM_ARROW(arrow::ExportSchema(*probe_schema_, &c_probe_schema)); PAIMON_RETURN_NOT_OK(inner_->SetReadSchema(&c_probe_schema, predicate_, selection_)); - // SetReadSchema may refresh the read ranges of the inner reader, so we set it again. - PAIMON_RETURN_NOT_OK(SetReadRanges(read_ranges_)); + // SetReadSchema may refresh the read ranges of the inner reader, so we set them again. + PAIMON_RETURN_NOT_OK(ReapplyReadRanges()); return Status::OK(); } @@ -259,7 +266,7 @@ Status LateMaterializingFileBatchReader::SetInnerPayloadSchema() { return Status::Invalid("late materialization bitmap is empty. Should return EOF."); } PAIMON_RETURN_NOT_OK(inner_->SetReadSchema(&c_payload_schema, /*predicate=*/nullptr, matched_bitmap_)); - PAIMON_RETURN_NOT_OK(SetReadRanges(read_ranges_)); + PAIMON_RETURN_NOT_OK(ReapplyReadRanges()); return Status::OK(); } @@ -267,7 +274,7 @@ Status LateMaterializingFileBatchReader::SetInnerFullSchema() { ::ArrowSchema c_full_schema; PAIMON_RETURN_NOT_OK_FROM_ARROW(arrow::ExportSchema(*full_schema_, &c_full_schema)); PAIMON_RETURN_NOT_OK(inner_->SetReadSchema(&c_full_schema, predicate_, selection_)); - PAIMON_RETURN_NOT_OK(SetReadRanges(read_ranges_)); + PAIMON_RETURN_NOT_OK(ReapplyReadRanges()); return Status::OK(); } diff --git a/src/paimon/common/reader/late_materializing_file_batch_reader.h b/src/paimon/common/reader/late_materializing_file_batch_reader.h index 0ca9fc733..7a1ab8a84 100644 --- a/src/paimon/common/reader/late_materializing_file_batch_reader.h +++ b/src/paimon/common/reader/late_materializing_file_batch_reader.h @@ -95,6 +95,7 @@ class LateMaterializingFileBatchReader : public PrefetchFileBatchReader { Status SetInnerProbeSchema(); Status SetInnerPayloadSchema(); Status SetInnerFullSchema(); + Status ReapplyReadRanges(); LatMatState state_ = kInit; std::unique_ptr inner_; From 428a567ac8de20dcf407157ec08d8d639bb32e6a Mon Sep 17 00:00:00 2001 From: zhouhongfeng Date: Mon, 24 Aug 2026 11:54:02 +0800 Subject: [PATCH 08/34] style: refractor code and add comments --- .../late_materializing_file_batch_reader.cpp | 65 +++++-------------- .../late_materializing_file_batch_reader.h | 28 ++++---- ...e_materializing_file_batch_reader_test.cpp | 48 +++++++------- 3 files changed, 54 insertions(+), 87 deletions(-) diff --git a/src/paimon/common/reader/late_materializing_file_batch_reader.cpp b/src/paimon/common/reader/late_materializing_file_batch_reader.cpp index 96a94b1ec..f7aa9c69a 100644 --- a/src/paimon/common/reader/late_materializing_file_batch_reader.cpp +++ b/src/paimon/common/reader/late_materializing_file_batch_reader.cpp @@ -54,13 +54,13 @@ Result LateMaterializingFileBatchReader::NextBatch() Result LateMaterializingFileBatchReader::NextBatchWithBitmap() { if (state_ == kProbing) { - // This calling will update matched_bitmap_ and probe_data_ PAIMON_RETURN_NOT_OK(ReadAndFilterProbeData()); if (matched_bitmap_.IsEmpty()) { state_ = kEOF; } else { - // inner_->SetReadSchema with matched_bitmap_ - PAIMON_RETURN_NOT_OK(SetInnerPayloadSchema()); + // payload pass reads only the matched rows (matched_bitmap_ is non-empty here). + PAIMON_RETURN_NOT_OK( + SetInnerReadSchema(payload_schema_, /*predicate=*/nullptr, matched_bitmap_)); state_ = kRunning; } } @@ -77,8 +77,6 @@ LateMaterializingFileBatchReader::NextBatchWithBitmap() { } Result> LateMaterializingFileBatchReader::BindProbeFilter() { - // non_partition_filter's field_index_ is relative to the full data field list, so the - // predicate must be rebound by name to the probe schema before evaluating it. std::map name_to_idx; for (int32_t i = 0; i < probe_schema_->num_fields(); ++i) { name_to_idx.emplace(probe_schema_->field(i)->name(), i); @@ -107,11 +105,9 @@ Result LateMaterializingFileBatchReader::FilterProbeBatch( if (!results[static_cast(i)]) { continue; } - // The format reader prunes row groups/pages, so a batch offset is not a file row id. + // map batch offset to file row id PAIMON_ASSIGN_OR_RAISE(uint64_t file_row, inner_->GetPreviousBatchFileRowId(static_cast(i))); - // keep matched_bitmap_ ⊆ selection_ (file index ∩ ¬DV), which the probe reader applies - // imprecisely if (selection_ && !selection_->Contains(file_row)) { continue; } @@ -174,9 +170,7 @@ Result LateMaterializingFileBatchReader::R PAIMON_ASSIGN_OR_RAISE_FROM_ARROW(std::shared_ptr payload_array, arrow::ImportArray(c_array.get(), c_schema.get())); - // Recompute the precise selection: keep only the candidate rows whose file row id is in - // matched_bitmap_. Also record their file row ids so GetPreviousBatchFileRowId stays - // correct after compaction/reassembly. + // Generate the valid bitmap and row_mapping_ RoaringBitmap32 valid; row_mapping_.clear(); for (auto it = bitmap.Begin(); it != bitmap.End(); ++it) { @@ -199,8 +193,6 @@ Result LateMaterializingFileBatchReader::R PAIMON_ASSIGN_OR_RAISE_FROM_ARROW(std::shared_ptr payload_compacted, arrow::Concatenate(payload_slices)); - // probe_data_ holds the matched probe rows in the same ascending file order, so a running - // cursor yields row-for-row alignment with the compacted payload. int64_t card = static_cast(valid.Cardinality()); if (probe_cursor_ + card > probe_data_->length()) { return Status::Invalid( @@ -243,38 +235,16 @@ Result LateMaterializingFileBatchReader::AssembleFul return std::make_pair(std::move(c_array), std::move(c_schema)); } -Status LateMaterializingFileBatchReader::ReapplyReadRanges() { - if (read_ranges_.empty()) { - return Status::OK(); +Status LateMaterializingFileBatchReader::SetInnerReadSchema( + const std::shared_ptr& read_schema, const std::shared_ptr& predicate, + const std::optional& selection) { + ::ArrowSchema c_read_schema; + PAIMON_RETURN_NOT_OK_FROM_ARROW(arrow::ExportSchema(*read_schema, &c_read_schema)); + /// Note: calling inner->SetReadSchema may refresh the read ranges of the inner reader. + PAIMON_RETURN_NOT_OK(inner_->SetReadSchema(&c_read_schema, predicate, selection)); + if (!read_ranges_.empty()) { + PAIMON_RETURN_NOT_OK(inner_->SetReadRanges(read_ranges_)); } - return inner_->SetReadRanges(read_ranges_); -} - -Status LateMaterializingFileBatchReader::SetInnerProbeSchema() { - ::ArrowSchema c_probe_schema; - PAIMON_RETURN_NOT_OK_FROM_ARROW(arrow::ExportSchema(*probe_schema_, &c_probe_schema)); - PAIMON_RETURN_NOT_OK(inner_->SetReadSchema(&c_probe_schema, predicate_, selection_)); - // SetReadSchema may refresh the read ranges of the inner reader, so we set them again. - PAIMON_RETURN_NOT_OK(ReapplyReadRanges()); - return Status::OK(); -} - -Status LateMaterializingFileBatchReader::SetInnerPayloadSchema() { - ::ArrowSchema c_payload_schema; - PAIMON_RETURN_NOT_OK_FROM_ARROW(arrow::ExportSchema(*payload_schema_, &c_payload_schema)); - if (matched_bitmap_.IsEmpty()) { - return Status::Invalid("late materialization bitmap is empty. Should return EOF."); - } - PAIMON_RETURN_NOT_OK(inner_->SetReadSchema(&c_payload_schema, /*predicate=*/nullptr, matched_bitmap_)); - PAIMON_RETURN_NOT_OK(ReapplyReadRanges()); - return Status::OK(); -} - -Status LateMaterializingFileBatchReader::SetInnerFullSchema() { - ::ArrowSchema c_full_schema; - PAIMON_RETURN_NOT_OK_FROM_ARROW(arrow::ExportSchema(*full_schema_, &c_full_schema)); - PAIMON_RETURN_NOT_OK(inner_->SetReadSchema(&c_full_schema, predicate_, selection_)); - PAIMON_RETURN_NOT_OK(ReapplyReadRanges()); return Status::OK(); } @@ -322,10 +292,10 @@ Status LateMaterializingFileBatchReader::SetReadSchema( } if (predicate_ == nullptr || probe_schema_ == nullptr) { - PAIMON_RETURN_NOT_OK(SetInnerFullSchema()); + PAIMON_RETURN_NOT_OK(SetInnerReadSchema(full_schema_, predicate_, selection_)); state_ = kNoLatMat; } else { - PAIMON_RETURN_NOT_OK(SetInnerProbeSchema()); + PAIMON_RETURN_NOT_OK(SetInnerReadSchema(probe_schema_, predicate_, selection_)); state_ = kProbing; } return Status::OK(); @@ -365,7 +335,8 @@ Status LateMaterializingFileBatchReader::SeekToRow(uint64_t row_number) { ++cursor; } probe_cursor_ = cursor; - state_ = kRunning; // a seek after EOF re-activates payload reading + // a seek after EOF re-activates payload reading + state_ = kRunning; } return Status::OK(); } diff --git a/src/paimon/common/reader/late_materializing_file_batch_reader.h b/src/paimon/common/reader/late_materializing_file_batch_reader.h index 7a1ab8a84..cc3aecfa3 100644 --- a/src/paimon/common/reader/late_materializing_file_batch_reader.h +++ b/src/paimon/common/reader/late_materializing_file_batch_reader.h @@ -33,6 +33,7 @@ namespace paimon { class PredicateFilter; // For convenience, we abbreviate `Later Materializing` as `LatMat`. +// TODO(zhouhongfeng.zhf): add this reader to the split read path. class LateMaterializingFileBatchReader : public PrefetchFileBatchReader { public: static Result> Create( @@ -68,22 +69,18 @@ class LateMaterializingFileBatchReader : public PrefetchFileBatchReader { kRunning, // Lat-mat is enable an is reading data kEOF }; - /// Scan the probe projection once, evaluating the predicate batch by batch: the matched file - /// row ids go into matched_bitmap_ and the matched probe values into probe_data_. + + /// Read the probe projection once (whole file) and evaluating the predicate batch by batch. + /// This function updates matched_bitmap_ and probe_data_. + /// TODO(zhouhongfeng.zhf): Read the probe data batch by batch to save memory. Status ReadAndFilterProbeData(); - /// Rebind the predicate to probe_schema_'s field indices, so that it can be evaluated over - /// the probe batches. Result> BindProbeFilter(); - /// Evaluate bound_filter over a single probe batch, adding the matched file row ids (that also - /// pass selection_) to matched_bitmap_. Returns the batch-local offsets of the matched rows so - /// the caller can compact the probe batch down to those rows. Result FilterProbeBatch(const std::shared_ptr& array, const std::shared_ptr& bound_filter); - /// In kRunning state, read one payload batch, map its matched rows back to the cached probe - /// rows by file row id, and reassemble the full read schema. + /// Read one payload batch with bitmap (matched rows only) Result ReadPayloadBatch(); /// Combine the compacted payload columns and the selected probe columns into a single struct @@ -92,10 +89,9 @@ class LateMaterializingFileBatchReader : public PrefetchFileBatchReader { const std::shared_ptr& payload_array, const std::shared_ptr& probe_array); - Status SetInnerProbeSchema(); - Status SetInnerPayloadSchema(); - Status SetInnerFullSchema(); - Status ReapplyReadRanges(); + Status SetInnerReadSchema(const std::shared_ptr& read_schema, + const std::shared_ptr& predicate, + const std::optional& selection); LatMatState state_ = kInit; std::unique_ptr inner_; @@ -107,14 +103,12 @@ class LateMaterializingFileBatchReader : public PrefetchFileBatchReader { std::shared_ptr payload_schema_; std::shared_ptr predicate_; std::optional selection_; - // matched probe rows only, concatenated in ascending file row order (aligned 1:1 with - // matched_bitmap_) + // the probe_data_ is sliced and compacted with the matched_bitmap_ std::shared_ptr probe_data_; - // file-level row ids that pass the predicate (and selection_); drives the payload read RoaringBitmap32 matched_bitmap_; // read cursor into probe_data_ for the payload phase int64_t probe_cursor_ = 0; - // file row id of each row in the batch last emitted in kRunning state + // to support GetPreviousBatchFileRowId std::vector row_mapping_; }; diff --git a/src/paimon/common/reader/late_materializing_file_batch_reader_test.cpp b/src/paimon/common/reader/late_materializing_file_batch_reader_test.cpp index 92c11e9fc..3dd702176 100644 --- a/src/paimon/common/reader/late_materializing_file_batch_reader_test.cpp +++ b/src/paimon/common/reader/late_materializing_file_batch_reader_test.cpp @@ -162,10 +162,9 @@ class LateMaterializingFileBatchReaderTest : public ::testing::Test { if (BatchReader::IsEofBatch(batch_with_bitmap)) { break; } - PAIMON_ASSIGN_OR_RAISE( - BatchReader::ReadBatch batch, - ReaderUtils::ApplyBitmapToReadBatch(std::move(batch_with_bitmap), - arrow::default_memory_pool())); + PAIMON_ASSIGN_OR_RAISE(BatchReader::ReadBatch batch, + ReaderUtils::ApplyBitmapToReadBatch( + std::move(batch_with_bitmap), arrow::default_memory_pool())); auto& [c_array, c_schema] = batch; PAIMON_ASSIGN_OR_RAISE_FROM_ARROW(std::shared_ptr array, arrow::ImportArray(c_array.get(), c_schema.get())); @@ -182,11 +181,10 @@ class LateMaterializingFileBatchReaderTest : public ::testing::Test { // Build a struct with 5 columns [a:int64, b:utf8, c:int64, d:utf8, e:int64], each carrying a // distinct value pattern so any column reordering is detected. std::shared_ptr BuildMultiFieldData(int32_t n) { - auto type = arrow::struct_({arrow::field("a", arrow::int64()), - arrow::field("b", arrow::utf8()), - arrow::field("c", arrow::int64()), - arrow::field("d", arrow::utf8()), - arrow::field("e", arrow::int64())}); + auto type = + arrow::struct_({arrow::field("a", arrow::int64()), arrow::field("b", arrow::utf8()), + arrow::field("c", arrow::int64()), arrow::field("d", arrow::utf8()), + arrow::field("e", arrow::int64())}); arrow::StructBuilder builder( type, arrow::default_memory_pool(), {std::make_shared(), std::make_shared(), @@ -457,10 +455,10 @@ TEST_F(LateMaterializingFileBatchReaderTest, MultiFieldPreservesColumnOrder) { auto mock = std::make_unique(data, type, /*batch_size=*/3); ASSERT_OK_AND_ASSIGN(auto reader, LateMaterializingFileBatchReader::Create(std::move(mock))); // probe columns = {a (idx0), c (idx2)}; payload columns = {b, d, e} - auto pred_a = PredicateBuilder::GreaterOrEqual(/*field_index=*/0, "a", FieldType::BIGINT, - Literal(3l)); - auto pred_c = PredicateBuilder::LessThan(/*field_index=*/2, "c", FieldType::BIGINT, - Literal(700l)); + auto pred_a = + PredicateBuilder::GreaterOrEqual(/*field_index=*/0, "a", FieldType::BIGINT, Literal(3l)); + auto pred_c = + PredicateBuilder::LessThan(/*field_index=*/2, "c", FieldType::BIGINT, Literal(700l)); ASSERT_OK_AND_ASSIGN(auto predicate, PredicateBuilder::And({pred_a, pred_c})); ASSERT_OK(SetReadSchema(reader.get(), arrow::schema(type->fields()), predicate, std::nullopt)); @@ -499,15 +497,16 @@ TEST_F(LateMaterializingFileBatchReaderTest, NestedPayloadColumn) { auto mock = std::make_unique(data, type, /*batch_size=*/3); ASSERT_OK_AND_ASSIGN(auto reader, LateMaterializingFileBatchReader::Create(std::move(mock))); // probe = {k}; payload = {arr (list), tag} - auto predicate = PredicateBuilder::GreaterOrEqual(/*field_index=*/0, "k", FieldType::BIGINT, - Literal(5l)); + auto predicate = + PredicateBuilder::GreaterOrEqual(/*field_index=*/0, "k", FieldType::BIGINT, Literal(5l)); ASSERT_OK(SetReadSchema(reader.get(), arrow::schema(type->fields()), predicate, std::nullopt)); ASSERT_OK_AND_ASSIGN(std::shared_ptr result, CollectStruct(reader.get())); ASSERT_TRUE(result); ASSERT_EQ(result->length(), 3); // k = 5,6,7 auto k = arrow::internal::checked_pointer_cast(result->GetFieldByName("k")); - auto arr = arrow::internal::checked_pointer_cast(result->GetFieldByName("arr")); + auto arr = + arrow::internal::checked_pointer_cast(result->GetFieldByName("arr")); auto tag = arrow::internal::checked_pointer_cast(result->GetFieldByName("tag")); ASSERT_TRUE(k && arr && tag); @@ -537,8 +536,8 @@ TEST_F(LateMaterializingFileBatchReaderTest, WorksAsInnerOfPrefetchReader) { /*enable_adaptive_prefetch_strategy=*/false, executor, /*initialize_read_ranges=*/false, /*read_ahead_cache_enabled=*/false, CacheConfig(), GetDefaultPool())); - auto predicate = PredicateBuilder::GreaterOrEqual(/*field_index=*/0, "k", FieldType::BIGINT, - Literal(4l)); + auto predicate = + PredicateBuilder::GreaterOrEqual(/*field_index=*/0, "k", FieldType::BIGINT, Literal(4l)); ::ArrowSchema c_schema; ASSERT_TRUE(arrow::ExportSchema(*arrow::schema(full_fields_), &c_schema).ok()); ASSERT_OK(impl->SetReadSchema(&c_schema, predicate, std::nullopt)); @@ -581,7 +580,8 @@ TEST_F(LateMaterializingFileBatchReaderTest, PrefetchInnerReentrantSetReadSchema ASSERT_OK_AND_ASSIGN(std::shared_ptr result1, CollectStruct(impl.get())); ASSERT_TRUE(result1); ASSERT_EQ(result1->length(), 3); // k = 7,8,9 - auto k1 = arrow::internal::checked_pointer_cast(result1->GetFieldByName("k")); + auto k1 = + arrow::internal::checked_pointer_cast(result1->GetFieldByName("k")); for (int64_t j = 0; j < result1->length(); ++j) { EXPECT_EQ(k1->Value(j), 7 + j); } @@ -594,8 +594,10 @@ TEST_F(LateMaterializingFileBatchReaderTest, PrefetchInnerReentrantSetReadSchema ASSERT_OK_AND_ASSIGN(std::shared_ptr result2, CollectStruct(impl.get())); ASSERT_TRUE(result2); ASSERT_EQ(result2->length(), 3); // k = 0,1,2 - auto k2 = arrow::internal::checked_pointer_cast(result2->GetFieldByName("k")); - auto v2 = arrow::internal::checked_pointer_cast(result2->GetFieldByName("v")); + auto k2 = + arrow::internal::checked_pointer_cast(result2->GetFieldByName("k")); + auto v2 = + arrow::internal::checked_pointer_cast(result2->GetFieldByName("v")); for (int64_t j = 0; j < result2->length(); ++j) { EXPECT_EQ(k2->Value(j), j); EXPECT_EQ(v2->GetString(j), "v_" + std::to_string(j)); @@ -625,8 +627,8 @@ TEST_F(LateMaterializingFileBatchReaderTest, PrefetchInnerParallelReadersWithSee /*enable_adaptive_prefetch_strategy=*/false, executor, /*initialize_read_ranges=*/false, /*read_ahead_cache_enabled=*/false, CacheConfig(), GetDefaultPool())); - auto predicate = PredicateBuilder::GreaterOrEqual(/*field_index=*/0, "k", FieldType::BIGINT, - Literal(5l)); + auto predicate = + PredicateBuilder::GreaterOrEqual(/*field_index=*/0, "k", FieldType::BIGINT, Literal(5l)); ::ArrowSchema c_schema; ASSERT_TRUE(arrow::ExportSchema(*arrow::schema(full_fields_), &c_schema).ok()); ASSERT_OK(impl->SetReadSchema(&c_schema, predicate, std::nullopt)); From 2d81627b5a3b60645c8e7816e62482019d86c1c0 Mon Sep 17 00:00:00 2001 From: zhouhongfeng Date: Mon, 24 Aug 2026 14:12:15 +0800 Subject: [PATCH 09/34] fix: support memory pool --- .../late_materializing_file_batch_reader.cpp | 21 ++++++++-- .../late_materializing_file_batch_reader.h | 14 +++++-- ...e_materializing_file_batch_reader_test.cpp | 41 ++++++++++++------- 3 files changed, 54 insertions(+), 22 deletions(-) diff --git a/src/paimon/common/reader/late_materializing_file_batch_reader.cpp b/src/paimon/common/reader/late_materializing_file_batch_reader.cpp index f7aa9c69a..19c1e0a32 100644 --- a/src/paimon/common/reader/late_materializing_file_batch_reader.cpp +++ b/src/paimon/common/reader/late_materializing_file_batch_reader.cpp @@ -28,11 +28,13 @@ #include "arrow/array/concatenate.h" #include "arrow/array/util.h" #include "arrow/c/bridge.h" +#include "arrow/memory_pool.h" #include "arrow/type.h" #include "arrow/util/checked_cast.h" #include "fmt/format.h" #include "paimon/common/predicate/predicate_filter.h" #include "paimon/common/reader/reader_utils.h" +#include "paimon/common/utils/arrow/mem_utils.h" #include "paimon/common/utils/arrow/status_utils.h" #include "paimon/predicate/predicate_utils.h" #include "paimon/status.h" @@ -40,12 +42,22 @@ namespace paimon { Result> LateMaterializingFileBatchReader::Create( - std::unique_ptr inner) { + std::unique_ptr inner, std::shared_ptr pool) { + // The reader's own compaction allocations go through an arrow pool; bridge the paimon pool + // once here so the accounting matches the rest of the read path. + std::shared_ptr arrow_pool; + if (pool != nullptr) { + arrow_pool = GetArrowPool(pool); + } auto reader = std::unique_ptr( - new LateMaterializingFileBatchReader(std::move(inner))); + new LateMaterializingFileBatchReader(std::move(inner), std::move(arrow_pool))); return reader; } +arrow::MemoryPool* LateMaterializingFileBatchReader::ArrowPool() const { + return arrow_pool_ ? arrow_pool_.get() : arrow::default_memory_pool(); +} + Result LateMaterializingFileBatchReader::NextBatch() { return Status::Invalid( "paimon inner reader PrefetchFileBatchReader should use NextBatchWithBitmap"); @@ -147,7 +159,8 @@ Status LateMaterializingFileBatchReader::ReadAndFilterProbeData() { PAIMON_ASSIGN_OR_RAISE_FROM_ARROW( probe_array, arrow::MakeEmptyArray(arrow::struct_(probe_schema_->fields()))); } else { - PAIMON_ASSIGN_OR_RAISE_FROM_ARROW(probe_array, arrow::Concatenate(probe_arrays)); + PAIMON_ASSIGN_OR_RAISE_FROM_ARROW(probe_array, + arrow::Concatenate(probe_arrays, ArrowPool())); } probe_data_ = arrow::internal::checked_pointer_cast(probe_array); return Status::OK(); @@ -191,7 +204,7 @@ Result LateMaterializingFileBatchReader::R PAIMON_ASSIGN_OR_RAISE(arrow::ArrayVector payload_slices, ReaderUtils::GenerateFilteredArrayVector(payload_array, valid)); PAIMON_ASSIGN_OR_RAISE_FROM_ARROW(std::shared_ptr payload_compacted, - arrow::Concatenate(payload_slices)); + arrow::Concatenate(payload_slices, ArrowPool())); int64_t card = static_cast(valid.Cardinality()); if (probe_cursor_ + card > probe_data_->length()) { diff --git a/src/paimon/common/reader/late_materializing_file_batch_reader.h b/src/paimon/common/reader/late_materializing_file_batch_reader.h index cc3aecfa3..2c7548f75 100644 --- a/src/paimon/common/reader/late_materializing_file_batch_reader.h +++ b/src/paimon/common/reader/late_materializing_file_batch_reader.h @@ -37,7 +37,7 @@ class PredicateFilter; class LateMaterializingFileBatchReader : public PrefetchFileBatchReader { public: static Result> Create( - std::unique_ptr inner); + std::unique_ptr inner, std::shared_ptr pool); Result NextBatch() override; Result NextBatchWithBitmap() override; @@ -59,8 +59,9 @@ class LateMaterializingFileBatchReader : public PrefetchFileBatchReader { Status SetReadRanges(const std::vector>& read_ranges) override; private: - explicit LateMaterializingFileBatchReader(std::unique_ptr inner) - : inner_(std::move(inner)) {} + explicit LateMaterializingFileBatchReader(std::unique_ptr inner, + std::shared_ptr arrow_pool) + : inner_(std::move(inner)), arrow_pool_(std::move(arrow_pool)) {} enum LatMatState { kInit, @@ -93,8 +94,13 @@ class LateMaterializingFileBatchReader : public PrefetchFileBatchReader { const std::shared_ptr& predicate, const std::optional& selection); - LatMatState state_ = kInit; + // Arrow pool for this reader's own allocations (probe/payload compaction). Falls back to the + // arrow default pool when no pool was provided. + arrow::MemoryPool* ArrowPool() const; + std::unique_ptr inner_; + std::shared_ptr arrow_pool_; + LatMatState state_ = kInit; std::vector> read_ranges_; std::shared_ptr full_schema_; // projection holding only the predicate fields; nullptr when probing is not applicable diff --git a/src/paimon/common/reader/late_materializing_file_batch_reader_test.cpp b/src/paimon/common/reader/late_materializing_file_batch_reader_test.cpp index 3dd702176..e6f577139 100644 --- a/src/paimon/common/reader/late_materializing_file_batch_reader_test.cpp +++ b/src/paimon/common/reader/late_materializing_file_batch_reader_test.cpp @@ -68,8 +68,9 @@ class LmReaderBuilder : public ReaderBuilder { Result> Build( const std::shared_ptr& /*stream*/) const override { auto mock = std::make_unique(data_, type_, batch_size_); - PAIMON_ASSIGN_OR_RAISE(std::unique_ptr reader, - LateMaterializingFileBatchReader::Create(std::move(mock))); + PAIMON_ASSIGN_OR_RAISE( + std::unique_ptr reader, + LateMaterializingFileBatchReader::Create(std::move(mock), GetDefaultPool())); return std::unique_ptr(std::move(reader)); } @@ -247,7 +248,8 @@ class LateMaterializingFileBatchReaderTest : public ::testing::Test { TEST_F(LateMaterializingFileBatchReaderTest, PassThroughWhenNoPredicate) { auto data = BuildData({0, 1, 2, 3, 4}); auto mock = std::make_unique(data, full_type_, /*batch_size=*/2); - ASSERT_OK_AND_ASSIGN(auto reader, LateMaterializingFileBatchReader::Create(std::move(mock))); + ASSERT_OK_AND_ASSIGN( + auto reader, LateMaterializingFileBatchReader::Create(std::move(mock), GetDefaultPool())); ASSERT_OK(SetReadSchema(reader.get(), arrow::schema(full_fields_), /*predicate=*/nullptr, std::nullopt)); @@ -265,7 +267,8 @@ TEST_F(LateMaterializingFileBatchReaderTest, PassThroughWhenNoPredicate) { TEST_F(LateMaterializingFileBatchReaderTest, PassThroughWhenPayloadEmpty) { auto data = BuildData({0, 1, 2, 3, 4}); auto mock = std::make_unique(data, full_type_, /*batch_size=*/2); - ASSERT_OK_AND_ASSIGN(auto reader, LateMaterializingFileBatchReader::Create(std::move(mock))); + ASSERT_OK_AND_ASSIGN( + auto reader, LateMaterializingFileBatchReader::Create(std::move(mock), GetDefaultPool())); // read schema is just {k}; the predicate on k covers all columns -> payload empty auto predicate = PredicateBuilder::GreaterOrEqual(/*field_index=*/0, /*field_name=*/"k", FieldType::BIGINT, Literal(0l)); @@ -279,7 +282,8 @@ TEST_F(LateMaterializingFileBatchReaderTest, PassThroughWhenPayloadEmpty) { TEST_F(LateMaterializingFileBatchReaderTest, ContiguousSubsetAcrossBatches) { auto data = BuildData({0, 1, 2, 3, 4, 5, 6, 7, 8, 9}); auto mock = std::make_unique(data, full_type_, /*batch_size=*/3); - ASSERT_OK_AND_ASSIGN(auto reader, LateMaterializingFileBatchReader::Create(std::move(mock))); + ASSERT_OK_AND_ASSIGN( + auto reader, LateMaterializingFileBatchReader::Create(std::move(mock), GetDefaultPool())); auto predicate = PredicateBuilder::GreaterThan(/*field_index=*/0, /*field_name=*/"k", FieldType::BIGINT, Literal(4l)); ASSERT_OK(SetReadSchema(reader.get(), arrow::schema(full_fields_), predicate, std::nullopt)); @@ -303,7 +307,8 @@ TEST_F(LateMaterializingFileBatchReaderTest, ScatteredAlternatingMatch) { } auto data = BuildData(ks); auto mock = std::make_unique(data, full_type_, /*batch_size=*/3); - ASSERT_OK_AND_ASSIGN(auto reader, LateMaterializingFileBatchReader::Create(std::move(mock))); + ASSERT_OK_AND_ASSIGN( + auto reader, LateMaterializingFileBatchReader::Create(std::move(mock), GetDefaultPool())); auto predicate = PredicateBuilder::Equal(/*field_index=*/0, /*field_name=*/"k", FieldType::BIGINT, Literal(1l)); ASSERT_OK(SetReadSchema(reader.get(), arrow::schema(full_fields_), predicate, std::nullopt)); @@ -326,7 +331,8 @@ TEST_F(LateMaterializingFileBatchReaderTest, MatchedIntersectsSelection) { } auto data = BuildData(ks); auto mock = std::make_unique(data, full_type_, /*batch_size=*/4); - ASSERT_OK_AND_ASSIGN(auto reader, LateMaterializingFileBatchReader::Create(std::move(mock))); + ASSERT_OK_AND_ASSIGN( + auto reader, LateMaterializingFileBatchReader::Create(std::move(mock), GetDefaultPool())); auto predicate = PredicateBuilder::Equal(/*field_index=*/0, /*field_name=*/"k", FieldType::BIGINT, Literal(1l)); // predicate hits {1,3,5,7,9,11}; selection keeps only {1,5,9} @@ -351,7 +357,8 @@ TEST_F(LateMaterializingFileBatchReaderTest, MatchedIntersectsSelection) { TEST_F(LateMaterializingFileBatchReaderTest, EmptyMatchReturnsEof) { auto data = BuildData({0, 1, 2, 3, 4}); auto mock = std::make_unique(data, full_type_, /*batch_size=*/2); - ASSERT_OK_AND_ASSIGN(auto reader, LateMaterializingFileBatchReader::Create(std::move(mock))); + ASSERT_OK_AND_ASSIGN( + auto reader, LateMaterializingFileBatchReader::Create(std::move(mock), GetDefaultPool())); auto predicate = PredicateBuilder::GreaterThan(/*field_index=*/0, /*field_name=*/"k", FieldType::BIGINT, Literal(100l)); ASSERT_OK(SetReadSchema(reader.get(), arrow::schema(full_fields_), predicate, std::nullopt)); @@ -364,7 +371,8 @@ TEST_F(LateMaterializingFileBatchReaderTest, EmptyMatchReturnsEof) { TEST_F(LateMaterializingFileBatchReaderTest, SeekToRowRealignsProbeCursor) { auto data = BuildData({0, 1, 2, 3, 4, 5, 6, 7, 8, 9}); auto mock = std::make_unique(data, full_type_, /*batch_size=*/4); - ASSERT_OK_AND_ASSIGN(auto reader, LateMaterializingFileBatchReader::Create(std::move(mock))); + ASSERT_OK_AND_ASSIGN( + auto reader, LateMaterializingFileBatchReader::Create(std::move(mock), GetDefaultPool())); auto predicate = PredicateBuilder::GreaterOrEqual(/*field_index=*/0, /*field_name=*/"k", FieldType::BIGINT, Literal(5l)); ASSERT_OK(SetReadSchema(reader.get(), arrow::schema(full_fields_), predicate, std::nullopt)); @@ -392,7 +400,8 @@ TEST_F(LateMaterializingFileBatchReaderTest, ReadRangesForwardedAcrossPhases) { auto data = BuildData({0, 1, 2, 3, 4, 5, 6, 7}); auto mock = std::make_unique(data, full_type_, /*batch_size=*/4); auto* mock_ptr = mock.get(); - ASSERT_OK_AND_ASSIGN(auto reader, LateMaterializingFileBatchReader::Create(std::move(mock))); + ASSERT_OK_AND_ASSIGN( + auto reader, LateMaterializingFileBatchReader::Create(std::move(mock), GetDefaultPool())); auto predicate = PredicateBuilder::GreaterOrEqual(/*field_index=*/0, /*field_name=*/"k", FieldType::BIGINT, Literal(2l)); ASSERT_OK(SetReadSchema(reader.get(), arrow::schema(full_fields_), predicate, std::nullopt)); @@ -414,7 +423,8 @@ TEST_F(LateMaterializingFileBatchReaderTest, ReadRangesForwardedAcrossPhases) { TEST_F(LateMaterializingFileBatchReaderTest, ReentrantSetReadSchema) { auto data = BuildData({0, 1, 2, 3, 4, 5, 6, 7, 8, 9}); auto mock = std::make_unique(data, full_type_, /*batch_size=*/3); - ASSERT_OK_AND_ASSIGN(auto reader, LateMaterializingFileBatchReader::Create(std::move(mock))); + ASSERT_OK_AND_ASSIGN( + auto reader, LateMaterializingFileBatchReader::Create(std::move(mock), GetDefaultPool())); auto predicate1 = PredicateBuilder::GreaterThan(/*field_index=*/0, /*field_name=*/"k", FieldType::BIGINT, Literal(7l)); @@ -437,7 +447,8 @@ TEST_F(LateMaterializingFileBatchReaderTest, ReentrantSetReadSchema) { TEST_F(LateMaterializingFileBatchReaderTest, ForwardsRowCountAndFileSchema) { auto data = BuildData({0, 1, 2, 3}); auto mock = std::make_unique(data, full_type_, /*batch_size=*/2); - ASSERT_OK_AND_ASSIGN(auto reader, LateMaterializingFileBatchReader::Create(std::move(mock))); + ASSERT_OK_AND_ASSIGN( + auto reader, LateMaterializingFileBatchReader::Create(std::move(mock), GetDefaultPool())); ASSERT_OK_AND_ASSIGN(uint64_t num_rows, reader->GetNumberOfRows()); EXPECT_EQ(num_rows, 4u); @@ -453,7 +464,8 @@ TEST_F(LateMaterializingFileBatchReaderTest, MultiFieldPreservesColumnOrder) { auto data = BuildMultiFieldData(10); auto type = data->type(); auto mock = std::make_unique(data, type, /*batch_size=*/3); - ASSERT_OK_AND_ASSIGN(auto reader, LateMaterializingFileBatchReader::Create(std::move(mock))); + ASSERT_OK_AND_ASSIGN( + auto reader, LateMaterializingFileBatchReader::Create(std::move(mock), GetDefaultPool())); // probe columns = {a (idx0), c (idx2)}; payload columns = {b, d, e} auto pred_a = PredicateBuilder::GreaterOrEqual(/*field_index=*/0, "a", FieldType::BIGINT, Literal(3l)); @@ -495,7 +507,8 @@ TEST_F(LateMaterializingFileBatchReaderTest, NestedPayloadColumn) { auto data = BuildNestedData(8); auto type = data->type(); auto mock = std::make_unique(data, type, /*batch_size=*/3); - ASSERT_OK_AND_ASSIGN(auto reader, LateMaterializingFileBatchReader::Create(std::move(mock))); + ASSERT_OK_AND_ASSIGN( + auto reader, LateMaterializingFileBatchReader::Create(std::move(mock), GetDefaultPool())); // probe = {k}; payload = {arr (list), tag} auto predicate = PredicateBuilder::GreaterOrEqual(/*field_index=*/0, "k", FieldType::BIGINT, Literal(5l)); From 750769211cb232b8b53f5fc9435d550bb3acdad0 Mon Sep 17 00:00:00 2001 From: zhouhongfeng Date: Mon, 24 Aug 2026 14:17:17 +0800 Subject: [PATCH 10/34] fix: support calling NextBatch --- .../late_materializing_file_batch_reader.cpp | 16 +++++----------- .../late_materializing_file_batch_reader.h | 3 +-- 2 files changed, 6 insertions(+), 13 deletions(-) diff --git a/src/paimon/common/reader/late_materializing_file_batch_reader.cpp b/src/paimon/common/reader/late_materializing_file_batch_reader.cpp index 19c1e0a32..9a1399a54 100644 --- a/src/paimon/common/reader/late_materializing_file_batch_reader.cpp +++ b/src/paimon/common/reader/late_materializing_file_batch_reader.cpp @@ -59,12 +59,6 @@ arrow::MemoryPool* LateMaterializingFileBatchReader::ArrowPool() const { } Result LateMaterializingFileBatchReader::NextBatch() { - return Status::Invalid( - "paimon inner reader PrefetchFileBatchReader should use NextBatchWithBitmap"); -} - -Result -LateMaterializingFileBatchReader::NextBatchWithBitmap() { if (state_ == kProbing) { PAIMON_RETURN_NOT_OK(ReadAndFilterProbeData()); if (matched_bitmap_.IsEmpty()) { @@ -78,11 +72,11 @@ LateMaterializingFileBatchReader::NextBatchWithBitmap() { } if (state_ == kNoLatMat) { - return inner_->NextBatchWithBitmap(); + return inner_->NextBatch(); } else if (state_ == kRunning) { return ReadPayloadBatch(); } else if (state_ == kEOF) { - return MakeEofBatchWithBitmap(); + return MakeEofBatch(); } return Status::Invalid("invalid state when calling NextBatchWithBitmap: " + std::to_string(state_)); @@ -166,13 +160,13 @@ Status LateMaterializingFileBatchReader::ReadAndFilterProbeData() { return Status::OK(); } -Result LateMaterializingFileBatchReader::ReadPayloadBatch() { +Result LateMaterializingFileBatchReader::ReadPayloadBatch() { while (true) { PAIMON_ASSIGN_OR_RAISE(FileBatchReader::ReadBatchWithBitmap batch_with_bitmap, inner_->NextBatchWithBitmap()); if (BatchReader::IsEofBatch(batch_with_bitmap)) { state_ = kEOF; - return MakeEofBatchWithBitmap(); + return MakeEofBatch(); } auto& [batch, bitmap] = batch_with_bitmap; if (bitmap.IsEmpty()) { @@ -217,7 +211,7 @@ Result LateMaterializingFileBatchReader::R PAIMON_ASSIGN_OR_RAISE(FileBatchReader::ReadBatch assembled, AssembleFullBatch(payload_compacted, probe_selected)); - return ReaderUtils::AddAllValidBitmap(std::move(assembled)); + return assembled; } } diff --git a/src/paimon/common/reader/late_materializing_file_batch_reader.h b/src/paimon/common/reader/late_materializing_file_batch_reader.h index 2c7548f75..9c77cf407 100644 --- a/src/paimon/common/reader/late_materializing_file_batch_reader.h +++ b/src/paimon/common/reader/late_materializing_file_batch_reader.h @@ -40,7 +40,6 @@ class LateMaterializingFileBatchReader : public PrefetchFileBatchReader { std::unique_ptr inner, std::shared_ptr pool); Result NextBatch() override; - Result NextBatchWithBitmap() override; std::shared_ptr GetReaderMetrics() const override; void Close() override; @@ -82,7 +81,7 @@ class LateMaterializingFileBatchReader : public PrefetchFileBatchReader { const std::shared_ptr& bound_filter); /// Read one payload batch with bitmap (matched rows only) - Result ReadPayloadBatch(); + Result ReadPayloadBatch(); /// Combine the compacted payload columns and the selected probe columns into a single struct /// array following full_schema_'s field order. From 0325f657b2e4e0346e950d33ef0b519934e408f1 Mon Sep 17 00:00:00 2001 From: zhouhongfeng Date: Mon, 24 Aug 2026 14:26:11 +0800 Subject: [PATCH 11/34] style: add TODOs --- .../common/reader/late_materializing_file_batch_reader.cpp | 1 + 1 file changed, 1 insertion(+) diff --git a/src/paimon/common/reader/late_materializing_file_batch_reader.cpp b/src/paimon/common/reader/late_materializing_file_batch_reader.cpp index 9a1399a54..b7cb71aed 100644 --- a/src/paimon/common/reader/late_materializing_file_batch_reader.cpp +++ b/src/paimon/common/reader/late_materializing_file_batch_reader.cpp @@ -99,6 +99,7 @@ Result> LateMaterializingFileBatchReader::BindP Result LateMaterializingFileBatchReader::FilterProbeBatch( const std::shared_ptr& array, const std::shared_ptr& bound_filter) { + // TODO(zhouhonfeng.zhf): use arrow::compute::Filter instead of PredicateFilter PAIMON_ASSIGN_OR_RAISE(std::vector results, bound_filter->Test(*array)); if (results.size() != static_cast(array->length())) { return Status::Invalid( From e5e2a6a4d77a7c994c2b9005e69c43c6d9f28b63 Mon Sep 17 00:00:00 2001 From: zhouhongfeng Date: Mon, 24 Aug 2026 15:11:49 +0800 Subject: [PATCH 12/34] feat: add support LM builder --- ...e_materializing_file_batch_reader_test.cpp | 42 +++-------- .../late_materializing_reader_builder.h | 75 +++++++++++++++++++ 2 files changed, 86 insertions(+), 31 deletions(-) create mode 100644 src/paimon/common/reader/late_materializing_reader_builder.h diff --git a/src/paimon/common/reader/late_materializing_file_batch_reader_test.cpp b/src/paimon/common/reader/late_materializing_file_batch_reader_test.cpp index e6f577139..27c84e848 100644 --- a/src/paimon/common/reader/late_materializing_file_batch_reader_test.cpp +++ b/src/paimon/common/reader/late_materializing_file_batch_reader_test.cpp @@ -32,6 +32,7 @@ #include "arrow/array/builder_nested.h" #include "arrow/c/bridge.h" #include "gtest/gtest.h" +#include "paimon/common/reader/late_materializing_reader_builder.h" #include "paimon/common/reader/prefetch_file_batch_reader_impl.h" #include "paimon/common/reader/reader_utils.h" #include "paimon/common/utils/arrow/status_utils.h" @@ -47,39 +48,12 @@ #include "paimon/status.h" #include "paimon/testing/mock/mock_file_batch_reader.h" #include "paimon/testing/mock/mock_file_system.h" +#include "paimon/testing/mock/mock_format_reader_builder.h" #include "paimon/testing/utils/testharness.h" #include "paimon/utils/roaring_bitmap32.h" namespace paimon::test { -// A ReaderBuilder whose Build() wraps a projecting, range-honoring mock in a -// LateMaterializingFileBatchReader, so the late-materialization reader can be exercised as an -// inner reader of PrefetchFileBatchReaderImpl. -class LmReaderBuilder : public ReaderBuilder { - public: - LmReaderBuilder(std::shared_ptr data, std::shared_ptr type, - int32_t batch_size) - : data_(std::move(data)), type_(std::move(type)), batch_size_(batch_size) {} - - ReaderBuilder* WithMemoryPool(const std::shared_ptr& /*pool*/) override { - return this; - } - - Result> Build( - const std::shared_ptr& /*stream*/) const override { - auto mock = std::make_unique(data_, type_, batch_size_); - PAIMON_ASSIGN_OR_RAISE( - std::unique_ptr reader, - LateMaterializingFileBatchReader::Create(std::move(mock), GetDefaultPool())); - return std::unique_ptr(std::move(reader)); - } - - private: - std::shared_ptr data_; - std::shared_ptr type_; - int32_t batch_size_ = 0; -}; - class LateMaterializingFileBatchReaderTest : public ::testing::Test { public: void SetUp() override { @@ -538,7 +512,9 @@ TEST_F(LateMaterializingFileBatchReaderTest, NestedPayloadColumn) { // PrefetchFileBatchReaderImpl (schema broadcast, range dispatch, seek, row-id tracking). TEST_F(LateMaterializingFileBatchReaderTest, WorksAsInnerOfPrefetchReader) { auto data = BuildData({0, 1, 2, 3, 4, 5, 6, 7, 8, 9}); - LmReaderBuilder builder(data, full_type_, /*batch_size=*/3); + LateMaterializingReaderBuilder builder( + std::make_unique(data, full_type_, /*batch_size=*/3), + GetDefaultPool()); auto mock_fs = std::make_shared(); ASSERT_OK_AND_ASSIGN(std::shared_ptr executor, CreateDefaultExecutor(2)); ASSERT_OK_AND_ASSIGN( @@ -572,7 +548,9 @@ TEST_F(LateMaterializingFileBatchReaderTest, WorksAsInnerOfPrefetchReader) { // re-plans ranges) must reset the probe state and produce correct results for the new predicate. TEST_F(LateMaterializingFileBatchReaderTest, PrefetchInnerReentrantSetReadSchema) { auto data = BuildData({0, 1, 2, 3, 4, 5, 6, 7, 8, 9}); - LmReaderBuilder builder(data, full_type_, /*batch_size=*/3); + LateMaterializingReaderBuilder builder( + std::make_unique(data, full_type_, /*batch_size=*/3), + GetDefaultPool()); auto mock_fs = std::make_shared(); ASSERT_OK_AND_ASSIGN(std::shared_ptr executor, CreateDefaultExecutor(2)); ASSERT_OK_AND_ASSIGN( @@ -629,7 +607,9 @@ TEST_F(LateMaterializingFileBatchReaderTest, PrefetchInnerParallelReadersWithSee auto data = BuildData(ks); // Per-batch ranges (the mock's default) let the impl split work across the parallel readers, // and each range-honoring reader only reads its assigned slice. - LmReaderBuilder builder(data, full_type_, /*batch_size=*/3); + LateMaterializingReaderBuilder builder( + std::make_unique(data, full_type_, /*batch_size=*/3), + GetDefaultPool()); auto mock_fs = std::make_shared(); ASSERT_OK_AND_ASSIGN(std::shared_ptr executor, CreateDefaultExecutor(3)); ASSERT_OK_AND_ASSIGN( diff --git a/src/paimon/common/reader/late_materializing_reader_builder.h b/src/paimon/common/reader/late_materializing_reader_builder.h new file mode 100644 index 000000000..07317924d --- /dev/null +++ b/src/paimon/common/reader/late_materializing_reader_builder.h @@ -0,0 +1,75 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you under the Apache License, Version 2.0 (the + * "License"); you may not use this file except in compliance + * with the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, + * software distributed under the License is distributed on an + * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY + * KIND, either express or implied. See the License for the + * specific language governing permissions and limitations + * under the License. + */ + +#pragma once + +#include +#include + +#include "paimon/common/reader/late_materializing_file_batch_reader.h" +#include "paimon/format/reader_builder.h" +#include "paimon/memory/memory_pool.h" +#include "paimon/reader/prefetch_file_batch_reader.h" +#include "paimon/result.h" + +namespace paimon { + +class LateMaterializingReaderBuilder : public ReaderBuilder { + public: + LateMaterializingReaderBuilder(std::unique_ptr inner, + std::shared_ptr pool) + : inner_(std::move(inner)), pool_(std::move(pool)) {} + + ReaderBuilder* WithMemoryPool(const std::shared_ptr& pool) override { + pool_ = pool; + inner_->WithMemoryPool(pool); + return this; + } + + ReaderBuilder* WithCache(const std::shared_ptr& cache) override { + inner_->WithCache(cache); + return this; + } + + ReaderBuilder* WithReadHints(const std::optional& hints) override { + inner_->WithReadHints(hints); + return this; + } + + Result> Build( + const std::shared_ptr& stream) const override { + PAIMON_ASSIGN_OR_RAISE(std::unique_ptr base, inner_->Build(stream)); + auto* prefetch = dynamic_cast(base.get()); + if (prefetch == nullptr) { + return Status::Invalid("Late materialization requires prefetch interface"); + } + base.release(); + PAIMON_ASSIGN_OR_RAISE( + std::unique_ptr reader, + LateMaterializingFileBatchReader::Create( + std::unique_ptr(prefetch), pool_)); + return std::unique_ptr(std::move(reader)); + } + + private: + std::unique_ptr inner_; + std::shared_ptr pool_; +}; + +} // namespace paimon From 1ab58053e906934891583707f88aa9254f6b0208 Mon Sep 17 00:00:00 2001 From: zhouhongfeng Date: Mon, 24 Aug 2026 15:52:59 +0800 Subject: [PATCH 13/34] feat: install LM reader below PrefetchFileBatchReader --- .../late_materializing_file_batch_reader.cpp | 12 ++++++++++++ .../reader/late_materializing_file_batch_reader.h | 4 +++- .../reader/late_materializing_reader_builder.h | 7 +++---- src/paimon/core/operation/abstract_split_read.cpp | 15 ++++++++++----- src/paimon/core/operation/abstract_split_read.h | 4 ++-- 5 files changed, 30 insertions(+), 12 deletions(-) diff --git a/src/paimon/common/reader/late_materializing_file_batch_reader.cpp b/src/paimon/common/reader/late_materializing_file_batch_reader.cpp index b7cb71aed..b7a157db6 100644 --- a/src/paimon/common/reader/late_materializing_file_batch_reader.cpp +++ b/src/paimon/common/reader/late_materializing_file_batch_reader.cpp @@ -59,6 +59,11 @@ arrow::MemoryPool* LateMaterializingFileBatchReader::ArrowPool() const { } Result LateMaterializingFileBatchReader::NextBatch() { + 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()) { @@ -296,6 +301,13 @@ Status LateMaterializingFileBatchReader::SetReadSchema( if (!probe_fields.empty() && !payload_fields.empty()) { probe_schema_ = arrow::schema(probe_fields, full_schema_->metadata()); payload_schema_ = arrow::schema(payload_fields, full_schema_->metadata()); + // Guard: probing must be able to bind the predicate to the probe projection. If + // binding fails (unexpected predicate/schema shapes, e.g. after a wrapper translated + // the schema), stay on the plain passthrough path instead of failing mid-read. + if (!BindProbeFilter().ok()) { + probe_schema_.reset(); + payload_schema_.reset(); + } } } diff --git a/src/paimon/common/reader/late_materializing_file_batch_reader.h b/src/paimon/common/reader/late_materializing_file_batch_reader.h index 9c77cf407..75a5ed207 100644 --- a/src/paimon/common/reader/late_materializing_file_batch_reader.h +++ b/src/paimon/common/reader/late_materializing_file_batch_reader.h @@ -33,7 +33,9 @@ namespace paimon { class PredicateFilter; // For convenience, we abbreviate `Later Materializing` as `LatMat`. -// TODO(zhouhongfeng.zhf): add this reader to the split read path. +// 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. class LateMaterializingFileBatchReader : public PrefetchFileBatchReader { public: static Result> Create( diff --git a/src/paimon/common/reader/late_materializing_reader_builder.h b/src/paimon/common/reader/late_materializing_reader_builder.h index 07317924d..0420a35cc 100644 --- a/src/paimon/common/reader/late_materializing_reader_builder.h +++ b/src/paimon/common/reader/late_materializing_reader_builder.h @@ -60,10 +60,9 @@ class LateMaterializingReaderBuilder : public ReaderBuilder { return Status::Invalid("Late materialization requires prefetch interface"); } base.release(); - PAIMON_ASSIGN_OR_RAISE( - std::unique_ptr reader, - LateMaterializingFileBatchReader::Create( - std::unique_ptr(prefetch), pool_)); + PAIMON_ASSIGN_OR_RAISE(std::unique_ptr reader, + LateMaterializingFileBatchReader::Create( + std::unique_ptr(prefetch), pool_)); return std::unique_ptr(std::move(reader)); } diff --git a/src/paimon/core/operation/abstract_split_read.cpp b/src/paimon/core/operation/abstract_split_read.cpp index bb82c5d82..de0b05582 100644 --- a/src/paimon/core/operation/abstract_split_read.cpp +++ b/src/paimon/core/operation/abstract_split_read.cpp @@ -34,6 +34,7 @@ #include "paimon/common/data/variant/variant_shredding_read_plan_factory.h" #include "paimon/common/data/variant/variant_type_utils.h" #include "paimon/common/reader/delegating_prefetch_reader.h" +#include "paimon/common/reader/late_materializing_reader_builder.h" #include "paimon/common/reader/predicate_batch_reader.h" #include "paimon/common/reader/prefetch_file_batch_reader_impl.h" #include "paimon/common/table/special_fields.h" @@ -97,7 +98,7 @@ Result>> AbstractSplitRead::CreateR PrepareReaderBuilder(data_file_identifier, extra_format_options)); PAIMON_ASSIGN_OR_RAISE( std::unique_ptr file_reader, - CreateFieldMappingReader(data_file_path, file, partition, reader_builder.get(), + CreateFieldMappingReader(data_file_path, file, partition, std::move(reader_builder), field_mapping_builder.get(), dv_factory, row_ranges, data_file_path_factory)); if (file_reader) { @@ -152,13 +153,17 @@ Result> AbstractSplitRead::PrepareReaderBuilder( Result> AbstractSplitRead::CreateFileBatchReader( const std::string& file_format_identifier, const std::string& data_file_path, - int64_t data_file_size, const ReaderBuilder* reader_builder) const { + int64_t data_file_size, std::unique_ptr reader_builder) const { if (context_->EnablePrefetch() && file_format_identifier != "blob" && file_format_identifier != "avro") { + // Wrap the format builder so each parallel reader under the prefetch layer performs + // probe/payload two-phase reads when a predicate is pushed down; without a predicate + // the late-materializing reader degrades to a plain passthrough. + LateMaterializingReaderBuilder lm_builder(std::move(reader_builder), pool_); PAIMON_ASSIGN_OR_RAISE( std::unique_ptr prefetch_reader, PrefetchFileBatchReaderImpl::Create( - data_file_path, data_file_size, reader_builder, options_.GetFileSystem(), + data_file_path, data_file_size, &lm_builder, options_.GetFileSystem(), context_->GetPrefetchMaxParallelNum(), options_.GetReadBatchSize(), context_->GetPrefetchBatchCount(), options_.EnableAdaptivePrefetchStrategy(), executor_, @@ -175,7 +180,7 @@ Result> AbstractSplitRead::CreateFileBatchReade Result> AbstractSplitRead::CreateFieldMappingReader( const std::string& data_file_path, const std::shared_ptr& file_meta, - const BinaryRow& partition, const ReaderBuilder* reader_builder, + const BinaryRow& partition, std::unique_ptr reader_builder, const FieldMappingBuilder* field_mapping_builder, DeletionVector::Factory dv_factory, const std::optional>& row_ranges, const std::shared_ptr& data_file_path_factory) const { @@ -215,7 +220,7 @@ Result> AbstractSplitRead::CreateFieldMappingRe PAIMON_ASSIGN_OR_RAISE(std::string file_format_identifier, file_meta->FileFormat()); PAIMON_ASSIGN_OR_RAISE(std::unique_ptr file_reader, CreateFileBatchReader(file_format_identifier, data_file_path, - file_meta->file_size, reader_builder)); + file_meta->file_size, std::move(reader_builder))); if (VectorFileBatchReader::ContainsVector(read_schema)) { file_reader = std::make_unique(std::move(file_reader), pool_); } diff --git a/src/paimon/core/operation/abstract_split_read.h b/src/paimon/core/operation/abstract_split_read.h index 27349fec1..612936c23 100644 --- a/src/paimon/core/operation/abstract_split_read.h +++ b/src/paimon/core/operation/abstract_split_read.h @@ -107,12 +107,12 @@ class AbstractSplitRead : public SplitRead { Result> CreateFileBatchReader( const std::string& file_format_identifier, const std::string& data_file_path, - int64_t data_file_size, const ReaderBuilder* reader_builder) const; + int64_t data_file_size, std::unique_ptr reader_builder) const; // return nullptr if data file is skipped by index or dv Result> CreateFieldMappingReader( const std::string& data_file_path, const std::shared_ptr& file_meta, - const BinaryRow& partition, const ReaderBuilder* reader_builder, + const BinaryRow& partition, std::unique_ptr reader_builder, const FieldMappingBuilder* field_mapping_builder, DeletionVector::Factory dv_factory, const std::optional>& row_ranges, const std::shared_ptr& data_file_path_factory) const; From 913f222016b4d14ea9054b0c7575ed1f560e624c Mon Sep 17 00:00:00 2001 From: zhouhongfeng Date: Mon, 24 Aug 2026 17:35:18 +0800 Subject: [PATCH 14/34] fix: normalize arrow array after slicing --- .../reader/late_materializing_file_batch_reader.cpp | 10 +++------- src/paimon/common/utils/arrow/arrow_utils.cpp | 11 ++++++++--- src/paimon/common/utils/arrow/arrow_utils.h | 3 +++ 3 files changed, 14 insertions(+), 10 deletions(-) diff --git a/src/paimon/common/reader/late_materializing_file_batch_reader.cpp b/src/paimon/common/reader/late_materializing_file_batch_reader.cpp index b7a157db6..75bc87376 100644 --- a/src/paimon/common/reader/late_materializing_file_batch_reader.cpp +++ b/src/paimon/common/reader/late_materializing_file_batch_reader.cpp @@ -36,6 +36,7 @@ #include "paimon/common/reader/reader_utils.h" #include "paimon/common/utils/arrow/mem_utils.h" #include "paimon/common/utils/arrow/status_utils.h" +#include "paimon/common/utils/arrow/arrow_utils.h" #include "paimon/predicate/predicate_utils.h" #include "paimon/status.h" @@ -213,6 +214,7 @@ Result LateMaterializingFileBatchReader::ReadPayload probe_cursor_, card, probe_data_->length())); } std::shared_ptr probe_selected = probe_data_->Slice(probe_cursor_, card); + PAIMON_ASSIGN_OR_RAISE(probe_selected, ArrowUtils::NormalizeArrayOffsets(probe_selected, arrow_pool_.get())); probe_cursor_ += card; PAIMON_ASSIGN_OR_RAISE(FileBatchReader::ReadBatch assembled, @@ -237,6 +239,7 @@ Result LateMaterializingFileBatchReader::AssembleFul return Status::Invalid( fmt::format("field {} missing in both payload and probe columns", field->name())); } + PAIMON_ASSIGN_OR_RAISE(col, ArrowUtils::NormalizeArrayOffsets(col, arrow_pool_.get())); children.push_back(std::move(col)); } PAIMON_ASSIGN_OR_RAISE_FROM_ARROW(std::shared_ptr full_struct, @@ -301,13 +304,6 @@ Status LateMaterializingFileBatchReader::SetReadSchema( if (!probe_fields.empty() && !payload_fields.empty()) { probe_schema_ = arrow::schema(probe_fields, full_schema_->metadata()); payload_schema_ = arrow::schema(payload_fields, full_schema_->metadata()); - // Guard: probing must be able to bind the predicate to the probe projection. If - // binding fails (unexpected predicate/schema shapes, e.g. after a wrapper translated - // the schema), stay on the plain passthrough path instead of failing mid-read. - if (!BindProbeFilter().ok()) { - probe_schema_.reset(); - payload_schema_.reset(); - } } } diff --git a/src/paimon/common/utils/arrow/arrow_utils.cpp b/src/paimon/common/utils/arrow/arrow_utils.cpp index f29e1d11e..6cbc62afe 100644 --- a/src/paimon/common/utils/arrow/arrow_utils.cpp +++ b/src/paimon/common/utils/arrow/arrow_utils.cpp @@ -449,9 +449,7 @@ Result> ArrowUtils::NormalizeRecordBatchOffs if (normalized_columns.empty()) { normalized_columns = record_batch->columns(); } - PAIMON_ASSIGN_OR_RAISE(std::shared_ptr normalized_data, - RebaseToZeroOffset(column->data(), pool)); - normalized_columns[i] = arrow::MakeArray(normalized_data); + PAIMON_ASSIGN_OR_RAISE(normalized_columns[i], NormalizeArrayOffsets(column, pool)); } if (normalized_columns.empty()) { return record_batch; @@ -460,6 +458,13 @@ Result> ArrowUtils::NormalizeRecordBatchOffs std::move(normalized_columns)); } +Result> ArrowUtils::NormalizeArrayOffsets( + const std::shared_ptr& array, arrow::MemoryPool* pool) { + PAIMON_ASSIGN_OR_RAISE(std::shared_ptr normalized_data, + RebaseToZeroOffset(array->data(), pool)); + return arrow::MakeArray(normalized_data); +} + Result ArrowUtils::GetCompressionType(const std::string& compression) { std::string normalized = StringUtils::ToLowerCase(compression); if (normalized.empty() || normalized == "none") { diff --git a/src/paimon/common/utils/arrow/arrow_utils.h b/src/paimon/common/utils/arrow/arrow_utils.h index 326b3889e..13bd81549 100644 --- a/src/paimon/common/utils/arrow/arrow_utils.h +++ b/src/paimon/common/utils/arrow/arrow_utils.h @@ -57,6 +57,9 @@ class PAIMON_EXPORT ArrowUtils { static Result> NormalizeRecordBatchOffsets( const std::shared_ptr& record_batch, arrow::MemoryPool* pool); + static Result> NormalizeArrayOffsets( + const std::shared_ptr& array, arrow::MemoryPool* pool); + static bool EqualsIgnoreNullable(const std::shared_ptr& type, const std::shared_ptr& other_type); From 19bb1069bfb3fab3c661a90ce6699f58d7a82102 Mon Sep 17 00:00:00 2001 From: zhouhongfeng Date: Mon, 24 Aug 2026 17:56:02 +0800 Subject: [PATCH 15/34] feat: enable context options for LatMat --- include/paimon/read_context.h | 17 ++++++++++++++++- .../core/operation/abstract_split_read.cpp | 9 ++++----- .../core/operation/internal_read_context.h | 3 +++ src/paimon/core/operation/read_context.cpp | 10 ++++++++++ 4 files changed, 33 insertions(+), 6 deletions(-) diff --git a/include/paimon/read_context.h b/include/paimon/read_context.h index 3e58b1c45..0d6ea58d5 100644 --- a/include/paimon/read_context.h +++ b/include/paimon/read_context.h @@ -51,7 +51,7 @@ class PAIMON_EXPORT ReadContext { const std::vector& read_field_names, const std::vector& read_field_ids, const std::shared_ptr& predicate, bool enable_predicate_filter, - bool enable_prefetch, uint32_t prefetch_batch_count, + bool enable_prefetch, bool enable_late_materializing, uint32_t prefetch_batch_count, uint32_t prefetch_max_parallel_num, bool enable_multi_thread_row_to_batch, uint32_t row_to_batch_thread_number, const std::optional& table_schema, const std::shared_ptr& memory_pool, @@ -97,6 +97,11 @@ class PAIMON_EXPORT ReadContext { bool EnablePrefetch() const { return enable_prefetch_; } + /// Whether late materialization (probe/payload two-phase reads) is enabled for the + /// prefetch read path. Defaults to false. + bool EnableLateMaterializing() const { + return enable_late_materializing_; + } uint32_t GetPrefetchBatchCount() const { return prefetch_batch_count_; } @@ -163,6 +168,7 @@ class PAIMON_EXPORT ReadContext { std::shared_ptr predicate_; bool enable_predicate_filter_; bool enable_prefetch_; + bool enable_late_materializing_; uint32_t prefetch_batch_count_; uint32_t prefetch_max_parallel_num_; bool enable_multi_thread_row_to_batch_; @@ -306,6 +312,15 @@ class PAIMON_EXPORT ReadContextBuilder { /// @return Reference to this builder for method chaining. ReadContextBuilder& EnablePrefetch(bool enabled); + /// Enable or disable late materialization (probe/payload two-phase reads) for the prefetch + /// read path. When enabled, each parallel reader under the prefetch layer performs a probe + /// read of predicate columns first and only materializes payload columns for matched rows. + /// @param enabled Whether to enable late materialization (default: false) + /// @return Reference to this builder for method chaining. + /// @note Only takes effect when prefetch is enabled; without a pushed-down predicate the + /// late-materializing reader degrades to a plain passthrough. + ReadContextBuilder& EnableLateMaterializing(bool enabled); + /// Enable or disable the read-ahead cache for read operations. /// /// A read-ahead cache is used to prebuffer data ranges before they are needed, diff --git a/src/paimon/core/operation/abstract_split_read.cpp b/src/paimon/core/operation/abstract_split_read.cpp index de0b05582..a7efa9b19 100644 --- a/src/paimon/core/operation/abstract_split_read.cpp +++ b/src/paimon/core/operation/abstract_split_read.cpp @@ -156,14 +156,13 @@ Result> AbstractSplitRead::CreateFileBatchReade int64_t data_file_size, std::unique_ptr reader_builder) const { if (context_->EnablePrefetch() && file_format_identifier != "blob" && file_format_identifier != "avro") { - // Wrap the format builder so each parallel reader under the prefetch layer performs - // probe/payload two-phase reads when a predicate is pushed down; without a predicate - // the late-materializing reader degrades to a plain passthrough. - LateMaterializingReaderBuilder lm_builder(std::move(reader_builder), pool_); + if (context_->EnableLateMaterializing()) { + reader_builder = std::make_unique(std::move(reader_builder), pool_); + } PAIMON_ASSIGN_OR_RAISE( std::unique_ptr prefetch_reader, PrefetchFileBatchReaderImpl::Create( - data_file_path, data_file_size, &lm_builder, options_.GetFileSystem(), + data_file_path, data_file_size, reader_builder.get(), options_.GetFileSystem(), context_->GetPrefetchMaxParallelNum(), options_.GetReadBatchSize(), context_->GetPrefetchBatchCount(), options_.EnableAdaptivePrefetchStrategy(), executor_, diff --git a/src/paimon/core/operation/internal_read_context.h b/src/paimon/core/operation/internal_read_context.h index 8ef9f2d26..8e773cdcd 100644 --- a/src/paimon/core/operation/internal_read_context.h +++ b/src/paimon/core/operation/internal_read_context.h @@ -74,6 +74,9 @@ class InternalReadContext { bool EnablePrefetch() const { return read_context_->EnablePrefetch(); } + bool EnableLateMaterializing() const { + return read_context_->EnableLateMaterializing(); + } uint32_t GetPrefetchBatchCount() const { return read_context_->GetPrefetchBatchCount(); } diff --git a/src/paimon/core/operation/read_context.cpp b/src/paimon/core/operation/read_context.cpp index 08a854d84..11eadddef 100644 --- a/src/paimon/core/operation/read_context.cpp +++ b/src/paimon/core/operation/read_context.cpp @@ -35,6 +35,7 @@ ReadContext::ReadContext( const std::string& path, const std::string& branch, const std::vector& read_field_names, const std::vector& read_field_ids, const std::shared_ptr& predicate, bool enable_predicate_filter, bool enable_prefetch, + bool enable_late_materializing, uint32_t prefetch_batch_count, uint32_t prefetch_max_parallel_num, bool enable_multi_thread_row_to_batch, uint32_t row_to_batch_thread_number, const std::optional& table_schema, const std::shared_ptr& memory_pool, @@ -51,6 +52,7 @@ ReadContext::ReadContext( predicate_(predicate), enable_predicate_filter_(enable_predicate_filter), enable_prefetch_(enable_prefetch), + enable_late_materializing_(enable_late_materializing), prefetch_batch_count_(prefetch_batch_count), prefetch_max_parallel_num_(prefetch_max_parallel_num), enable_multi_thread_row_to_batch_(enable_multi_thread_row_to_batch), @@ -97,6 +99,7 @@ class ReadContextBuilder::Impl { predicate_.reset(); enable_predicate_filter_ = false; enable_prefetch_ = false; + enable_late_materializing_ = false; read_ahead_cache_enabled_ = true; prefetch_batch_count_ = 600; prefetch_max_parallel_num_ = 3; @@ -122,6 +125,7 @@ class ReadContextBuilder::Impl { std::shared_ptr predicate_; bool enable_predicate_filter_ = false; bool enable_prefetch_ = false; + bool enable_late_materializing_ = false; uint32_t prefetch_batch_count_ = 600; uint32_t prefetch_max_parallel_num_ = 3; bool enable_multi_thread_row_to_batch_ = false; @@ -191,6 +195,11 @@ ReadContextBuilder& ReadContextBuilder::EnablePrefetch(bool enabled) { return *this; } +ReadContextBuilder& ReadContextBuilder::EnableLateMaterializing(bool enabled) { + impl_->enable_late_materializing_ = enabled; + return *this; +} + ReadContextBuilder& ReadContextBuilder::SetPrefetchBatchCount(uint32_t batch_count) { impl_->prefetch_batch_count_ = batch_count; return *this; @@ -297,6 +306,7 @@ Result> ReadContextBuilder::Finish() { auto ctx = std::make_unique( impl_->path_, impl_->branch_, impl_->read_field_names_, impl_->read_field_ids_, impl_->predicate_, impl_->enable_predicate_filter_, impl_->enable_prefetch_, + impl_->enable_late_materializing_, impl_->prefetch_batch_count_, impl_->prefetch_max_parallel_num_, impl_->enable_multi_thread_row_to_batch_, impl_->row_to_batch_thread_number_, impl_->table_schema_, impl_->memory_pool_, impl_->executor_, impl_->specific_file_system_, From 7f44acd393a060e01eea96f30d584d641053b9d0 Mon Sep 17 00:00:00 2001 From: zhouhongfeng Date: Mon, 24 Aug 2026 18:07:23 +0800 Subject: [PATCH 16/34] pre-commit --- .../late_materializing_file_batch_reader.cpp | 5 +++-- src/paimon/common/utils/arrow/arrow_utils.cpp | 2 +- .../core/operation/abstract_split_read.cpp | 3 ++- src/paimon/core/operation/read_context.cpp | 21 +++++++++---------- 4 files changed, 16 insertions(+), 15 deletions(-) diff --git a/src/paimon/common/reader/late_materializing_file_batch_reader.cpp b/src/paimon/common/reader/late_materializing_file_batch_reader.cpp index 75bc87376..1b0d03c85 100644 --- a/src/paimon/common/reader/late_materializing_file_batch_reader.cpp +++ b/src/paimon/common/reader/late_materializing_file_batch_reader.cpp @@ -34,9 +34,9 @@ #include "fmt/format.h" #include "paimon/common/predicate/predicate_filter.h" #include "paimon/common/reader/reader_utils.h" +#include "paimon/common/utils/arrow/arrow_utils.h" #include "paimon/common/utils/arrow/mem_utils.h" #include "paimon/common/utils/arrow/status_utils.h" -#include "paimon/common/utils/arrow/arrow_utils.h" #include "paimon/predicate/predicate_utils.h" #include "paimon/status.h" @@ -214,7 +214,8 @@ Result LateMaterializingFileBatchReader::ReadPayload probe_cursor_, card, probe_data_->length())); } std::shared_ptr probe_selected = probe_data_->Slice(probe_cursor_, card); - PAIMON_ASSIGN_OR_RAISE(probe_selected, ArrowUtils::NormalizeArrayOffsets(probe_selected, arrow_pool_.get())); + PAIMON_ASSIGN_OR_RAISE( + probe_selected, ArrowUtils::NormalizeArrayOffsets(probe_selected, arrow_pool_.get())); probe_cursor_ += card; PAIMON_ASSIGN_OR_RAISE(FileBatchReader::ReadBatch assembled, diff --git a/src/paimon/common/utils/arrow/arrow_utils.cpp b/src/paimon/common/utils/arrow/arrow_utils.cpp index 6cbc62afe..c6a07b2bf 100644 --- a/src/paimon/common/utils/arrow/arrow_utils.cpp +++ b/src/paimon/common/utils/arrow/arrow_utils.cpp @@ -461,7 +461,7 @@ Result> ArrowUtils::NormalizeRecordBatchOffs Result> ArrowUtils::NormalizeArrayOffsets( const std::shared_ptr& array, arrow::MemoryPool* pool) { PAIMON_ASSIGN_OR_RAISE(std::shared_ptr normalized_data, - RebaseToZeroOffset(array->data(), pool)); + RebaseToZeroOffset(array->data(), pool)); return arrow::MakeArray(normalized_data); } diff --git a/src/paimon/core/operation/abstract_split_read.cpp b/src/paimon/core/operation/abstract_split_read.cpp index a7efa9b19..c8d33c9f2 100644 --- a/src/paimon/core/operation/abstract_split_read.cpp +++ b/src/paimon/core/operation/abstract_split_read.cpp @@ -157,7 +157,8 @@ Result> AbstractSplitRead::CreateFileBatchReade if (context_->EnablePrefetch() && file_format_identifier != "blob" && file_format_identifier != "avro") { if (context_->EnableLateMaterializing()) { - reader_builder = std::make_unique(std::move(reader_builder), pool_); + reader_builder = + std::make_unique(std::move(reader_builder), pool_); } PAIMON_ASSIGN_OR_RAISE( std::unique_ptr prefetch_reader, diff --git a/src/paimon/core/operation/read_context.cpp b/src/paimon/core/operation/read_context.cpp index 11eadddef..deacfa78b 100644 --- a/src/paimon/core/operation/read_context.cpp +++ b/src/paimon/core/operation/read_context.cpp @@ -35,11 +35,10 @@ ReadContext::ReadContext( const std::string& path, const std::string& branch, const std::vector& read_field_names, const std::vector& read_field_ids, const std::shared_ptr& predicate, bool enable_predicate_filter, bool enable_prefetch, - bool enable_late_materializing, - uint32_t prefetch_batch_count, uint32_t prefetch_max_parallel_num, - bool enable_multi_thread_row_to_batch, uint32_t row_to_batch_thread_number, - const std::optional& table_schema, const std::shared_ptr& memory_pool, - const std::shared_ptr& executor, + bool enable_late_materializing, uint32_t prefetch_batch_count, + uint32_t prefetch_max_parallel_num, bool enable_multi_thread_row_to_batch, + uint32_t row_to_batch_thread_number, const std::optional& table_schema, + const std::shared_ptr& memory_pool, const std::shared_ptr& executor, const std::shared_ptr& specific_file_system, const std::map& fs_scheme_to_identifier_map, const std::shared_ptr& realtime_context, @@ -306,12 +305,12 @@ Result> ReadContextBuilder::Finish() { auto ctx = std::make_unique( impl_->path_, impl_->branch_, impl_->read_field_names_, impl_->read_field_ids_, impl_->predicate_, impl_->enable_predicate_filter_, impl_->enable_prefetch_, - impl_->enable_late_materializing_, - impl_->prefetch_batch_count_, impl_->prefetch_max_parallel_num_, - impl_->enable_multi_thread_row_to_batch_, impl_->row_to_batch_thread_number_, - impl_->table_schema_, impl_->memory_pool_, impl_->executor_, impl_->specific_file_system_, - impl_->fs_scheme_to_identifier_map_, impl_->realtime_context_, impl_->options_, - impl_->read_ahead_cache_enabled_, impl_->cache_config_, impl_->cache_); + impl_->enable_late_materializing_, impl_->prefetch_batch_count_, + impl_->prefetch_max_parallel_num_, impl_->enable_multi_thread_row_to_batch_, + impl_->row_to_batch_thread_number_, impl_->table_schema_, impl_->memory_pool_, + impl_->executor_, impl_->specific_file_system_, impl_->fs_scheme_to_identifier_map_, + impl_->realtime_context_, impl_->options_, impl_->read_ahead_cache_enabled_, + impl_->cache_config_, impl_->cache_); if (impl_->read_schema_ && impl_->read_schema_->release) { ctx->SetReadSchema(std::move(impl_->read_schema_)); } From 95fe9584a03ab13030fe735af584095781a4cf36 Mon Sep 17 00:00:00 2001 From: zhouhongfeng Date: Tue, 25 Aug 2026 11:38:04 +0800 Subject: [PATCH 17/34] fix: forward PrefetchFileBatchReader::PreBufferRange to fix cache issues --- .../common/reader/late_materializing_file_batch_reader.h | 3 +++ 1 file changed, 3 insertions(+) diff --git a/src/paimon/common/reader/late_materializing_file_batch_reader.h b/src/paimon/common/reader/late_materializing_file_batch_reader.h index 75a5ed207..809dc6a36 100644 --- a/src/paimon/common/reader/late_materializing_file_batch_reader.h +++ b/src/paimon/common/reader/late_materializing_file_batch_reader.h @@ -59,6 +59,9 @@ class LateMaterializingFileBatchReader : public PrefetchFileBatchReader { bool* need_prefetch) const override; Status SetReadRanges(const std::vector>& read_ranges) override; + Result>> PreBufferRange() override { + return inner_->PreBufferRange(); + } private: explicit LateMaterializingFileBatchReader(std::unique_ptr inner, std::shared_ptr arrow_pool) From dc124e2f5665187608c5bb7710b2c6cf87dcf375 Mon Sep 17 00:00:00 2001 From: zhouhongfeng Date: Tue, 25 Aug 2026 15:14:38 +0800 Subject: [PATCH 18/34] fix: diable lat-mat in previout tests and add new tests for lat-mat --- .../operation/merge_file_split_read_test.cpp | 63 +++++++ test/inte/read_inte_test.cpp | 171 ++++++++++++++++++ test/inte/read_inte_with_index_test.cpp | 74 +++++++- test/inte/scan_and_read_inte_test.cpp | 101 ++++++++++- 4 files changed, 406 insertions(+), 3 deletions(-) diff --git a/src/paimon/core/operation/merge_file_split_read_test.cpp b/src/paimon/core/operation/merge_file_split_read_test.cpp index 911cf7961..ee689851e 100644 --- a/src/paimon/core/operation/merge_file_split_read_test.cpp +++ b/src/paimon/core/operation/merge_file_split_read_test.cpp @@ -99,6 +99,8 @@ class MergeFileSplitReadTest : public ::testing::Test, void AddOptions(ReadContextBuilder* context_builder) const { auto [use_min_heap, enable_io_prefetch, enable_multi_thread_row_to_batch] = GetParam(); + // disable late materializing by default + context_builder->EnableLateMaterializing(false); if (use_min_heap) { context_builder->AddOption(Options::SORT_ENGINE, "min-heap"); } else { @@ -842,6 +844,67 @@ TEST_P(MergeFileSplitReadTest, TestReadWithPredicate) { CheckResult(result_array, expected_array, read_schema); } +// Late materialization reads the predicate columns first and only materializes the remaining +// columns for matched rows. Combined with the top-level predicate filter the read path returns +// the exact user-predicate match set: in a section with several sorted runs only the key part of +// the predicate reaches the data files, and the merge result is filtered afterwards. +TEST_P(MergeFileSplitReadTest, TestReadWithPredicateAndLateMaterializing) { + std::string path = + paimon::test::GetDataDir() + "/parquet/pk_table_with_mor.db/pk_table_with_mor"; + ReadContextBuilder context_builder(path); + + std::vector raw_read_fields = {DataField(1, arrow::field("k1", arrow::int32())), + DataField(3, arrow::field("p1", arrow::int32())), + DataField(5, arrow::field("s1", arrow::utf8())), + DataField(4, arrow::field("s0", arrow::utf8())), + DataField(6, arrow::field("v0", arrow::float64())), + DataField(7, arrow::field("v1", arrow::boolean()))}; + auto read_schema = DataField::ConvertDataFieldsToArrowSchema(raw_read_fields); + ASSERT_TRUE(read_schema); + + context_builder.SetReadFieldNames({"k1", "p1", "s1", "s0", "v0", "v1"}); + context_builder.SetOptions({{Options::SEQUENCE_FIELD, "s0,s1"}, + {Options::MERGE_ENGINE, "deduplicate"}, + {Options::IGNORE_DELETE, "true"}}); + AddOptions(&context_builder); + context_builder.EnableLateMaterializing(true); + // key predicate, always pushed down into the data files + auto greater_or_equal = PredicateBuilder::GreaterOrEqual(/*field_index=*/0, /*field_name=*/"k1", + FieldType::INT, Literal(1)); + // value predicate, only pushed down when a section holds a single sorted run + auto greater_than = PredicateBuilder::GreaterThan(/*field_index=*/4, /*field_name=*/"v0", + FieldType::DOUBLE, Literal(12.0)); + ASSERT_OK_AND_ASSIGN(std::shared_ptr predicate_result, + PredicateBuilder::And({greater_or_equal, greater_than})); + context_builder.SetPredicate(predicate_result); + context_builder.EnablePredicateFilter(true).EnableLateMaterializing(true); + ASSERT_OK_AND_ASSIGN(std::shared_ptr read_context, context_builder.Finish()); + + auto internal_context = CreateInternalReadContext(read_context); + ASSERT_OK_AND_ASSIGN(auto batch_reader, CreateReader(internal_context, PrepareDataSplit())); + + ASSERT_OK_AND_ASSIGN(std::shared_ptr result_array, + ReadResultCollector::CollectResult(batch_reader.get())); + + auto fields_with_row_kind = read_schema->fields(); + fields_with_row_kind.insert(fields_with_row_kind.begin(), + arrow::field("_VALUE_KIND", arrow::int8())); + + // Only the merged rows with k1 >= 1 and v0 > 12.0 remain. + std::shared_ptr expected_array; + auto array_status = + arrow::ipc::internal::json::ChunkedArrayFromJSON(arrow::struct_(fields_with_row_kind), {R"([ + [0, 1, 0, "!", "driver", 13.3, false], + [0, 2, 0, "!", "driver", 13.3, false], + [0, 200, 0, "number", "max", 140.4, false], + [0, 1, 1, "you", "zoo", 130.0, false] + + ])"}, + &expected_array); + ASSERT_TRUE(array_status.ok()); + CheckResult(result_array, expected_array, read_schema); +} + TEST_P(MergeFileSplitReadTest, TestReadWithAlterTable) { std::string path = paimon::test::GetDataDir() + "/parquet/pk_table_with_mor.db/pk_table_with_mor"; diff --git a/test/inte/read_inte_test.cpp b/test/inte/read_inte_test.cpp index 1e0952e08..99d9f0983 100644 --- a/test/inte/read_inte_test.cpp +++ b/test/inte/read_inte_test.cpp @@ -2209,6 +2209,7 @@ TEST_P(ReadInteTest, TestAppendReadWithPredicateOnlyPushdown) { .AddOption("test.enable-adaptive-prefetch-strategy", param.enable_adaptive_prefetch_strategy) .SetPredicate(predicate) + .EnableLateMaterializing(false) .EnablePrefetch(param.enable_prefetch); ASSERT_OK_AND_ASSIGN(auto read_context, context_builder.Finish()); @@ -2265,6 +2266,90 @@ TEST_P(ReadInteTest, TestAppendReadWithPredicateOnlyPushdown) { ASSERT_TRUE(result_array->Equals(*expected_array)) << result_array->ToString(); } +// Late materialization reads the predicate columns first and only materializes the remaining +// columns for matched rows. Combined with the top-level predicate filter, the read path returns +// the exact user-predicate match set. +TEST_P(ReadInteTest, TestAppendReadWithLateMaterializing) { + std::vector read_fields = {DataField(3, arrow::field("f3", arrow::float64())), + DataField(0, arrow::field("f0", arrow::utf8())), + DataField(1, arrow::field("f1", arrow::int32()))}; + ASSERT_OK_AND_ASSIGN( + auto predicate, + PredicateBuilder::Or( + {PredicateBuilder::GreaterThan(/*field_index=*/0, /*field_name=*/"f3", + FieldType::DOUBLE, Literal(static_cast(15.0))), + PredicateBuilder::IsNull(/*field_index=*/0, /*field_name=*/"f3", FieldType::DOUBLE)})); + + auto param = GetParam(); + std::string path = + paimon::test::GetDataDir() + "/" + param.file_format + "/append_09.db/append_09"; + + ReadContextBuilder context_builder(path); + context_builder.SetReadAheadCacheEnabled(param.read_ahead_cache_enabled); + context_builder.SetReadFieldNames({"f3", "f0", "f1"}); + context_builder.AddOption(Options::FILE_FORMAT, param.file_format) + .AddOption("read.batch-size", "2") + .AddOption("test.enable-adaptive-prefetch-strategy", + param.enable_adaptive_prefetch_strategy) + .SetPredicate(predicate) + .EnablePredicateFilter(true) + .EnableLateMaterializing(true) + .EnablePrefetch(param.enable_prefetch); + ASSERT_OK_AND_ASSIGN(auto read_context, context_builder.Finish()); + ASSERT_OK_AND_ASSIGN(auto table_read, TableRead::Create(std::move(read_context))); + + std::vector file_list_0; + std::vector file_list_1; + std::vector file_list_2; + if (param.file_format == "orc") { + file_list_0 = {"data-d41fd7d1-b3e4-4905-aad9-b20a780e90a2-0.orc"}; + file_list_1 = {"data-4e30d6c0-f109-4300-a010-4ba03047dd9d-0.orc", + "data-10b9eea8-241d-4e4b-8ab8-2a82d72d79a2-0.orc", + "data-e2bb59ee-ae25-4e5b-9bcc-257250bc5fdd-0.orc", + "data-2d5ea1ea-77c1-47ff-bb87-19a509962a37-0.orc"}; + file_list_2 = {"data-db2b44c0-0d73-449d-82a0-4075bd2cb6e3-0.orc", + "data-b913a160-a4d1-4084-af2a-18333c35668e-0.orc"}; + } else if (param.file_format == "parquet") { + file_list_0 = {"data-46e27d5b-4850-4d1e-abb6-b3aabbbc08cb-0.parquet"}; + file_list_1 = {"data-864a052b-a938-4e04-b32c-6c72699a0c92-0.parquet", + "data-c0401350-64a3-4a54-a143-dd125ad9a8e5-0.parquet", + "data-7a912f84-04b7-4bbb-8dc6-53f4a292ea25-0.parquet", + "data-bb891df7-ea12-4b7e-9017-41aabe08c8ec-0.parquet"}; + file_list_2 = {"data-b446f78a-2cfb-4b3b-add8-31295d24a277-0.parquet", + "data-fd72a479-53ae-42f7-aec0-e982ee555928-0.parquet"}; + } + + DataSplitsSimple input_data_splits = { + {paimon::test::GetDataDir() + "/" + param.file_format + + "/append_09.db/append_09/f1=10/bucket-0", + BinaryRowGenerator::GenerateRow({10}, pool_.get()), file_list_0}, + {paimon::test::GetDataDir() + "/" + param.file_format + + "/append_09.db/append_09/f1=10/bucket-1", + BinaryRowGenerator::GenerateRow({10}, pool_.get()), file_list_1}, + {paimon::test::GetDataDir() + "/" + param.file_format + + "/append_09.db/append_09/f1=20/bucket-0", + BinaryRowGenerator::GenerateRow({20}, pool_.get()), file_list_2}}; + + auto data_splits = CreateDataSplits(input_data_splits, /*snapshot_id=*/4); + ASSERT_OK_AND_ASSIGN(auto batch_reader, table_read->CreateReader(data_splits)); + ASSERT_OK_AND_ASSIGN(auto result_array, ReadResultCollector::CollectResult(batch_reader.get())); + + auto fields_with_row_kind = read_fields; + fields_with_row_kind.insert(fields_with_row_kind.begin(), SpecialFields::ValueKind()); + std::shared_ptr arrow_data_type = + DataField::ConvertDataFieldsToArrowStructType(fields_with_row_kind); + + // "Bob" (f3 = 12.1) is the only row that does not match the predicate. + std::shared_ptr expected_array; + auto array_status = arrow::ipc::internal::json::ChunkedArrayFromJSON(arrow_data_type, {R"([ + [0, 15.1, "Emily", 10], [0, 16.1, "Alex", 10], [0, 17.1, "David", 10], + [0, 17.1, "Lily", 10], [0, null, "Paul", 20] + ])"}, + &expected_array); + ASSERT_TRUE(array_status.ok()); + ASSERT_TRUE(result_array->Equals(*expected_array)) << result_array->ToString(); +} + TEST_P(ReadInteTest, TestAppendReadWithPredicateAllFiltered) { std::vector read_fields = {DataField(3, arrow::field("f3", arrow::float64())), DataField(0, arrow::field("f0", arrow::utf8())), @@ -3109,6 +3194,7 @@ TEST_P(ReadInteTest, TestPkReadSnapshot6WithSchemaEvolutionWithPredicateOnlyPush context_builder.SetReadAheadCacheEnabled(param.read_ahead_cache_enabled); context_builder.SetPredicate(predicate); context_builder.EnablePrefetch(param.enable_prefetch) + .EnableLateMaterializing(false) .AddOption("test.enable-adaptive-prefetch-strategy", param.enable_adaptive_prefetch_strategy); ASSERT_OK_AND_ASSIGN(auto read_context, context_builder.Finish()); @@ -3166,6 +3252,91 @@ TEST_P(ReadInteTest, TestPkReadSnapshot6WithSchemaEvolutionWithPredicateOnlyPush ASSERT_TRUE(result_array->Equals(*expected_array)); } +// Same coverage for the primary-key path with schema evolution and a deletion vector: late +// materialization runs below the prefetch layer per data file, so the deletion vector and the +// schema-evolution field mapping keep working. +TEST_P(ReadInteTest, TestPkReadSnapshot6WithSchemaEvolutionWithLateMaterializing) { + std::vector read_fields = {DataField(1, arrow::field("key1", arrow::int32())), + DataField(7, arrow::field("k", arrow::utf8())), + DataField(2, arrow::field("key_2", arrow::int32())), + DataField(4, arrow::field("c", arrow::int32())), + DataField(8, arrow::field("d", arrow::int32())), + DataField(6, arrow::field("a", arrow::int32())), + DataField(0, arrow::field("key0", arrow::int32())), + DataField(9, arrow::field("e", arrow::int32()))}; + auto param = GetParam(); + std::string path = paimon::test::GetDataDir() + "/" + param.file_format + + "/pk_table_with_alter_table.db/pk_table_with_alter_table/"; + // equal is a partition filter and is not pushed into the data files; less_than is pushed down + // and only matches the column added by schema evolution, where the older files yield nulls. + auto equal = PredicateBuilder::Equal(/*field_index=*/6, /*field_name=*/"key0", FieldType::INT, + Literal(0)); + auto less_than = PredicateBuilder::LessThan(/*field_index=*/7, /*field_name=*/"e", + FieldType::INT, Literal(510)); + ASSERT_OK_AND_ASSIGN(auto predicate, PredicateBuilder::And({equal, less_than})); + + ReadContextBuilder context_builder(path); + context_builder.SetReadFieldNames({{"key1", "k", "key_2", "c", "d", "a", "key0", "e"}}); + context_builder.AddOption(Options::FILE_FORMAT, param.file_format) + .AddOption("read.batch-size", "2"); + context_builder.SetReadAheadCacheEnabled(param.read_ahead_cache_enabled); + context_builder.SetPredicate(predicate); + context_builder.EnablePredicateFilter(true) + .EnableLateMaterializing(true) + .EnablePrefetch(param.enable_prefetch) + .AddOption("test.enable-adaptive-prefetch-strategy", + param.enable_adaptive_prefetch_strategy); + ASSERT_OK_AND_ASSIGN(auto read_context, context_builder.Finish()); + ASSERT_OK_AND_ASSIGN(auto table_read, TableRead::Create(std::move(read_context))); + + std::vector file_list_0; + std::vector file_list_1; + std::string deletion_file; + if (param.file_format == "orc") { + file_list_0 = {"data-3842c1d6-6b34-4b2c-a648-9e95b4fb941b-0.orc", + "data-d6d370f3-242b-45c9-8739-44bf31b2b449-0.orc"}; + file_list_1 = {"data-7b538b91-5dbb-4e16-a639-1b5c0696db8c-0.orc"}; + deletion_file = "index-51804749-ed6c-4e7b-b3e9-337cfe38499c-1"; + } else if (param.file_format == "parquet") { + file_list_0 = {"data-8969384c-d715-4113-b663-2248c9a8c8d9-0.parquet", + "data-f2f38e80-7d28-4d51-90b3-c28951e5cdc0-0.parquet"}; + file_list_1 = {"data-d7a33230-223e-4d65-8e39-bc7ed26bdd32-0.parquet"}; + deletion_file = "index-c93829f3-1a72-4d88-8401-70663ce46426-1"; + } + + DataSplitsSchemaDv input_data_splits = { + {path + "key0=1/key1=1/bucket-0", + BinaryRowGenerator::GenerateRow({1, 1}, pool_.get()), + file_list_0, + /*schema ids*/ {0, 1}, + /*deletion file*/ + {DeletionFile(path + "index/" + deletion_file, + /*offset=*/1, /*length=*/26, /*cardinality=*/std::nullopt), + std::nullopt}}, + {path + "key0=0/key1=1/bucket-0", BinaryRowGenerator::GenerateRow({0, 1}, pool_.get()), + file_list_1, + /*schema ids*/ {1}, + /*deletion file*/ {std::nullopt}}}; + + auto data_splits = CreateDataSplits(input_data_splits, /*snapshot_id=*/6); + ASSERT_OK_AND_ASSIGN(auto batch_reader, table_read->CreateReader(data_splits)); + ASSERT_OK_AND_ASSIGN(auto result_array, ReadResultCollector::CollectResult(batch_reader.get())); + + auto fields_with_row_kind = read_fields; + fields_with_row_kind.insert(fields_with_row_kind.begin(), SpecialFields::ValueKind()); + std::shared_ptr arrow_data_type = + DataField::ConvertDataFieldsToArrowStructType(fields_with_row_kind); + + // "Paul" is the only row in partition key0 = 0 whose e is not null and matches e < 510. + std::shared_ptr expected_array; + auto array_status = arrow::ipc::internal::json::ChunkedArrayFromJSON(arrow_data_type, {R"([ + [0, 1, "Paul", 502, 504, 508, 506, 0, 509] +])"}, + &expected_array); + ASSERT_TRUE(array_status.ok()); + ASSERT_TRUE(result_array->Equals(*expected_array)) << result_array->ToString(); +} + TEST_P(ReadInteTest, TestPkReadSnapshot6WithSchemaEvolutionWithPredicateFilter) { std::vector read_fields = {DataField(1, arrow::field("key1", arrow::int32())), DataField(7, arrow::field("k", arrow::utf8())), diff --git a/test/inte/read_inte_with_index_test.cpp b/test/inte/read_inte_with_index_test.cpp index f1316c6ce..ac9add2c0 100644 --- a/test/inte/read_inte_with_index_test.cpp +++ b/test/inte/read_inte_with_index_test.cpp @@ -84,7 +84,8 @@ class ReadInteWithIndexTest : public testing::Test, ReadContextBuilder context_builder(table_path); context_builder.AddOption("read.batch-size", "2") .AddOption("test.enable-adaptive-prefetch-strategy", "false") - .SetPredicate(predicate); + .SetPredicate(predicate) + .EnableLateMaterializing(false); if (enable_prefetch) { context_builder.EnablePrefetch(true).SetPrefetchBatchCount(3); } @@ -1232,6 +1233,77 @@ TEST_P(ReadInteWithIndexTest, TestNoEmbeddingBitmapIndex) { CheckResultForBitmapWithSingleRowGroup(path, arrow_data_type, split); } +// Late materialization reads the predicate columns first and only materializes the remaining +// columns for matched rows. The bitmap index selection is pushed into the same reader, so +// combined with the top-level predicate filter the read path returns the exact match set. +TEST_P(ReadInteWithIndexTest, TestBitmapIndexWithLateMaterializing) { + auto [file_format, enable_prefetch] = GetParam(); + std::string path = GetDataDir() + "/" + file_format + + "/append_with_bitmap_no_embedding.db/append_with_bitmap_no_embedding/"; + std::string file_name; + if (file_format == "orc") { + file_name = "data-414509f5-e40c-4245-b992-bbf486778ac9-0.orc"; + } else if (file_format == "parquet") { + file_name = "data-783929b2-49d4-4006-a898-194a62e3278d-0.parquet"; + } + + std::vector read_fields = {SpecialFields::ValueKind(), + DataField(0, arrow::field("f0", arrow::utf8())), + DataField(1, arrow::field("f1", arrow::int32())), + DataField(2, arrow::field("f2", arrow::int32())), + DataField(3, arrow::field("f3", arrow::float64()))}; + std::shared_ptr arrow_data_type = + DataField::ConvertDataFieldsToArrowStructType(read_fields); + + auto data_file_meta = std::make_shared( + file_name, /*file_size=*/689, + /*row_count=*/8, /*min_key=*/BinaryRow::EmptyRow(), + /*max_key=*/BinaryRow::EmptyRow(), /*key_stats=*/SimpleStats::EmptyStats(), + /*value_stats=*/SimpleStats::EmptyStats(), /*min_sequence_number=*/0, + /*max_sequence_number=*/7, /*schema_id=*/0, + /*level=*/0, + /*extra_files=*/ + std::vector>({file_name + ".index"}), + /*creation_time=*/Timestamp(0ll, 0), /*delete_row_count=*/0, + /*embedded_index=*/nullptr, FileSource::Append(), + /*value_stats_cols=*/std::nullopt, + /*external_path=*/std::nullopt, /*first_row_id=*/std::nullopt, /*write_cols=*/std::nullopt); + DataSplitImpl::Builder builder(BinaryRow::EmptyRow(), /*bucket=*/0, + /*bucket_path=*/path + "bucket-0/", {data_file_meta}); + ASSERT_OK_AND_ASSIGN(auto split, + builder.WithSnapshot(1).IsStreaming(false).RawConvertible(true).Build()); + + std::string literal_str = "Bob"; + auto predicate = PredicateBuilder::Equal( + /*field_index=*/0, /*field_name=*/"f0", FieldType::STRING, + Literal(FieldType::STRING, literal_str.data(), literal_str.size())); + + ReadContextBuilder context_builder(path); + context_builder.AddOption("read.batch-size", "2") + .AddOption("test.enable-adaptive-prefetch-strategy", "false") + .SetPredicate(predicate) + .EnablePredicateFilter(true) + .EnableLateMaterializing(true); + if (enable_prefetch) { + context_builder.EnablePrefetch(true).SetPrefetchBatchCount(3); + } + ASSERT_OK_AND_ASSIGN(auto read_context, context_builder.Finish()); + ASSERT_OK_AND_ASSIGN(auto table_read, TableRead::Create(std::move(read_context))); + ASSERT_OK_AND_ASSIGN(auto batch_reader, + table_read->CreateReader(std::vector>{split})); + ASSERT_OK_AND_ASSIGN(auto result_array, ReadResultCollector::CollectResult(batch_reader.get())); + + // Only the two "Bob" rows match the predicate. + std::shared_ptr expected_array; + auto array_status = arrow::ipc::internal::json::ChunkedArrayFromJSON(arrow_data_type, {R"([ +[0, "Bob", 10, 1, 12.1], +[0, "Bob", 10, 1, 16.1] + ])"}, + &expected_array); + ASSERT_TRUE(array_status.ok()); + ASSERT_TRUE(result_array->Equals(*expected_array)) << result_array->ToString(); +} + TEST_P(ReadInteWithIndexTest, TestNoEmbeddingBitmapIndexWithExternalPath) { auto [file_format, enable_prefetch] = GetParam(); std::string path = GetDataDir() + "/" + file_format + diff --git a/test/inte/scan_and_read_inte_test.cpp b/test/inte/scan_and_read_inte_test.cpp index 96a7c9a19..4d3569086 100644 --- a/test/inte/scan_and_read_inte_test.cpp +++ b/test/inte/scan_and_read_inte_test.cpp @@ -724,7 +724,8 @@ TEST_P(ScanAndReadInteTest, TestWithPKWithDvBatchScanSnapshot6WithPredicate) { ReadContextBuilder read_context_builder(table_path); AddReadOptionsForPrefetch(&read_context_builder); - read_context_builder.SetPredicate(predicate); + read_context_builder.SetPredicate(predicate) + .EnableLateMaterializing(false); ASSERT_OK_AND_ASSIGN(auto read_context, read_context_builder.Finish()); ASSERT_OK_AND_ASSIGN(auto table_read, TableRead::Create(std::move(read_context))); @@ -744,6 +745,50 @@ TEST_P(ScanAndReadInteTest, TestWithPKWithDvBatchScanSnapshot6WithPredicate) { ASSERT_TRUE(expected->Equals(read_result)) << read_result->ToString(); } +// Late materialization reads the predicate columns first and only materializes the remaining +// columns for matched rows. Combined with the top-level predicate filter, the read path returns +// the exact user-predicate match set (deletion vectors stay applied per data file). +TEST_P(ScanAndReadInteTest, TestWithPKWithDvBatchScanSnapshot6WithLateMaterializing) { + auto file_format = FileFormat(); + std::string table_path = paimon::test::GetDataDir() + file_format + + "/pk_table_scan_and_read_dv.db/pk_table_scan_and_read_dv/"; + ScanContextBuilder scan_context_builder(table_path); + scan_context_builder.AddOption(Options::SCAN_SNAPSHOT_ID, "6"); + + std::string literal_str = "Alice"; + auto not_equal = PredicateBuilder::NotEqual( + /*field_index=*/0, /*field_name=*/"f0", FieldType::STRING, + Literal(FieldType::STRING, literal_str.data(), literal_str.size())); + auto greater_than = PredicateBuilder::GreaterThan(/*field_index=*/3, /*field_name=*/"f3", + FieldType::DOUBLE, Literal(18.0)); + ASSERT_OK_AND_ASSIGN(auto predicate, PredicateBuilder::And({not_equal, greater_than})); + scan_context_builder.SetPredicate(predicate); + ASSERT_OK_AND_ASSIGN(auto scan_context, FinishScanContext(scan_context_builder)); + ASSERT_OK_AND_ASSIGN(auto table_scan, TableScan::Create(std::move(scan_context))); + + ReadContextBuilder read_context_builder(table_path); + AddReadOptionsForPrefetch(&read_context_builder); + read_context_builder.SetPredicate(predicate) + .EnablePredicateFilter(true) + .EnableLateMaterializing(true); + ASSERT_OK_AND_ASSIGN(auto read_context, read_context_builder.Finish()); + ASSERT_OK_AND_ASSIGN(auto table_read, TableRead::Create(std::move(read_context))); + + ASSERT_OK_AND_ASSIGN(auto result_plan, table_scan->CreatePlan()); + ASSERT_EQ(result_plan->SnapshotId().value(), 6); + ASSERT_OK_AND_ASSIGN(auto batch_reader, table_read->CreateReader(result_plan->Splits())); + ASSERT_OK_AND_ASSIGN(auto read_result, ReadResultCollector::CollectResult(batch_reader.get())); + + // check result: "Lucy" (f3 = 14.1) does not match f3 > 18 and is filtered out. + auto expected = std::make_shared( + arrow::ipc::internal::json::ArrayFromJSON(arrow_data_type_, R"([ +[0, "Paul", 20, 1, 18.1] + ])") + .ValueOrDie()); + ASSERT_TRUE(expected); + ASSERT_TRUE(expected->Equals(read_result)) << read_result->ToString(); +} + TEST_P(ScanAndReadInteTest, TestWithPKWithDvBatchScanSnapshot4WithPredicate) { auto file_format = FileFormat(); std::string table_path = paimon::test::GetDataDir() + file_format + @@ -1251,7 +1296,8 @@ TEST_P(ScanAndReadInteTest, TestWithPKWithMorBatchScanSnapshot5WithPredicate) { ReadContextBuilder read_context_builder(table_path); AddReadOptionsForPrefetch(&read_context_builder); - read_context_builder.SetPredicate(predicate); + read_context_builder.SetPredicate(predicate) + .EnableLateMaterializing(false); ASSERT_OK_AND_ASSIGN(auto read_context, read_context_builder.Finish()); ASSERT_OK_AND_ASSIGN(auto table_read, TableRead::Create(std::move(read_context))); @@ -1279,6 +1325,57 @@ TEST_P(ScanAndReadInteTest, TestWithPKWithMorBatchScanSnapshot5WithPredicate) { ASSERT_TRUE(expected->Equals(read_result)) << read_result->ToString(); } +// Same coverage as the deletion-vector case above, for the merge-on-read path where only the +// key part of the predicate is pushed down into the data files. +TEST_P(ScanAndReadInteTest, TestWithPKWithMorBatchScanSnapshot5WithLateMaterializing) { + auto file_format = FileFormat(); + std::string table_path = paimon::test::GetDataDir() + file_format + + "/pk_table_scan_and_read_mor.db/pk_table_scan_and_read_mor/"; + + ScanContextBuilder scan_context_builder(table_path); + scan_context_builder.AddOption(Options::SCAN_SNAPSHOT_ID, "5"); + + std::string literal_str = "Alice"; + auto not_equal = PredicateBuilder::NotEqual( + /*field_index=*/0, /*field_name=*/"f0", FieldType::STRING, + Literal(FieldType::STRING, literal_str.data(), literal_str.size())); + std::string literal_str2 = "Lucy"; + auto less_than = PredicateBuilder::LessThan( + /*field_index=*/0, /*field_name=*/"f0", FieldType::STRING, + Literal(FieldType::STRING, literal_str2.data(), literal_str2.size())); + auto less_or_equal = PredicateBuilder::LessOrEqual(/*field_index=*/3, /*field_name=*/"f3", + FieldType::DOUBLE, Literal(30.0)); + ASSERT_OK_AND_ASSIGN(auto predicate, + PredicateBuilder::And({not_equal, less_than, less_or_equal})); + scan_context_builder.SetPredicate(predicate); + ASSERT_OK_AND_ASSIGN(auto scan_context, FinishScanContext(scan_context_builder)); + ASSERT_OK_AND_ASSIGN(auto table_scan, TableScan::Create(std::move(scan_context))); + + ReadContextBuilder read_context_builder(table_path); + AddReadOptionsForPrefetch(&read_context_builder); + read_context_builder.SetPredicate(predicate) + .EnablePredicateFilter(true) + .EnableLateMaterializing(true); + ASSERT_OK_AND_ASSIGN(auto read_context, read_context_builder.Finish()); + ASSERT_OK_AND_ASSIGN(auto table_read, TableRead::Create(std::move(read_context))); + + ASSERT_OK_AND_ASSIGN(auto result_plan, table_scan->CreatePlan()); + ASSERT_EQ(result_plan->SnapshotId().value(), 5); + ASSERT_OK_AND_ASSIGN(auto batch_reader, table_read->CreateReader(result_plan->Splits())); + ASSERT_OK_AND_ASSIGN(auto read_result, ReadResultCollector::CollectResult(batch_reader.get())); + + // check result: only the rows before "Lucy" with f3 <= 30.0 remain. + auto expected = std::make_shared( + arrow::ipc::internal::json::ArrayFromJSON(arrow_data_type_, R"([ +[0, "Bob", 10, 0, 12.1], +[0, "David", 10, 0, 17.1], +[0, "Emily", 10, 0, 13.1] + ])") + .ValueOrDie()); + ASSERT_TRUE(expected); + ASSERT_TRUE(expected->Equals(read_result)) << read_result->ToString(); +} + TEST_P(ScanAndReadInteTest, TestWithPKWithMorBatchScanSnapshot3WithPredicate) { auto file_format = FileFormat(); std::string table_path = paimon::test::GetDataDir() + file_format + From 369d97fbed55656d9134eacb5e110fde47d3519e Mon Sep 17 00:00:00 2001 From: zhouhongfeng Date: Tue, 25 Aug 2026 15:16:23 +0800 Subject: [PATCH 19/34] fix: lat-mat not controlled by enable-prefetch --- src/paimon/core/operation/abstract_split_read.cpp | 8 ++++---- 1 file changed, 4 insertions(+), 4 deletions(-) diff --git a/src/paimon/core/operation/abstract_split_read.cpp b/src/paimon/core/operation/abstract_split_read.cpp index c8d33c9f2..2759f17a5 100644 --- a/src/paimon/core/operation/abstract_split_read.cpp +++ b/src/paimon/core/operation/abstract_split_read.cpp @@ -154,12 +154,12 @@ Result> AbstractSplitRead::PrepareReaderBuilder( Result> AbstractSplitRead::CreateFileBatchReader( const std::string& file_format_identifier, const std::string& data_file_path, int64_t data_file_size, std::unique_ptr reader_builder) const { + if (context_->EnableLateMaterializing()) { + reader_builder = + std::make_unique(std::move(reader_builder), pool_); + } if (context_->EnablePrefetch() && file_format_identifier != "blob" && file_format_identifier != "avro") { - if (context_->EnableLateMaterializing()) { - reader_builder = - std::make_unique(std::move(reader_builder), pool_); - } PAIMON_ASSIGN_OR_RAISE( std::unique_ptr prefetch_reader, PrefetchFileBatchReaderImpl::Create( From f0a7f75560298d2c6caf306388f1e5c902fd2999 Mon Sep 17 00:00:00 2001 From: zhouhongfeng Date: Tue, 25 Aug 2026 15:26:10 +0800 Subject: [PATCH 20/34] refractor --- .../late_materializing_file_batch_reader.cpp | 47 +++---------------- .../late_materializing_file_batch_reader.h | 41 +++++++++++----- .../operation/merge_file_split_read_test.cpp | 4 +- test/inte/read_inte_with_index_test.cpp | 2 +- test/inte/scan_and_read_inte_test.cpp | 6 +-- 5 files changed, 42 insertions(+), 58 deletions(-) diff --git a/src/paimon/common/reader/late_materializing_file_batch_reader.cpp b/src/paimon/common/reader/late_materializing_file_batch_reader.cpp index 1b0d03c85..89d286ac6 100644 --- a/src/paimon/common/reader/late_materializing_file_batch_reader.cpp +++ b/src/paimon/common/reader/late_materializing_file_batch_reader.cpp @@ -46,19 +46,15 @@ Result> LateMaterializingFileB std::unique_ptr inner, std::shared_ptr pool) { // The reader's own compaction allocations go through an arrow pool; bridge the paimon pool // once here so the accounting matches the rest of the read path. - std::shared_ptr arrow_pool; - if (pool != nullptr) { - arrow_pool = GetArrowPool(pool); + if (pool == nullptr) { + return Status::Invalid("pool could not be nullptr."); } + std::shared_ptr arrow_pool = GetArrowPool(pool); auto reader = std::unique_ptr( new LateMaterializingFileBatchReader(std::move(inner), std::move(arrow_pool))); return reader; } -arrow::MemoryPool* LateMaterializingFileBatchReader::ArrowPool() const { - return arrow_pool_ ? arrow_pool_.get() : arrow::default_memory_pool(); -} - Result LateMaterializingFileBatchReader::NextBatch() { if (state_ == kInit) { // SetReadSchema has not been called: read with the file schema, matching the @@ -161,7 +157,7 @@ Status LateMaterializingFileBatchReader::ReadAndFilterProbeData() { probe_array, arrow::MakeEmptyArray(arrow::struct_(probe_schema_->fields()))); } else { PAIMON_ASSIGN_OR_RAISE_FROM_ARROW(probe_array, - arrow::Concatenate(probe_arrays, ArrowPool())); + arrow::Concatenate(probe_arrays, arrow_pool_.get())); } probe_data_ = arrow::internal::checked_pointer_cast(probe_array); return Status::OK(); @@ -188,7 +184,7 @@ Result LateMaterializingFileBatchReader::ReadPayload RoaringBitmap32 valid; row_mapping_.clear(); for (auto it = bitmap.Begin(); it != bitmap.End(); ++it) { - uint64_t offset = static_cast(*it); + auto offset = static_cast(*it); PAIMON_ASSIGN_OR_RAISE(uint64_t file_row, inner_->GetPreviousBatchFileRowId(offset)); if (!matched_bitmap_.Contains(file_row)) { continue; @@ -205,9 +201,9 @@ Result LateMaterializingFileBatchReader::ReadPayload PAIMON_ASSIGN_OR_RAISE(arrow::ArrayVector payload_slices, ReaderUtils::GenerateFilteredArrayVector(payload_array, valid)); PAIMON_ASSIGN_OR_RAISE_FROM_ARROW(std::shared_ptr payload_compacted, - arrow::Concatenate(payload_slices, ArrowPool())); + arrow::Concatenate(payload_slices, arrow_pool_.get())); - int64_t card = static_cast(valid.Cardinality()); + auto card = static_cast(valid.Cardinality()); if (probe_cursor_ + card > probe_data_->length()) { return Status::Invalid( fmt::format("probe cache underflow: cursor {} + {} exceeds probe rows {}", @@ -265,18 +261,6 @@ Status LateMaterializingFileBatchReader::SetInnerReadSchema( return Status::OK(); } -std::shared_ptr LateMaterializingFileBatchReader::GetReaderMetrics() const { - return inner_->GetReaderMetrics(); -} - -void LateMaterializingFileBatchReader::Close() { - inner_->Close(); -} - -Result> LateMaterializingFileBatchReader::GetFileSchema() const { - return inner_->GetFileSchema(); -} - Status LateMaterializingFileBatchReader::SetReadSchema( ::ArrowSchema* read_schema, const std::shared_ptr& predicate, const std::optional& selection_bitmap) { @@ -333,14 +317,6 @@ Result LateMaterializingFileBatchReader::GetPreviousBatchFileRowId( return row_mapping_[batch_row_id]; } -Result LateMaterializingFileBatchReader::GetNumberOfRows() const { - return inner_->GetNumberOfRows(); -} - -bool LateMaterializingFileBatchReader::SupportPreciseBitmapSelection() const { - return inner_->SupportPreciseBitmapSelection(); -} - Status LateMaterializingFileBatchReader::SeekToRow(uint64_t row_number) { PAIMON_RETURN_NOT_OK(inner_->SeekToRow(row_number)); if (state_ == kRunning || state_ == kEOF) { @@ -358,15 +334,6 @@ Status LateMaterializingFileBatchReader::SeekToRow(uint64_t row_number) { return Status::OK(); } -uint64_t LateMaterializingFileBatchReader::GetNextRowToRead() const { - return inner_->GetNextRowToRead(); -} - -Result>> LateMaterializingFileBatchReader::GenReadRanges( - bool* need_prefetch) const { - return inner_->GenReadRanges(need_prefetch); -} - Status LateMaterializingFileBatchReader::SetReadRanges( const std::vector>& read_ranges) { read_ranges_ = read_ranges; diff --git a/src/paimon/common/reader/late_materializing_file_batch_reader.h b/src/paimon/common/reader/late_materializing_file_batch_reader.h index 809dc6a36..5d67bb61b 100644 --- a/src/paimon/common/reader/late_materializing_file_batch_reader.h +++ b/src/paimon/common/reader/late_materializing_file_batch_reader.h @@ -43,25 +43,48 @@ class LateMaterializingFileBatchReader : public PrefetchFileBatchReader { Result NextBatch() override; - std::shared_ptr GetReaderMetrics() const override; - void Close() override; + std::shared_ptr GetReaderMetrics() const override { + return inner_->GetReaderMetrics(); + }; + + void Close() override { + inner_->Close(); + } + + Result> GetFileSchema() const override { + return inner_->GetFileSchema(); + } - Result> GetFileSchema() const override; Status SetReadSchema(::ArrowSchema* read_schema, const std::shared_ptr& predicate, const std::optional& selection_bitmap) override; + Result GetPreviousBatchFileRowId(uint64_t batch_row_id) const override; - Result GetNumberOfRows() const override; - bool SupportPreciseBitmapSelection() const override; + + Result GetNumberOfRows() const override { + return inner_->GetNumberOfRows(); + } + + bool SupportPreciseBitmapSelection() const override { + return inner_->SupportPreciseBitmapSelection(); + } Status SeekToRow(uint64_t row_number) override; - uint64_t GetNextRowToRead() const override; + + uint64_t GetNextRowToRead() const override { + return inner_->GetNextRowToRead(); + } + Result>> GenReadRanges( - bool* need_prefetch) const override; + bool* need_prefetch) const override { + return inner_->GenReadRanges(need_prefetch); + } + Status SetReadRanges(const std::vector>& read_ranges) override; Result>> PreBufferRange() override { return inner_->PreBufferRange(); } + private: explicit LateMaterializingFileBatchReader(std::unique_ptr inner, std::shared_ptr arrow_pool) @@ -98,10 +121,6 @@ class LateMaterializingFileBatchReader : public PrefetchFileBatchReader { const std::shared_ptr& predicate, const std::optional& selection); - // Arrow pool for this reader's own allocations (probe/payload compaction). Falls back to the - // arrow default pool when no pool was provided. - arrow::MemoryPool* ArrowPool() const; - std::unique_ptr inner_; std::shared_ptr arrow_pool_; LatMatState state_ = kInit; diff --git a/src/paimon/core/operation/merge_file_split_read_test.cpp b/src/paimon/core/operation/merge_file_split_read_test.cpp index ee689851e..fad88774a 100644 --- a/src/paimon/core/operation/merge_file_split_read_test.cpp +++ b/src/paimon/core/operation/merge_file_split_read_test.cpp @@ -866,7 +866,7 @@ TEST_P(MergeFileSplitReadTest, TestReadWithPredicateAndLateMaterializing) { context_builder.SetOptions({{Options::SEQUENCE_FIELD, "s0,s1"}, {Options::MERGE_ENGINE, "deduplicate"}, {Options::IGNORE_DELETE, "true"}}); - AddOptions(&context_builder); + AddOptions(&context_builder); context_builder.EnableLateMaterializing(true); // key predicate, always pushed down into the data files auto greater_or_equal = PredicateBuilder::GreaterOrEqual(/*field_index=*/0, /*field_name=*/"k1", @@ -900,7 +900,7 @@ TEST_P(MergeFileSplitReadTest, TestReadWithPredicateAndLateMaterializing) { [0, 1, 1, "you", "zoo", 130.0, false] ])"}, - &expected_array); + &expected_array); ASSERT_TRUE(array_status.ok()); CheckResult(result_array, expected_array, read_schema); } diff --git a/test/inte/read_inte_with_index_test.cpp b/test/inte/read_inte_with_index_test.cpp index ac9add2c0..f241229cb 100644 --- a/test/inte/read_inte_with_index_test.cpp +++ b/test/inte/read_inte_with_index_test.cpp @@ -1299,7 +1299,7 @@ TEST_P(ReadInteWithIndexTest, TestBitmapIndexWithLateMaterializing) { [0, "Bob", 10, 1, 12.1], [0, "Bob", 10, 1, 16.1] ])"}, - &expected_array); + &expected_array); ASSERT_TRUE(array_status.ok()); ASSERT_TRUE(result_array->Equals(*expected_array)) << result_array->ToString(); } diff --git a/test/inte/scan_and_read_inte_test.cpp b/test/inte/scan_and_read_inte_test.cpp index 4d3569086..35ef8f7e3 100644 --- a/test/inte/scan_and_read_inte_test.cpp +++ b/test/inte/scan_and_read_inte_test.cpp @@ -724,8 +724,7 @@ TEST_P(ScanAndReadInteTest, TestWithPKWithDvBatchScanSnapshot6WithPredicate) { ReadContextBuilder read_context_builder(table_path); AddReadOptionsForPrefetch(&read_context_builder); - read_context_builder.SetPredicate(predicate) - .EnableLateMaterializing(false); + read_context_builder.SetPredicate(predicate).EnableLateMaterializing(false); ASSERT_OK_AND_ASSIGN(auto read_context, read_context_builder.Finish()); ASSERT_OK_AND_ASSIGN(auto table_read, TableRead::Create(std::move(read_context))); @@ -1296,8 +1295,7 @@ TEST_P(ScanAndReadInteTest, TestWithPKWithMorBatchScanSnapshot5WithPredicate) { ReadContextBuilder read_context_builder(table_path); AddReadOptionsForPrefetch(&read_context_builder); - read_context_builder.SetPredicate(predicate) - .EnableLateMaterializing(false); + read_context_builder.SetPredicate(predicate).EnableLateMaterializing(false); ASSERT_OK_AND_ASSIGN(auto read_context, read_context_builder.Finish()); ASSERT_OK_AND_ASSIGN(auto table_read, TableRead::Create(std::move(read_context))); From 4925c4e8f18173271478ce65869402a185e1e97b Mon Sep 17 00:00:00 2001 From: zhouhongfeng Date: Tue, 25 Aug 2026 15:40:15 +0800 Subject: [PATCH 21/34] style: add comments --- .../common/reader/late_materializing_file_batch_reader.h | 4 ++-- src/paimon/core/operation/data_evolution_split_read.h | 3 ++- src/paimon/core/operation/merge_file_split_read.h | 3 ++- src/paimon/core/operation/merge_file_split_read_test.cpp | 4 ---- src/paimon/core/operation/raw_file_split_read.h | 3 ++- test/inte/read_inte_test.cpp | 3 --- test/inte/read_inte_with_index_test.cpp | 3 --- test/inte/scan_and_read_inte_test.cpp | 3 --- 8 files changed, 8 insertions(+), 18 deletions(-) diff --git a/src/paimon/common/reader/late_materializing_file_batch_reader.h b/src/paimon/common/reader/late_materializing_file_batch_reader.h index 5d67bb61b..75f165641 100644 --- a/src/paimon/common/reader/late_materializing_file_batch_reader.h +++ b/src/paimon/common/reader/late_materializing_file_batch_reader.h @@ -92,9 +92,9 @@ class LateMaterializingFileBatchReader : public PrefetchFileBatchReader { enum LatMatState { kInit, - kProbing, // schema is set + kProbing, // schema is set, probing is in progress kNoLatMat, // no need to late materialization - kRunning, // Lat-mat is enable an is reading data + kRunning, // Lat-mat is enabled and the payload reader is reading data kEOF }; diff --git a/src/paimon/core/operation/data_evolution_split_read.h b/src/paimon/core/operation/data_evolution_split_read.h index b4f80adb6..7ac851d4b 100644 --- a/src/paimon/core/operation/data_evolution_split_read.h +++ b/src/paimon/core/operation/data_evolution_split_read.h @@ -65,7 +65,8 @@ struct DeletionFile; /// ->FieldMappingReader->(ApplyDeletionVectorBatchReader)->(ApplyBitmapIndexBatchReader) /// ->(CompleteRowTrackingFieldsBatchReader)->(ShreddingFileReader) /// ->(MapSharedShreddingFileReader) -/// ->(VectorFileBatchReader)->(DelegatingPrefetchReader)->(PrefetchFileBatchReader)->FormatReader +/// ->(VectorFileBatchReader)->(DelegatingPrefetchReader)->(PrefetchFileBatchReader) +/// ->(LateMaterializingFileBatchReader)->FormatReader /// /// /// A union `SplitRead` to read multiple inner files to merge columns. A single-file row range diff --git a/src/paimon/core/operation/merge_file_split_read.h b/src/paimon/core/operation/merge_file_split_read.h index d4bfa727c..ac0a8246d 100644 --- a/src/paimon/core/operation/merge_file_split_read.h +++ b/src/paimon/core/operation/merge_file_split_read.h @@ -75,7 +75,8 @@ class MergeFunctionWrapper; /// ->DropDeleteReader->SortMergeReader->ConcatKeyValueRecordReader->KeyValueDataFileRecordReader /// ->FieldMappingReader->(ApplyDeletionVectorBatchReader)->(ShreddingFileReader) /// ->(MapSharedShreddingFileReader) -/// ->(DelegatingPrefetchReader)->(PrefetchFileBatchReader)->FormatReader +/// ->(DelegatingPrefetchReader)->(PrefetchFileBatchReader) +/// ->(LateMaterializingFileBatchReader)->FormatReader class MergeFileSplitRead : public AbstractSplitRead { public: static Result> Create( diff --git a/src/paimon/core/operation/merge_file_split_read_test.cpp b/src/paimon/core/operation/merge_file_split_read_test.cpp index fad88774a..b3ddad65e 100644 --- a/src/paimon/core/operation/merge_file_split_read_test.cpp +++ b/src/paimon/core/operation/merge_file_split_read_test.cpp @@ -844,10 +844,6 @@ TEST_P(MergeFileSplitReadTest, TestReadWithPredicate) { CheckResult(result_array, expected_array, read_schema); } -// Late materialization reads the predicate columns first and only materializes the remaining -// columns for matched rows. Combined with the top-level predicate filter the read path returns -// the exact user-predicate match set: in a section with several sorted runs only the key part of -// the predicate reaches the data files, and the merge result is filtered afterwards. TEST_P(MergeFileSplitReadTest, TestReadWithPredicateAndLateMaterializing) { std::string path = paimon::test::GetDataDir() + "/parquet/pk_table_with_mor.db/pk_table_with_mor"; diff --git a/src/paimon/core/operation/raw_file_split_read.h b/src/paimon/core/operation/raw_file_split_read.h index 6a97b9b37..064b96eb6 100644 --- a/src/paimon/core/operation/raw_file_split_read.h +++ b/src/paimon/core/operation/raw_file_split_read.h @@ -55,7 +55,8 @@ struct DeletionFile; /// ->ConcatBatchReader across /// files->FieldMappingReader->(ApplyBitmapIndexBatchReader)->(CompleteRowTrackingFieldsBatchReader) /// ->(ShreddingFileReader)->(MapSharedShreddingFileReader)->(VectorFileBatchReader) -/// ->(DelegatingPrefetchReader)->(PrefetchFileBatchReader)->FormatReader +/// ->(DelegatingPrefetchReader)->(PrefetchFileBatchReader)->(LateMaterializingFileBatchReader) +/// ->FormatReader class RawFileSplitRead : public AbstractSplitRead { public: diff --git a/test/inte/read_inte_test.cpp b/test/inte/read_inte_test.cpp index 99d9f0983..b218d133b 100644 --- a/test/inte/read_inte_test.cpp +++ b/test/inte/read_inte_test.cpp @@ -3252,9 +3252,6 @@ TEST_P(ReadInteTest, TestPkReadSnapshot6WithSchemaEvolutionWithPredicateOnlyPush ASSERT_TRUE(result_array->Equals(*expected_array)); } -// Same coverage for the primary-key path with schema evolution and a deletion vector: late -// materialization runs below the prefetch layer per data file, so the deletion vector and the -// schema-evolution field mapping keep working. TEST_P(ReadInteTest, TestPkReadSnapshot6WithSchemaEvolutionWithLateMaterializing) { std::vector read_fields = {DataField(1, arrow::field("key1", arrow::int32())), DataField(7, arrow::field("k", arrow::utf8())), diff --git a/test/inte/read_inte_with_index_test.cpp b/test/inte/read_inte_with_index_test.cpp index f241229cb..7a4439734 100644 --- a/test/inte/read_inte_with_index_test.cpp +++ b/test/inte/read_inte_with_index_test.cpp @@ -1233,9 +1233,6 @@ TEST_P(ReadInteWithIndexTest, TestNoEmbeddingBitmapIndex) { CheckResultForBitmapWithSingleRowGroup(path, arrow_data_type, split); } -// Late materialization reads the predicate columns first and only materializes the remaining -// columns for matched rows. The bitmap index selection is pushed into the same reader, so -// combined with the top-level predicate filter the read path returns the exact match set. TEST_P(ReadInteWithIndexTest, TestBitmapIndexWithLateMaterializing) { auto [file_format, enable_prefetch] = GetParam(); std::string path = GetDataDir() + "/" + file_format + diff --git a/test/inte/scan_and_read_inte_test.cpp b/test/inte/scan_and_read_inte_test.cpp index 35ef8f7e3..12d5ea8ea 100644 --- a/test/inte/scan_and_read_inte_test.cpp +++ b/test/inte/scan_and_read_inte_test.cpp @@ -744,9 +744,6 @@ TEST_P(ScanAndReadInteTest, TestWithPKWithDvBatchScanSnapshot6WithPredicate) { ASSERT_TRUE(expected->Equals(read_result)) << read_result->ToString(); } -// Late materialization reads the predicate columns first and only materializes the remaining -// columns for matched rows. Combined with the top-level predicate filter, the read path returns -// the exact user-predicate match set (deletion vectors stay applied per data file). TEST_P(ScanAndReadInteTest, TestWithPKWithDvBatchScanSnapshot6WithLateMaterializing) { auto file_format = FileFormat(); std::string table_path = paimon::test::GetDataDir() + file_format + From f1f1c765b067cd1a1e9b3f2fa17a284518da05ec Mon Sep 17 00:00:00 2001 From: zhouhongfeng Date: Tue, 25 Aug 2026 16:22:31 +0800 Subject: [PATCH 22/34] fix: lat-mat silently fails when format is avro or blob --- include/paimon/read_context.h | 12 +++--- .../core/operation/abstract_split_read.cpp | 43 ++++++++++--------- 2 files changed, 27 insertions(+), 28 deletions(-) diff --git a/include/paimon/read_context.h b/include/paimon/read_context.h index 0d6ea58d5..e645645ba 100644 --- a/include/paimon/read_context.h +++ b/include/paimon/read_context.h @@ -97,8 +97,6 @@ class PAIMON_EXPORT ReadContext { bool EnablePrefetch() const { return enable_prefetch_; } - /// Whether late materialization (probe/payload two-phase reads) is enabled for the - /// prefetch read path. Defaults to false. bool EnableLateMaterializing() const { return enable_late_materializing_; } @@ -312,13 +310,13 @@ class PAIMON_EXPORT ReadContextBuilder { /// @return Reference to this builder for method chaining. ReadContextBuilder& EnablePrefetch(bool enabled); - /// Enable or disable late materialization (probe/payload two-phase reads) for the prefetch - /// read path. When enabled, each parallel reader under the prefetch layer performs a probe - /// read of predicate columns first and only materializes payload columns for matched rows. + /// Enable or disable late materialization (probe/payload two-phase reads). When enabled, + /// each parallel reader under the prefetch layer performs a probe read of predicate + /// columns first and only materializes payload columns for matched rows. /// @param enabled Whether to enable late materialization (default: false) /// @return Reference to this builder for method chaining. - /// @note Only takes effect when prefetch is enabled; without a pushed-down predicate the - /// late-materializing reader degrades to a plain passthrough. + /// @note Without a pushed-down predicate the late-materializing reader degrades to a + /// plain passthrough. ReadContextBuilder& EnableLateMaterializing(bool enabled); /// Enable or disable the read-ahead cache for read operations. diff --git a/src/paimon/core/operation/abstract_split_read.cpp b/src/paimon/core/operation/abstract_split_read.cpp index 2759f17a5..60bed5f8a 100644 --- a/src/paimon/core/operation/abstract_split_read.cpp +++ b/src/paimon/core/operation/abstract_split_read.cpp @@ -154,28 +154,29 @@ Result> AbstractSplitRead::PrepareReaderBuilder( Result> AbstractSplitRead::CreateFileBatchReader( const std::string& file_format_identifier, const std::string& data_file_path, int64_t data_file_size, std::unique_ptr reader_builder) const { - if (context_->EnableLateMaterializing()) { - reader_builder = - std::make_unique(std::move(reader_builder), pool_); - } - if (context_->EnablePrefetch() && file_format_identifier != "blob" && - file_format_identifier != "avro") { - PAIMON_ASSIGN_OR_RAISE( - std::unique_ptr prefetch_reader, - PrefetchFileBatchReaderImpl::Create( - data_file_path, data_file_size, reader_builder.get(), options_.GetFileSystem(), - context_->GetPrefetchMaxParallelNum(), options_.GetReadBatchSize(), - context_->GetPrefetchBatchCount(), options_.EnableAdaptivePrefetchStrategy(), - executor_, - /*initialize_read_ranges=*/false, context_->ReadAheadCacheEnabled(), - context_->GetCacheConfig(), pool_)); - return std::make_unique(std::move(prefetch_reader)); - } else { - PAIMON_ASSIGN_OR_RAISE( - std::shared_ptr input_stream, - options_.GetFileSystem()->Open(FileStatus(data_file_path, data_file_size))); - return reader_builder->Build(input_stream); + // blob and avro do not support prefetch or late materializing + if (file_format_identifier != "blob" && file_format_identifier != "avro") { + if (context_->EnableLateMaterializing()) { + reader_builder = + std::make_unique(std::move(reader_builder), pool_); + } + if (context_->EnablePrefetch()) { + PAIMON_ASSIGN_OR_RAISE( + std::unique_ptr prefetch_reader, + PrefetchFileBatchReaderImpl::Create( + data_file_path, data_file_size, reader_builder.get(), options_.GetFileSystem(), + context_->GetPrefetchMaxParallelNum(), options_.GetReadBatchSize(), + context_->GetPrefetchBatchCount(), options_.EnableAdaptivePrefetchStrategy(), + executor_, + /*initialize_read_ranges=*/false, context_->ReadAheadCacheEnabled(), + context_->GetCacheConfig(), pool_)); + return std::make_unique(std::move(prefetch_reader)); + } } + PAIMON_ASSIGN_OR_RAISE( + std::shared_ptr input_stream, + options_.GetFileSystem()->Open(FileStatus(data_file_path, data_file_size))); + return reader_builder->Build(input_stream); } Result> AbstractSplitRead::CreateFieldMappingReader( From 708dadfe06069ccb30dd1aa5060af7b62b2dd79b Mon Sep 17 00:00:00 2001 From: zhouhongfeng Date: Tue, 25 Aug 2026 16:34:17 +0800 Subject: [PATCH 23/34] fix: minor issues in LM reader --- .../late_materializing_file_batch_reader.cpp | 14 +++++++++++--- .../reader/late_materializing_file_batch_reader.h | 2 ++ 2 files changed, 13 insertions(+), 3 deletions(-) diff --git a/src/paimon/common/reader/late_materializing_file_batch_reader.cpp b/src/paimon/common/reader/late_materializing_file_batch_reader.cpp index 89d286ac6..c0370b077 100644 --- a/src/paimon/common/reader/late_materializing_file_batch_reader.cpp +++ b/src/paimon/common/reader/late_materializing_file_batch_reader.cpp @@ -80,8 +80,7 @@ Result LateMaterializingFileBatchReader::NextBatch() } else if (state_ == kEOF) { return MakeEofBatch(); } - return Status::Invalid("invalid state when calling NextBatchWithBitmap: " + - std::to_string(state_)); + return Status::Invalid("invalid state when calling NextBatch: " + std::to_string(state_)); } Result> LateMaterializingFileBatchReader::BindProbeFilter() { @@ -117,7 +116,7 @@ Result LateMaterializingFileBatchReader::FilterProbeBatch( // map batch offset to file row id PAIMON_ASSIGN_OR_RAISE(uint64_t file_row, inner_->GetPreviousBatchFileRowId(static_cast(i))); - if (selection_ && !selection_->Contains(file_row)) { + if (selection_ && !selection_->Contains(static_cast(file_row))) { continue; } batch_matched.Add(static_cast(i)); @@ -169,6 +168,11 @@ Result LateMaterializingFileBatchReader::ReadPayload inner_->NextBatchWithBitmap()); if (BatchReader::IsEofBatch(batch_with_bitmap)) { state_ = kEOF; + if (probe_cursor_ != probe_data_->length()) { + return Status::Invalid( + fmt::format("probe cursor {} does not match probe data length {}", + probe_cursor_, probe_data_->length())); + } return MakeEofBatch(); } auto& [batch, bitmap] = batch_with_bitmap; @@ -320,6 +324,10 @@ Result LateMaterializingFileBatchReader::GetPreviousBatchFileRowId( Status LateMaterializingFileBatchReader::SeekToRow(uint64_t row_number) { PAIMON_RETURN_NOT_OK(inner_->SeekToRow(row_number)); if (state_ == kRunning || state_ == kEOF) { + if (matched_bitmap_.IsEmpty()) { + state_ = kEOF; + return Status::OK(); + } int64_t cursor = 0; for (auto it = matched_bitmap_.Begin(); it != matched_bitmap_.End(); ++it) { if (static_cast(*it) >= row_number) { diff --git a/src/paimon/common/reader/late_materializing_file_batch_reader.h b/src/paimon/common/reader/late_materializing_file_batch_reader.h index 75f165641..76e41d2cf 100644 --- a/src/paimon/common/reader/late_materializing_file_batch_reader.h +++ b/src/paimon/common/reader/late_materializing_file_batch_reader.h @@ -65,6 +65,8 @@ class LateMaterializingFileBatchReader : public PrefetchFileBatchReader { } bool SupportPreciseBitmapSelection() const override { + // When probe_schema_ or payload_schema_ is null, lat-mat does not take effect. + // Here we simply pass through the inner reader's support. return inner_->SupportPreciseBitmapSelection(); } From 029e5e1e62693ce588c81e88d75a81eb3b233df0 Mon Sep 17 00:00:00 2001 From: zhouhongfeng Date: Tue, 25 Aug 2026 17:12:55 +0800 Subject: [PATCH 24/34] fix: predicate issues --- .../late_materializing_file_batch_reader.cpp | 32 +++++++++---------- .../late_materializing_file_batch_reader.h | 4 +-- ...e_materializing_file_batch_reader_test.cpp | 16 ++++++++++ 3 files changed, 34 insertions(+), 18 deletions(-) diff --git a/src/paimon/common/reader/late_materializing_file_batch_reader.cpp b/src/paimon/common/reader/late_materializing_file_batch_reader.cpp index c0370b077..c262d0a1b 100644 --- a/src/paimon/common/reader/late_materializing_file_batch_reader.cpp +++ b/src/paimon/common/reader/late_materializing_file_batch_reader.cpp @@ -33,6 +33,7 @@ #include "arrow/util/checked_cast.h" #include "fmt/format.h" #include "paimon/common/predicate/predicate_filter.h" +#include "paimon/common/predicate/predicate_validator.h" #include "paimon/common/reader/reader_utils.h" #include "paimon/common/utils/arrow/arrow_utils.h" #include "paimon/common/utils/arrow/mem_utils.h" @@ -83,20 +84,6 @@ Result LateMaterializingFileBatchReader::NextBatch() return Status::Invalid("invalid state when calling NextBatch: " + std::to_string(state_)); } -Result> LateMaterializingFileBatchReader::BindProbeFilter() { - std::map 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 bound_predicate, - PredicateUtils::CreatePickedFieldFilter(predicate_, name_to_idx)); - auto bound_filter = std::dynamic_pointer_cast(bound_predicate); - if (!bound_filter) { - return Status::Invalid("failed to bind predicate to probe schema"); - } - return bound_filter; -} - Result LateMaterializingFileBatchReader::FilterProbeBatch( const std::shared_ptr& array, const std::shared_ptr& bound_filter) { @@ -126,7 +113,6 @@ Result LateMaterializingFileBatchReader::FilterProbeBatch( } Status LateMaterializingFileBatchReader::ReadAndFilterProbeData() { - PAIMON_ASSIGN_OR_RAISE(std::shared_ptr bound_filter, BindProbeFilter()); matched_bitmap_ = RoaringBitmap32(); probe_cursor_ = 0; arrow::ArrayVector probe_arrays; @@ -139,7 +125,7 @@ Status LateMaterializingFileBatchReader::ReadAndFilterProbeData() { PAIMON_ASSIGN_OR_RAISE_FROM_ARROW(std::shared_ptr array, arrow::ImportArray(c_array.get(), c_schema.get())); PAIMON_ASSIGN_OR_RAISE(RoaringBitmap32 batch_matched, - FilterProbeBatch(array, bound_filter)); + FilterProbeBatch(array, probe_filter_)); // 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()) { @@ -277,6 +263,7 @@ Status LateMaterializingFileBatchReader::SetReadSchema( row_mapping_.clear(); probe_schema_.reset(); payload_schema_.reset(); + probe_filter_.reset(); if (predicate_ != nullptr) { std::set probe_names; PAIMON_RETURN_NOT_OK(PredicateUtils::GetAllNames(predicate_, &probe_names)); @@ -293,6 +280,19 @@ Status LateMaterializingFileBatchReader::SetReadSchema( if (!probe_fields.empty() && !payload_fields.empty()) { 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 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 bound_predicate, + PredicateUtils::CreatePickedFieldFilter(predicate_, name_to_idx)); + probe_filter_ = std::dynamic_pointer_cast(bound_predicate); + if (!probe_filter_) { + return Status::Invalid("failed to bind predicate to probe schema"); + } } } diff --git a/src/paimon/common/reader/late_materializing_file_batch_reader.h b/src/paimon/common/reader/late_materializing_file_batch_reader.h index 76e41d2cf..f232b68d0 100644 --- a/src/paimon/common/reader/late_materializing_file_batch_reader.h +++ b/src/paimon/common/reader/late_materializing_file_batch_reader.h @@ -105,8 +105,6 @@ class LateMaterializingFileBatchReader : public PrefetchFileBatchReader { /// TODO(zhouhongfeng.zhf): Read the probe data batch by batch to save memory. Status ReadAndFilterProbeData(); - Result> BindProbeFilter(); - Result FilterProbeBatch(const std::shared_ptr& array, const std::shared_ptr& bound_filter); @@ -133,6 +131,8 @@ class LateMaterializingFileBatchReader : public PrefetchFileBatchReader { // projection holding the payload (non-probe) fields; nullptr when probing is not applicable std::shared_ptr payload_schema_; std::shared_ptr predicate_; + // predicate bound to probe_schema_'s field indices; null when probing is not applicable + std::shared_ptr probe_filter_; std::optional selection_; // the probe_data_ is sliced and compacted with the matched_bitmap_ std::shared_ptr probe_data_; diff --git a/src/paimon/common/reader/late_materializing_file_batch_reader_test.cpp b/src/paimon/common/reader/late_materializing_file_batch_reader_test.cpp index 27c84e848..f948a702e 100644 --- a/src/paimon/common/reader/late_materializing_file_batch_reader_test.cpp +++ b/src/paimon/common/reader/late_materializing_file_batch_reader_test.cpp @@ -639,4 +639,20 @@ TEST_F(LateMaterializingFileBatchReaderTest, PrefetchInnerParallelReadersWithSee impl->Close(); } +// When the predicate's field type does not match the probe schema, the +// ValidatePredicateWithSchema check must fail with a clear error instead +// of silently producing incorrect results. +TEST_F(LateMaterializingFileBatchReaderTest, FailsOnPredicateTypeMismatch) { + auto data = BuildData({0, 1, 2, 3, 4}); + auto mock = std::make_unique(data, full_type_, /*batch_size=*/2); + ASSERT_OK_AND_ASSIGN( + auto reader, LateMaterializingFileBatchReader::Create(std::move(mock), GetDefaultPool())); + // k is int64 in the schema, but the predicate claims FieldType::INT (int32). + auto predicate = + PredicateBuilder::Equal(/*field_index=*/0, /*field_name=*/"k", FieldType::INT, Literal(10)); + ASSERT_NOK_WITH_MSG( + SetReadSchema(reader.get(), arrow::schema(full_fields_), predicate, std::nullopt), + "mismatches"); +} + } // namespace paimon::test From 660be343a467cd61158e04db080a1657112b0e78 Mon Sep 17 00:00:00 2001 From: zhouhongfeng Date: Tue, 25 Aug 2026 17:19:20 +0800 Subject: [PATCH 25/34] fix: clang-tidy --- src/paimon/common/reader/late_materializing_file_batch_reader.h | 1 + 1 file changed, 1 insertion(+) diff --git a/src/paimon/common/reader/late_materializing_file_batch_reader.h b/src/paimon/common/reader/late_materializing_file_batch_reader.h index f232b68d0..f1de45dbd 100644 --- a/src/paimon/common/reader/late_materializing_file_batch_reader.h +++ b/src/paimon/common/reader/late_materializing_file_batch_reader.h @@ -20,6 +20,7 @@ #pragma once #include +#include #include #include From 66916f31a54f96d62bede27adf06f52ded687e8d Mon Sep 17 00:00:00 2001 From: zhouhongfeng Date: Tue, 25 Aug 2026 17:22:48 +0800 Subject: [PATCH 26/34] add TODOs --- src/paimon/common/reader/late_materializing_file_batch_reader.h | 2 ++ 1 file changed, 2 insertions(+) diff --git a/src/paimon/common/reader/late_materializing_file_batch_reader.h b/src/paimon/common/reader/late_materializing_file_batch_reader.h index f1de45dbd..2d7f264cc 100644 --- a/src/paimon/common/reader/late_materializing_file_batch_reader.h +++ b/src/paimon/common/reader/late_materializing_file_batch_reader.h @@ -85,6 +85,8 @@ class LateMaterializingFileBatchReader : public PrefetchFileBatchReader { Status SetReadRanges(const std::vector>& read_ranges) override; Result>> PreBufferRange() override { + // TODO(zhouhongfeng.zhf): PrebufferRange (called by PrefetchFileBatchReader) only read the probe data, + // consider read the payload data as well. return inner_->PreBufferRange(); } From d366c120ac6c2d22528935298dc6d1f3234744f3 Mon Sep 17 00:00:00 2001 From: zhouhongfeng Date: Tue, 25 Aug 2026 18:23:03 +0800 Subject: [PATCH 27/34] test: restore behavior of merge_file_split_read_test --- src/paimon/core/operation/merge_file_split_read_test.cpp | 3 +-- 1 file changed, 1 insertion(+), 2 deletions(-) diff --git a/src/paimon/core/operation/merge_file_split_read_test.cpp b/src/paimon/core/operation/merge_file_split_read_test.cpp index b3ddad65e..19cd96d67 100644 --- a/src/paimon/core/operation/merge_file_split_read_test.cpp +++ b/src/paimon/core/operation/merge_file_split_read_test.cpp @@ -99,8 +99,6 @@ class MergeFileSplitReadTest : public ::testing::Test, void AddOptions(ReadContextBuilder* context_builder) const { auto [use_min_heap, enable_io_prefetch, enable_multi_thread_row_to_batch] = GetParam(); - // disable late materializing by default - context_builder->EnableLateMaterializing(false); if (use_min_heap) { context_builder->AddOption(Options::SORT_ENGINE, "min-heap"); } else { @@ -800,6 +798,7 @@ TEST_P(MergeFileSplitReadTest, TestReadWithPredicate) { context_builder.SetOptions({{Options::SEQUENCE_FIELD, "s0,s1"}, {Options::MERGE_ENGINE, "deduplicate"}, {Options::IGNORE_DELETE, "true"}}); + context_builder.EnableLateMaterializing(false); AddOptions(&context_builder); // less_than will be ignore as it is partition predicate From 9215628f27cf698b242097665c40f81fcad077e9 Mon Sep 17 00:00:00 2001 From: zhouhongfeng Date: Tue, 25 Aug 2026 18:23:49 +0800 Subject: [PATCH 28/34] pre-commit --- .../common/reader/late_materializing_file_batch_reader.h | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/src/paimon/common/reader/late_materializing_file_batch_reader.h b/src/paimon/common/reader/late_materializing_file_batch_reader.h index 2d7f264cc..c1b91b086 100644 --- a/src/paimon/common/reader/late_materializing_file_batch_reader.h +++ b/src/paimon/common/reader/late_materializing_file_batch_reader.h @@ -85,8 +85,8 @@ class LateMaterializingFileBatchReader : public PrefetchFileBatchReader { Status SetReadRanges(const std::vector>& read_ranges) override; Result>> PreBufferRange() override { - // TODO(zhouhongfeng.zhf): PrebufferRange (called by PrefetchFileBatchReader) only read the probe data, - // consider read the payload data as well. + // TODO(zhouhongfeng.zhf): PrebufferRange (called by PrefetchFileBatchReader) only read the + // probe data, consider read the payload data as well. return inner_->PreBufferRange(); } From 86a0e996ca41a7494f8d2a01f8373ca3084d872a Mon Sep 17 00:00:00 2001 From: zhouhongfeng Date: Wed, 26 Aug 2026 11:25:47 +0800 Subject: [PATCH 29/34] fix: reset LM reader state on Close --- .../late_materializing_file_batch_reader.cpp | 25 +++++++++++++------ .../late_materializing_file_batch_reader.h | 4 +++ 2 files changed, 22 insertions(+), 7 deletions(-) diff --git a/src/paimon/common/reader/late_materializing_file_batch_reader.cpp b/src/paimon/common/reader/late_materializing_file_batch_reader.cpp index c262d0a1b..e399d9edf 100644 --- a/src/paimon/common/reader/late_materializing_file_batch_reader.cpp +++ b/src/paimon/common/reader/late_materializing_file_batch_reader.cpp @@ -254,16 +254,10 @@ Status LateMaterializingFileBatchReader::SetInnerReadSchema( Status LateMaterializingFileBatchReader::SetReadSchema( ::ArrowSchema* read_schema, const std::shared_ptr& predicate, const std::optional& selection_bitmap) { + Reset(); PAIMON_ASSIGN_OR_RAISE_FROM_ARROW(full_schema_, arrow::ImportSchema(read_schema)); predicate_ = predicate; selection_ = selection_bitmap; - matched_bitmap_ = RoaringBitmap32(); - probe_data_.reset(); - probe_cursor_ = 0; - row_mapping_.clear(); - probe_schema_.reset(); - payload_schema_.reset(); - probe_filter_.reset(); if (predicate_ != nullptr) { std::set probe_names; PAIMON_RETURN_NOT_OK(PredicateUtils::GetAllNames(predicate_, &probe_names)); @@ -349,4 +343,21 @@ Status LateMaterializingFileBatchReader::SetReadRanges( return Status::OK(); } +void LateMaterializingFileBatchReader::Reset() { + state_ = kInit; + read_ranges_.clear(); + matched_bitmap_ = RoaringBitmap32(); + probe_data_.reset(); + probe_cursor_ = 0; + row_mapping_.clear(); + probe_schema_.reset(); + payload_schema_.reset(); + full_schema_.reset(); + probe_filter_.reset(); + predicate_.reset(); + selection_.reset(); + probe_cursor_ = 0; + row_mapping_.clear(); +} + } // namespace paimon diff --git a/src/paimon/common/reader/late_materializing_file_batch_reader.h b/src/paimon/common/reader/late_materializing_file_batch_reader.h index c1b91b086..3a2c7917e 100644 --- a/src/paimon/common/reader/late_materializing_file_batch_reader.h +++ b/src/paimon/common/reader/late_materializing_file_batch_reader.h @@ -49,6 +49,7 @@ class LateMaterializingFileBatchReader : public PrefetchFileBatchReader { }; void Close() override { + Reset(); inner_->Close(); } @@ -95,6 +96,9 @@ class LateMaterializingFileBatchReader : public PrefetchFileBatchReader { std::shared_ptr arrow_pool) : inner_(std::move(inner)), arrow_pool_(std::move(arrow_pool)) {} + /// Reset the state of the late materializing reader, does not close inner reader. + void Reset(); + enum LatMatState { kInit, kProbing, // schema is set, probing is in progress From 7f4712742145176358ecb1269ce017ef3033a328 Mon Sep 17 00:00:00 2001 From: zhouhongfeng Date: Wed, 26 Aug 2026 11:27:50 +0800 Subject: [PATCH 30/34] fix: return invalid when bitmap is empty --- .../common/reader/late_materializing_file_batch_reader.cpp | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/paimon/common/reader/late_materializing_file_batch_reader.cpp b/src/paimon/common/reader/late_materializing_file_batch_reader.cpp index e399d9edf..0ce7bda42 100644 --- a/src/paimon/common/reader/late_materializing_file_batch_reader.cpp +++ b/src/paimon/common/reader/late_materializing_file_batch_reader.cpp @@ -164,7 +164,7 @@ Result LateMaterializingFileBatchReader::ReadPayload auto& [batch, bitmap] = batch_with_bitmap; if (bitmap.IsEmpty()) { ReaderUtils::ReleaseReadBatch(std::move(batch)); - continue; + return Status::Invalid("inner read bitmap is empty."); } auto& [c_array, c_schema] = batch; PAIMON_ASSIGN_OR_RAISE_FROM_ARROW(std::shared_ptr payload_array, From 68ad0e898fa1d5e660c4a52a21aa952643c329a8 Mon Sep 17 00:00:00 2001 From: zhouhongfeng Date: Wed, 26 Aug 2026 14:13:24 +0800 Subject: [PATCH 31/34] test: update test cases --- ...e_materializing_file_batch_reader_test.cpp | 25 ++++++------------- test/inte/read_inte_test.cpp | 8 +++--- test/inte/scan_and_read_inte_test.cpp | 8 ++---- 3 files changed, 14 insertions(+), 27 deletions(-) diff --git a/src/paimon/common/reader/late_materializing_file_batch_reader_test.cpp b/src/paimon/common/reader/late_materializing_file_batch_reader_test.cpp index f948a702e..611e50f92 100644 --- a/src/paimon/common/reader/late_materializing_file_batch_reader_test.cpp +++ b/src/paimon/common/reader/late_materializing_file_batch_reader_test.cpp @@ -49,6 +49,7 @@ #include "paimon/testing/mock/mock_file_batch_reader.h" #include "paimon/testing/mock/mock_file_system.h" #include "paimon/testing/mock/mock_format_reader_builder.h" +#include "paimon/testing/utils/read_result_collector.h" #include "paimon/testing/utils/testharness.h" #include "paimon/utils/roaring_bitmap32.h" @@ -128,28 +129,16 @@ class LateMaterializingFileBatchReaderTest : public ::testing::Test { return reader->SetReadSchema(&c_schema, predicate, selection); } - // Collect all output rows as a single concatenated struct array (for schema/nested checks). + // Collect all output rows as a single concatenated struct array (for schema/nested checks), + // reusing the shared collector so the batch-offset and bitmap contracts are checked too. Result> CollectStruct(FileBatchReader* reader) { - arrow::ArrayVector chunks; - while (true) { - PAIMON_ASSIGN_OR_RAISE(BatchReader::ReadBatchWithBitmap batch_with_bitmap, - reader->NextBatchWithBitmap()); - if (BatchReader::IsEofBatch(batch_with_bitmap)) { - break; - } - PAIMON_ASSIGN_OR_RAISE(BatchReader::ReadBatch batch, - ReaderUtils::ApplyBitmapToReadBatch( - std::move(batch_with_bitmap), arrow::default_memory_pool())); - auto& [c_array, c_schema] = batch; - PAIMON_ASSIGN_OR_RAISE_FROM_ARROW(std::shared_ptr array, - arrow::ImportArray(c_array.get(), c_schema.get())); - chunks.push_back(array); - } - if (chunks.empty()) { + PAIMON_ASSIGN_OR_RAISE(std::shared_ptr chunked, + ReadResultCollector::CollectResult(reader)); + if (chunked == nullptr) { return std::shared_ptr(); } PAIMON_ASSIGN_OR_RAISE_FROM_ARROW(std::shared_ptr combined, - arrow::Concatenate(chunks)); + arrow::Concatenate(chunked->chunks())); return arrow::internal::checked_pointer_cast(combined); } diff --git a/test/inte/read_inte_test.cpp b/test/inte/read_inte_test.cpp index b218d133b..a2a52d343 100644 --- a/test/inte/read_inte_test.cpp +++ b/test/inte/read_inte_test.cpp @@ -2292,7 +2292,6 @@ TEST_P(ReadInteTest, TestAppendReadWithLateMaterializing) { .AddOption("test.enable-adaptive-prefetch-strategy", param.enable_adaptive_prefetch_strategy) .SetPredicate(predicate) - .EnablePredicateFilter(true) .EnableLateMaterializing(true) .EnablePrefetch(param.enable_prefetch); ASSERT_OK_AND_ASSIGN(auto read_context, context_builder.Finish()); @@ -3278,8 +3277,7 @@ TEST_P(ReadInteTest, TestPkReadSnapshot6WithSchemaEvolutionWithLateMaterializing .AddOption("read.batch-size", "2"); context_builder.SetReadAheadCacheEnabled(param.read_ahead_cache_enabled); context_builder.SetPredicate(predicate); - context_builder.EnablePredicateFilter(true) - .EnableLateMaterializing(true) + context_builder.EnableLateMaterializing(true) .EnablePrefetch(param.enable_prefetch) .AddOption("test.enable-adaptive-prefetch-strategy", param.enable_adaptive_prefetch_strategy); @@ -3327,6 +3325,10 @@ TEST_P(ReadInteTest, TestPkReadSnapshot6WithSchemaEvolutionWithLateMaterializing // "Paul" is the only row in partition key0 = 0 whose e is not null and matches e < 510. std::shared_ptr expected_array; auto array_status = arrow::ipc::internal::json::ChunkedArrayFromJSON(arrow_data_type, {R"([ + [0, 1, "Bob", 22, 24, null, 26, 1, null], + [0, 1, "Emily", 32, 34, null, 36, 1, null], + [0, 1, "David", 62, 64, null, 66, 1, null], + [0, 1, "Whether I shall turn out to be the hero of my own life.", 72, 74, null, 76, 1, null], [0, 1, "Paul", 502, 504, 508, 506, 0, 509] ])"}, &expected_array); diff --git a/test/inte/scan_and_read_inte_test.cpp b/test/inte/scan_and_read_inte_test.cpp index 12d5ea8ea..537a4cb36 100644 --- a/test/inte/scan_and_read_inte_test.cpp +++ b/test/inte/scan_and_read_inte_test.cpp @@ -764,9 +764,7 @@ TEST_P(ScanAndReadInteTest, TestWithPKWithDvBatchScanSnapshot6WithLateMaterializ ReadContextBuilder read_context_builder(table_path); AddReadOptionsForPrefetch(&read_context_builder); - read_context_builder.SetPredicate(predicate) - .EnablePredicateFilter(true) - .EnableLateMaterializing(true); + read_context_builder.SetPredicate(predicate).EnableLateMaterializing(true); ASSERT_OK_AND_ASSIGN(auto read_context, read_context_builder.Finish()); ASSERT_OK_AND_ASSIGN(auto table_read, TableRead::Create(std::move(read_context))); @@ -1348,9 +1346,7 @@ TEST_P(ScanAndReadInteTest, TestWithPKWithMorBatchScanSnapshot5WithLateMateriali ReadContextBuilder read_context_builder(table_path); AddReadOptionsForPrefetch(&read_context_builder); - read_context_builder.SetPredicate(predicate) - .EnablePredicateFilter(true) - .EnableLateMaterializing(true); + read_context_builder.SetPredicate(predicate).EnableLateMaterializing(true); ASSERT_OK_AND_ASSIGN(auto read_context, read_context_builder.Finish()); ASSERT_OK_AND_ASSIGN(auto table_read, TableRead::Create(std::move(read_context))); From 585114393de29062c86a90ec50c37e28e2bb6a24 Mon Sep 17 00:00:00 2001 From: zhouhongfeng Date: Wed, 26 Aug 2026 14:57:37 +0800 Subject: [PATCH 32/34] test: diable lat-mat in paimon-global-index-test and paimon-blob-table-inte-test --- test/inte/blob_table_inte_test.cpp | 1 + test/inte/global_index_test.cpp | 3 ++- 2 files changed, 3 insertions(+), 1 deletion(-) diff --git a/test/inte/blob_table_inte_test.cpp b/test/inte/blob_table_inte_test.cpp index ea951a311..b3c5cb237 100644 --- a/test/inte/blob_table_inte_test.cpp +++ b/test/inte/blob_table_inte_test.cpp @@ -299,6 +299,7 @@ class BlobTableInteTest : public testing::Test, public ::testing::WithParamInter auto splits = plan->Splits(); ReadContextBuilder read_context_builder(table_path); read_context_builder.SetReadFieldNames(read_schema).SetPredicate(predicate); + read_context_builder.EnableLateMaterializing(false); if (!options.empty()) { read_context_builder.SetOptions(options); } diff --git a/test/inte/global_index_test.cpp b/test/inte/global_index_test.cpp index ccb9ce324..84ab22df8 100644 --- a/test/inte/global_index_test.cpp +++ b/test/inte/global_index_test.cpp @@ -209,7 +209,8 @@ class GlobalIndexTest : public ::testing::Test, public ::testing::WithParamInter ReadContextBuilder read_context_builder(table_path); read_context_builder.SetReadFieldNames(read_schema) .SetPredicate(predicate) - .WithFileSystem(fs_); + .WithFileSystem(fs_) + .EnableLateMaterializing(false); PAIMON_ASSIGN_OR_RAISE(std::unique_ptr read_context, read_context_builder.Finish()); PAIMON_ASSIGN_OR_RAISE(auto table_read, TableRead::Create(std::move(read_context))); From 3f654c7280ffdb786d8029b3e9b0daeed5fd249e Mon Sep 17 00:00:00 2001 From: zhouhongfeng Date: Wed, 26 Aug 2026 18:15:11 +0800 Subject: [PATCH 33/34] feat: support non-prefetch format (avro, blob, etc.) --- .../reader/prefetch_file_batch_reader.h | 2 +- .../late_materializing_file_batch_reader.cpp | 30 ++++++----- .../late_materializing_file_batch_reader.h | 51 +++++++++++++++---- .../late_materializing_reader_builder.h | 14 ++--- .../prefetch_file_batch_reader_impl.cpp | 9 ++-- .../reader/prefetch_file_batch_reader_impl.h | 2 +- .../core/operation/abstract_split_read.cpp | 43 ++++++++-------- src/paimon/format/orc/orc_file_batch_reader.h | 2 +- .../parquet/parquet_file_batch_reader.cpp | 2 + .../parquet/parquet_file_batch_reader.h | 5 +- .../testing/mock/mock_file_batch_reader.h | 3 +- 11 files changed, 101 insertions(+), 62 deletions(-) diff --git a/include/paimon/reader/prefetch_file_batch_reader.h b/include/paimon/reader/prefetch_file_batch_reader.h index acc7d0bbd..5e6313e8f 100644 --- a/include/paimon/reader/prefetch_file_batch_reader.h +++ b/include/paimon/reader/prefetch_file_batch_reader.h @@ -40,7 +40,7 @@ class PAIMON_EXPORT PrefetchFileBatchReader : public FileBatchReader { /// Retrieves the row number of the next row to be read. /// This method indicates the current read position within the file. /// @return The row number of the next row to read. - virtual uint64_t GetNextRowToRead() const = 0; + virtual Result GetNextRowToRead() const = 0; /// Generates a list of row ranges to be read in batches. /// Each range specifies the start and end row numbers for a batch, diff --git a/src/paimon/common/reader/late_materializing_file_batch_reader.cpp b/src/paimon/common/reader/late_materializing_file_batch_reader.cpp index 0ce7bda42..ed7fa304d 100644 --- a/src/paimon/common/reader/late_materializing_file_batch_reader.cpp +++ b/src/paimon/common/reader/late_materializing_file_batch_reader.cpp @@ -19,6 +19,7 @@ #include "paimon/common/reader/late_materializing_file_batch_reader.h" +#include #include #include #include @@ -44,15 +45,20 @@ namespace paimon { Result> LateMaterializingFileBatchReader::Create( - std::unique_ptr inner, std::shared_ptr pool) { + std::unique_ptr inner, std::shared_ptr pool) { // The reader's own compaction allocations go through an arrow pool; bridge the paimon pool // once here so the accounting matches the rest of the read path. if (pool == nullptr) { return Status::Invalid("pool could not be nullptr."); } + if (inner == nullptr) { + return Status::Invalid("inner could not be nullptr."); + } + auto* prefetch_inner = dynamic_cast(inner.get()); std::shared_ptr arrow_pool = GetArrowPool(pool); - auto reader = std::unique_ptr( - new LateMaterializingFileBatchReader(std::move(inner), std::move(arrow_pool))); + auto reader = + std::unique_ptr(new LateMaterializingFileBatchReader( + std::move(inner), prefetch_inner, std::move(arrow_pool))); return reader; } @@ -243,11 +249,7 @@ Status LateMaterializingFileBatchReader::SetInnerReadSchema( const std::optional& selection) { ::ArrowSchema c_read_schema; PAIMON_RETURN_NOT_OK_FROM_ARROW(arrow::ExportSchema(*read_schema, &c_read_schema)); - /// Note: calling inner->SetReadSchema may refresh the read ranges of the inner reader. PAIMON_RETURN_NOT_OK(inner_->SetReadSchema(&c_read_schema, predicate, selection)); - if (!read_ranges_.empty()) { - PAIMON_RETURN_NOT_OK(inner_->SetReadRanges(read_ranges_)); - } return Status::OK(); } @@ -316,7 +318,9 @@ Result LateMaterializingFileBatchReader::GetPreviousBatchFileRowId( } Status LateMaterializingFileBatchReader::SeekToRow(uint64_t row_number) { - PAIMON_RETURN_NOT_OK(inner_->SeekToRow(row_number)); + PAIMON_ASSIGN_OR_RAISE(PrefetchFileBatchReader * prefetch_reader, + GetPrefetchReaderOrRaise("SeekToRow")); + PAIMON_RETURN_NOT_OK(prefetch_reader->SeekToRow(row_number)); if (state_ == kRunning || state_ == kEOF) { if (matched_bitmap_.IsEmpty()) { state_ = kEOF; @@ -338,14 +342,16 @@ Status LateMaterializingFileBatchReader::SeekToRow(uint64_t row_number) { Status LateMaterializingFileBatchReader::SetReadRanges( const std::vector>& read_ranges) { - read_ranges_ = read_ranges; - PAIMON_RETURN_NOT_OK(inner_->SetReadRanges(read_ranges_)); - return Status::OK(); + if (prefetch_inner_ == nullptr) { + // Only the format reader can act on this hint, and the PrefetchFileBatchReader contract + // lets an implementation that cannot honor it ignore the hint. + return Status::OK(); + } + return prefetch_inner_->SetReadRanges(read_ranges); } void LateMaterializingFileBatchReader::Reset() { state_ = kInit; - read_ranges_.clear(); matched_bitmap_ = RoaringBitmap32(); probe_data_.reset(); probe_cursor_ = 0; diff --git a/src/paimon/common/reader/late_materializing_file_batch_reader.h b/src/paimon/common/reader/late_materializing_file_batch_reader.h index 3a2c7917e..231625db6 100644 --- a/src/paimon/common/reader/late_materializing_file_batch_reader.h +++ b/src/paimon/common/reader/late_materializing_file_batch_reader.h @@ -24,9 +24,11 @@ #include #include +#include #include #include +#include "fmt/format.h" #include "paimon/reader/prefetch_file_batch_reader.h" namespace paimon { @@ -40,7 +42,7 @@ class PredicateFilter; class LateMaterializingFileBatchReader : public PrefetchFileBatchReader { public: static Result> Create( - std::unique_ptr inner, std::shared_ptr pool); + std::unique_ptr inner, std::shared_ptr pool); Result NextBatch() override; @@ -74,13 +76,17 @@ class LateMaterializingFileBatchReader : public PrefetchFileBatchReader { Status SeekToRow(uint64_t row_number) override; - uint64_t GetNextRowToRead() const override { - return inner_->GetNextRowToRead(); + Result GetNextRowToRead() const override { + PAIMON_ASSIGN_OR_RAISE(PrefetchFileBatchReader * prefetch_reader, + GetPrefetchReaderOrRaise("GetNextRowToRead")); + return prefetch_reader->GetNextRowToRead(); } Result>> GenReadRanges( bool* need_prefetch) const override { - return inner_->GenReadRanges(need_prefetch); + PAIMON_ASSIGN_OR_RAISE(PrefetchFileBatchReader * prefetch_reader, + GetPrefetchReaderOrRaise("GenReadRanges")); + return prefetch_reader->GenReadRanges(need_prefetch); } Status SetReadRanges(const std::vector>& read_ranges) override; @@ -88,13 +94,19 @@ class LateMaterializingFileBatchReader : public PrefetchFileBatchReader { Result>> PreBufferRange() override { // TODO(zhouhongfeng.zhf): PrebufferRange (called by PrefetchFileBatchReader) only read the // probe data, consider read the payload data as well. - return inner_->PreBufferRange(); + if (prefetch_inner_ == nullptr) { + return std::vector>{}; + } + return prefetch_inner_->PreBufferRange(); } private: - explicit LateMaterializingFileBatchReader(std::unique_ptr inner, - std::shared_ptr arrow_pool) - : inner_(std::move(inner)), arrow_pool_(std::move(arrow_pool)) {} + LateMaterializingFileBatchReader(std::unique_ptr inner, + PrefetchFileBatchReader* prefetch_inner, + std::shared_ptr arrow_pool) + : inner_(std::move(inner)), + prefetch_inner_(prefetch_inner), + arrow_pool_(std::move(arrow_pool)) {} /// Reset the state of the late materializing reader, does not close inner reader. void Reset(); @@ -128,10 +140,29 @@ class LateMaterializingFileBatchReader : public PrefetchFileBatchReader { const std::shared_ptr& predicate, const std::optional& selection); - std::unique_ptr inner_; + /// Returns the inner reader's prefetch interface, or an error when the format reader does not + /// implement it (avro and blob do not). + Result GetPrefetchReaderOrRaise( + std::string_view function_name) const { + if (prefetch_inner_ == nullptr) { + return Status::NotImplemented( + fmt::format("format reader is not a prefetch reader, function {} not supported", + function_name)); + } + return prefetch_inner_; + } + + /// The probe/payload logic needs nothing beyond FileBatchReader, so the inner reader is held as + /// the base type: parquet and orc readers implement PrefetchFileBatchReader, while avro and + /// blob readers only implement FileBatchReader. The prefetch-only methods are rejected for the + /// latter; AbstractSplitRead::CreateFileBatchReader keeps those formats out of the prefetch + /// layer so nothing calls them. + std::unique_ptr inner_; + /// Non-owning view of inner_ when it implements the prefetch interface, nullptr otherwise. + /// inner_ is never reassigned, so the cast is resolved once in Create(). + PrefetchFileBatchReader* prefetch_inner_ = nullptr; std::shared_ptr arrow_pool_; LatMatState state_ = kInit; - std::vector> read_ranges_; std::shared_ptr full_schema_; // projection holding only the predicate fields; nullptr when probing is not applicable std::shared_ptr probe_schema_; diff --git a/src/paimon/common/reader/late_materializing_reader_builder.h b/src/paimon/common/reader/late_materializing_reader_builder.h index 0420a35cc..5c7be5077 100644 --- a/src/paimon/common/reader/late_materializing_reader_builder.h +++ b/src/paimon/common/reader/late_materializing_reader_builder.h @@ -54,15 +54,11 @@ class LateMaterializingReaderBuilder : public ReaderBuilder { Result> Build( const std::shared_ptr& stream) const override { - PAIMON_ASSIGN_OR_RAISE(std::unique_ptr base, inner_->Build(stream)); - auto* prefetch = dynamic_cast(base.get()); - if (prefetch == nullptr) { - return Status::Invalid("Late materialization requires prefetch interface"); - } - base.release(); - PAIMON_ASSIGN_OR_RAISE(std::unique_ptr reader, - LateMaterializingFileBatchReader::Create( - std::unique_ptr(prefetch), pool_)); + PAIMON_ASSIGN_OR_RAISE(std::unique_ptr format_reader, + inner_->Build(stream)); + PAIMON_ASSIGN_OR_RAISE( + std::unique_ptr reader, + LateMaterializingFileBatchReader::Create(std::move(format_reader), pool_)); return std::unique_ptr(std::move(reader)); } diff --git a/src/paimon/common/reader/prefetch_file_batch_reader_impl.cpp b/src/paimon/common/reader/prefetch_file_batch_reader_impl.cpp index 12e38966d..369bb58fe 100644 --- a/src/paimon/common/reader/prefetch_file_batch_reader_impl.cpp +++ b/src/paimon/common/reader/prefetch_file_batch_reader_impl.cpp @@ -411,7 +411,8 @@ std::optional> PrefetchFileBatchReaderImpl::GetCur Status PrefetchFileBatchReaderImpl::EnsureReaderPosition( size_t reader_idx, const std::pair& current_read_range) const { uint64_t pos = std::max(readers_pos_[reader_idx]->load(), current_read_range.first); - if (readers_[reader_idx]->GetNextRowToRead() != pos) { + PAIMON_ASSIGN_OR_RAISE(uint64_t next_row_to_read, readers_[reader_idx]->GetNextRowToRead()); + if (next_row_to_read != pos) { return readers_[reader_idx]->SeekToRow(pos); } return Status::OK(); @@ -480,7 +481,9 @@ Status PrefetchFileBatchReaderImpl::HandleReadResult( } else { // all within the range, data before readers_[reader_idx]->GetNextRowToRead() has been // effectively consumed - readers_pos_[reader_idx]->store(readers_[reader_idx]->GetNextRowToRead()); + PAIMON_ASSIGN_OR_RAISE(uint64_t next_row_to_read, + readers_[reader_idx]->GetNextRowToRead()); + readers_pos_[reader_idx]->store(next_row_to_read); } if (bitmap.IsEmpty()) { ReaderUtils::ReleaseReadBatch(std::move(read_batch)); @@ -646,7 +649,7 @@ Result PrefetchFileBatchReaderImpl::GetNumberOfRows() const { return readers_[0]->GetNumberOfRows(); } -uint64_t PrefetchFileBatchReaderImpl::GetNextRowToRead() const { +Result PrefetchFileBatchReaderImpl::GetNextRowToRead() const { assert(false); return -1; } diff --git a/src/paimon/common/reader/prefetch_file_batch_reader_impl.h b/src/paimon/common/reader/prefetch_file_batch_reader_impl.h index c21856d09..4750501e2 100644 --- a/src/paimon/common/reader/prefetch_file_batch_reader_impl.h +++ b/src/paimon/common/reader/prefetch_file_batch_reader_impl.h @@ -80,7 +80,7 @@ class PrefetchFileBatchReaderImpl : public PrefetchFileBatchReader { Status SeekToRow(uint64_t row_number) override; Result GetPreviousBatchFileRowId(uint64_t batch_row_id) const override; Result GetNumberOfRows() const override; - uint64_t GetNextRowToRead() const override; + Result GetNextRowToRead() const override; void Close() override; Status SetReadRanges(const std::vector>& read_ranges) override; diff --git a/src/paimon/core/operation/abstract_split_read.cpp b/src/paimon/core/operation/abstract_split_read.cpp index 6723971c5..89f5071e9 100644 --- a/src/paimon/core/operation/abstract_split_read.cpp +++ b/src/paimon/core/operation/abstract_split_read.cpp @@ -153,29 +153,28 @@ Result> AbstractSplitRead::PrepareReaderBuilder( Result> AbstractSplitRead::CreateFileBatchReader( const std::string& file_format_identifier, const std::string& data_file_path, int64_t data_file_size, std::unique_ptr reader_builder) const { - // blob and avro do not support prefetch or late materializing - if (file_format_identifier != "blob" && file_format_identifier != "avro") { - if (context_->EnableLateMaterializing()) { - reader_builder = - std::make_unique(std::move(reader_builder), pool_); - } - if (context_->EnablePrefetch()) { - PAIMON_ASSIGN_OR_RAISE( - std::unique_ptr prefetch_reader, - PrefetchFileBatchReaderImpl::Create( - data_file_path, data_file_size, reader_builder.get(), options_.GetFileSystem(), - context_->GetPrefetchMaxParallelNum(), options_.GetReadBatchSize(), - context_->GetPrefetchBatchCount(), options_.EnableAdaptivePrefetchStrategy(), - executor_, - /*initialize_read_ranges=*/false, context_->ReadAheadCacheEnabled(), - context_->GetCacheConfig(), pool_)); - return std::make_unique(std::move(prefetch_reader)); - } + if (context_->EnableLateMaterializing()) { + reader_builder = + std::make_unique(std::move(reader_builder), pool_); + } + if (context_->EnablePrefetch() && file_format_identifier != "blob" && + file_format_identifier != "avro") { + PAIMON_ASSIGN_OR_RAISE( + std::unique_ptr prefetch_reader, + PrefetchFileBatchReaderImpl::Create( + data_file_path, data_file_size, reader_builder.get(), options_.GetFileSystem(), + context_->GetPrefetchMaxParallelNum(), options_.GetReadBatchSize(), + context_->GetPrefetchBatchCount(), options_.EnableAdaptivePrefetchStrategy(), + executor_, + /*initialize_read_ranges=*/false, context_->ReadAheadCacheEnabled(), + context_->GetCacheConfig(), pool_)); + return std::make_unique(std::move(prefetch_reader)); + } else { + PAIMON_ASSIGN_OR_RAISE( + std::shared_ptr input_stream, + options_.GetFileSystem()->Open(FileStatus(data_file_path, data_file_size))); + return reader_builder->Build(input_stream); } - PAIMON_ASSIGN_OR_RAISE( - std::shared_ptr input_stream, - options_.GetFileSystem()->Open(FileStatus(data_file_path, data_file_size))); - return reader_builder->Build(input_stream); } Result> AbstractSplitRead::CreateFieldMappingReader( diff --git a/src/paimon/format/orc/orc_file_batch_reader.h b/src/paimon/format/orc/orc_file_batch_reader.h index 85673a93e..b48f7cdb0 100644 --- a/src/paimon/format/orc/orc_file_batch_reader.h +++ b/src/paimon/format/orc/orc_file_batch_reader.h @@ -85,7 +85,7 @@ class OrcFileBatchReader : public PrefetchFileBatchReader { return reader_->GetNumberOfRows(); } - uint64_t GetNextRowToRead() const override { + Result GetNextRowToRead() const override { return reader_->GetNextRowToRead(); } diff --git a/src/paimon/format/parquet/parquet_file_batch_reader.cpp b/src/paimon/format/parquet/parquet_file_batch_reader.cpp index 0c9d065e5..7605c4242 100644 --- a/src/paimon/format/parquet/parquet_file_batch_reader.cpp +++ b/src/paimon/format/parquet/parquet_file_batch_reader.cpp @@ -154,6 +154,7 @@ ParquetFileBatchReader::ParquetFileBatchReader( arrow_pool_(arrow_pool), input_stream_(std::move(input_stream)), reader_(std::move(reader)), + read_ranges_(reader_->GetAllRowGroupRanges()), metrics_(std::make_shared()), storage_read_bytes_(std::move(storage_read_bytes)), logger_(Logger::GetLogger("ParquetFileBatchReader")) {} @@ -309,6 +310,7 @@ Status ParquetFileBatchReader::SetReadSchema( PAIMON_RETURN_NOT_OK(UpdateAllTargetRowRanges(target_row_groups)); PAIMON_RETURN_NOT_OK(reader_->PrepareForReadingLazy(target_row_groups, column_indices)); + PAIMON_RETURN_NOT_OK(reader_->ApplyReadRanges(read_ranges_)); } PAIMON_PARQUET_CATCH_AND_RETURN_STATUS("ParquetFileBatchReader::SetReadSchema") return Status::OK(); diff --git a/src/paimon/format/parquet/parquet_file_batch_reader.h b/src/paimon/format/parquet/parquet_file_batch_reader.h index daa18f040..2b1097fb4 100644 --- a/src/paimon/format/parquet/parquet_file_batch_reader.h +++ b/src/paimon/format/parquet/parquet_file_batch_reader.h @@ -128,12 +128,13 @@ class ParquetFileBatchReader : public PrefetchFileBatchReader { return reader_->GetNumberOfRows(); } - uint64_t GetNextRowToRead() const override { + Result GetNextRowToRead() const override { assert(reader_); return reader_->GetNextRowToRead(); } Status SetReadRanges(const std::vector>& read_ranges) override { + read_ranges_ = read_ranges; return reader_->ApplyReadRanges(read_ranges); } @@ -261,6 +262,8 @@ class ParquetFileBatchReader : public PrefetchFileBatchReader { std::shared_ptr read_data_type_; + std::vector> read_ranges_; + std::shared_ptr metrics_; // storageReadBytes counter shared with the underlying ArrowInputStreamAdapter. std::shared_ptr> storage_read_bytes_; diff --git a/src/paimon/testing/mock/mock_file_batch_reader.h b/src/paimon/testing/mock/mock_file_batch_reader.h index fb3cc3ada..4566289f4 100644 --- a/src/paimon/testing/mock/mock_file_batch_reader.h +++ b/src/paimon/testing/mock/mock_file_batch_reader.h @@ -87,7 +87,6 @@ class MockFileBatchReader : public PrefetchFileBatchReader { // late-materialization reader moving from its probe pass to its payload pass, rely on it. current_pos_ = 0; previous_batch_first_row_num_ = std::numeric_limits::max(); - read_ranges_.clear(); return Status::OK(); } @@ -196,7 +195,7 @@ class MockFileBatchReader : public PrefetchFileBatchReader { Result GetNumberOfRows() const override { return ToReaderRowNumber(read_end_pos_); } - uint64_t GetNextRowToRead() const override { + Result GetNextRowToRead() const override { return ToReaderRowNumber(current_pos_); } void Close() override {} From a02e7f57258430a4993edbc92cf0ec65b6626976 Mon Sep 17 00:00:00 2001 From: zhouhongfeng Date: Wed, 26 Aug 2026 18:26:42 +0800 Subject: [PATCH 34/34] test: update lat-mat tests --- ...e_materializing_file_batch_reader_test.cpp | 35 +++++++++++++++---- 1 file changed, 29 insertions(+), 6 deletions(-) diff --git a/src/paimon/common/reader/late_materializing_file_batch_reader_test.cpp b/src/paimon/common/reader/late_materializing_file_batch_reader_test.cpp index 611e50f92..ff571a284 100644 --- a/src/paimon/common/reader/late_materializing_file_batch_reader_test.cpp +++ b/src/paimon/common/reader/late_materializing_file_batch_reader_test.cpp @@ -239,6 +239,12 @@ TEST_F(LateMaterializingFileBatchReaderTest, PassThroughWhenPayloadEmpty) { ASSERT_OK_AND_ASSIGN(std::vector rows, Collect(reader.get())); ASSERT_EQ(rows.size(), 5u); + for (size_t idx = 0; idx < rows.size(); ++idx) { + EXPECT_EQ(rows[idx].k, static_cast(idx)); + // v is outside the read schema, so the pass-through output must not carry it + EXPECT_EQ(rows[idx].v, ""); + EXPECT_EQ(rows[idx].file_row, static_cast(idx)); + } } // Contiguous matched subset spanning multiple batches. @@ -308,12 +314,12 @@ TEST_F(LateMaterializingFileBatchReaderTest, MatchedIntersectsSelection) { ASSERT_OK_AND_ASSIGN(std::vector rows, Collect(reader.get())); ASSERT_EQ(rows.size(), 3u); - EXPECT_EQ(rows[0].file_row, 1u); - EXPECT_EQ(rows[1].file_row, 5u); - EXPECT_EQ(rows[2].file_row, 9u); - EXPECT_EQ(rows[0].v, "v_1"); - EXPECT_EQ(rows[1].v, "v_5"); - EXPECT_EQ(rows[2].v, "v_9"); + const std::vector expected_rows = {1u, 5u, 9u}; + for (size_t idx = 0; idx < rows.size(); ++idx) { + EXPECT_EQ(rows[idx].k, 1); + EXPECT_EQ(rows[idx].v, "v_" + std::to_string(expected_rows[idx])); + EXPECT_EQ(rows[idx].file_row, expected_rows[idx]); + } } // No matched rows: the reader returns EOF immediately. @@ -375,6 +381,12 @@ TEST_F(LateMaterializingFileBatchReaderTest, ReadRangesForwardedAcrossPhases) { // Drive to EOF; this performs the probe pass and the payload schema switch. ASSERT_OK_AND_ASSIGN(std::vector rows, Collect(reader.get())); ASSERT_EQ(rows.size(), 6u); // k = 2..7 + for (size_t idx = 0; idx < rows.size(); ++idx) { + int64_t expected = 2 + static_cast(idx); + EXPECT_EQ(rows[idx].k, expected); + EXPECT_EQ(rows[idx].v, "v_" + std::to_string(expected)); + EXPECT_EQ(rows[idx].file_row, static_cast(expected)); + } // The inner reader must have received the cached ranges again after the payload switch. // SetReadSchema (invoked on the payload switch) clears the inner reader's ranges, so the @@ -394,6 +406,12 @@ TEST_F(LateMaterializingFileBatchReaderTest, ReentrantSetReadSchema) { ASSERT_OK(SetReadSchema(reader.get(), arrow::schema(full_fields_), predicate1, std::nullopt)); ASSERT_OK_AND_ASSIGN(std::vector rows1, Collect(reader.get())); ASSERT_EQ(rows1.size(), 2u); // k = 8,9 + for (size_t idx = 0; idx < rows1.size(); ++idx) { + int64_t expected = 8 + static_cast(idx); + EXPECT_EQ(rows1[idx].k, expected); + EXPECT_EQ(rows1[idx].v, "v_" + std::to_string(expected)); + EXPECT_EQ(rows1[idx].file_row, static_cast(expected)); + } auto predicate2 = PredicateBuilder::LessThan(/*field_index=*/0, /*field_name=*/"k", FieldType::BIGINT, Literal(3l)); @@ -402,6 +420,7 @@ TEST_F(LateMaterializingFileBatchReaderTest, ReentrantSetReadSchema) { ASSERT_EQ(rows2.size(), 3u); // k = 0,1,2 for (size_t idx = 0; idx < rows2.size(); ++idx) { EXPECT_EQ(rows2[idx].k, static_cast(idx)); + EXPECT_EQ(rows2[idx].v, "v_" + std::to_string(idx)); EXPECT_EQ(rows2[idx].file_row, static_cast(idx)); } } @@ -562,8 +581,12 @@ TEST_F(LateMaterializingFileBatchReaderTest, PrefetchInnerReentrantSetReadSchema ASSERT_EQ(result1->length(), 3); // k = 7,8,9 auto k1 = arrow::internal::checked_pointer_cast(result1->GetFieldByName("k")); + auto v1 = + arrow::internal::checked_pointer_cast(result1->GetFieldByName("v")); + ASSERT_TRUE(k1 && v1); for (int64_t j = 0; j < result1->length(); ++j) { EXPECT_EQ(k1->Value(j), 7 + j); + EXPECT_EQ(v1->GetString(j), "v_" + std::to_string(7 + j)); } auto predicate2 =