fix: properly handle larger final IPC chunk in mmap store (#645)

src/db/index/storage/mmap_forward_store.cc

@@ -118,11 +118,13 @@ arrow::Status MmapForwardStore::OpenIPC(
     chunk_index_map_.emplace_back(num_rows_, num_rows_ + chunk->length() - 1);
     num_rows_ += chunk->length();
 
-    // Check if all chunks have the same size except possibly the last one
+    // All non-last chunks must have the same size. The last chunk may be
+    // smaller, but a larger last chunk also requires the general lookup path.
     if (fixed_batch_size_ == -1) {
       fixed_batch_size_ = chunk->length();
     } else if (fixed_batch_size_ != chunk->length()) {
-      if (i != chunked_array->num_chunks() - 1) {
+      if (i != chunked_array->num_chunks() - 1 ||
+          chunk->length() > fixed_batch_size_) {
         is_fixed_batch_size_ = false;
       }
     }
@@ -568,4 +570,4 @@ RecordBatchReaderPtr MmapForwardStore::scan(
   }
 }
 
-}  // namespace zvec
+}  // namespace zvec

tests/db/index/storage/mmap_store_test.cc

@@ -22,6 +22,7 @@
 #include <arrow/table.h>
 #include <gtest/gtest.h>
 #include "db/common/constants.h"
+#include "db/index/storage/chunked_file_writer.h"
 #define private public
 #define protected public
 #include "db/index/storage/mmap_forward_store.h"
@@ -31,6 +32,31 @@
 
 using namespace zvec;
 
+namespace {
+
+arrow::Status WriteUnevenBatchIPC(const std::string &path) {
+  auto schema = arrow::schema({arrow::field("id", arrow::int32())});
+  auto writer = ChunkedFileWriter::Open(path, schema, FileFormat::IPC);
+  if (!writer) {
+    return arrow::Status::IOError("failed to open IPC writer");
+  }
+
+  int32_t next_id = 0;
+  for (int64_t batch_size : {3, 4}) {
+    arrow::Int32Builder builder;
+    for (int64_t i = 0; i < batch_size; ++i) {
+      ARROW_RETURN_NOT_OK(builder.Append(next_id++));
+    }
+    std::shared_ptr<arrow::Array> ids;
+    ARROW_RETURN_NOT_OK(builder.Finish(&ids));
+    auto batch = arrow::RecordBatch::Make(schema, batch_size, {ids});
+    ARROW_RETURN_NOT_OK(writer->Write(*batch));
+  }
+  return writer->Close();
+}
+
+}  // namespace
+
 class MmapStoreTest : public testing::Test {
  protected:
   void SetUp() override {
@@ -53,10 +79,14 @@ class MmapStoreTest : public testing::Test {
     if (std::filesystem::exists(parquet_path)) {
       std::filesystem::remove(parquet_path);
     }
+    if (std::filesystem::exists(uneven_ipc_path)) {
+      std::filesystem::remove(uneven_ipc_path);
+    }
   }
 
   std::string ipc_path = "test.ipc";
   std::string parquet_path = "test.parquet";
+  std::string uneven_ipc_path = "uneven.ipc";
 };
 
 
@@ -516,6 +546,20 @@ TEST_F(MmapStoreTest, IPCFetchSingleRow) {
   }
 }
 
+TEST_F(MmapStoreTest, IPCFetchFromLargerLastChunk) {
+  ASSERT_TRUE(WriteUnevenBatchIPC(uneven_ipc_path).ok());
+  auto ipc_store = std::make_shared<MmapForwardStore>(uneven_ipc_path);
+  ASSERT_TRUE(ipc_store->Open().ok());
+
+  auto result = ipc_store->fetch({"id"}, std::vector<int>{6});
+  ASSERT_NE(result, nullptr);
+  ASSERT_EQ(result->num_rows(), 1);
+  auto ids =
+      std::dynamic_pointer_cast<arrow::Int32Array>(result->column(0)->chunk(0));
+  ASSERT_NE(ids, nullptr);
+  EXPECT_EQ(ids->Value(0), 6);
+}
+
 TEST_F(MmapStoreTest, ParquetFetchSingleRow) {
   auto parquet_store = std::make_shared<MmapForwardStore>(parquet_path);
   ASSERT_TRUE(parquet_store->Open().ok());
@@ -673,4 +717,4 @@ TEST_F(MmapStoreTest, ConstructorAndPhysicSchema) {
 TEST_F(MmapStoreTest, DeleteDestructs) {
   MmapForwardStore *store = new MmapForwardStore(ipc_path);
   delete store;
-}
+}