diff --git a/src/db/index/storage/mmap_forward_store.cc b/src/db/index/storage/mmap_forward_store.cc index dcaff40..f32765a 100644 --- a/src/db/index/storage/mmap_forward_store.cc +++ b/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 \ No newline at end of file +} // namespace zvec diff --git a/tests/db/index/storage/mmap_store_test.cc b/tests/db/index/storage/mmap_store_test.cc index c2d9bb1..e127dc4 100644 --- a/tests/db/index/storage/mmap_store_test.cc +++ b/tests/db/index/storage/mmap_store_test.cc @@ -22,6 +22,7 @@ #include #include #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 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(uneven_ipc_path); + ASSERT_TRUE(ipc_store->Open().ok()); + + auto result = ipc_store->fetch({"id"}, std::vector{6}); + ASSERT_NE(result, nullptr); + ASSERT_EQ(result->num_rows(), 1); + auto ids = + std::dynamic_pointer_cast(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(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; -} \ No newline at end of file +}