diff --git a/include/paimon/read_context.h b/include/paimon/read_context.h index 3e58b1c45..e645645ba 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,9 @@ class PAIMON_EXPORT ReadContext { bool EnablePrefetch() const { return enable_prefetch_; } + bool EnableLateMaterializing() const { + return enable_late_materializing_; + } uint32_t GetPrefetchBatchCount() const { return prefetch_batch_count_; } @@ -163,6 +166,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 +310,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). 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 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/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/CMakeLists.txt b/src/paimon/CMakeLists.txt index 61fd7c96d..8fb6c17d1 100644 --- a/src/paimon/CMakeLists.txt +++ b/src/paimon/CMakeLists.txt @@ -134,6 +134,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 @@ -606,6 +607,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.cpp b/src/paimon/common/reader/late_materializing_file_batch_reader.cpp new file mode 100644 index 000000000..ed7fa304d --- /dev/null +++ b/src/paimon/common/reader/late_materializing_file_batch_reader.cpp @@ -0,0 +1,369 @@ +/* + * 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 "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/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" +#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, 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), prefetch_inner, std::move(arrow_pool))); + return reader; +} + +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()) { + state_ = kEOF; + } else { + // 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; + } + } + + if (state_ == kNoLatMat) { + return inner_->NextBatch(); + } else if (state_ == kRunning) { + return ReadPayloadBatch(); + } else if (state_ == kEOF) { + return MakeEofBatch(); + } + return Status::Invalid("invalid state when calling NextBatch: " + std::to_string(state_)); +} + +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( + 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; + } + // map batch offset to file row id + PAIMON_ASSIGN_OR_RAISE(uint64_t file_row, + inner_->GetPreviousBatchFileRowId(static_cast(i))); + if (selection_ && !selection_->Contains(static_cast(file_row))) { + continue; + } + batch_matched.Add(static_cast(i)); + matched_bitmap_.Add(file_row); + } + return batch_matched; +} + +Status LateMaterializingFileBatchReader::ReadAndFilterProbeData() { + 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, 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()) { + 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, arrow_pool_.get())); + } + 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; + 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; + if (bitmap.IsEmpty()) { + ReaderUtils::ReleaseReadBatch(std::move(batch)); + 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, + arrow::ImportArray(c_array.get(), c_schema.get())); + + // Generate the valid bitmap and row_mapping_ + RoaringBitmap32 valid; + row_mapping_.clear(); + for (auto it = bitmap.Begin(); it != bitmap.End(); ++it) { + auto 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, valid)); + PAIMON_ASSIGN_OR_RAISE_FROM_ARROW(std::shared_ptr payload_compacted, + arrow::Concatenate(payload_slices, arrow_pool_.get())); + + 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 {}", + 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, + AssembleFullBatch(payload_compacted, probe_selected)); + return 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())); + } + 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, + 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::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)); + PAIMON_RETURN_NOT_OK(inner_->SetReadSchema(&c_read_schema, predicate, selection)); + return Status::OK(); +} + +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; + 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()); + 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"); + } + } + } + + if (predicate_ == nullptr || probe_schema_ == nullptr) { + PAIMON_RETURN_NOT_OK(SetInnerReadSchema(full_schema_, predicate_, selection_)); + state_ = kNoLatMat; + } else { + PAIMON_RETURN_NOT_OK(SetInnerReadSchema(probe_schema_, predicate_, selection_)); + 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]; +} + +Status LateMaterializingFileBatchReader::SeekToRow(uint64_t row_number) { + PAIMON_ASSIGN_OR_RAISE(PrefetchFileBatchReader * prefetch_reader, + GetPrefetchReaderOrRaise("SeekToRow")); + PAIMON_RETURN_NOT_OK(prefetch_reader->SeekToRow(row_number)); + if (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) { + break; + } + ++cursor; + } + probe_cursor_ = cursor; + // a seek after EOF re-activates payload reading + state_ = kRunning; + } + return Status::OK(); +} + +Status LateMaterializingFileBatchReader::SetReadRanges( + const std::vector>& read_ranges) { + 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; + 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 new file mode 100644 index 000000000..231625db6 --- /dev/null +++ b/src/paimon/common/reader/late_materializing_file_batch_reader.h @@ -0,0 +1,184 @@ +/* + * 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 +#include + +#include "fmt/format.h" +#include "paimon/reader/prefetch_file_batch_reader.h" + +namespace paimon { + +class PredicateFilter; + +// For convenience, we abbreviate `Later Materializing` as `LatMat`. +// This reader is installed below the prefetch layer (see +// AbstractSplitRead::CreateFileBatchReader) and performs probe/payload two-phase reads when a +// predicate is pushed down through SetReadSchema; without a predicate it is a plain passthrough. +class LateMaterializingFileBatchReader : public PrefetchFileBatchReader { + public: + static Result> Create( + std::unique_ptr inner, std::shared_ptr pool); + + Result NextBatch() override; + + std::shared_ptr GetReaderMetrics() const override { + return inner_->GetReaderMetrics(); + }; + + void Close() override { + Reset(); + inner_->Close(); + } + + Result> GetFileSchema() const override { + return inner_->GetFileSchema(); + } + + 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 { + return inner_->GetNumberOfRows(); + } + + 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(); + } + + Status SeekToRow(uint64_t row_number) override; + + Result GetNextRowToRead() const override { + PAIMON_ASSIGN_OR_RAISE(PrefetchFileBatchReader * prefetch_reader, + GetPrefetchReaderOrRaise("GetNextRowToRead")); + return prefetch_reader->GetNextRowToRead(); + } + + Result>> GenReadRanges( + bool* need_prefetch) const override { + PAIMON_ASSIGN_OR_RAISE(PrefetchFileBatchReader * prefetch_reader, + GetPrefetchReaderOrRaise("GenReadRanges")); + return prefetch_reader->GenReadRanges(need_prefetch); + } + + 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. + if (prefetch_inner_ == nullptr) { + return std::vector>{}; + } + return prefetch_inner_->PreBufferRange(); + } + + private: + 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(); + + enum LatMatState { + kInit, + kProbing, // schema is set, probing is in progress + kNoLatMat, // no need to late materialization + kRunning, // Lat-mat is enabled and the payload reader is reading data + kEOF + }; + + /// 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(); + + Result FilterProbeBatch(const std::shared_ptr& array, + const std::shared_ptr& bound_filter); + + /// 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 + /// array following full_schema_'s field order. + Result AssembleFullBatch( + const std::shared_ptr& payload_array, + const std::shared_ptr& probe_array); + + Status SetInnerReadSchema(const std::shared_ptr& read_schema, + const std::shared_ptr& predicate, + const std::optional& selection); + + /// 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::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_; + // 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_; + RoaringBitmap32 matched_bitmap_; + // read cursor into probe_data_ for the payload phase + int64_t probe_cursor_ = 0; + // to support GetPreviousBatchFileRowId + std::vector row_mapping_; +}; + +} // namespace paimon 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..ff571a284 --- /dev/null +++ b/src/paimon/common/reader/late_materializing_file_batch_reader_test.cpp @@ -0,0 +1,670 @@ +/* + * 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/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" +#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/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" + +namespace paimon::test { + +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), + // reusing the shared collector so the batch-offset and bitmap contracts are checked too. + Result> CollectStruct(FileBatchReader* reader) { + 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(chunked->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), GetDefaultPool())); + 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), 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)); + 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); + 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. +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), 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)); + + 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), 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)); + + 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), 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} + 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); + 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. +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), 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)); + + 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), 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)); + + // 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), 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)); + + 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 + 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 + // 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), GetDefaultPool())); + + 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 + 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)); + 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].v, "v_" + std::to_string(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), GetDefaultPool())); + + 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), 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)); + 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), GetDefaultPool())); + // 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}); + 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( + 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}); + 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( + 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")); + 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 = + 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. + 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( + 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(); +} + +// 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 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..5c7be5077 --- /dev/null +++ b/src/paimon/common/reader/late_materializing_reader_builder.h @@ -0,0 +1,70 @@ +/* + * 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 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)); + } + + private: + std::unique_ptr inner_; + std::shared_ptr pool_; +}; + +} // namespace paimon 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/common/utils/arrow/arrow_utils.cpp b/src/paimon/common/utils/arrow/arrow_utils.cpp index f29e1d11e..c6a07b2bf 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); diff --git a/src/paimon/core/operation/abstract_split_read.cpp b/src/paimon/core/operation/abstract_split_read.cpp index d057e1d72..89f5071e9 100644 --- a/src/paimon/core/operation/abstract_split_read.cpp +++ b/src/paimon/core/operation/abstract_split_read.cpp @@ -33,6 +33,7 @@ #include "paimon/common/data/shredding/shredding_file_reader.h" #include "paimon/common/data/variant/variant_shredding_read_plan_factory.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" @@ -96,7 +97,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) { @@ -151,13 +152,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_->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, options_.GetFileSystem(), + data_file_path, data_file_size, reader_builder.get(), options_.GetFileSystem(), context_->GetPrefetchMaxParallelNum(), options_.GetReadBatchSize(), context_->GetPrefetchBatchCount(), options_.EnableAdaptivePrefetchStrategy(), executor_, @@ -174,7 +179,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 { @@ -214,7 +219,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 a56b48fdf..a02ed5fb2 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; diff --git a/src/paimon/core/operation/data_evolution_split_read.h b/src/paimon/core/operation/data_evolution_split_read.h index fad0e6742..2568eae46 100644 --- a/src/paimon/core/operation/data_evolution_split_read.h +++ b/src/paimon/core/operation/data_evolution_split_read.h @@ -64,7 +64,8 @@ struct DeletionFile; /// ->(ConcatBatchReader across blob files | BlobFallbackBatchReader across blob sequence layers) /// ->FieldMappingReader->(ApplyDeletionVectorBatchReader)->(ApplyBitmapIndexBatchReader) /// ->(CompleteRowTrackingFieldsBatchReader)->(ShreddingFileReader) -/// ->(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/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/merge_file_split_read.h b/src/paimon/core/operation/merge_file_split_read.h index 5003cb55a..11dcd0b37 100644 --- a/src/paimon/core/operation/merge_file_split_read.h +++ b/src/paimon/core/operation/merge_file_split_read.h @@ -74,7 +74,8 @@ class MergeFunctionWrapper; /// files->KeyValueProjectionReader/AsyncKeyValueProjectionReader /// ->DropDeleteReader->SortMergeReader->ConcatKeyValueRecordReader->KeyValueDataFileRecordReader /// ->FieldMappingReader->(ApplyDeletionVectorBatchReader)->(ShreddingFileReader) -/// ->(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 911cf7961..19cd96d67 100644 --- a/src/paimon/core/operation/merge_file_split_read_test.cpp +++ b/src/paimon/core/operation/merge_file_split_read_test.cpp @@ -798,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 @@ -842,6 +843,63 @@ TEST_P(MergeFileSplitReadTest, TestReadWithPredicate) { CheckResult(result_array, expected_array, read_schema); } +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/src/paimon/core/operation/raw_file_split_read.h b/src/paimon/core/operation/raw_file_split_read.h index 93eab5509..646f24ac7 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)->(VectorFileBatchReader) -/// ->(DelegatingPrefetchReader)->(PrefetchFileBatchReader)->FormatReader +/// ->(DelegatingPrefetchReader)->(PrefetchFileBatchReader) +/// ->(LateMaterializingFileBatchReader)->FormatReader class RawFileSplitRead : public AbstractSplitRead { public: diff --git a/src/paimon/core/operation/read_context.cpp b/src/paimon/core/operation/read_context.cpp index 08a854d84..deacfa78b 100644 --- a/src/paimon/core/operation/read_context.cpp +++ b/src/paimon/core/operation/read_context.cpp @@ -35,10 +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, - 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, @@ -51,6 +51,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 +98,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 +124,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 +194,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,11 +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_->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_)); } 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 f05a2347b..4566289f4 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" @@ -78,11 +79,14 @@ 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(); return Status::OK(); } @@ -119,8 +123,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_; @@ -128,21 +153,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( @@ -168,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 {} @@ -181,7 +208,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(); @@ -189,6 +216,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_; 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))); diff --git a/test/inte/read_inte_test.cpp b/test/inte/read_inte_test.cpp index 1e0952e08..a2a52d343 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,89 @@ 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) + .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 +3193,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 +3251,91 @@ TEST_P(ReadInteTest, TestPkReadSnapshot6WithSchemaEvolutionWithPredicateOnlyPush ASSERT_TRUE(result_array->Equals(*expected_array)); } +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.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, "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); + 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..7a4439734 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,74 @@ TEST_P(ReadInteWithIndexTest, TestNoEmbeddingBitmapIndex) { CheckResultForBitmapWithSingleRowGroup(path, arrow_data_type, split); } +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..537a4cb36 100644 --- a/test/inte/scan_and_read_inte_test.cpp +++ b/test/inte/scan_and_read_inte_test.cpp @@ -724,7 +724,7 @@ 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 +744,45 @@ TEST_P(ScanAndReadInteTest, TestWithPKWithDvBatchScanSnapshot6WithPredicate) { ASSERT_TRUE(expected->Equals(read_result)) << read_result->ToString(); } +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).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 +1290,7 @@ 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 +1318,55 @@ 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).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 +