fix: handle non-contiguous RowRanges when resolving global row IDs (#383)

Author: zhf999

include/paimon/reader/file_batch_reader.h

@@ -46,8 +46,14 @@ class PAIMON_EXPORT FileBatchReader : public BatchReader {
     using BatchReader::NextBatch;
     using BatchReader::NextBatchWithBitmap;
 
-    /// Get the row number of the first row in the previously read batch.
-    virtual Result<uint64_t> GetPreviousBatchFirstRowNumber() const = 0;
+    /// Get the file-level row ID for a given batch-relative row index
+    /// in the previously read batch.
+    ///
+    /// @param batch_row_id  Zero-based index within the current batch.
+    /// @return The corresponding file-level row ID, or Status::Invalid
+    ///         if no batch has been read yet, the last batch was EOF,
+    ///         or batch_row_id is out of range.
+    virtual Result<uint64_t> GetPreviousBatchFileRowId(uint64_t batch_row_id) const = 0;
 
     /// Get the number of rows in the file.
     virtual Result<uint64_t> GetNumberOfRows() const = 0;

src/paimon/common/data/shredding/map_shared_shredding_file_reader.cpp

@@ -391,8 +391,9 @@ void MapSharedShreddingFileReader::Close() {
     reader_->Close();
 }
 
-Result<uint64_t> MapSharedShreddingFileReader::GetPreviousBatchFirstRowNumber() const {
-    return reader_->GetPreviousBatchFirstRowNumber();
+Result<uint64_t> MapSharedShreddingFileReader::GetPreviousBatchFileRowId(
+    uint64_t batch_row_id) const {
+    return reader_->GetPreviousBatchFileRowId(batch_row_id);
 }
 
 Result<uint64_t> MapSharedShreddingFileReader::GetNumberOfRows() const {

src/paimon/common/data/shredding/map_shared_shredding_file_reader.h

@@ -60,7 +60,7 @@ class MapSharedShreddingFileReader : public FileBatchReader {
 
     void Close() override;
 
-    Result<uint64_t> GetPreviousBatchFirstRowNumber() const override;
+    Result<uint64_t> GetPreviousBatchFileRowId(uint64_t batch_row_id) const override;
 
     Result<uint64_t> GetNumberOfRows() const override;
 

src/paimon/common/file_index/bitmap/apply_bitmap_index_batch_reader.h

@@ -80,8 +80,8 @@ class ApplyBitmapIndexBatchReader : public FileBatchReader {
         return Status::Invalid("ApplyBitmapIndexBatchReader does not support SetReadSchema");
     }
 
-    Result<uint64_t> GetPreviousBatchFirstRowNumber() const override {
-        return reader_->GetPreviousBatchFirstRowNumber();
+    Result<uint64_t> GetPreviousBatchFileRowId(uint64_t batch_row_id) const override {
+        return reader_->GetPreviousBatchFileRowId(batch_row_id);
     }
 
     Result<uint64_t> GetNumberOfRows() const override {
@@ -94,14 +94,23 @@ class ApplyBitmapIndexBatchReader : public FileBatchReader {
 
  private:
     Result<RoaringBitmap32> Filter(int32_t batch_size) const {
-        RoaringBitmap32 is_valid;
-        PAIMON_ASSIGN_OR_RAISE(int32_t start_pos, reader_->GetPreviousBatchFirstRowNumber());
-        int32_t length = batch_size;
-        for (auto iter = bitmap_.EqualOrLarger(start_pos);
-             iter != bitmap_.End() && *iter < start_pos + length; ++iter) {
-            is_valid.Add(*iter - start_pos);
+        RoaringBitmap32 result;
+        auto bitmap_iter = bitmap_.Begin();
+        auto bitmap_end = bitmap_.End();
+
+        for (int32_t i = 0; i < batch_size; ++i) {
+            PAIMON_ASSIGN_OR_RAISE(uint64_t file_row_id, reader_->GetPreviousBatchFileRowId(i));
+            while (bitmap_iter != bitmap_end && static_cast<uint64_t>(*bitmap_iter) < file_row_id) {
+                ++bitmap_iter;
+            }
+            if (bitmap_iter == bitmap_end) {
+                break;
+            }
+            if (static_cast<uint64_t>(*bitmap_iter) == file_row_id) {
+                result.Add(i);
+            }
         }
-        return is_valid;
+        return result;
     }
 
  private:

src/paimon/common/reader/delegating_prefetch_reader.h

@@ -54,8 +54,8 @@ class DelegatingPrefetchReader : public FileBatchReader {
         return prefetch_reader_->SetReadSchema(read_schema, predicate, selection_bitmap);
     }
 
-    Result<uint64_t> GetPreviousBatchFirstRowNumber() const override {
-        return GetReader()->GetPreviousBatchFirstRowNumber();
+    Result<uint64_t> GetPreviousBatchFileRowId(uint64_t batch_row_id) const override {
+        return GetReader()->GetPreviousBatchFileRowId(batch_row_id);
     }
 
     Result<uint64_t> GetNumberOfRows() const override {

src/paimon/common/reader/prefetch_file_batch_reader_impl.cpp

@@ -40,6 +40,19 @@ class Schema;
 
 namespace paimon {
 
+namespace {
+
+std::pair<int64_t, int64_t> ComputeBatchSliceByReadRange(
+    const std::vector<uint64_t>& global_row_ids, const std::pair<uint64_t, uint64_t>& read_range) {
+    auto begin_it =
+        std::lower_bound(global_row_ids.begin(), global_row_ids.end(), read_range.first);
+    auto end_it = std::lower_bound(global_row_ids.begin(), global_row_ids.end(), read_range.second);
+    return {static_cast<int64_t>(std::distance(global_row_ids.begin(), begin_it)),
+            static_cast<int64_t>(std::distance(global_row_ids.begin(), end_it))};
+}
+
+}  // namespace
+
 Result<std::unique_ptr<PrefetchFileBatchReaderImpl>> PrefetchFileBatchReaderImpl::Create(
     const std::string& data_file_path, const ReaderBuilder* reader_builder,
     const std::shared_ptr<FileSystem>& fs, uint32_t prefetch_max_parallel_num, int32_t batch_size,
@@ -265,6 +278,7 @@ Status PrefetchFileBatchReaderImpl::CleanUp() {
 
     read_ranges_.clear();
     read_ranges_in_group_.clear();
+    current_batch_global_row_ids_.clear();
     clean_prefetch_queue();
     for (size_t i = 0; i < readers_pos_.size(); i++) {
         readers_pos_[i]->store(0);
@@ -409,28 +423,44 @@ Status PrefetchFileBatchReaderImpl::EnsureReaderPosition(
 Status PrefetchFileBatchReaderImpl::HandleReadResult(
     size_t reader_idx, const std::pair<uint64_t, uint64_t>& read_range,
     ReadBatchWithBitmap&& read_batch_with_bitmap) {
-    PAIMON_ASSIGN_OR_RAISE(uint64_t first_row_number,
-                           readers_[reader_idx]->GetPreviousBatchFirstRowNumber());
     auto& prefetch_queue = prefetch_queues_[reader_idx];
     if (!BatchReader::IsEofBatch(read_batch_with_bitmap)) {
         auto& [read_batch, bitmap] = read_batch_with_bitmap;
         auto& [c_array, c_schema] = read_batch;
+        std::vector<uint64_t> global_row_ids;
+        global_row_ids.reserve(c_array->length);
+        for (int64_t i = 0; i < c_array->length; ++i) {
+            PAIMON_ASSIGN_OR_RAISE(uint64_t global_row_id,
+                                   readers_[reader_idx]->GetPreviousBatchFileRowId(i));
+            global_row_ids.push_back(global_row_id);
+        }
+        if (global_row_ids.empty()) {
+            ReaderUtils::ReleaseReadBatch(std::move(read_batch));
+            return Status::OK();
+        }
+        auto [slice_begin, slice_end] = ComputeBatchSliceByReadRange(global_row_ids, read_range);
+        // slice_begin should always be 0, records before read_range.first have been consumed or
+        // filtered out.
+        if (slice_begin != 0) {
+            return Status::Invalid(fmt::format("Slice begin is {}, which is not 0.", slice_begin));
+        }
 
-        if (first_row_number >= read_range.second) {
-            // fully out of range, data before first_row_number has been filtered out
-            readers_pos_[reader_idx]->store(first_row_number);
+        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();
-        } else if (first_row_number + c_array->length > read_range.second) {
+        } 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);
             PAIMON_ASSIGN_OR_RAISE_FROM_ARROW(std::shared_ptr<arrow::Array> src_array,
                                               arrow::ImportArray(c_array.get(), c_schema.get()));
-            int32_t target_length = read_range.second - first_row_number;
-            auto array = src_array->Slice(/*offset=*/0, target_length);
+            auto array = src_array->Slice(0, slice_end);
             PAIMON_RETURN_NOT_OK_FROM_ARROW(
                 arrow::ExportArray(*array, c_array.get(), c_schema.get()));
-            bitmap.RemoveRange(target_length, src_array->length());
+            bitmap.RemoveRange(slice_end, src_array->length());
+            global_row_ids =
+                std::vector<uint64_t>(global_row_ids.begin(), global_row_ids.begin() + slice_end);
         } else {
             // all within the range, data before readers_[reader_idx]->GetNextRowToRead() has been
             // effectively consumed
@@ -440,11 +470,12 @@ Status PrefetchFileBatchReaderImpl::HandleReadResult(
             ReaderUtils::ReleaseReadBatch(std::move(read_batch));
             return Status::OK();
         }
-        prefetch_queue->push({read_range, std::move(read_batch_with_bitmap), first_row_number});
+        prefetch_queue->push(
+            {read_range, std::move(read_batch_with_bitmap), std::move(global_row_ids)});
     } else {
         std::pair<uint64_t, uint64_t> eof_range;
         PAIMON_ASSIGN_OR_RAISE(eof_range, EofRange());
-        prefetch_queue->push({eof_range, std::move(read_batch_with_bitmap), first_row_number});
+        prefetch_queue->push({eof_range, std::move(read_batch_with_bitmap), {}});
         readers_pos_[reader_idx]->store(std::numeric_limits<uint64_t>::max());
     }
     return Status::OK();
@@ -527,7 +558,7 @@ Result<BatchReader::ReadBatchWithBitmap> PrefetchFileBatchReaderImpl::NextBatchW
                     std::unique_lock<std::mutex> lock(working_mutex_);
                     cv_.notify_one();
                 }
-                previous_batch_first_row_num_ = prefetch_batch.value().previous_batch_first_row_num;
+                current_batch_global_row_ids_ = std::move(prefetch_batch.value().global_row_ids);
                 return std::move(prefetch_batch).value().batch;
             }
         }
@@ -537,7 +568,7 @@ Result<BatchReader::ReadBatchWithBitmap> PrefetchFileBatchReaderImpl::NextBatchW
                 assert(false);
                 return Status::Invalid("peek batch not suppose to be nullptr");
             }
-            previous_batch_first_row_num_ = peek_batch->previous_batch_first_row_num;
+            current_batch_global_row_ids_.clear();
             return BatchReader::MakeEofBatchWithBitmap();
         }
         if (value_count == prefetch_queues_.size()) {
@@ -571,8 +602,19 @@ Result<std::unique_ptr<::ArrowSchema>> PrefetchFileBatchReaderImpl::GetFileSchem
     return readers_[0]->GetFileSchema();
 }
 
-Result<uint64_t> PrefetchFileBatchReaderImpl::GetPreviousBatchFirstRowNumber() const {
-    return previous_batch_first_row_num_;
+Result<uint64_t> PrefetchFileBatchReaderImpl::GetPreviousBatchFileRowId(
+    uint64_t batch_row_id) const {
+    if (current_batch_global_row_ids_.empty()) {
+        return Status::Invalid(
+            "Last batch is not read or last batch is empty, cannot get previous batch global row "
+            "id");
+    }
+    if (batch_row_id >= current_batch_global_row_ids_.size()) {
+        return Status::Invalid(
+            fmt::format("batch_row_id {} is out of range, last batch row count is {}", batch_row_id,
+                        current_batch_global_row_ids_.size()));
+    }
+    return current_batch_global_row_ids_[batch_row_id];
 }
 
 Result<uint64_t> PrefetchFileBatchReaderImpl::GetNumberOfRows() const {

src/paimon/common/reader/prefetch_file_batch_reader_impl.h

@@ -76,7 +76,7 @@ class PrefetchFileBatchReaderImpl : public PrefetchFileBatchReader {
                          const std::optional<RoaringBitmap32>& selection_bitmap) override;
 
     Status SeekToRow(uint64_t row_number) override;
-    Result<uint64_t> GetPreviousBatchFirstRowNumber() const override;
+    Result<uint64_t> GetPreviousBatchFileRowId(uint64_t batch_row_id) const override;
     Result<uint64_t> GetNumberOfRows() const override;
     uint64_t GetNextRowToRead() const override;
     void Close() override;
@@ -105,7 +105,7 @@ class PrefetchFileBatchReaderImpl : public PrefetchFileBatchReader {
     struct PrefetchBatch {
         std::pair<uint64_t, uint64_t> read_range;
         BatchReader::ReadBatchWithBitmap batch;
-        uint64_t previous_batch_first_row_num;
+        std::vector<uint64_t> global_row_ids;
     };
 
     PrefetchFileBatchReaderImpl(
@@ -160,7 +160,7 @@ class PrefetchFileBatchReaderImpl : public PrefetchFileBatchReader {
     std::unique_ptr<std::thread> background_thread_;
     Status read_status_;
     std::atomic<bool> is_shutdown_ = false;
-    uint64_t previous_batch_first_row_num_ = std::numeric_limits<uint64_t>::max();
+    std::vector<uint64_t> current_batch_global_row_ids_;
     bool need_prefetch_ = false;
     bool read_ranges_freshed_ = false;
     const uint32_t prefetch_queue_capacity_;