wehub/alibaba--zvec
1
0
Fork 0
Code Issues Pull Requests Actions 2 Deployments Models Agent Notes Packages Projects Releases Wiki Activity
Files
main
alibaba--zvec/tests/db/index/segment/segment_test.cc
T
wehub-resource-sync be3ef883e1 chore: import upstream snapshot with attribution
2026-07-13 12:47:42 +08:00

2299 lines
80 KiB
C++
Raw Permalink Blame History

// Copyright 2025-present the zvec project
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
// http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.
#include <iostream>
#include <sstream> // NOLINT(misc-include-cleaner): protects <sstream> from the private/public test macro below.
#define private public
#define protected public
#include "db/index/segment/segment.h"
#undef private
#undef protected
#include <cstdint>
#include <memory>
#include <thread>
#include <arrow/array/array_binary.h>
#include <arrow/io/file.h>
#include <arrow/ipc/reader.h>
#include <arrow/pretty_print.h>
#include <arrow/result.h>
#include <arrow/table.h>
#include <gtest/gtest.h>
#include <zvec/ailego/buffer/block_eviction_queue.h>
#include "db/common/file_helper.h"
#include "db/index/common/delete_store.h"
#include "db/index/common/id_map.h"
#include "db/index/common/version_manager.h"
#include "db/index/storage/wal/wal_file.h"
#include "segment_test_fixture.h"
#include "utils/utils.h"
#include "zvec/db/options.h"
using namespace zvec;
TEST_P(SegmentTest, EmptySchema) {
auto segment = test::TestHelper::CreateSegmentWithDoc(
col_path_, *schema_, 0, 0, id_map_, delete_store_, version_manager_, options_,
0, 0);
ASSERT_TRUE(segment != nullptr);
EXPECT_EQ(segment->id(), 0);
segment.reset();
}
TEST_P(SegmentTest, General) {
options_.max_buffer_size_ = 1 * 1024;
auto segment = test::TestHelper::CreateSegmentWithDoc(
col_path_, *schema_, 0, 0, id_map_, delete_store_, version_manager_, options_,
0, 25);
ASSERT_TRUE(segment != nullptr);
auto combined_reader = segment->scan({LOCAL_ROW_ID, "id", "name", "age"});
ASSERT_TRUE(combined_reader != nullptr);
EXPECT_TRUE(combined_reader->schema() != nullptr);
std::shared_ptr<arrow::RecordBatch> batch;
uint32_t total_doc = 0;
while (true) {
auto status = combined_reader->ReadNext(&batch);
if (status.ok() == false) break;
if (batch == nullptr) break;
EXPECT_EQ(batch->num_columns(), 4);
total_doc += batch->num_rows();
}
EXPECT_EQ(total_doc, 25);
std::vector<int> segment_doc_ids = {0, 3, 6, 1, 0, 14, 12, 21};
auto combined_table = segment->fetch(
{LOCAL_ROW_ID, "id", "name", "age", "binary", "array_binary"},
segment_doc_ids);
ASSERT_TRUE(combined_table != nullptr);
EXPECT_EQ(combined_table->num_columns(), 6);
EXPECT_EQ(combined_table->num_rows(), 8);
auto field = combined_table->schema()->field(0);
EXPECT_EQ(field->name(), LOCAL_ROW_ID);
// Get data from the LOCAL_ROW_ID column for each row
auto id_column = combined_table->column(0);
auto id_array =
std::dynamic_pointer_cast<arrow::UInt64Array>(id_column->chunk(0));
std::vector<int32_t> &expected_ids = segment_doc_ids;
std::vector<int32_t> actual_ids;
for (int i = 0; i < id_array->length(); ++i) {
actual_ids.push_back(id_array->Value(i));
}
EXPECT_EQ(actual_ids, expected_ids)
<< "ID column values don't match expected order";
}
TEST_P(SegmentTest, InsertMoreData) {
auto segment = test::TestHelper::CreateSegmentWithDoc(
col_path_, *schema_, 0, 0, id_map_, delete_store_, version_manager_, options_,
0, 0);
ASSERT_TRUE(segment != nullptr);
uint64_t MAX_DOC = 1000;
auto start = std::chrono::system_clock::now();
test::TestHelper::SegmentInsertDoc(segment, *schema_, 0, MAX_DOC);
auto end = std::chrono::system_clock::now();
auto cost = std::chrono::duration_cast<std::chrono::milliseconds>(end - start)
.count();
std::cout << "insert cost " << cost << "ms" << std::endl;
auto combined_reader = segment->scan({"id", "name", "age"});
std::shared_ptr<arrow::RecordBatch> batch;
uint32_t total_doc = 0;
while (true) {
auto status = combined_reader->ReadNext(&batch);
if (status.ok() == false) break;
if (batch == nullptr) break;
total_doc += batch->num_rows();
}
EXPECT_EQ(total_doc, MAX_DOC);
}
TEST_P(SegmentTest, InsertScalarTypes) {
auto tmp_schema =
test::TestHelper::CreateSchemaWithScalarIndex(true, true, col_name_);
auto invert_params = std::make_shared<InvertIndexParams>(false);
schema_->add_field(std::make_shared<FieldSchema>("binary", DataType::BINARY,
false, invert_params));
schema_->add_field(std::make_shared<FieldSchema>(
"array_binary", DataType::ARRAY_BINARY, false, invert_params));
auto segment = test::TestHelper::CreateSegmentWithDoc(
col_path_, *schema_, 0, 0, id_map_, delete_store_, version_manager_, options_,
0, 10);
ASSERT_TRUE(segment != nullptr);
}
TEST_P(SegmentTest, InsertVectorTypes) {
auto tmp_schema = test::TestHelper::CreateSchemaWithVectorIndex(
false, col_name_,
std::make_shared<HnswIndexParams>(MetricType::IP, 16, 20,
QuantizeType::FP16));
// first insert 100 doc
int doc_count = 100;
{
auto segment = test::TestHelper::CreateSegmentWithDoc(
col_path_, *tmp_schema, 0, 0, id_map_, delete_store_, version_manager_,
options_, 0, doc_count);
ASSERT_TRUE(segment != nullptr);
}
// Open
{
Version v = version_manager_->get_current_version();
auto result =
Segment::Open(col_path_, *tmp_schema, *v.writing_segment_meta(), id_map_,
delete_store_, version_manager_, options_);
ASSERT_TRUE(result.has_value());
auto segment = result.value();
EXPECT_GT(segment->get_vector_indexer("dense_fp32").size(), 0);
EXPECT_GT(segment->get_quant_vector_indexer("dense_fp32").size(), 0);
}
}
TEST_P(SegmentTest, FetchByGlobalDocID) {
auto segment = test::TestHelper::CreateSegmentWithDoc(
col_path_, *schema_, 0, 0, id_map_, delete_store_, version_manager_, options_,
0, 1);
ASSERT_TRUE(segment != nullptr);
auto ret_doc = segment->Fetch(0);
EXPECT_TRUE(ret_doc != nullptr);
EXPECT_EQ(ret_doc->doc_id(), 0);
EXPECT_EQ(ret_doc->pk(), "pk_0");
}
TEST_P(SegmentTest, FetchSingleRow) {
int doc_count = 10;
auto segment = test::TestHelper::CreateSegmentWithDoc(
col_path_, *schema_, 0, 0, id_map_, delete_store_, version_manager_, options_,
0, doc_count);
ASSERT_TRUE(segment != nullptr);
auto func = [&](int index) -> void {
ExecBatchPtr batch = segment->fetch({"id", "name", "age"}, index);
ASSERT_TRUE(batch != nullptr);
EXPECT_EQ(batch->length, 1);
EXPECT_EQ(batch->values.size(), 3);
auto id_scalar = batch->values[0].scalar();
ASSERT_TRUE(id_scalar != nullptr);
auto id_value = std::dynamic_pointer_cast<arrow::Int32Scalar>(id_scalar);
ASSERT_TRUE(id_value != nullptr);
EXPECT_EQ(id_value->value, index);
};
for (int i = 0; i < doc_count; ++i) {
func(i);
}
}
TEST_P(SegmentTest, FetchSingleRowWithPersistStore) {
// first insert 1000 doc
int doc_count = 1000;
{
auto segment = test::TestHelper::CreateSegmentWithDoc(
col_path_, *schema_, 0, 0, id_map_, delete_store_, version_manager_, options_,
0, doc_count);
ASSERT_TRUE(segment != nullptr);
}
// Open
{
Version v = version_manager_->get_current_version();
SegmentOptions open_options;
open_options.read_only_ = false;
auto result = Segment::Open(col_path_, *schema_, *v.writing_segment_meta(),
id_map_, delete_store_, version_manager_,
open_options);
ASSERT_TRUE(result.has_value());
auto segment = result.value();
test::TestHelper::SegmentInsertDoc(segment, *schema_, doc_count,
doc_count * 2);
auto func = [&](int index) -> void {
ExecBatchPtr batch = segment->fetch({"id", "name", "age"}, index);
ASSERT_TRUE(batch != nullptr);
EXPECT_EQ(batch->length, 1);
EXPECT_EQ(batch->values.size(), 3);
auto id_scalar = batch->values[0].scalar();
ASSERT_TRUE(id_scalar != nullptr);
auto id_value = std::dynamic_pointer_cast<arrow::Int32Scalar>(id_scalar);
ASSERT_TRUE(id_value != nullptr);
EXPECT_EQ(id_value->value, index);
};
for (int i = 0; i < doc_count * 2; ++i) {
func(i);
}
}
}
TEST_P(SegmentTest, FetchSingleRowWithUserID) {
auto segment = test::TestHelper::CreateSegmentWithDoc(
col_path_, *schema_, 0, 0, id_map_, delete_store_, version_manager_, options_,
0, 10);
ASSERT_TRUE(segment != nullptr);
ExecBatchPtr batch = segment->fetch({USER_ID, "id", "name"}, 2);
ASSERT_TRUE(batch != nullptr);
EXPECT_EQ(batch->length, 1);
EXPECT_EQ(batch->values.size(), 3);
auto user_id_scalar = batch->values[0].scalar();
ASSERT_TRUE(user_id_scalar != nullptr);
EXPECT_TRUE(std::dynamic_pointer_cast<arrow::StringScalar>(user_id_scalar) !=
nullptr);
}
TEST_P(SegmentTest, FetchSingleRowWithGlobalDocID) {
auto segment = test::TestHelper::CreateSegmentWithDoc(
col_path_, *schema_, 0, 0, id_map_, delete_store_, version_manager_, options_,
0, 10);
ASSERT_TRUE(segment != nullptr);
ExecBatchPtr batch = segment->fetch({GLOBAL_DOC_ID, "id", "name"}, 4);
ASSERT_TRUE(batch != nullptr);
EXPECT_EQ(batch->length, 1);
EXPECT_EQ(batch->values.size(), 3);
auto global_doc_id_scalar = batch->values[0].scalar();
ASSERT_TRUE(global_doc_id_scalar != nullptr);
EXPECT_TRUE(std::dynamic_pointer_cast<arrow::UInt64Scalar>(
global_doc_id_scalar) != nullptr);
}
TEST_P(SegmentTest, FetchSingleRowWithNegativeIndex) {
auto segment = test::TestHelper::CreateSegmentWithDoc(
col_path_, *schema_, 0, 0, id_map_, delete_store_, version_manager_, options_,
0, 10);
ASSERT_TRUE(segment != nullptr);
ExecBatchPtr batch = segment->fetch({"id", "name"}, -1);
EXPECT_EQ(batch, nullptr);
}
TEST_P(SegmentTest, FetchSingleRowWithOutOfRangeIndex) {
auto segment = test::TestHelper::CreateSegmentWithDoc(
col_path_, *schema_, 0, 0, id_map_, delete_store_, version_manager_, options_,
0, 10);
ASSERT_TRUE(segment != nullptr);
ExecBatchPtr batch = segment->fetch({"id", "name"}, 15);
EXPECT_EQ(batch, nullptr);
}
TEST_P(SegmentTest, FetchSingleRowWithInvalidColumn) {
auto segment = test::TestHelper::CreateSegmentWithDoc(
col_path_, *schema_, 0, 0, id_map_, delete_store_, version_manager_, options_,
0, 10);
ASSERT_TRUE(segment != nullptr);
ExecBatchPtr batch = segment->fetch({"id", "invalid_column"}, 0);
EXPECT_EQ(batch, nullptr);
}
TEST_P(SegmentTest, FetchSingleRowWithEmptyColumns) {
auto segment = test::TestHelper::CreateSegmentWithDoc(
col_path_, *schema_, 0, 0, id_map_, delete_store_, version_manager_, options_,
0, 10);
ASSERT_TRUE(segment != nullptr);
ExecBatchPtr batch = segment->fetch({}, 0);
EXPECT_EQ(batch, nullptr);
}
TEST_P(SegmentTest, FetchSingleRowFromEmptySegment) {
auto segment = test::TestHelper::CreateSegmentWithDoc(
col_path_, *schema_, 0, 0, id_map_, delete_store_, version_manager_, options_,
0, 0);
ASSERT_TRUE(segment != nullptr);
ExecBatchPtr batch = segment->fetch({"id", "name"}, 0);
EXPECT_EQ(batch, nullptr);
}
TEST_P(SegmentTest, FetchSingleRowWithBinaryFields) {
auto segment = test::TestHelper::CreateSegmentWithDoc(
col_path_, *schema_, 0, 0, id_map_, delete_store_, version_manager_, options_,
0, 10);
ASSERT_TRUE(segment != nullptr);
ExecBatchPtr batch = segment->fetch({"binary", "array_binary"}, 1);
ASSERT_TRUE(batch != nullptr);
EXPECT_EQ(batch->length, 1);
EXPECT_EQ(batch->values.size(), 2);
auto binary_scalar = batch->values[0].scalar();
ASSERT_TRUE(binary_scalar != nullptr);
EXPECT_TRUE(std::dynamic_pointer_cast<arrow::BinaryScalar>(binary_scalar) !=
nullptr);
auto array_binary_scalar = batch->values[1].scalar();
ASSERT_TRUE(array_binary_scalar != nullptr);
EXPECT_TRUE(std::dynamic_pointer_cast<arrow::ListScalar>(
array_binary_scalar) != nullptr);
}
TEST_P(SegmentTest, Recover) {
// first insert 100 doc
int doc_count = 100;
{
auto segment = test::TestHelper::CreateSegmentWithDoc(
col_path_, *schema_, 0, 0, id_map_, delete_store_, version_manager_, options_,
0, doc_count);
ASSERT_TRUE(segment != nullptr);
}
// simulate wal file
{
Version v = version_manager_->get_current_version();
auto writing_block_id =
v.writing_segment_meta()->writing_forward_block_->id();
auto wal_file = FileHelper::MakeWalPath(col_path_, 0, writing_block_id);
WalOptions wal_option{0, true};
WalFilePtr wal_file_;
WalFile::CreateAndOpen(wal_file, wal_option, &wal_file_);
ASSERT_TRUE(wal_file_ != nullptr);
for (int i = doc_count; i < doc_count + 100; i++) {
Doc doc = test::TestHelper::CreateDoc(i, *schema_);
doc.set_operator(Operator::INSERT);
std::vector<uint8_t> buf = doc.serialize();
auto ret = wal_file_->append(std::string(buf.begin(), buf.end()));
ASSERT_EQ(ret, 0);
}
for (int i = 0; i < doc_count; i++) {
Doc doc = test::TestHelper::CreateDoc(i, *schema_);
doc.set_doc_id(i); // global doc id
doc.set_operator(Operator::UPDATE);
std::vector<uint8_t> buf = doc.serialize();
auto ret = wal_file_->append(std::string(buf.begin(), buf.end()));
ASSERT_EQ(ret, 0);
}
for (int i = 0; i < doc_count; i++) {
Doc doc = test::TestHelper::CreateDoc(i, *schema_);
doc.set_operator(Operator::UPSERT);
std::vector<uint8_t> buf = doc.serialize();
auto ret = wal_file_->append(std::string(buf.begin(), buf.end()));
ASSERT_EQ(ret, 0);
}
for (int i = 0; i < doc_count; i++) {
Doc doc = test::TestHelper::CreateDoc(i, *schema_);
doc.set_doc_id(i + 300); // global doc id
doc.set_operator(Operator::DELETE);
std::vector<uint8_t> buf = doc.serialize();
auto ret = wal_file_->append(std::string(buf.begin(), buf.end()));
ASSERT_EQ(ret, 0);
}
}
// recover
{
Version v = version_manager_->get_current_version();
SegmentOptions open_options;
open_options.read_only_ = false;
auto result = Segment::Open(col_path_, *schema_, *v.writing_segment_meta(),
id_map_, delete_store_, version_manager_,
open_options);
ASSERT_TRUE(result.has_value());
auto segment = result.value();
auto combined_reader = segment->scan({"id"});
std::shared_ptr<arrow::RecordBatch> batch;
uint32_t total_doc = 0;
while (true) {
auto status = combined_reader->ReadNext(&batch);
if (status.ok() == false) break;
if (batch == nullptr) break;
total_doc += batch->num_rows();
EXPECT_EQ(batch->num_columns(), 1);
}
// Why 400 ? because in segment we just mark deleted doc
EXPECT_EQ(total_doc, 400);
auto filter = delete_store_->make_filter();
auto actual_doc_count = segment->doc_count(filter);
EXPECT_EQ(actual_doc_count, 100);
}
}
TEST_P(SegmentTest, UpdateDoc) {
auto segment = test::TestHelper::CreateSegmentWithDoc(
col_path_, *schema_, 0, 0, id_map_, delete_store_, version_manager_, options_,
0, 10);
ASSERT_TRUE(segment != nullptr);
// before update
uint64_t count = segment->doc_count(delete_store_->make_filter());
EXPECT_EQ(count, 10);
// Create a new document to update
Doc update_doc = test::TestHelper::CreateDoc(5, *schema_);
update_doc.set<std::string>("name", "updated_name");
update_doc.set<uint32_t>("age", 99);
// Update the document
auto status = segment->Update(update_doc);
EXPECT_TRUE(status.ok()) << "Update failed: " << status.message();
// after update
count = segment->doc_count(delete_store_->make_filter());
EXPECT_EQ(count, 10);
// Fetch the updated document and verify changes
// Note: The parameter here is the internal global_doc_id, not user-specified
auto ret_doc = segment->Fetch(10);
EXPECT_TRUE(ret_doc != nullptr);
EXPECT_EQ(ret_doc->get<std::string>("name"), "updated_name");
EXPECT_EQ(ret_doc->get<uint32_t>("age"), 99);
}
TEST_P(SegmentTest, UpdateDocBatch) {
int doc_count = 10;
auto segment = test::TestHelper::CreateSegmentWithDoc(
col_path_, *schema_, 0, 0, id_map_, delete_store_, version_manager_, options_,
0, doc_count);
ASSERT_TRUE(segment != nullptr);
// before update
uint64_t count = segment->doc_count(delete_store_->make_filter());
EXPECT_EQ(count, doc_count);
// Create a new document to update
for (int i = 0; i < doc_count; i++) {
Doc update_doc = test::TestHelper::CreateDoc(i, *schema_);
// Update the document
auto status = segment->Update(update_doc);
EXPECT_TRUE(status.ok()) << "Update failed: " << status.message();
}
// after update
count = segment->doc_count(delete_store_->make_filter());
EXPECT_EQ(count, doc_count);
// Fetch the updated document and verify changes
// Note: The parameter here is the internal global_doc_id, not user-specified
auto ret_doc = segment->Fetch(doc_count * 2 - 1);
EXPECT_TRUE(ret_doc != nullptr);
EXPECT_EQ(ret_doc->get<std::string>("name"),
"value_" + std::to_string(doc_count - 1));
}
TEST_P(SegmentTest, DeleteDoc) {
auto segment = test::TestHelper::CreateSegmentWithDoc(
col_path_, *schema_, 0, 0, id_map_, delete_store_, version_manager_, options_,
0, 10);
ASSERT_TRUE(segment != nullptr);
// before update
uint64_t count = segment->doc_count(delete_store_->make_filter());
EXPECT_EQ(count, 10);
// Delete a document by primary key
auto status = segment->Delete("pk_5");
EXPECT_TRUE(status.ok()) << "Delete by pk failed: " << status.message();
// after delete
count = segment->doc_count(delete_store_->make_filter());
EXPECT_EQ(count, 9);
// Delete a document by global doc id
status = segment->Delete(3);
EXPECT_TRUE(status.ok()) << "Delete by global doc id failed: "
<< status.message();
count = segment->doc_count(delete_store_->make_filter());
EXPECT_EQ(count, 8);
}
TEST_P(SegmentTest, DeleteBatch) {
int doc_count = 10;
auto segment = test::TestHelper::CreateSegmentWithDoc(
col_path_, *schema_, 0, 0, id_map_, delete_store_, version_manager_, options_,
0, doc_count);
ASSERT_TRUE(segment != nullptr);
// before update
uint64_t count = segment->doc_count(delete_store_->make_filter());
EXPECT_EQ(count, doc_count);
for (int i = 0; i < doc_count; i++) {
auto status = segment->Delete("pk_" + std::to_string(i));
EXPECT_TRUE(status.ok()) << "Delete by pk failed: " << status.message();
}
// after delete
count = segment->doc_count(delete_store_->make_filter());
EXPECT_EQ(count, 0);
}
TEST_P(SegmentTest, UpsertDoc) {
auto segment = test::TestHelper::CreateSegmentWithDoc(
col_path_, *schema_, 0, 0, id_map_, delete_store_, version_manager_, options_,
0, 5);
ASSERT_TRUE(segment != nullptr);
// before update
uint64_t count = segment->doc_count(delete_store_->make_filter());
EXPECT_EQ(count, 5);
// Upsert an existing document
Doc upsert_doc1 = test::TestHelper::CreateDoc(3, *schema_);
upsert_doc1.set<std::string>("name", "upserted_name");
auto status = segment->Upsert(upsert_doc1);
EXPECT_TRUE(status.ok()) << "Upsert existing doc failed: "
<< status.message();
count = segment->doc_count(delete_store_->make_filter());
EXPECT_EQ(count, 5);
// Verify the update
auto ret_doc = segment->Fetch(5);
EXPECT_TRUE(ret_doc != nullptr);
EXPECT_EQ(ret_doc->get<std::string>("name"), "upserted_name");
// Upsert a new document
Doc upsert_doc2 = test::TestHelper::CreateDoc(6, *schema_);
upsert_doc2.set<std::string>("name", "new_upserted_doc");
status = segment->Upsert(upsert_doc2);
EXPECT_TRUE(status.ok()) << "Upsert new doc failed: " << status.message();
count = segment->doc_count(delete_store_->make_filter());
EXPECT_EQ(count, 6);
// Verify the new document was inserted
ret_doc = segment->Fetch(6);
EXPECT_TRUE(ret_doc != nullptr);
EXPECT_EQ(ret_doc->get<std::string>("name"), "new_upserted_doc");
}
TEST_P(SegmentTest, UpsertDocBatch) {
int doc_count = 10;
auto segment = test::TestHelper::CreateSegmentWithDoc(
col_path_, *schema_, 0, 0, id_map_, delete_store_, version_manager_, options_,
0, doc_count);
ASSERT_TRUE(segment != nullptr);
// before update
uint64_t count = segment->doc_count(delete_store_->make_filter());
EXPECT_EQ(count, doc_count);
for (int i = 0; i < doc_count; i++) {
// Upsert existing document
Doc upsert_doc1 = test::TestHelper::CreateDoc(i, *schema_);
upsert_doc1.set<std::string>("name", "upserted_name" + std::to_string(i));
auto status = segment->Upsert(upsert_doc1);
EXPECT_TRUE(status.ok())
<< "Upsert existing doc failed: " << status.message();
// Upsert new document
Doc upsert_doc2 = test::TestHelper::CreateDoc(doc_count + i, *schema_);
upsert_doc2.set<std::string>("name",
"new_upserted_doc" + std::to_string(i));
status = segment->Upsert(upsert_doc2);
EXPECT_TRUE(status.ok()) << "Upsert new doc failed: " << status.message();
}
count = segment->doc_count(delete_store_->make_filter());
EXPECT_EQ(count, doc_count * 2);
int incr_idx = 0;
for (int i = doc_count; i < doc_count + doc_count * 2; i += 2) {
// Verify the update
auto ret_doc = segment->Fetch(i);
EXPECT_TRUE(ret_doc != nullptr);
EXPECT_EQ(ret_doc->get<std::string>("name"),
"upserted_name" + std::to_string(incr_idx));
// Verify the new document was inserted
ret_doc = segment->Fetch(i + 1);
EXPECT_EQ(ret_doc->get<std::string>("name"),
"new_upserted_doc" + std::to_string(incr_idx));
incr_idx++;
}
}
TEST_P(SegmentTest, Flush) {
auto segment = test::TestHelper::CreateSegmentWithDoc(
col_path_, *schema_, 0, 0, id_map_, delete_store_, version_manager_, options_,
0, 100);
ASSERT_TRUE(segment != nullptr);
// Flush the segment
auto status = segment->flush();
EXPECT_TRUE(status.ok()) << "Flush failed: " << status.message();
}
TEST_P(SegmentTest, FlushAfterInsert) {
auto segment = test::TestHelper::CreateSegmentWithDoc(
col_path_, *schema_, 0, 0, id_map_, delete_store_, version_manager_, options_,
0, 100);
ASSERT_TRUE(segment != nullptr);
// Flush the segment
auto status = segment->flush();
EXPECT_TRUE(status.ok()) << "Flush failed: " << status.message();
test::TestHelper::SegmentInsertDoc(segment, *schema_, 100, 150);
ASSERT_EQ(segment->doc_count(), 150);
for (int i = 0; i < 150; i++) {
auto ret_doc = segment->Fetch(i);
EXPECT_TRUE(ret_doc != nullptr);
Doc verify_doc = test::TestHelper::CreateDoc(i, *schema_);
auto vv = verify_doc.get<std::vector<float>>("dense_fp32").value();
auto v = ret_doc->get<std::vector<float>>("dense_fp32").value();
for (uint32_t j = 0; j < vv.size(); j++) {
ASSERT_FLOAT_EQ(v[j], vv[j]);
}
}
}
TEST_P(SegmentTest, Dump) {
auto segment = test::TestHelper::CreateSegmentWithDoc(
col_path_, *schema_, 0, 0, id_map_, delete_store_, version_manager_, options_,
0, 100);
ASSERT_TRUE(segment != nullptr);
// Dump the segment
auto status = segment->dump();
EXPECT_TRUE(status.ok()) << "Flush failed: " << status.message();
status = segment->dump();
EXPECT_FALSE(status.ok());
EXPECT_EQ(status.code(), StatusCode::NOT_SUPPORTED);
}
TEST_P(SegmentTest, DocCount) {
auto segment = test::TestHelper::CreateSegmentWithDoc(
col_path_, *schema_, 0, 0, id_map_, delete_store_, version_manager_, options_,
0, 50);
ASSERT_TRUE(segment != nullptr);
// Get document count
uint64_t count = segment->doc_count();
EXPECT_EQ(count, 50);
// Delete some documents
segment->Delete("pk_10");
segment->Delete("pk_20");
segment->Delete("pk_30");
// Get document count again
count = segment->doc_count(delete_store_->make_filter());
EXPECT_EQ(count, 47);
}
// TEST_P(SegmentTest, Insert100WData) {
// options_.max_buffer_size_ = 8 * 1024 * 1024;
// auto segment = test::TestHelper::CreateSegmentWithDoc(
// col_path_, *schema_, 0, 0, id_map_, delete_store_, version_manager_,
// options_, 0, 0);
// ASSERT_TRUE(segment != nullptr);
// uint64_t MAX_DOC = 1000000;
// auto start = std::chrono::system_clock::now();
// test::TestHelper::SegmentInsertDoc(segment, *schema_, 0, MAX_DOC);
// auto end = std::chrono::system_clock::now();
// auto cost = std::chrono::duration_cast<std::chrono::milliseconds>(end -
// start)
// .count();
// std::cout << "insert cost " << cost << "ms" << std::endl;
// start = std::chrono::system_clock::now();
// ;
// auto combined_reader = segment->scan(
// {"id", "name", "age", USER_ID, GLOBAL_DOC_ID, LOCAL_ROW_ID});
// std::shared_ptr<arrow::RecordBatch> batch;
// uint32_t total_doc = 0;
// while (true) {
// auto status = combined_reader->ReadNext(&batch);
// if (status.ok() == false) break;
// if (batch == nullptr) break;
// total_doc += batch->num_rows();
// }
// end = std::chrono::system_clock::now();
// cost = std::chrono::duration_cast<std::chrono::milliseconds>(end - start)
// .count();
// std::cout << "scan cost " << cost << "ms" << std::endl;
// EXPECT_EQ(total_doc, MAX_DOC);
// }
TEST_P(SegmentTest, CombinedVectorColumnIndexer) {
options_.max_buffer_size_ = 10 * 1024;
auto segment = test::TestHelper::CreateSegmentWithDoc(
col_path_, *schema_, 0, 0, id_map_, delete_store_, version_manager_, options_,
0, 0);
ASSERT_TRUE(segment != nullptr);
uint64_t MAX_DOC = 1000;
test::TestHelper::SegmentInsertDoc(segment, *schema_, 0, MAX_DOC);
Doc new_doc = test::TestHelper::CreateDoc(1000, *schema_);
auto status = segment->Insert(new_doc);
ASSERT_TRUE(status.ok());
auto combined_indexer = segment->get_combined_vector_indexer("dense_fp32");
ASSERT_TRUE(combined_indexer != nullptr);
// fetch
auto fetched_data = combined_indexer->Fetch(1000);
ASSERT_TRUE(fetched_data);
const float *dense_vector = reinterpret_cast<const float *>(
std::get<vector_column_params::DenseVectorBuffer>(
fetched_data->vector_buffer)
.data.data());
auto vv = new_doc.get<std::vector<float>>("dense_fp32").value();
for (uint32_t i = 0; i < vv.size(); i++) {
ASSERT_FLOAT_EQ(dense_vector[i], vv[i]);
}
// query
auto dense_fp32_field = schema_->get_field("dense_fp32");
auto query_vector = new_doc.get<std::vector<float>>("dense_fp32").value();
auto query = vector_column_params::VectorData{
vector_column_params::DenseVector{query_vector.data()}};
vector_column_params::QueryParams query_params;
query_params.dimension = dense_fp32_field->dimension();
query_params.topk = 10;
query_params.filter = nullptr;
query_params.fetch_vector = false;
auto results = combined_indexer->Search(query, query_params);
ASSERT_TRUE(results.has_value());
auto vector_results =
dynamic_cast<VectorIndexResults *>(results.value().get());
ASSERT_TRUE(vector_results);
ASSERT_EQ(vector_results->count(), 10);
int count = 0;
auto iter = vector_results->create_iterator();
while (iter->valid()) {
count++;
iter->next();
}
ASSERT_EQ(count, 10);
}
TEST_P(SegmentTest, CombinedVectorColumnIndexerWithQuantVectorIndex) {
options_.max_buffer_size_ = 10 * 1024;
auto tmp_schema = test::TestHelper::CreateSchemaWithVectorIndex(
false, "demo",
std::make_shared<HnswIndexParams>(MetricType::IP, 16, 20,
QuantizeType::FP16));
auto segment = test::TestHelper::CreateSegmentWithDoc(
col_path_, *tmp_schema, 0, 0, id_map_, delete_store_, version_manager_,
options_, 0, 0);
ASSERT_TRUE(segment != nullptr);
uint64_t MAX_DOC = 1000;
test::TestHelper::SegmentInsertDoc(segment, *schema_, 0, MAX_DOC);
Doc new_doc = test::TestHelper::CreateDoc(1000, *schema_);
auto status = segment->Insert(new_doc);
ASSERT_TRUE(status.ok());
auto combined_indexer =
segment->get_quant_combined_vector_indexer("dense_fp32");
ASSERT_TRUE(combined_indexer != nullptr);
// fetch
auto fetched_data = combined_indexer->Fetch(1000);
ASSERT_TRUE(fetched_data);
const float *dense_vector = reinterpret_cast<const float *>(
std::get<vector_column_params::DenseVectorBuffer>(
fetched_data->vector_buffer)
.data.data());
auto vv = new_doc.get<std::vector<float>>("dense_fp32").value();
for (uint32_t i = 0; i < vv.size(); i++) {
EXPECT_NEAR(dense_vector[i], vv[i], 0.1);
}
// query
auto dense_fp32_field = schema_->get_field("dense_fp32");
auto query_vector = new_doc.get<std::vector<float>>("dense_fp32").value();
auto query = vector_column_params::VectorData{
vector_column_params::DenseVector{query_vector.data()}};
vector_column_params::QueryParams query_params;
query_params.dimension = dense_fp32_field->dimension();
query_params.topk = 10;
query_params.filter = nullptr;
query_params.fetch_vector = false;
query_params.query_params =
std::make_shared<zvec::QueryParams>(IndexType::HNSW);
query_params.query_params->set_is_using_refiner(true);
auto results = combined_indexer->Search(query, query_params);
ASSERT_TRUE(results.has_value());
auto vector_results =
dynamic_cast<VectorIndexResults *>(results.value().get());
ASSERT_TRUE(vector_results);
ASSERT_EQ(vector_results->count(), 10);
int count = 0;
auto iter = vector_results->create_iterator();
while (iter->valid()) {
count++;
iter->next();
}
ASSERT_EQ(count, 10);
}
TEST_P(SegmentTest, CombinedVectorColumnIndexerQueryWithPks) {
options_.max_buffer_size_ = 10 * 1024;
auto tmp_schema = test::TestHelper::CreateSchemaWithVectorIndex(
false, "demo", std::make_shared<HnswIndexParams>(MetricType::IP));
auto segment = test::TestHelper::CreateSegmentWithDoc(
col_path_, *tmp_schema, 0, 0, id_map_, delete_store_, version_manager_,
options_, 0, 0);
ASSERT_TRUE(segment != nullptr);
uint64_t MAX_DOC = 1000;
test::TestHelper::SegmentInsertDoc(segment, *schema_, 0, MAX_DOC);
auto combined_indexer = segment->get_combined_vector_indexer("dense_fp32");
ASSERT_TRUE(combined_indexer != nullptr);
Doc verify_doc = test::TestHelper::CreateDoc(999, *schema_);
std::vector<std::vector<uint64_t>> bf_pks = {
{10, 20, 30, 40, 50, 60, 70, 80, 90, 999}};
// query
auto dense_fp32_field = schema_->get_field("dense_fp32");
auto query_vector = verify_doc.get<std::vector<float>>("dense_fp32").value();
auto query = vector_column_params::VectorData{
vector_column_params::DenseVector{query_vector.data()}};
vector_column_params::QueryParams query_params;
query_params.data_type = dense_fp32_field->data_type();
query_params.dimension = dense_fp32_field->dimension();
query_params.topk = 10;
query_params.filter = nullptr;
query_params.fetch_vector = false;
query_params.query_params =
std::make_shared<zvec::QueryParams>(IndexType::HNSW);
query_params.bf_pks = bf_pks;
auto results = combined_indexer->Search(query, query_params);
ASSERT_TRUE(results.has_value());
auto vector_results =
dynamic_cast<VectorIndexResults *>(results.value().get());
ASSERT_TRUE(vector_results);
ASSERT_EQ(vector_results->count(), 10);
int count = 0;
std::vector<uint64_t> result_doc_ids;
auto iter = vector_results->create_iterator();
while (iter->valid()) {
count++;
result_doc_ids.push_back(iter->doc_id());
iter->next();
}
ASSERT_EQ(count, 10);
// need reverse result_doc_ids
std::reverse(result_doc_ids.begin(), result_doc_ids.end());
ASSERT_EQ(result_doc_ids, bf_pks[0]);
}
TEST_P(SegmentTest, ConcurrentInsertOperations) {
auto segment = test::TestHelper::CreateSegmentWithDoc(
col_path_, *schema_, 0, 0, id_map_, delete_store_, version_manager_, options_,
0, 0);
ASSERT_TRUE(segment != nullptr);
const int num_threads = 4;
const int docs_per_thread = 50;
std::vector<std::thread> threads;
// Launch multiple threads to insert documents concurrently
for (int t = 0; t < num_threads; ++t) {
threads.emplace_back([&, t]() {
for (int i = 0; i < docs_per_thread; ++i) {
int doc_id = t * docs_per_thread + i;
Doc doc = test::TestHelper::CreateDoc(doc_id, *schema_);
auto status = segment->Insert(doc);
EXPECT_TRUE(status.ok())
<< "Thread " << t << " insert failed for doc " << doc_id;
}
});
}
// Wait for all threads to complete
for (auto &thread : threads) {
thread.join();
}
// Verify total document count
uint64_t count = segment->doc_count();
EXPECT_EQ(count, num_threads * docs_per_thread);
}
TEST_P(SegmentTest, ConcurrentMixedOperations) {
auto segment = test::TestHelper::CreateSegmentWithDoc(
col_path_, *schema_, 0, 0, id_map_, delete_store_, version_manager_, options_,
0, 100);
ASSERT_TRUE(segment != nullptr);
std::vector<std::thread> threads;
// Thread 1: Insert new documents
threads.emplace_back([&]() {
for (int i = 100; i < 120; ++i) {
Doc doc = test::TestHelper::CreateDoc(i, *schema_);
auto status = segment->Insert(doc);
EXPECT_TRUE(status.ok() || status.code() == StatusCode::ALREADY_EXISTS);
}
});
// Thread 2: Update existing documents
threads.emplace_back([&]() {
for (int i = 0; i < 50; i += 5) {
Doc doc = test::TestHelper::CreateDoc(i, *schema_);
doc.set<std::string>("name", "updated_concurrent_" + std::to_string(i));
auto status = segment->Update(doc);
EXPECT_TRUE(status.ok() || status.code() == StatusCode::NOT_FOUND);
}
});
// Thread 3: Delete documents
threads.emplace_back([&]() {
for (int i = 50; i < 100; i += 10) {
auto status = segment->Delete("pk_" + std::to_string(i));
EXPECT_TRUE(status.ok() || status.code() == StatusCode::NOT_FOUND);
}
});
// Wait for all threads to complete
for (auto &thread : threads) {
thread.join();
}
}
// corner cases
TEST_P(SegmentTest, DuplicateInsert) {
auto segment = test::TestHelper::CreateSegmentWithDoc(
col_path_, *schema_, 0, 0, id_map_, delete_store_, version_manager_, options_,
0, 0);
ASSERT_TRUE(segment != nullptr);
Doc doc1 = test::TestHelper::CreateDoc(0, *schema_);
auto status1 = segment->Insert(doc1);
EXPECT_TRUE(status1.ok()) << "First insert failed: " << status1.message();
auto meta = segment->meta();
ASSERT_TRUE(meta != nullptr);
auto &mem_block = meta->writing_forward_block().value();
EXPECT_EQ(mem_block.doc_count_, 1);
EXPECT_EQ(mem_block.min_doc_id_, 0);
EXPECT_EQ(mem_block.max_doc_id_, 0);
auto doc = segment->Fetch(0);
EXPECT_TRUE(doc != nullptr);
EXPECT_EQ(*doc, doc1);
auto status2 = segment->Insert(doc1);
EXPECT_FALSE(status2.ok()) << "Duplicate insert should fail";
auto fetched_doc = segment->Fetch(0);
ASSERT_TRUE(fetched_doc != nullptr);
EXPECT_NE(fetched_doc->get<std::string>("name").value(), "duplicate_name");
}
TEST_P(SegmentTest, DuplicateDelete) {
auto segment = test::TestHelper::CreateSegmentWithDoc(
col_path_, *schema_, 0, 0, id_map_, delete_store_, version_manager_, options_,
0, 5);
ASSERT_TRUE(segment != nullptr);
auto status1 = segment->Delete("pk_2");
EXPECT_TRUE(status1.ok()) << "First delete failed: " << status1.message();
auto status2 = segment->Delete("pk_2");
EXPECT_FALSE(status2.ok()) << "Duplicate delete should fail";
auto status3 = segment->Delete(2);
EXPECT_FALSE(status3.ok())
<< "Delete by doc_id of already deleted doc should fail";
}
TEST_P(SegmentTest, DeleteNonExistentDoc) {
auto segment = test::TestHelper::CreateSegmentWithDoc(
col_path_, *schema_, 0, 0, id_map_, delete_store_, version_manager_, options_,
0, 5);
ASSERT_TRUE(segment != nullptr);
auto status1 = segment->Delete("pk_999");
EXPECT_FALSE(status1.ok()) << "Delete non-existent pk should fail";
}
TEST_P(SegmentTest, UpdateNonExistentDoc) {
auto segment = test::TestHelper::CreateSegmentWithDoc(
col_path_, *schema_, 0, 0, id_map_, delete_store_, version_manager_, options_,
0, 5);
ASSERT_TRUE(segment != nullptr);
Doc doc = test::TestHelper::CreateDoc(999, *schema_);
doc.set<std::string>("name", "non_existent_doc");
auto status = segment->Update(doc);
EXPECT_FALSE(status.ok()) << "Update non-existent doc should fail";
}
TEST_P(SegmentTest, UpsertNonExistentDoc) {
auto segment = test::TestHelper::CreateSegmentWithDoc(
col_path_, *schema_, 0, 0, id_map_, delete_store_, version_manager_, options_,
0, 5);
ASSERT_TRUE(segment != nullptr);
Doc doc = test::TestHelper::CreateDoc(999, *schema_);
doc.set<std::string>("name", "new_upserted_doc");
auto status = segment->Upsert(doc);
EXPECT_TRUE(status.ok()) << "Upsert non-existent doc should succeed: "
<< status.message();
auto filter = delete_store_->make_filter();
uint64_t count = segment->doc_count(filter);
EXPECT_EQ(count, 6);
}
TEST_P(SegmentTest, ScanWithEmptyColumns) {
auto segment = test::TestHelper::CreateSegmentWithDoc(
col_path_, *schema_, 0, 0, id_map_, delete_store_, version_manager_, options_,
0, 5);
ASSERT_TRUE(segment != nullptr);
auto reader = segment->scan({});
ASSERT_TRUE(reader == nullptr);
}
TEST_P(SegmentTest, ScanWithInvalidColumns) {
auto segment = test::TestHelper::CreateSegmentWithDoc(
col_path_, *schema_, 0, 0, id_map_, delete_store_, version_manager_, options_,
0, 10);
ASSERT_TRUE(segment != nullptr);
// Try to scan with invalid column name
auto reader = segment->scan({"invalid_column"});
EXPECT_TRUE(reader == nullptr);
}
TEST_P(SegmentTest, FetchNonExistentDoc) {
auto segment = test::TestHelper::CreateSegmentWithDoc(
col_path_, *schema_, 0, 0, id_map_, delete_store_, version_manager_, options_,
0, 5);
ASSERT_TRUE(segment != nullptr);
auto doc = segment->Fetch(999);
EXPECT_TRUE(doc == nullptr) << "Fetch non-existent doc should return nullptr";
}
TEST_P(SegmentTest, FetchWithInvalidSegmentDocIDs) {
auto segment = test::TestHelper::CreateSegmentWithDoc(
col_path_, *schema_, 0, 0, id_map_, delete_store_, version_manager_, options_,
0, 5);
ASSERT_TRUE(segment != nullptr);
std::vector<int> invalid_segment_doc_ids = {999, 1000};
auto table = segment->fetch({"id", "name"}, invalid_segment_doc_ids);
ASSERT_TRUE(table == nullptr);
}
TEST_P(SegmentTest, FetchWithInvalidColumns) {
auto segment = test::TestHelper::CreateSegmentWithDoc(
col_path_, *schema_, 0, 0, id_map_, delete_store_, version_manager_, options_,
0, 10);
ASSERT_TRUE(segment != nullptr);
// Try to fetch with invalid column name
std::vector<int> segment_doc_ids = {0, 1, 2};
auto table = segment->fetch({"invalid_column"}, segment_doc_ids);
EXPECT_TRUE(table == nullptr);
}
TEST_P(SegmentTest, InsertEmptyDocWithNullableSchema) {
auto nullable_schema = test::TestHelper::CreateNormalSchema(true, col_name_);
auto segment = test::TestHelper::CreateSegmentWithDoc(
col_path_, *nullable_schema, 0, 0, id_map_, delete_store_, version_manager_,
options_, 0, 0);
ASSERT_TRUE(segment != nullptr);
Doc empty_doc;
empty_doc.set_pk("pk_empty");
auto status = segment->Insert(empty_doc);
EXPECT_TRUE(status.ok());
}
TEST_P(SegmentTest, MultipleDuplicateDeletes) {
auto segment = test::TestHelper::CreateSegmentWithDoc(
col_path_, *schema_, 0, 0, id_map_, delete_store_, version_manager_, options_,
0, 5);
ASSERT_TRUE(segment != nullptr);
auto status1 = segment->Delete("pk_1");
EXPECT_TRUE(status1.ok());
for (int i = 0; i < 10; ++i) {
auto status = segment->Delete("pk_1");
EXPECT_FALSE(status.ok()) << "Delete iteration " << i << " should fail";
}
auto filter = delete_store_->make_filter();
uint64_t count = segment->doc_count(filter);
EXPECT_EQ(count, 4);
}
TEST_P(SegmentTest, FetchWithTwoVectorFields) {
schema_->add_field(std::make_shared<FieldSchema>(
"dense2_fp32", DataType::VECTOR_FP32, 128, false,
std::make_shared<FlatIndexParams>(MetricType::IP)));
int doc_count = 1000;
auto segment = test::TestHelper::CreateSegmentWithDoc(
col_path_, *schema_, 0, 0, id_map_, delete_store_, version_manager_, options_,
0, doc_count);
ASSERT_TRUE(segment != nullptr);
segment.reset();
version_manager_.reset();
id_map_->flush();
id_map_.reset();
std::string delete_store_path =
FileHelper::MakeFilePath(col_path_, FileID::DELETE_FILE, 0);
delete_store_->flush(delete_store_path);
delete_store_.reset();
auto recover_version_manager = VersionManager::Recovery(col_path_);
auto recover_version_mgr = recover_version_manager.value();
ASSERT_TRUE(recover_version_mgr != nullptr);
auto v = recover_version_mgr->get_current_version();
// idmap
std::string idmap_path = FileHelper::MakeFilePath(col_path_, FileID::ID_FILE,
v.id_map_path_suffix());
IDMap::Ptr recover_id_map = std::make_shared<IDMap>(col_name_);
auto status = recover_id_map->open(idmap_path, false, false);
ASSERT_TRUE(status.ok());
delete_store_path = FileHelper::MakeFilePath(col_path_, FileID::DELETE_FILE,
v.delete_snapshot_path_suffix());
auto recover_delete_store =
DeleteStore::CreateAndLoad(col_name_, delete_store_path);
ASSERT_TRUE(recover_delete_store != nullptr);
int incr_doc_count = 1000;
auto result = Segment::Open(col_path_, *schema_, *v.writing_segment_meta(),
recover_id_map, recover_delete_store,
recover_version_mgr, options_);
ASSERT_TRUE(result.has_value());
segment = std::move(result).value();
ASSERT_TRUE(segment != nullptr);
auto s = test::TestHelper::SegmentInsertDoc(
segment, *schema_, doc_count, doc_count + incr_doc_count, false);
ASSERT_TRUE(s.ok());
for (int i = 0; i < doc_count + incr_doc_count; i++) {
auto expect_doc = test::TestHelper::CreateDoc(i, *schema_);
auto ret_doc = segment->Fetch(i);
if (*ret_doc != expect_doc) {
std::cout << " ret_doc: " << ret_doc->to_string() << std::endl;
std::cout << "expect_doc: " << expect_doc.to_string() << std::endl;
}
ASSERT_EQ(*ret_doc, expect_doc);
}
}
TEST_P(SegmentTest, FetchPerf) {
// create segment
int doc_count = 1000;
options_.max_buffer_size_ = 100 * 1024;
auto segment = test::TestHelper::CreateSegmentWithDoc(
col_path_, *schema_, 0, 0, id_map_, delete_store_, version_manager_, options_,
0, doc_count);
ASSERT_TRUE(segment != nullptr);
segment->dump();
auto writing_segment_meta = segment->meta();
// convert writing segment meta to persisted segment meta
Version version = version_manager_->get_current_version();
writing_segment_meta->remove_writing_forward_block();
auto s = version.add_persisted_segment_meta(writing_segment_meta);
ASSERT_TRUE(s.ok());
s = version_manager_->apply(version);
ASSERT_TRUE(s.ok());
s = version_manager_->flush();
ASSERT_TRUE(s.ok());
segment.reset();
version_manager_.reset();
id_map_->flush();
id_map_.reset();
std::string delete_store_path =
FileHelper::MakeFilePath(col_path_, FileID::DELETE_FILE, 0);
delete_store_->flush(delete_store_path);
delete_store_.reset();
auto recover_version_manager = VersionManager::Recovery(col_path_);
auto recover_version_mgr = recover_version_manager.value();
ASSERT_TRUE(recover_version_mgr != nullptr);
Version v = recover_version_mgr->get_current_version();
const auto &persist_metas = v.persisted_segment_metas();
// idmap
std::string idmap_path = FileHelper::MakeFilePath(col_path_, FileID::ID_FILE,
v.id_map_path_suffix());
IDMap::Ptr recover_id_map = std::make_shared<IDMap>(col_name_);
auto status = recover_id_map->open(idmap_path, false, false);
ASSERT_TRUE(status.ok());
delete_store_path = FileHelper::MakeFilePath(col_path_, FileID::DELETE_FILE,
v.delete_snapshot_path_suffix());
auto recover_delete_store =
DeleteStore::CreateAndLoad(col_name_, delete_store_path);
ASSERT_TRUE(recover_delete_store != nullptr);
// open persist segment
options_.read_only_ = true;
auto result =
Segment::Open(col_path_, *schema_, *persist_metas[0], recover_id_map,
recover_delete_store, recover_version_mgr, options_);
ASSERT_TRUE(result.has_value());
segment = std::move(result).value();
ASSERT_TRUE(segment != nullptr);
s = segment->add_column(
std::make_shared<FieldSchema>("add_int32", DataType::INT32, false),
"int32 + 1", AddColumnOptions());
EXPECT_TRUE(s.ok());
std::vector<int> segment_doc_ids = {0, 3, 6, 1, 0, 501, 999};
auto func = [&](const std::vector<std::string> columns,
int local_row_id_idx) -> void {
auto combined_table = segment->fetch(columns, segment_doc_ids);
ASSERT_TRUE(combined_table != nullptr);
EXPECT_EQ(combined_table->num_columns(), columns.size());
EXPECT_EQ(combined_table->num_rows(), segment_doc_ids.size());
auto field = combined_table->schema()->field(local_row_id_idx);
EXPECT_EQ(field->name(), LOCAL_ROW_ID);
// Get data from the LOCAL_ROW_ID column for each row
auto id_column = combined_table->column(local_row_id_idx);
auto id_array =
std::dynamic_pointer_cast<arrow::UInt64Array>(id_column->chunk(0));
std::vector<int32_t> &expected_ids = segment_doc_ids;
std::vector<int32_t> actual_ids;
for (int i = 0; i < id_array->length(); ++i) {
actual_ids.push_back(id_array->Value(i));
}
EXPECT_EQ(actual_ids, expected_ids)
<< "ID column values don't match expected order";
};
func({LOCAL_ROW_ID, "id", "name", "add_int32"}, 0);
func(
{
"id",
LOCAL_ROW_ID,
"name",
"add_int32",
},
1);
func({"id", "name", "add_int32", LOCAL_ROW_ID}, 3);
}
TEST_P(SegmentTest, AddColumn) {
// create segment
options_.max_buffer_size_ = 10 * 1024 * 1024;
int doc_count = 1000;
auto segment = test::TestHelper::CreateSegmentWithDoc(
col_path_, *schema_, 0, 0, id_map_, delete_store_, version_manager_, options_,
0, doc_count);
ASSERT_TRUE(segment != nullptr);
auto s = segment->add_column(
std::make_shared<FieldSchema>("add_int32", DataType::INT32, false),
"int32 + 1", AddColumnOptions());
EXPECT_FALSE(s.ok());
segment->dump();
auto writing_segment_meta = segment->meta();
// convert writing segment meta to persisted segment meta
Version version = version_manager_->get_current_version();
writing_segment_meta->remove_writing_forward_block();
s = version.add_persisted_segment_meta(writing_segment_meta);
ASSERT_TRUE(s.ok());
s = version_manager_->apply(version);
ASSERT_TRUE(s.ok());
s = version_manager_->flush();
ASSERT_TRUE(s.ok());
segment.reset();
version_manager_.reset();
id_map_->flush();
id_map_.reset();
std::string delete_store_path =
FileHelper::MakeFilePath(col_path_, FileID::DELETE_FILE, 0);
delete_store_->flush(delete_store_path);
delete_store_.reset();
auto recover_version_manager = VersionManager::Recovery(col_path_);
auto recover_version_mgr = recover_version_manager.value();
ASSERT_TRUE(recover_version_mgr != nullptr);
Version v = recover_version_mgr->get_current_version();
const auto &persist_metas = v.persisted_segment_metas();
// idmap
std::string idmap_path = FileHelper::MakeFilePath(col_path_, FileID::ID_FILE,
v.id_map_path_suffix());
IDMap::Ptr recover_id_map = std::make_shared<IDMap>(col_name_);
auto status = recover_id_map->open(idmap_path, false, false);
ASSERT_TRUE(status.ok());
delete_store_path = FileHelper::MakeFilePath(col_path_, FileID::DELETE_FILE,
v.delete_snapshot_path_suffix());
auto recover_delete_store =
DeleteStore::CreateAndLoad(col_name_, delete_store_path);
ASSERT_TRUE(recover_delete_store != nullptr);
// open persist segment
options_.read_only_ = true;
auto result =
Segment::Open(col_path_, *schema_, *persist_metas[0], recover_id_map,
recover_delete_store, recover_version_mgr, options_);
ASSERT_TRUE(result.has_value());
segment = std::move(result).value();
ASSERT_TRUE(segment != nullptr);
s = segment->add_column(
std::make_shared<FieldSchema>("add_int32", DataType::INT32, false), "",
AddColumnOptions());
EXPECT_FALSE(s.ok());
s = segment->add_column(std::make_shared<FieldSchema>(
"add_undefined", DataType::UNDEFINED, false),
"", AddColumnOptions());
EXPECT_FALSE(s.ok());
// before add column
auto meta = segment->meta();
auto &persist_blocks = meta->persisted_blocks();
int old_scalar_blocks_cnt = 0;
for (auto &block : persist_blocks) {
if (block.type() == BlockType::SCALAR) {
old_scalar_blocks_cnt++;
}
}
int add_column_cnt = 0;
auto func = [&](const std::shared_ptr<FieldSchema> &field_schema,
const std::string &expression) {
auto &column_name = field_schema->name();
AddColumnOptions add_options;
status = segment->add_column(field_schema, expression, add_options);
EXPECT_TRUE(status.ok());
// after add column
int new_scalar_blocks_cnt = 0;
for (auto &block : persist_blocks) {
if (block.type() == BlockType::SCALAR) {
new_scalar_blocks_cnt++;
}
}
EXPECT_EQ(
new_scalar_blocks_cnt,
old_scalar_blocks_cnt + old_scalar_blocks_cnt * (++add_column_cnt));
auto combined_reader = segment->scan({"id", "name", "age", column_name});
ASSERT_TRUE(combined_reader != nullptr);
std::shared_ptr<arrow::RecordBatch> batch;
uint32_t total_doc = 0;
while (true) {
auto status = combined_reader->ReadNext(&batch);
if (status.ok() == false) break;
if (batch == nullptr) break;
EXPECT_EQ(batch->num_columns(), 4);
total_doc += batch->num_rows();
}
EXPECT_EQ(total_doc, doc_count);
auto new_schema = *schema_;
new_schema.add_field(field_schema);
auto check_doc = [&](int doc_count) {
for (int i = 0; i < doc_count; i++) {
auto expect_doc = test::TestHelper::CreateDoc(i, new_schema);
auto doc = segment->Fetch(i);
ASSERT_EQ(doc->pk(), expect_doc.pk());
// column in same persist block
{
ExecBatchPtr exec_batch = segment->fetch({"id", "name", "age"}, i);
ASSERT_TRUE(exec_batch != nullptr);
EXPECT_EQ(exec_batch->length, 1);
EXPECT_EQ(exec_batch->values.size(), 3);
auto id_scalar = exec_batch->values[0].scalar();
ASSERT_TRUE(id_scalar != nullptr);
auto id_value =
std::dynamic_pointer_cast<arrow::Int32Scalar>(id_scalar);
ASSERT_TRUE(id_value != nullptr);
EXPECT_EQ(id_value->value, i);
}
{
ExecBatchPtr exec_batch = segment->fetch({column_name}, i);
ASSERT_TRUE(exec_batch != nullptr);
EXPECT_EQ(exec_batch->length, 1);
EXPECT_EQ(exec_batch->values.size(), 1);
auto id_scalar = exec_batch->values[0].scalar();
ASSERT_TRUE(id_scalar != nullptr);
}
// column in different persist block
{
ExecBatchPtr exec_batch =
segment->fetch({"id", "name", "age", column_name}, i);
ASSERT_TRUE(exec_batch != nullptr);
EXPECT_EQ(exec_batch->length, 1);
EXPECT_EQ(exec_batch->values.size(), 4);
auto id_scalar = exec_batch->values[0].scalar();
ASSERT_TRUE(id_scalar != nullptr);
auto id_value =
std::dynamic_pointer_cast<arrow::Int32Scalar>(id_scalar);
ASSERT_TRUE(id_value != nullptr);
EXPECT_EQ(id_value->value, i);
}
}
};
check_doc(doc_count);
};
auto index_param = std::make_shared<InvertIndexParams>();
std::vector<std::pair<std::string, std::shared_ptr<FieldSchema>>>
test_column_schemas = {
{"add_int32", std::make_shared<FieldSchema>("", DataType::INT32,
false, index_param)},
{"add_int64", std::make_shared<FieldSchema>("", DataType::INT64,
false, index_param)},
{"add_uint32", std::make_shared<FieldSchema>("", DataType::UINT32,
false, index_param)},
{"add_uint64", std::make_shared<FieldSchema>("", DataType::UINT64,
false, index_param)},
{"add_float", std::make_shared<FieldSchema>("", DataType::FLOAT,
false, index_param)},
{"add_double", std::make_shared<FieldSchema>("", DataType::DOUBLE,
false, index_param)},
{"add_int32_nullable", std::make_shared<FieldSchema>(
"", DataType::INT32, true, index_param)},
{"add_int64_nullable", std::make_shared<FieldSchema>(
"", DataType::INT64, true, index_param)},
{"add_uint32_nullable", std::make_shared<FieldSchema>(
"", DataType::UINT32, true, index_param)},
{"add_uint64_nullable", std::make_shared<FieldSchema>(
"", DataType::UINT64, true, index_param)},
{"add_float_nullable", std::make_shared<FieldSchema>(
"", DataType::FLOAT, true, index_param)},
{"add_double_nullable", std::make_shared<FieldSchema>(
"", DataType::DOUBLE, true, index_param)},
};
std::unordered_map<std::string, std::vector<std::string>> test_expressions = {
{"add_int32", {"int32 + 1", "-int32", "+int32", "1", "-1"}},
{"add_int64", {"int64 + 1", "-int64", "+int64", "1", "-1"}},
{"add_uint32", {"uint32 + 1", "-uint32", "+int32", "1", "0"}},
{"add_uint64", {"uint64 + 1", "-uint64", "+uint64", "1", "0"}},
{"add_float", {"float + 1.0", "-float", "+float", "0.1", "-0.1"}},
{"add_double", {"double + 1.0", "-double", "+double", "0.1", "-0.1"}},
{"add_int32_nullable", {""}},
{"add_int64_nullable", {""}},
{"add_uint32_nullable", {""}},
{"add_uint64_nullable", {""}},
{"add_float_nullable", {""}},
{"add_double_nullable", {""}},
};
for (auto &[column_name, field_schema] : test_column_schemas) {
auto expressions = test_expressions[column_name];
for (auto &expression : expressions) {
std::string generated_col_name =
column_name + "_" +
std::to_string(
ailego::Crc32c::Hash(expression.data(), expression.size()));
auto new_field_schema = std::make_shared<FieldSchema>(
field_schema->name(), field_schema->data_type(),
field_schema->nullable(), field_schema->index_params());
new_field_schema->set_name(generated_col_name);
func(new_field_schema, expression);
}
}
}
TEST_P(SegmentTest, AddNullableColumnWithoutExpressionMultiBlock) {
// Use small buffer to force multiple scalar blocks within a single segment
options_.max_buffer_size_ = 1 * 1024;
int doc_count = 100;
auto segment = test::TestHelper::CreateSegmentWithDoc(
col_path_, *schema_, 0, 0, id_map_, delete_store_, version_manager_, options_,
0, doc_count);
ASSERT_TRUE(segment != nullptr);
segment->dump();
auto writing_segment_meta = segment->meta();
Version version = version_manager_->get_current_version();
writing_segment_meta->remove_writing_forward_block();
auto s = version.add_persisted_segment_meta(writing_segment_meta);
ASSERT_TRUE(s.ok());
s = version_manager_->apply(version);
ASSERT_TRUE(s.ok());
s = version_manager_->flush();
ASSERT_TRUE(s.ok());
segment.reset();
version_manager_.reset();
id_map_->flush();
id_map_.reset();
std::string delete_store_path =
FileHelper::MakeFilePath(col_path_, FileID::DELETE_FILE, 0);
delete_store_->flush(delete_store_path);
delete_store_.reset();
auto recover_version_manager = VersionManager::Recovery(col_path_);
auto recover_version_mgr = recover_version_manager.value();
ASSERT_TRUE(recover_version_mgr != nullptr);
Version v = recover_version_mgr->get_current_version();
const auto &persist_metas = v.persisted_segment_metas();
std::string idmap_path = FileHelper::MakeFilePath(col_path_, FileID::ID_FILE,
v.id_map_path_suffix());
IDMap::Ptr recover_id_map = std::make_shared<IDMap>(col_name_);
auto status = recover_id_map->open(idmap_path, false, false);
ASSERT_TRUE(status.ok());
delete_store_path = FileHelper::MakeFilePath(col_path_, FileID::DELETE_FILE,
v.delete_snapshot_path_suffix());
auto recover_delete_store =
DeleteStore::CreateAndLoad(col_name_, delete_store_path);
ASSERT_TRUE(recover_delete_store != nullptr);
// Verify we have multiple scalar blocks
int scalar_block_count = 0;
for (auto &block : persist_metas[0]->persisted_blocks()) {
if (block.type() == BlockType::SCALAR) {
scalar_block_count++;
}
}
ASSERT_GT(scalar_block_count, 1);
options_.read_only_ = true;
auto result =
Segment::Open(col_path_, *schema_, *persist_metas[0], recover_id_map,
recover_delete_store, recover_version_mgr, options_);
ASSERT_TRUE(result.has_value());
segment = std::move(result).value();
ASSERT_TRUE(segment != nullptr);
// Add nullable columns without expression — this used to crash on
// multi-block segments
std::vector<std::pair<std::string, DataType>> nullable_types = {
{"add_int32_null", DataType::INT32},
{"add_uint32_null", DataType::UINT32},
{"add_int64_null", DataType::INT64},
};
for (auto &[nullable_col_name, data_type] : nullable_types) {
auto field_schema =
std::make_shared<FieldSchema>(nullable_col_name, data_type, true);
s = segment->add_column(field_schema, "", AddColumnOptions());
ASSERT_TRUE(s.ok())
<< "Failed to add nullable column " << nullable_col_name << ": "
<< s.message();
auto combined_reader =
segment->scan({"id", "name", "age", nullable_col_name});
ASSERT_TRUE(combined_reader != nullptr);
std::shared_ptr<arrow::RecordBatch> batch;
uint32_t total_doc = 0;
while (true) {
auto st = combined_reader->ReadNext(&batch);
if (!st.ok()) break;
if (batch == nullptr) break;
EXPECT_EQ(batch->num_columns(), 4);
total_doc += batch->num_rows();
}
EXPECT_EQ(total_doc, doc_count);
}
}
TEST_P(SegmentTest, AddColumnWithExpressionMultiBlock) {
// Use small buffer to force multiple scalar blocks within a single segment
options_.max_buffer_size_ = 1 * 1024;
int doc_count = 100;
auto segment = test::TestHelper::CreateSegmentWithDoc(
col_path_, *schema_, 0, 0, id_map_, delete_store_, version_manager_, options_,
0, doc_count);
ASSERT_TRUE(segment != nullptr);
segment->dump();
auto writing_segment_meta = segment->meta();
Version version = version_manager_->get_current_version();
writing_segment_meta->remove_writing_forward_block();
auto s = version.add_persisted_segment_meta(writing_segment_meta);
ASSERT_TRUE(s.ok());
s = version_manager_->apply(version);
ASSERT_TRUE(s.ok());
s = version_manager_->flush();
ASSERT_TRUE(s.ok());
segment.reset();
version_manager_.reset();
id_map_->flush();
id_map_.reset();
std::string delete_store_path =
FileHelper::MakeFilePath(col_path_, FileID::DELETE_FILE, 0);
delete_store_->flush(delete_store_path);
delete_store_.reset();
auto recover_version_manager = VersionManager::Recovery(col_path_);
auto recover_version_mgr = recover_version_manager.value();
ASSERT_TRUE(recover_version_mgr != nullptr);
Version v = recover_version_mgr->get_current_version();
const auto &persist_metas = v.persisted_segment_metas();
std::string idmap_path = FileHelper::MakeFilePath(col_path_, FileID::ID_FILE,
v.id_map_path_suffix());
IDMap::Ptr recover_id_map = std::make_shared<IDMap>(col_name_);
auto status = recover_id_map->open(idmap_path, false, false);
ASSERT_TRUE(status.ok());
delete_store_path = FileHelper::MakeFilePath(col_path_, FileID::DELETE_FILE,
v.delete_snapshot_path_suffix());
auto recover_delete_store =
DeleteStore::CreateAndLoad(col_name_, delete_store_path);
ASSERT_TRUE(recover_delete_store != nullptr);
// Verify we have multiple scalar blocks
int scalar_block_count = 0;
for (auto &block : persist_metas[0]->persisted_blocks()) {
if (block.type() == BlockType::SCALAR) {
scalar_block_count++;
}
}
ASSERT_GT(scalar_block_count, 1);
options_.read_only_ = true;
auto result =
Segment::Open(col_path_, *schema_, *persist_metas[0], recover_id_map,
recover_delete_store, recover_version_mgr, options_);
ASSERT_TRUE(result.has_value());
segment = std::move(result).value();
ASSERT_TRUE(segment != nullptr);
// Add column with expression on multi-block segment
auto field_schema =
std::make_shared<FieldSchema>("add_expr_col", DataType::INT32, false);
s = segment->add_column(field_schema, "id + 1", AddColumnOptions());
ASSERT_TRUE(s.ok()) << "Failed to add column with expression: "
<< s.message();
auto combined_reader = segment->scan({"id", "name", "age", "add_expr_col"});
ASSERT_TRUE(combined_reader != nullptr);
std::shared_ptr<arrow::RecordBatch> batch;
uint32_t total_doc = 0;
while (true) {
auto st = combined_reader->ReadNext(&batch);
if (!st.ok()) break;
if (batch == nullptr) break;
EXPECT_EQ(batch->num_columns(), 4);
total_doc += batch->num_rows();
}
EXPECT_EQ(total_doc, doc_count);
// Add nullable column without expression on the same multi-block segment
auto nullable_field =
std::make_shared<FieldSchema>("add_null_col", DataType::INT64, true);
s = segment->add_column(nullable_field, "", AddColumnOptions());
ASSERT_TRUE(s.ok())
<< "Failed to add nullable column after expression column: "
<< s.message();
combined_reader = segment->scan({"id", "add_expr_col", "add_null_col"});
ASSERT_TRUE(combined_reader != nullptr);
total_doc = 0;
while (true) {
auto st = combined_reader->ReadNext(&batch);
if (!st.ok()) break;
if (batch == nullptr) break;
EXPECT_EQ(batch->num_columns(), 3);
total_doc += batch->num_rows();
}
EXPECT_EQ(total_doc, doc_count);
}
TEST_P(SegmentTest, AlterColumnMultiBlock) {
options_.max_buffer_size_ = 1 * 1024;
int doc_count = 100;
auto segment = test::TestHelper::CreateSegmentWithDoc(
col_path_, *schema_, 0, 0, id_map_, delete_store_, version_manager_, options_,
0, doc_count);
ASSERT_TRUE(segment != nullptr);
segment->dump();
auto writing_segment_meta = segment->meta();
Version version = version_manager_->get_current_version();
writing_segment_meta->remove_writing_forward_block();
auto s = version.add_persisted_segment_meta(writing_segment_meta);
ASSERT_TRUE(s.ok());
s = version_manager_->apply(version);
ASSERT_TRUE(s.ok());
s = version_manager_->flush();
ASSERT_TRUE(s.ok());
segment.reset();
version_manager_.reset();
id_map_->flush();
id_map_.reset();
std::string delete_store_path =
FileHelper::MakeFilePath(col_path_, FileID::DELETE_FILE, 0);
delete_store_->flush(delete_store_path);
delete_store_.reset();
auto recover_version_manager = VersionManager::Recovery(col_path_);
auto recover_version_mgr = recover_version_manager.value();
ASSERT_TRUE(recover_version_mgr != nullptr);
Version v = recover_version_mgr->get_current_version();
const auto &persist_metas = v.persisted_segment_metas();
std::string idmap_path = FileHelper::MakeFilePath(col_path_, FileID::ID_FILE,
v.id_map_path_suffix());
IDMap::Ptr recover_id_map = std::make_shared<IDMap>(col_name_);
auto status = recover_id_map->open(idmap_path, false, false);
ASSERT_TRUE(status.ok());
delete_store_path = FileHelper::MakeFilePath(col_path_, FileID::DELETE_FILE,
v.delete_snapshot_path_suffix());
auto recover_delete_store =
DeleteStore::CreateAndLoad(col_name_, delete_store_path);
ASSERT_TRUE(recover_delete_store != nullptr);
int scalar_block_count = 0;
for (auto &block : persist_metas[0]->persisted_blocks()) {
if (block.type() == BlockType::SCALAR) {
scalar_block_count++;
}
}
ASSERT_GT(scalar_block_count, 1);
options_.read_only_ = true;
auto result =
Segment::Open(col_path_, *schema_, *persist_metas[0], recover_id_map,
recover_delete_store, recover_version_mgr, options_);
ASSERT_TRUE(result.has_value());
segment = std::move(result).value();
ASSERT_TRUE(segment != nullptr);
// Alter column type: int32 -> int64 on multi-block segment
auto new_field = std::make_shared<FieldSchema>("id", DataType::INT64, false);
s = segment->alter_column("id", new_field, AlterColumnOptions());
ASSERT_TRUE(s.ok()) << "alter_column failed: " << s.message();
auto combined_reader = segment->scan({"id", "name", "age"});
ASSERT_TRUE(combined_reader != nullptr);
std::shared_ptr<arrow::RecordBatch> batch;
uint32_t total_doc = 0;
while (true) {
auto st = combined_reader->ReadNext(&batch);
if (!st.ok()) break;
if (batch == nullptr) break;
EXPECT_EQ(batch->num_columns(), 3);
total_doc += batch->num_rows();
}
EXPECT_EQ(total_doc, doc_count);
}
TEST_P(SegmentTest, DropColumnMultiBlock) {
options_.max_buffer_size_ = 1 * 1024;
int doc_count = 100;
auto segment = test::TestHelper::CreateSegmentWithDoc(
col_path_, *schema_, 0, 0, id_map_, delete_store_, version_manager_, options_,
0, doc_count);
ASSERT_TRUE(segment != nullptr);
segment->dump();
auto writing_segment_meta = segment->meta();
Version version = version_manager_->get_current_version();
writing_segment_meta->remove_writing_forward_block();
auto s = version.add_persisted_segment_meta(writing_segment_meta);
ASSERT_TRUE(s.ok());
s = version_manager_->apply(version);
ASSERT_TRUE(s.ok());
s = version_manager_->flush();
ASSERT_TRUE(s.ok());
segment.reset();
version_manager_.reset();
id_map_->flush();
id_map_.reset();
std::string delete_store_path =
FileHelper::MakeFilePath(col_path_, FileID::DELETE_FILE, 0);
delete_store_->flush(delete_store_path);
delete_store_.reset();
auto recover_version_manager = VersionManager::Recovery(col_path_);
auto recover_version_mgr = recover_version_manager.value();
ASSERT_TRUE(recover_version_mgr != nullptr);
Version v = recover_version_mgr->get_current_version();
const auto &persist_metas = v.persisted_segment_metas();
std::string idmap_path = FileHelper::MakeFilePath(col_path_, FileID::ID_FILE,
v.id_map_path_suffix());
IDMap::Ptr recover_id_map = std::make_shared<IDMap>(col_name_);
auto status = recover_id_map->open(idmap_path, false, false);
ASSERT_TRUE(status.ok());
delete_store_path = FileHelper::MakeFilePath(col_path_, FileID::DELETE_FILE,
v.delete_snapshot_path_suffix());
auto recover_delete_store =
DeleteStore::CreateAndLoad(col_name_, delete_store_path);
ASSERT_TRUE(recover_delete_store != nullptr);
int scalar_block_count = 0;
for (auto &block : persist_metas[0]->persisted_blocks()) {
if (block.type() == BlockType::SCALAR) {
scalar_block_count++;
}
}
ASSERT_GT(scalar_block_count, 1);
options_.read_only_ = true;
auto result =
Segment::Open(col_path_, *schema_, *persist_metas[0], recover_id_map,
recover_delete_store, recover_version_mgr, options_);
ASSERT_TRUE(result.has_value());
segment = std::move(result).value();
ASSERT_TRUE(segment != nullptr);
// Drop column on multi-block segment
s = segment->drop_column("id");
ASSERT_TRUE(s.ok()) << "drop_column failed: " << s.message();
auto combined_reader = segment->scan({"id"});
ASSERT_TRUE(combined_reader == nullptr);
// Remaining columns should still be scannable
combined_reader = segment->scan({"name", "age"});
ASSERT_TRUE(combined_reader != nullptr);
std::shared_ptr<arrow::RecordBatch> batch;
uint32_t total_doc = 0;
while (true) {
auto st = combined_reader->ReadNext(&batch);
if (!st.ok()) break;
if (batch == nullptr) break;
EXPECT_EQ(batch->num_columns(), 2);
total_doc += batch->num_rows();
}
EXPECT_EQ(total_doc, doc_count);
}
TEST_P(SegmentTest, AddNullableThenAlterDropMultiBlock) {
options_.max_buffer_size_ = 1 * 1024;
int doc_count = 100;
auto segment = test::TestHelper::CreateSegmentWithDoc(
col_path_, *schema_, 0, 0, id_map_, delete_store_, version_manager_, options_,
0, doc_count);
ASSERT_TRUE(segment != nullptr);
segment->dump();
auto writing_segment_meta = segment->meta();
Version version = version_manager_->get_current_version();
writing_segment_meta->remove_writing_forward_block();
auto s = version.add_persisted_segment_meta(writing_segment_meta);
ASSERT_TRUE(s.ok());
s = version_manager_->apply(version);
ASSERT_TRUE(s.ok());
s = version_manager_->flush();
ASSERT_TRUE(s.ok());
segment.reset();
version_manager_.reset();
id_map_->flush();
id_map_.reset();
std::string delete_store_path =
FileHelper::MakeFilePath(col_path_, FileID::DELETE_FILE, 0);
delete_store_->flush(delete_store_path);
delete_store_.reset();
auto recover_version_manager = VersionManager::Recovery(col_path_);
auto recover_version_mgr = recover_version_manager.value();
ASSERT_TRUE(recover_version_mgr != nullptr);
Version v = recover_version_mgr->get_current_version();
const auto &persist_metas = v.persisted_segment_metas();
std::string idmap_path = FileHelper::MakeFilePath(col_path_, FileID::ID_FILE,
v.id_map_path_suffix());
IDMap::Ptr recover_id_map = std::make_shared<IDMap>(col_name_);
auto status = recover_id_map->open(idmap_path, false, false);
ASSERT_TRUE(status.ok());
delete_store_path = FileHelper::MakeFilePath(col_path_, FileID::DELETE_FILE,
v.delete_snapshot_path_suffix());
auto recover_delete_store =
DeleteStore::CreateAndLoad(col_name_, delete_store_path);
ASSERT_TRUE(recover_delete_store != nullptr);
options_.read_only_ = true;
auto result =
Segment::Open(col_path_, *schema_, *persist_metas[0], recover_id_map,
recover_delete_store, recover_version_mgr, options_);
ASSERT_TRUE(result.has_value());
segment = std::move(result).value();
ASSERT_TRUE(segment != nullptr);
// 1) Add nullable column without expression
auto nullable_field =
std::make_shared<FieldSchema>("nullable_col", DataType::INT32, true);
s = segment->add_column(nullable_field, "", AddColumnOptions());
ASSERT_TRUE(s.ok()) << "add nullable column failed: " << s.message();
// 2) Alter the nullable column type: INT32 -> INT64
auto altered_field =
std::make_shared<FieldSchema>("nullable_col", DataType::INT64, true);
s = segment->alter_column("nullable_col", altered_field,
AlterColumnOptions());
ASSERT_TRUE(s.ok()) << "alter nullable column failed: " << s.message();
auto combined_reader = segment->scan({"id", "nullable_col"});
ASSERT_TRUE(combined_reader != nullptr);
std::shared_ptr<arrow::RecordBatch> batch;
uint32_t total_doc = 0;
while (true) {
auto st = combined_reader->ReadNext(&batch);
if (!st.ok()) break;
if (batch == nullptr) break;
EXPECT_EQ(batch->num_columns(), 2);
total_doc += batch->num_rows();
}
EXPECT_EQ(total_doc, doc_count);
// 3) Drop the altered nullable column
s = segment->drop_column("nullable_col");
ASSERT_TRUE(s.ok()) << "drop nullable column failed: " << s.message();
combined_reader = segment->scan({"nullable_col"});
ASSERT_TRUE(combined_reader == nullptr);
// 4) Add it back again to verify structure is still consistent
auto re_add_field =
std::make_shared<FieldSchema>("nullable_col_v2", DataType::FLOAT, true);
s = segment->add_column(re_add_field, "", AddColumnOptions());
ASSERT_TRUE(s.ok()) << "re-add nullable column failed: " << s.message();
combined_reader = segment->scan({"id", "name", "nullable_col_v2"});
ASSERT_TRUE(combined_reader != nullptr);
total_doc = 0;
while (true) {
auto st = combined_reader->ReadNext(&batch);
if (!st.ok()) break;
if (batch == nullptr) break;
EXPECT_EQ(batch->num_columns(), 3);
total_doc += batch->num_rows();
}
EXPECT_EQ(total_doc, doc_count);
}
TEST_P(SegmentTest, AlterColumn) {
// create segment
int doc_count = 1000;
auto segment = test::TestHelper::CreateSegmentWithDoc(
col_path_, *schema_, 0, 0, id_map_, delete_store_, version_manager_, options_,
0, doc_count);
ASSERT_TRUE(segment != nullptr);
auto s = segment->alter_column(
"alter_int32",
std::make_shared<FieldSchema>("alter_int32", DataType::INT32, false),
AlterColumnOptions());
EXPECT_FALSE(s.ok());
segment->dump();
auto writing_segment_meta = segment->meta();
// convert writing segment meta to persisted segment meta
Version version = version_manager_->get_current_version();
writing_segment_meta->remove_writing_forward_block();
s = version.add_persisted_segment_meta(writing_segment_meta);
ASSERT_TRUE(s.ok());
s = version_manager_->apply(version);
ASSERT_TRUE(s.ok());
s = version_manager_->flush();
ASSERT_TRUE(s.ok());
segment.reset();
version_manager_.reset();
id_map_->flush();
id_map_.reset();
std::string delete_store_path =
FileHelper::MakeFilePath(col_path_, FileID::DELETE_FILE, 0);
delete_store_->flush(delete_store_path);
delete_store_.reset();
auto recover_version_manager = VersionManager::Recovery(col_path_);
auto recover_version_mgr = recover_version_manager.value();
ASSERT_TRUE(recover_version_mgr != nullptr);
Version v = recover_version_mgr->get_current_version();
const auto &persist_metas = v.persisted_segment_metas();
// idmap
std::string idmap_path = FileHelper::MakeFilePath(col_path_, FileID::ID_FILE,
v.id_map_path_suffix());
IDMap::Ptr recover_id_map = std::make_shared<IDMap>(col_name_);
auto status = recover_id_map->open(idmap_path, false, false);
ASSERT_TRUE(status.ok());
delete_store_path = FileHelper::MakeFilePath(col_path_, FileID::DELETE_FILE,
v.delete_snapshot_path_suffix());
auto recover_delete_store =
DeleteStore::CreateAndLoad(col_name_, delete_store_path);
ASSERT_TRUE(recover_delete_store != nullptr);
// open persist segment
options_.read_only_ = true;
auto result =
Segment::Open(col_path_, *schema_, *persist_metas[0], recover_id_map,
recover_delete_store, recover_version_mgr, options_);
ASSERT_TRUE(result.has_value());
segment = std::move(result).value();
ASSERT_TRUE(segment != nullptr);
s = segment->alter_column(
"alter_int32",
std::make_shared<FieldSchema>("alter_int32", DataType::INT32, false),
AlterColumnOptions());
EXPECT_FALSE(s.ok()); // not found
s = segment->alter_column(
"int32",
std::make_shared<FieldSchema>("int32", DataType::UNDEFINED, false),
AlterColumnOptions());
EXPECT_FALSE(s.ok()); // undefined type
auto func = [&](const std::string &column_name,
const std::shared_ptr<FieldSchema> &field_schema) {
AlterColumnOptions alter_options;
status = segment->alter_column(column_name, field_schema, alter_options);
EXPECT_TRUE(status.ok());
auto combined_reader = segment->scan({"id", "name", "age", column_name});
ASSERT_TRUE(combined_reader != nullptr);
std::shared_ptr<arrow::RecordBatch> batch;
uint32_t total_doc = 0;
while (true) {
auto status = combined_reader->ReadNext(&batch);
if (status.ok() == false) break;
if (batch == nullptr) break;
EXPECT_EQ(batch->num_columns(), 4);
total_doc += batch->num_rows();
}
EXPECT_EQ(total_doc, doc_count);
};
std::vector<std::string> test_alter_columns = {"int32", "int64", "uint32",
"uint64", "float", "double"};
for (auto &column_name : test_alter_columns) {
// std::string column_name = "int32";
for (auto &dest_column : test_alter_columns) {
if (column_name == dest_column) continue;
auto field_schema = schema_->get_field(dest_column);
auto new_field_schema = std::make_shared<FieldSchema>(*field_schema);
new_field_schema->set_name(column_name);
func(column_name, new_field_schema);
}
}
}
TEST_P(SegmentTest, DropColumn) {
// create segment
int doc_count = 1000;
auto segment = test::TestHelper::CreateSegmentWithDoc(
col_path_, *schema_, 0, 0, id_map_, delete_store_, version_manager_, options_,
0, doc_count);
ASSERT_TRUE(segment != nullptr);
auto s = segment->drop_column("int32");
EXPECT_FALSE(s.ok());
segment->dump();
auto writing_segment_meta = segment->meta();
// convert writing segment meta to persisted segment meta
Version version = version_manager_->get_current_version();
writing_segment_meta->remove_writing_forward_block();
s = version.add_persisted_segment_meta(writing_segment_meta);
ASSERT_TRUE(s.ok());
s = version_manager_->apply(version);
ASSERT_TRUE(s.ok());
s = version_manager_->flush();
ASSERT_TRUE(s.ok());
segment.reset();
version_manager_.reset();
id_map_->flush();
id_map_.reset();
std::string delete_store_path =
FileHelper::MakeFilePath(col_path_, FileID::DELETE_FILE, 0);
delete_store_->flush(delete_store_path);
delete_store_.reset();
auto recover_version_manager = VersionManager::Recovery(col_path_);
auto recover_version_mgr = recover_version_manager.value();
ASSERT_TRUE(recover_version_mgr != nullptr);
Version v = recover_version_mgr->get_current_version();
const auto &persist_metas = v.persisted_segment_metas();
// idmap
std::string idmap_path = FileHelper::MakeFilePath(col_path_, FileID::ID_FILE,
v.id_map_path_suffix());
IDMap::Ptr recover_id_map = std::make_shared<IDMap>(col_name_);
auto status = recover_id_map->open(idmap_path, false, false);
ASSERT_TRUE(status.ok());
delete_store_path = FileHelper::MakeFilePath(col_path_, FileID::DELETE_FILE,
v.delete_snapshot_path_suffix());
auto recover_delete_store =
DeleteStore::CreateAndLoad(col_name_, delete_store_path);
ASSERT_TRUE(recover_delete_store != nullptr);
// open persist segment
options_.read_only_ = true;
auto result =
Segment::Open(col_path_, *schema_, *persist_metas[0], recover_id_map,
recover_delete_store, recover_version_mgr, options_);
ASSERT_TRUE(result.has_value());
segment = std::move(result).value();
ASSERT_TRUE(segment != nullptr);
auto meta = segment->meta();
auto &persist_blocks = meta->persisted_blocks();
auto func = [&](const std::string &column_name) {
status = segment->drop_column(column_name);
EXPECT_TRUE(status.ok());
// after drop column
bool col_exit = false;
for (auto &block : persist_blocks) {
if (block.type() == BlockType::SCALAR) {
if (block.contain_column(column_name)) {
col_exit = true;
break;
}
}
}
EXPECT_EQ(col_exit, false);
auto combined_reader = segment->scan({column_name});
ASSERT_TRUE(combined_reader == nullptr);
};
std::vector<std::string> test_drop_columns = {"int32", "int64", "uint32",
"uint64", "float", "double"};
for (auto &column_name : test_drop_columns) {
func(column_name);
}
}
INSTANTIATE_TEST_SUITE_P(MMapTest, SegmentTest, testing::Values(true, false));
Reference in New Issue View Git Blame Copy Permalink