fix(parquet): avoid seeking into the middle of a row group when prefetch and page index filtering are both enabled (#185)
* test: add test cases to show problem
* fix: avoid seek to the middle of a RowGroup when reading a parquet with prefetch=on and page-inex-filter=on
* test: update test cases
* fix: partially-matched path do not push next_row_to_read
* test: update test comments and enable multi-thread reading
* style: update comments
* clang-format
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 3e8737a..defdd66 100644
--- a/src/paimon/common/reader/prefetch_file_batch_reader_impl.cpp
+++ b/src/paimon/common/reader/prefetch_file_batch_reader_impl.cpp
@@ -427,6 +427,16 @@
return Status::OK();
}
+std::optional<std::pair<uint64_t, uint64_t>> PrefetchFileBatchReaderImpl::FindReadRangeContaining(
+ size_t reader_idx, uint64_t row_id) const {
+ for (const auto& range : read_ranges_in_group_[reader_idx]) {
+ if (row_id >= range.first && row_id < range.second) {
+ return range;
+ }
+ }
+ return std::nullopt;
+}
+
Status PrefetchFileBatchReaderImpl::HandleReadResult(
size_t reader_idx, const std::pair<uint64_t, uint64_t>& read_range,
ReadBatchWithBitmap&& read_batch_with_bitmap) {
@@ -454,9 +464,18 @@
if (0 == slice_end) {
// fully out of range, data before global_row_ids has been filtered out
- readers_pos_[reader_idx]->store(global_row_ids[0]);
- ReaderUtils::ReleaseReadBatch(std::move(read_batch));
- return Status::OK();
+ // find the read range that contains the first row id and put it into queue in advance.
+ std::optional<std::pair<uint64_t, uint64_t>> owner_range =
+ FindReadRangeContaining(reader_idx, global_row_ids[0]);
+ if (owner_range == std::nullopt) {
+ readers_pos_[reader_idx]->store(global_row_ids[0]);
+ ReaderUtils::ReleaseReadBatch(std::move(read_batch));
+ return Status::OK();
+ }
+ // Recurses at most once: global_row_ids[0] is within owner_range, so the recursive
+ // call cannot compute a zero slice end again.
+ return HandleReadResult(reader_idx, owner_range.value(),
+ std::move(read_batch_with_bitmap));
} else if (slice_end < c_array->length) {
// partially out of range, data before read_range.second has been effectively consumed
readers_pos_[reader_idx]->store(read_range.second);
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 36673e8..f0c302e 100644
--- a/src/paimon/common/reader/prefetch_file_batch_reader_impl.h
+++ b/src/paimon/common/reader/prefetch_file_batch_reader_impl.h
@@ -134,6 +134,11 @@
Status RefreshReadRangesAfterCleanUp();
Result<std::pair<uint64_t, uint64_t>> EofRange() const;
std::optional<std::pair<uint64_t, uint64_t>> GetCurrentReadRange(size_t reader_idx) const;
+
+ /// Find the read range assigned to the given reader that contains the given file row id.
+ /// Returns nullopt when no assigned range contains it.
+ std::optional<std::pair<uint64_t, uint64_t>> FindReadRangeContaining(size_t reader_idx,
+ uint64_t row_id) const;
Status EnsureReaderPosition(size_t reader_idx,
const std::pair<uint64_t, uint64_t>& read_range) const;
Status HandleReadResult(size_t reader_idx, const std::pair<uint64_t, uint64_t>& read_range,
diff --git a/src/paimon/format/parquet/file_reader_wrapper.cpp b/src/paimon/format/parquet/file_reader_wrapper.cpp
index 4c90b95..48a4430 100644
--- a/src/paimon/format/parquet/file_reader_wrapper.cpp
+++ b/src/paimon/format/parquet/file_reader_wrapper.cpp
@@ -259,6 +259,14 @@
static_cast<uint64_t>(*original_row)
: current_filtered_rg_start_;
filtered_global_offset_ += record_batch->num_rows();
+ // Advance to the next row that survives filtering, or to the row group end when the
+ // filtered ranges are exhausted, so that next_row_to_read_ tracks the streaming position.
+ auto next_original_row =
+ current_filtered_row_ranges_.MapFilteredIndexToOriginalRow(filtered_global_offset_);
+ next_row_to_read_ =
+ next_original_row.has_value()
+ ? current_filtered_rg_start_ + static_cast<uint64_t>(*next_original_row)
+ : all_row_group_ranges_[rg_id].second;
return record_batch;
}
diff --git a/src/paimon/format/parquet/file_reader_wrapper_test.cpp b/src/paimon/format/parquet/file_reader_wrapper_test.cpp
index 3ac7f62..de3b6fd 100644
--- a/src/paimon/format/parquet/file_reader_wrapper_test.cpp
+++ b/src/paimon/format/parquet/file_reader_wrapper_test.cpp
@@ -374,6 +374,59 @@
}
}
+/// While streaming a page-filtered row group, GetNextRowToRead() must report the next row that
+/// survives filtering, and the row group end once the filtered ranges are exhausted. Reporting the
+/// row group start for the whole row group makes callers believe the reader has not moved.
+TEST_F(FileReaderWrapperTest, PageFilteredAdvancesNextRowToRead) {
+ std::string file_path = PathUtil::JoinPath(dir_->Str(), "page_next_row.parquet");
+ // 2000 rows produces 2 row groups (max_row_group_length=1000) with page index enabled.
+ PrepareParquetFile(file_path, /*row_count=*/2000, /*enable_page_index=*/true);
+ ASSERT_OK_AND_ASSIGN(auto reader_wrapper,
+ PrepareReaderWrapper(file_path, /*wrapper_batch_size=*/7));
+ ASSERT_EQ(2, reader_wrapper->GetNumberOfRowGroups());
+
+ // RowRanges are RG-local. RG0 keeps two non-contiguous stretches so that a batch can span the
+ // gap between them; RG1 keeps its first 20 rows.
+ RowRanges rg0_ranges({RowRanges::Range(10, 49), RowRanges::Range(100, 149)});
+ RowRanges rg1_ranges(RowRanges::Range(0, 19));
+ ASSERT_OK(reader_wrapper->PrepareForReading(
+ {TargetRowGroup(/*rg_index=*/0, /*is_partially_matched=*/true, /*ranges=*/rg0_ranges),
+ TargetRowGroup(/*rg_index=*/1, /*is_partially_matched=*/true, /*ranges=*/rg1_ranges)},
+ /*column_indices=*/{0, 1, 2}));
+
+ // Absolute row numbers the reader is expected to produce, in order.
+ std::vector<uint64_t> expected_rows;
+ for (uint64_t row = 10; row <= 49; ++row) {
+ expected_rows.push_back(row);
+ }
+ for (uint64_t row = 100; row <= 149; ++row) {
+ expected_rows.push_back(row);
+ }
+ for (uint64_t row = 1000; row <= 1019; ++row) {
+ expected_rows.push_back(row);
+ }
+
+ size_t consumed = 0;
+ while (true) {
+ ASSERT_OK_AND_ASSIGN(std::shared_ptr<arrow::RecordBatch> record_batch,
+ reader_wrapper->Next());
+ if (!record_batch) {
+ break;
+ }
+ ASSERT_LT(consumed, expected_rows.size());
+ ASSERT_EQ(expected_rows[consumed],
+ reader_wrapper->GetPreviousBatchFirstRowNumber().value());
+ consumed += record_batch->num_rows();
+ ASSERT_LE(consumed, expected_rows.size());
+ // RG0 ends exactly where RG1 starts, so the row group boundary is also covered by
+ // expected_rows; only the very last batch leaves the cursor at the file end.
+ uint64_t expected_next_row =
+ consumed < expected_rows.size() ? expected_rows[consumed] : 2000;
+ ASSERT_EQ(expected_next_row, reader_wrapper->GetNextRowToRead());
+ }
+ ASSERT_EQ(expected_rows.size(), consumed);
+}
+
TEST_F(FileReaderWrapperTest, GetRowGroupRanges) {
std::string file_path = PathUtil::JoinPath(dir_->Str(), "test.parquet");
PrepareParquetFile(file_path, /*row_count=*/5500);
diff --git a/test/inte/write_and_read_inte_test.cpp b/test/inte/write_and_read_inte_test.cpp
index 53920ba..dd63ed6 100644
--- a/test/inte/write_and_read_inte_test.cpp
+++ b/test/inte/write_and_read_inte_test.cpp
@@ -1529,6 +1529,96 @@
ASSERT_TRUE(expected->Equals(read_result)) << read_result->ToString();
}
+/// Reproduces the prefetch + parquet page-index filter failure: the predicate keeps only the last
+/// page of RG1, RG2 and RG3, so the prefetch reader ends up seeking to a row in the middle of a row
+/// group, which FileReaderWrapper::SeekToRow rejects.
+TEST_P(WriteAndReadInteTest, TestAppendWithParquetPageIndexFilterAndPrefetch) {
+ auto [file_format, file_system] = GetParam();
+ if (file_format != "parquet" || file_system != "local") {
+ return;
+ }
+
+ auto test_dir = UniqueTestDirectory::Create("local");
+ arrow::FieldVector fields = {arrow::field("f0", arrow::int32()),
+ arrow::field("f1", arrow::utf8())};
+ auto schema = arrow::schema(fields);
+ std::map<std::string, std::string> options = {
+ {Options::MANIFEST_FORMAT, "orc"},
+ {Options::FILE_FORMAT, "parquet"},
+ {Options::TARGET_FILE_SIZE, "1048576"},
+ {Options::BUCKET, "-1"},
+ {Options::FILE_SYSTEM, "local"},
+ // One row per page (see TestAppendWithParquetPageIndexFilter for why these three
+ // options are needed together) and 4 rows per row group, so the 16 rows below end up
+ // in 4 row groups of 4 single-row pages.
+ {Options::WRITE_BATCH_SIZE, "1"},
+ {"parquet.page.size", "1"},
+ {"parquet.enable-dictionary", "false"},
+ {"parquet.write.enable-page-index", "true"},
+ {"parquet.write.max-row-group-length", "4"},
+ {"parquet.read.enable-page-index-filter", "true"},
+ };
+ ASSERT_OK_AND_ASSIGN(
+ auto helper, TestHelper::Create(test_dir->Str(), schema, /*partition_keys=*/{},
+ /*primary_keys=*/{}, options, /*is_streaming_mode=*/true));
+ std::string table_path = test_dir->Str() + "/foo.db/bar";
+
+ std::string data = R"([
+ [0, "v0"], [1, "v1"], [2, "v2"], [3, "v3"],
+ [4, "v4"], [5, "v5"], [6, "v6"], [7, "v7"],
+ [8, "v8"], [9, "v9"], [10, "v10"], [11, "v11"],
+ [12, "v12"], [13, "v13"], [14, "v14"], [15, "v15"]
+ ])";
+ ASSERT_OK_AND_ASSIGN(std::unique_ptr<RecordBatch> batch,
+ TestHelper::MakeRecordBatch(arrow::struct_(fields), data,
+ /*partition_map=*/{}, /*bucket=*/0, {}));
+ ASSERT_OK(helper->WriteAndCommit(std::move(batch), /*commit_identifier=*/0,
+ /*expected_commit_messages=*/std::nullopt));
+
+ // Keep only the last row of RG1, RG2 and RG3, so each row group is partially matched and its
+ // first selected row is 3 rows behind the row group start.
+ auto predicate = PredicateBuilder::In(/*field_index=*/0, /*field_name=*/"f0", FieldType::INT,
+ {Literal(7), Literal(11), Literal(15)});
+ ASSERT_TRUE(predicate);
+
+ ScanContextBuilder scan_context_builder(table_path);
+ scan_context_builder.SetOptions(options)
+ .AddOption(Options::SCAN_MODE, StartupMode::LatestFull().ToString())
+ .SetPredicate(predicate);
+ ASSERT_OK_AND_ASSIGN(auto scan_context, scan_context_builder.Finish());
+ ASSERT_OK_AND_ASSIGN(auto table_scan, TableScan::Create(std::move(scan_context)));
+ ASSERT_OK_AND_ASSIGN(auto result_plan, table_scan->CreatePlan());
+ ASSERT_FALSE(result_plan->Splits().empty());
+
+ // The row group aligned read ranges still cover RG0 although the predicate pruned that row
+ // group, and 2 sub readers take the ranges round robin. The reader0 owning RG0's range
+ // therefore finds no data for it and skips ahead into the next row group it owns, whose first
+ // selected row sits in the middle of that row group. Row level filtering stays off: the
+ // expected rows below are exactly what page-index filtering selects.
+ ReadContextBuilder read_context_builder(table_path);
+ read_context_builder.SetOptions(options)
+ .SetPredicate(predicate)
+ .EnablePrefetch(true)
+ .SetPrefetchMaxParallelNum(2)
+ .SetPrefetchBatchCount(3)
+ .AddOption("test.enable-adaptive-prefetch-strategy", "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)));
+ 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()));
+
+ arrow::FieldVector fields_with_row_kind = fields;
+ fields_with_row_kind.insert(fields_with_row_kind.begin(),
+ arrow::field("_VALUE_KIND", arrow::int8()));
+ auto expected_data_type = arrow::struct_(fields_with_row_kind);
+ auto expected = std::make_shared<arrow::ChunkedArray>(
+ arrow::ipc::internal::json::ArrayFromJSON(expected_data_type, R"([
+[0, 7, "v7"], [0, 11, "v11"], [0, 15, "v15"]
+])")
+ .ValueOrDie());
+ ASSERT_TRUE(expected->Equals(read_result)) << read_result->ToString();
+}
+
TEST_P(WriteAndReadInteTest, TestAppendWithParquetMetadataCache) {
auto [file_format, file_system] = GetParam();
if (file_format != "parquet" || file_system != "local") {
