Skip to content

Commit ca625d6

Browse files
gripleafgripleaf
committed
fix(read): use file-local late materialization selections
1 parent a04b10d commit ca625d6

10 files changed

Lines changed: 176 additions & 35 deletions

‎docs/source/user_guide/read.rst‎

Lines changed: 2 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -50,7 +50,8 @@ columns (the *payload* columns) only for matching row ranges.
5050

5151
The optimization is disabled by default. Enable it with
5252
``read.late-materialization.enabled=true``. It is applied only when predicate filtering is enabled,
53-
the probe and payload projections are both non-empty, and every input file has a first row ID.
53+
and the probe and payload projections are both non-empty. The selected file-local row IDs are
54+
pushed directly to each file reader, so row tracking and global row IDs are not required.
5455
Otherwise, the reader uses the normal single-pass path.
5556

5657
``read.late-materialization.max-match-rows`` limits the number of matching rows accumulated for a

‎src/paimon/core/operation/abstract_split_read.cpp‎

Lines changed: 10 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -79,7 +79,8 @@ Result<std::vector<std::unique_ptr<FileBatchReader>>> AbstractSplitRead::CreateR
7979
const std::shared_ptr<arrow::Schema>& read_schema, const std::shared_ptr<Predicate>& predicate,
8080
DeletionVector::Factory dv_factory, const std::optional<std::vector<Range>>& row_ranges,
8181
const std::shared_ptr<DataFilePathFactory>& data_file_path_factory,
82-
const std::map<std::string, std::string>& extra_format_options) const {
82+
const std::map<std::string, std::string>& extra_format_options,
83+
const std::optional<RoaringBitmap32>& file_selection) const {
8384
if (data_files.empty()) {
8485
return std::vector<std::unique_ptr<FileBatchReader>>();
8586
}
@@ -98,7 +99,7 @@ Result<std::vector<std::unique_ptr<FileBatchReader>>> AbstractSplitRead::CreateR
9899
std::unique_ptr<FileBatchReader> file_reader,
99100
CreateFieldMappingReader(data_file_path, file, partition, reader_builder.get(),
100101
field_mapping_builder.get(), dv_factory, row_ranges,
101-
data_file_path_factory));
102+
data_file_path_factory, file_selection));
102103
if (file_reader) {
103104
raw_file_readers.push_back(std::move(file_reader));
104105
}
@@ -169,7 +170,8 @@ Result<std::unique_ptr<FileBatchReader>> AbstractSplitRead::CreateFieldMappingRe
169170
const BinaryRow& partition, const ReaderBuilder* reader_builder,
170171
const FieldMappingBuilder* field_mapping_builder, DeletionVector::Factory dv_factory,
171172
const std::optional<std::vector<Range>>& row_ranges,
172-
const std::shared_ptr<DataFilePathFactory>& data_file_path_factory) const {
173+
const std::shared_ptr<DataFilePathFactory>& data_file_path_factory,
174+
const std::optional<RoaringBitmap32>& file_selection) const {
173175
std::shared_ptr<TableSchema> data_schema;
174176
if (file_meta->schema_id == context_->GetTableSchema()->Id()) {
175177
data_schema = context_->GetTableSchema();
@@ -230,10 +232,11 @@ Result<std::unique_ptr<FileBatchReader>> AbstractSplitRead::CreateFieldMappingRe
230232
}
231233
const auto& predicate = field_mapping->non_partition_info.non_partition_filter;
232234
auto all_data_schema = DataField::ConvertDataFieldsToArrowSchema(data_schema->Fields());
233-
PAIMON_ASSIGN_OR_RAISE(std::unique_ptr<FileBatchReader> final_reader,
234-
ApplyIndexAndDvReaderIfNeeded(
235-
std::move(file_reader), file_meta, all_data_schema, read_schema,
236-
predicate, dv_factory, row_ranges, data_file_path_factory));
235+
PAIMON_ASSIGN_OR_RAISE(
236+
std::unique_ptr<FileBatchReader> final_reader,
237+
ApplyIndexAndDvReaderIfNeeded(std::move(file_reader), file_meta, all_data_schema,
238+
read_schema, predicate, dv_factory, row_ranges,
239+
data_file_path_factory, file_selection));
237240
if (!final_reader) {
238241
// file is skipped by index or dv
239242
return std::unique_ptr<FileBatchReader>();

‎src/paimon/core/operation/abstract_split_read.h‎

Lines changed: 8 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -66,13 +66,16 @@ class AbstractSplitRead : public SplitRead {
6666
/// `extra_format_options` are merged over the table options when building the format
6767
/// reader, e.g. to switch the blob format reader into placeholder-aware mode for the
6868
/// data-evolution blob fallback read path.
69+
/// `file_selection`, when present, contains file-local row IDs and is applied independently
70+
/// to every file in `data_files`.
6971
Result<std::vector<std::unique_ptr<FileBatchReader>>> CreateRawFileReaders(
7072
const BinaryRow& partition, const std::vector<std::shared_ptr<DataFileMeta>>& data_files,
7173
const std::shared_ptr<arrow::Schema>& read_schema,
7274
const std::shared_ptr<Predicate>& predicate, DeletionVector::Factory dv_factory,
7375
const std::optional<std::vector<Range>>& row_ranges,
7476
const std::shared_ptr<DataFilePathFactory>& data_file_path_factory,
75-
const std::map<std::string, std::string>& extra_format_options) const;
77+
const std::map<std::string, std::string>& extra_format_options,
78+
const std::optional<RoaringBitmap32>& file_selection = std::nullopt) const;
7679

7780
protected:
7881
AbstractSplitRead(const std::shared_ptr<FileStorePathFactory>& path_factory,
@@ -91,7 +94,8 @@ class AbstractSplitRead : public SplitRead {
9194
const std::shared_ptr<arrow::Schema>& read_schema,
9295
const std::shared_ptr<Predicate>& predicate, DeletionVector::Factory dv_factory,
9396
const std::optional<std::vector<Range>>& row_ranges,
94-
const std::shared_ptr<DataFilePathFactory>& data_file_path_factory) const = 0;
97+
const std::shared_ptr<DataFilePathFactory>& data_file_path_factory,
98+
const std::optional<RoaringBitmap32>& file_selection) const = 0;
9599

96100
// 1. project write cols to data schema
97101
// 2. add partition fields (if write cols not contain)
@@ -115,7 +119,8 @@ class AbstractSplitRead : public SplitRead {
115119
const BinaryRow& partition, const ReaderBuilder* reader_builder,
116120
const FieldMappingBuilder* field_mapping_builder, DeletionVector::Factory dv_factory,
117121
const std::optional<std::vector<Range>>& row_ranges,
118-
const std::shared_ptr<DataFilePathFactory>& data_file_path_factory) const;
122+
const std::shared_ptr<DataFilePathFactory>& data_file_path_factory,
123+
const std::optional<RoaringBitmap32>& file_selection) const;
119124

120125
Result<std::pair<std::unique_ptr<FileBatchReader>, std::set<int32_t>>>
121126
ApplySharedShreddingReaderIfNeeded(std::unique_ptr<FileBatchReader>&& file_reader,

‎src/paimon/core/operation/data_evolution_split_read.cpp‎

Lines changed: 7 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -428,7 +428,8 @@ Result<std::unique_ptr<FileBatchReader>> DataEvolutionSplitRead::ApplyIndexAndDv
428428
const std::shared_ptr<arrow::Schema>& data_schema,
429429
const std::shared_ptr<arrow::Schema>& read_schema, const std::shared_ptr<Predicate>& predicate,
430430
DeletionVector::Factory dv_factory, const std::optional<std::vector<Range>>& row_ranges,
431-
const std::shared_ptr<DataFilePathFactory>& data_file_path_factory) const {
431+
const std::shared_ptr<DataFilePathFactory>& data_file_path_factory,
432+
const std::optional<RoaringBitmap32>& file_selection) const {
432433
if (predicate) {
433434
assert(false);
434435
// as DataEvolutionSplitRead will skip predicate
@@ -443,6 +444,11 @@ Result<std::unique_ptr<FileBatchReader>> DataEvolutionSplitRead::ApplyIndexAndDv
443444
}
444445
PAIMON_ASSIGN_OR_RAISE(std::optional<RoaringBitmap32> selection_row_ids,
445446
file->ToFileSelection(row_ranges));
447+
if (file_selection) {
448+
selection_row_ids = selection_row_ids ? RoaringBitmap32::And(selection_row_ids.value(),
449+
file_selection.value())
450+
: file_selection;
451+
}
446452
::ArrowSchema c_read_schema;
447453
PAIMON_RETURN_NOT_OK_FROM_ARROW(arrow::ExportSchema(*read_schema, &c_read_schema));
448454
PAIMON_RETURN_NOT_OK(

‎src/paimon/core/operation/data_evolution_split_read.h‎

Lines changed: 2 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -94,7 +94,8 @@ class DataEvolutionSplitRead : public AbstractSplitRead {
9494
const std::shared_ptr<arrow::Schema>& read_schema,
9595
const std::shared_ptr<Predicate>& predicate, DeletionVector::Factory dv_factory,
9696
const std::optional<std::vector<Range>>& row_ranges,
97-
const std::shared_ptr<DataFilePathFactory>& data_file_path_factory) const override;
97+
const std::shared_ptr<DataFilePathFactory>& data_file_path_factory,
98+
const std::optional<RoaringBitmap32>& file_selection) const override;
9899

99100
private:
100101
/// Files for partial field.

‎src/paimon/core/operation/merge_file_split_read.cpp‎

Lines changed: 7 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -197,7 +197,8 @@ Result<std::unique_ptr<FileBatchReader>> MergeFileSplitRead::ApplyIndexAndDvRead
197197
const std::shared_ptr<arrow::Schema>& data_schema,
198198
const std::shared_ptr<arrow::Schema>& read_schema, const std::shared_ptr<Predicate>& predicate,
199199
DeletionVector::Factory dv_factory, const std::optional<std::vector<Range>>& ranges,
200-
const std::shared_ptr<DataFilePathFactory>& data_file_path_factory) const {
200+
const std::shared_ptr<DataFilePathFactory>& data_file_path_factory,
201+
const std::optional<RoaringBitmap32>& file_selection) const {
201202
// merge read does not use index
202203
std::shared_ptr<DeletionVector> deletion_vector;
203204
if (dv_factory) {
@@ -215,6 +216,11 @@ Result<std::unique_ptr<FileBatchReader>> MergeFileSplitRead::ApplyIndexAndDvRead
215216
PAIMON_ASSIGN_OR_RAISE(uint64_t num_rows, file_reader->GetNumberOfRows());
216217
actual_selection.value().Flip(0, num_rows);
217218
}
219+
if (file_selection) {
220+
actual_selection = actual_selection ? RoaringBitmap32::And(actual_selection.value(),
221+
file_selection.value())
222+
: file_selection;
223+
}
218224

219225
::ArrowSchema c_read_schema;
220226
PAIMON_RETURN_NOT_OK_FROM_ARROW(arrow::ExportSchema(*read_schema, &c_read_schema));

‎src/paimon/core/operation/merge_file_split_read.h‎

Lines changed: 2 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -97,7 +97,8 @@ class MergeFileSplitRead : public AbstractSplitRead {
9797
const std::shared_ptr<arrow::Schema>& read_schema,
9898
const std::shared_ptr<Predicate>& predicate, DeletionVector::Factory dv_factory,
9999
const std::optional<std::vector<Range>>& ranges,
100-
const std::shared_ptr<DataFilePathFactory>& data_file_path_factory) const override;
100+
const std::shared_ptr<DataFilePathFactory>& data_file_path_factory,
101+
const std::optional<RoaringBitmap32>& file_selection) const override;
101102

102103
Result<std::unique_ptr<SortMergeReader>> CreateSortMergeReaderForSection(
103104
const std::vector<SortedRun>& section, const BinaryRow& partition,

‎src/paimon/core/operation/raw_file_split_read.cpp‎

Lines changed: 22 additions & 15 deletions
Original file line numberDiff line numberDiff line change
@@ -19,6 +19,7 @@
1919
#include "paimon/core/operation/raw_file_split_read.h"
2020

2121
#include <cstdint>
22+
#include <limits>
2223
#include <map>
2324
#include <optional>
2425
#include <set>
@@ -225,11 +226,6 @@ auto RawFileSplitRead::TryCreateLateMaterializedReader(
225226
if (!plan) {
226227
return LateMaterializationReadResult();
227228
}
228-
for (const auto& file : data_files) {
229-
if (!file->first_row_id) {
230-
return LateMaterializationReadResult();
231-
}
232-
}
233229

234230
std::map<std::string, int32_t> probe_field_name_to_idx;
235231
for (int32_t i = 0; i < plan->probe_schema->num_fields(); ++i) {
@@ -261,7 +257,7 @@ auto RawFileSplitRead::TryCreateLateMaterializedReader(
261257
}
262258

263259
std::vector<std::shared_ptr<arrow::Array>> probe_chunks;
264-
std::vector<Range> selected_global_ranges;
260+
RoaringBitmap32 selected_file_rows;
265261
while (true) {
266262
PAIMON_ASSIGN_OR_RAISE(BatchReader::ReadBatchWithBitmap batch_with_bitmap,
267263
probe_readers[0]->NextBatchWithBitmap());
@@ -297,9 +293,12 @@ auto RawFileSplitRead::TryCreateLateMaterializedReader(
297293
}
298294
PAIMON_ASSIGN_OR_RAISE(uint64_t file_row_id,
299295
probe_readers[0]->GetPreviousBatchFileRowId(batch_row_id));
300-
int64_t global_row_id =
301-
file->first_row_id.value() + static_cast<int64_t>(file_row_id);
302-
selected_global_ranges.emplace_back(global_row_id, global_row_id);
296+
if (file_row_id > std::numeric_limits<uint32_t>::max()) {
297+
return Status::Invalid(
298+
fmt::format("late materialization file row id {} exceeds bitmap capacity",
299+
file_row_id));
300+
}
301+
selected_file_rows.Add(static_cast<uint32_t>(file_row_id));
303302
selected_bitmap.Add(batch_row_id);
304303
}
305304

@@ -326,7 +325,7 @@ auto RawFileSplitRead::TryCreateLateMaterializedReader(
326325
completed_metrics->Merge(probe_readers[0]->GetReaderMetrics());
327326
probe_readers[0]->Close();
328327

329-
if (selected_global_ranges.empty()) {
328+
if (selected_file_rows.IsEmpty()) {
330329
continue;
331330
}
332331

@@ -335,13 +334,12 @@ auto RawFileSplitRead::TryCreateLateMaterializedReader(
335334
arrow::Concatenate(probe_chunks, probe_arrow_pool.get()));
336335
std::shared_ptr<arrow::StructArray> probe_struct =
337336
arrow::internal::checked_pointer_cast<arrow::StructArray>(probe_data);
338-
std::vector<Range> row_ranges =
339-
Range::SortAndMergeOverlap(selected_global_ranges, /*adjacent=*/true);
340337
PAIMON_ASSIGN_OR_RAISE(
341338
std::vector<std::unique_ptr<FileBatchReader>> payload_readers,
342339
CreateRawFileReaders(partition, {file}, plan->payload_schema,
343-
/*predicate=*/nullptr, dv_factory, row_ranges,
344-
data_file_path_factory, /*extra_format_options=*/{}));
340+
/*predicate=*/nullptr, dv_factory,
341+
/*row_ranges=*/std::nullopt, data_file_path_factory,
342+
/*extra_format_options=*/{}, selected_file_rows));
345343
if (payload_readers.empty()) {
346344
return Status::Invalid("late materialization payload reader was filtered out");
347345
}
@@ -367,7 +365,8 @@ Result<std::unique_ptr<FileBatchReader>> RawFileSplitRead::ApplyIndexAndDvReader
367365
const std::shared_ptr<arrow::Schema>& data_schema,
368366
const std::shared_ptr<arrow::Schema>& read_schema, const std::shared_ptr<Predicate>& predicate,
369367
DeletionVector::Factory dv_factory, const std::optional<std::vector<Range>>& ranges,
370-
const std::shared_ptr<DataFilePathFactory>& data_file_path_factory) const {
368+
const std::shared_ptr<DataFilePathFactory>& data_file_path_factory,
369+
const std::optional<RoaringBitmap32>& file_selection) const {
371370
std::shared_ptr<FileIndexResult> file_index_result;
372371
if (options_.FileIndexReadEnabled()) {
373372
PAIMON_ASSIGN_OR_RAISE(
@@ -420,6 +419,14 @@ Result<std::unique_ptr<FileBatchReader>> RawFileSplitRead::ApplyIndexAndDvReader
420419
}
421420
}
422421

422+
// Late materialization already operates on one FileBatchReader at a time, so its selection
423+
// can stay file-local instead of making a round trip through global row IDs.
424+
if (file_selection) {
425+
actual_selection = actual_selection ? RoaringBitmap32::And(actual_selection.value(),
426+
file_selection.value())
427+
: file_selection;
428+
}
429+
423430
if (actual_selection && actual_selection.value().IsEmpty()) {
424431
return std::unique_ptr<FileBatchReader>();
425432
}

‎src/paimon/core/operation/raw_file_split_read.h‎

Lines changed: 2 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -85,7 +85,8 @@ class RawFileSplitRead : public AbstractSplitRead {
8585
const std::shared_ptr<arrow::Schema>& read_schema,
8686
const std::shared_ptr<Predicate>& predicate, DeletionVector::Factory dv_factory,
8787
const std::optional<std::vector<Range>>& ranges,
88-
const std::shared_ptr<DataFilePathFactory>& data_file_path_factory) const override;
88+
const std::shared_ptr<DataFilePathFactory>& data_file_path_factory,
89+
const std::optional<RoaringBitmap32>& file_selection) const override;
8990

9091
private:
9192
struct LateMaterializationPlan;

0 commit comments

Comments
 (0)