fix lance GetPreviousBatchFirstRowNumber

Author: lszskye

src/paimon/format/lance/lance_file_batch_reader.cpp

@@ -70,8 +70,8 @@ Status LanceFileBatchReader::SetReadSchema(::ArrowSchema* read_schema,
                                       arrow::ImportSchema(read_schema));
     read_field_names_ = arrow_schema->field_names();
     assert(!read_field_names_.empty());
+    read_row_ids_.clear();
     if (selection_bitmap) {
-        read_row_ids_.clear();
         read_row_ids_.reserve(selection_bitmap.value().Cardinality());
         for (auto iter = selection_bitmap.value().Begin(); iter != selection_bitmap.value().End();
              ++iter) {
@@ -84,6 +84,8 @@ Status LanceFileBatchReader::SetReadSchema(::ArrowSchema* read_schema,
             release_stream_reader(stream_reader_, error_message_.data(), error_message_.size());
         PAIMON_RETURN_NOT_OK(LanceToPaimonStatus(err_code, error_message_));
         stream_reader_ = nullptr;
+        previous_batch_first_row_num_ = std::numeric_limits<uint64_t>::max();
+        last_batch_row_num_ = 0;
     }
     return Status::OK();
 }
@@ -102,6 +104,12 @@ Result<BatchReader::ReadBatch> LanceFileBatchReader::NextBatch() {
         PAIMON_RETURN_NOT_OK(LanceToPaimonStatus(err_code, error_message_));
         assert(stream_reader_);
     }
+    if (previous_batch_first_row_num_ == std::numeric_limits<uint64_t>::max()) {
+        // first read
+        previous_batch_first_row_num_ = 0;
+    } else {
+        previous_batch_first_row_num_ += last_batch_row_num_;
+    }
     auto c_array = std::make_unique<ArrowArray>();
     auto c_schema = std::make_unique<ArrowSchema>();
     bool is_eof = false;
@@ -111,6 +119,7 @@ Result<BatchReader::ReadBatch> LanceFileBatchReader::NextBatch() {
     if (is_eof) {
         return BatchReader::MakeEofBatch();
     }
+    last_batch_row_num_ = c_array->length;
     return std::make_pair(std::move(c_array), std::move(c_schema));
 }
 

src/paimon/format/lance/lance_file_batch_reader.h

@@ -42,11 +42,14 @@ class LanceFileBatchReader : public FileBatchReader {
     Result<ReadBatch> NextBatch() override;
 
     Result<uint64_t> GetPreviousBatchFirstRowNumber() const override {
-        // TODO(xinyu.lxy): support function
-        return Status::Invalid(
-            "Cannot call GetPreviousBatchFirstRowNumber in LanceFileBatchReader because, after "
-            "bitmap pushdown, rows in the array returned by NextBatch are no longer "
-            "contiguous.");
+        if (!read_row_ids_.empty() && read_row_ids_.size() != num_rows_) {
+            // TODO(xinyu.lxy): support function
+            return Status::Invalid(
+                "Cannot call GetPreviousBatchFirstRowNumber in LanceFileBatchReader because, after "
+                "bitmap pushdown, rows in the array returned by NextBatch are no longer "
+                "contiguous.");
+        }
+        return previous_batch_first_row_num_;
     }
 
     Result<uint64_t> GetNumberOfRows() const override {
@@ -76,6 +79,9 @@ class LanceFileBatchReader : public FileBatchReader {
     int32_t batch_size_ = -1;
     int32_t batch_readahead_ = -1;
     uint64_t num_rows_ = 0;
+    // only validate when there is no bitmap pushdown
+    uint64_t previous_batch_first_row_num_ = std::numeric_limits<uint64_t>::max();
+    uint64_t last_batch_row_num_ = 0;
     mutable std::string error_message_;
     LanceFileReader* file_reader_ = nullptr;
     LanceReaderAdapter* stream_reader_ = nullptr;

src/paimon/format/lance/lance_format_reader_writer_test.cpp

@@ -419,4 +419,59 @@ TEST_F(LanceFileReaderWriterTest, TestTimestampType) {
     CheckResult(src_chunk_array, schema, /*enable_tz=*/false);
 }
 
+TEST_F(LanceFileReaderWriterTest, TestPreviousBatchFirstRowNumber) {
+    arrow::FieldVector fields = {arrow::field("f1", arrow::int32()),
+                                 arrow::field("f2", arrow::utf8())};
+    auto schema = arrow::schema(fields);
+    auto array = std::dynamic_pointer_cast<arrow::StructArray>(
+        arrow::ipc::internal::json::ArrayFromJSON(arrow::struct_({fields}), R"([
+        [1, "Hello"],
+        [2, "World"],
+        [3, "apple"],
+        [4, "Alice"],
+        [5, "Bob"],
+        [6, "Lucy"]
+    ])")
+            .ValueOrDie());
+    auto src_chunk_array = std::make_shared<arrow::ChunkedArray>(arrow::ArrayVector({array}));
+
+    auto dir = paimon::test::UniqueTestDirectory::Create();
+    ASSERT_TRUE(dir);
+    std::string file_path = dir->Str() + "/test.lance";
+    WriteFile(file_path, src_chunk_array, schema);
+    ASSERT_OK_AND_ASSIGN(
+        std::unique_ptr<LanceFileBatchReader> reader,
+        LanceFileBatchReader::Create(file_path, /*batch_size=*/4, /*batch_readahead=*/2));
+    ASSERT_EQ(std::numeric_limits<uint64_t>::max(),
+              reader->GetPreviousBatchFirstRowNumber().value());
+
+    // first batch row 0-3
+    ASSERT_OK_AND_ASSIGN(auto read_batch, reader->NextBatch());
+    ASSERT_OK_AND_ASSIGN(auto read_array,
+                         paimon::test::ReadResultCollector::GetArray(std::move(read_batch)));
+    ASSERT_TRUE(read_array->Equals(array->Slice(0, 4)));
+    ASSERT_EQ(0, reader->GetPreviousBatchFirstRowNumber().value());
+
+    // second batch 4-5
+    ASSERT_OK_AND_ASSIGN(read_batch, reader->NextBatch());
+    ASSERT_OK_AND_ASSIGN(read_array,
+                         paimon::test::ReadResultCollector::GetArray(std::move(read_batch)));
+    ASSERT_TRUE(read_array->Equals(array->Slice(4, 2)));
+    ASSERT_EQ(4, reader->GetPreviousBatchFirstRowNumber().value());
+
+    // eof
+    ASSERT_OK_AND_ASSIGN(read_batch, reader->NextBatch());
+    ASSERT_TRUE(BatchReader::IsEofBatch(read_batch));
+    ASSERT_EQ(6, reader->GetPreviousBatchFirstRowNumber().value());
+
+    // test with bitmap pushdown
+    ArrowSchema c_read_schema;
+    ASSERT_TRUE(arrow::ExportSchema(*schema, &c_read_schema).ok());
+    ASSERT_OK(reader->SetReadSchema(&c_read_schema, /*predicate=*/nullptr,
+                                    /*selection_bitmap=*/RoaringBitmap32::From({0, 3})));
+    ASSERT_NOK_WITH_MSG(
+        reader->GetPreviousBatchFirstRowNumber(),
+        "Cannot call GetPreviousBatchFirstRowNumber in LanceFileBatchReader because, after bitmap "
+        "pushdown, rows in the array returned by NextBatch are no longer contiguous.");
+}
 }  // namespace paimon::lance::test