milvus/internal/core/unittest/test_sealed.cpp

// Copyright (C) 2019-2020 Zilliz. All rights reserved.
//
// 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 <fmt/core.h>
#include <folly/CancellationToken.h>
#include <folly/FBVector.h>
#include <nlohmann/json.hpp>
#include <stdlib.h>
#include <time.h>
#include <algorithm>
#include <array>
#include <atomic>
#include <chrono>
#include <cstdint>
#include <exception>
#include <functional>
#include <future>
#include <iostream>
#include <limits>
#include <map>
#include <memory>
#include <string>
#include <string_view>
#include <stdexcept>
#include <tuple>
#include <utility>
#include <vector>

#include "NamedType/named_type_impl.hpp"
#include "cachinglayer/CacheSlot.h"
#include "common/Common.h"
#include "common/Consts.h"
#include "common/EasyAssert.h"
#include "common/FieldData.h"
#include "common/FieldDataInterface.h"
#include "common/IndexMeta.h"
#include "common/LoadInfo.h"
#include "common/PrometheusClient.h"
#include "common/QueryInfo.h"
#include "common/QueryResult.h"
#include "common/Schema.h"
#include "common/Span.h"
#include "common/Types.h"
#include "common/VectorArray.h"
#include "common/VectorTrait.h"
#include "common/protobuf_utils.h"
#include "exec/expression/function/FunctionFactory.h"
#include "gtest/gtest.h"
#include "index/Index.h"
#include "index/IndexFactory.h"
#include "index/IndexInfo.h"
#include "index/Meta.h"
#include "index/SkipIndex.h"
#include "index/VectorIndex.h"
#include "knowhere/comp/index_param.h"
#include "knowhere/config.h"
#include "knowhere/dataset.h"
#include "knowhere/version.h"
#include "mmap/ChunkedColumnGroup.h"
#include "pb/common.pb.h"
#include "pb/schema.pb.h"
#include "query/Plan.h"
#include "query/PlanImpl.h"
#include "query/Utils.h"
#include "segcore/ChunkedSegmentSealedImpl.h"
#include "segcore/SegcoreConfig.h"
#include "segcore/SegmentGrowing.h"
#include "segcore/SegmentGrowingImpl.h"
#include "segcore/SegmentLoadInfo.h"
#include "segcore/SegmentSealed.h"
#include "segcore/Types.h"
#include "segcore/storagev2translator/SystemIndexTranslator.h"
#include "storage/FileManager.h"
#include "storage/InsertData.h"
#include "storage/PayloadReader.h"
#include "storage/ThreadPools.h"
#include "storage/RemoteChunkManagerSingleton.h"
#include "storage/Types.h"
#include "storage/Util.h"
#include "test_utils/Constants.h"
#include "test_utils/DataGen.h"
#include "test_utils/cachinglayer_test_utils.h"
#include "test_utils/indexbuilder_test_utils.h"
#include "test_utils/storage_test_utils.h"

using namespace milvus;
using namespace milvus::query;
using namespace milvus::segcore;

using milvus::segcore::LoadIndexInfo;

namespace {

bool
GetFieldBit(const BitsetType& bitset, FieldId field_id) {
    auto pos = field_id.get() - START_USER_FIELDID;
    AssertInfo(pos >= 0, "invalid field id");
    return bitset[pos];
}

constexpr int64_t kWarmupPkFieldId = START_USER_FIELDID;
constexpr int64_t kWarmupVectorFieldId = START_USER_FIELDID + 1;

void
AddWarmupProperty(milvus::proto::schema::CollectionSchema& schema_proto,
                  const std::string& key,
                  const std::string& value) {
    auto* prop = schema_proto.add_properties();
    prop->set_key(key);
    prop->set_value(value);
}

SchemaPtr
CreateWarmupPolicySchema(bool include_vector) {
    milvus::proto::schema::CollectionSchema schema_proto;
    auto* pk_field = schema_proto.add_fields();
    pk_field->set_fieldid(kWarmupPkFieldId);
    pk_field->set_name("pk");
    pk_field->set_data_type(milvus::proto::schema::DataType::Int64);
    pk_field->set_is_primary_key(true);

    if (include_vector) {
        auto* vector_field = schema_proto.add_fields();
        vector_field->set_fieldid(kWarmupVectorFieldId);
        vector_field->set_name("vec");
        vector_field->set_data_type(
            milvus::proto::schema::DataType::FloatVector);
        auto* dim = vector_field->add_type_params();
        dim->set_key("dim");
        dim->set_value("4");
    }

    AddWarmupProperty(schema_proto, "warmup.vectorField", "sync");
    AddWarmupProperty(schema_proto, "warmup.vectorIndex", "disable");
    auto schema = Schema::ParseFrom(schema_proto);
    schema->set_schema_version(include_vector ? 200 : 100);
    return schema;
}

std::shared_ptr<milvus_storage::api::ColumnGroups>
MakeWarmupTestColumnGroups() {
    auto column_groups = std::make_shared<milvus_storage::api::ColumnGroups>();
    auto column_group = std::make_shared<milvus_storage::api::ColumnGroup>();
    column_group->columns.emplace_back(std::to_string(kWarmupVectorFieldId));
    column_groups->emplace_back(std::move(column_group));
    return column_groups;
}

class WarmupTestChunkReader : public milvus_storage::api::ChunkReader {
 public:
    size_t
    total_number_of_chunks() const override {
        return 1;
    }

    arrow::Result<std::vector<int64_t>>
    get_chunk_indices(const std::vector<int64_t>& row_indices) override {
        return std::vector<int64_t>(row_indices.size(), 0);
    }

    arrow::Result<std::shared_ptr<arrow::RecordBatch>>
    get_chunk(int64_t) override {
        return arrow::Status::NotImplemented("warmup test chunk reader");
    }

    arrow::Result<std::vector<std::shared_ptr<arrow::RecordBatch>>>
    get_chunks(const std::vector<int64_t>&, size_t) override {
        return arrow::Status::NotImplemented("warmup test chunk reader");
    }

    arrow::Result<std::vector<uint64_t>>
    get_chunk_size() override {
        return std::vector<uint64_t>{1};
    }

    arrow::Result<std::vector<uint64_t>>
    get_chunk_rows() override {
        return std::vector<uint64_t>{1};
    }
};

class WarmupTestReader : public milvus_storage::api::Reader {
 public:
    explicit WarmupTestReader(
        std::shared_ptr<milvus_storage::api::ColumnGroups> column_groups)
        : column_groups_(std::move(column_groups)) {
    }

    std::shared_ptr<milvus_storage::api::ColumnGroups>
    get_column_groups() const override {
        return column_groups_;
    }

    arrow::Result<std::shared_ptr<arrow::RecordBatchReader>>
    get_record_batch_reader(const std::string&) const override {
        return arrow::Status::NotImplemented("warmup test reader");
    }

    arrow::Result<std::unique_ptr<milvus_storage::api::ChunkReader>>
    get_chunk_reader(int64_t, const std::shared_ptr<std::vector<std::string>>&)
        const override {
        return std::make_unique<WarmupTestChunkReader>();
    }

    arrow::Result<std::shared_ptr<arrow::Table>>
    take(const std::vector<int64_t>&,
         size_t,
         const std::shared_ptr<std::vector<std::string>>&) override {
        return arrow::Status::NotImplemented("warmup test reader");
    }

    void
    set_keyretriever(
        const std::function<std::string(const std::string&)>&) override {
    }

 private:
    std::shared_ptr<milvus_storage::api::ColumnGroups> column_groups_;
};

class CancellationObservingIndexTranslator
    : public cachinglayer::Translator<index::IndexBase> {
 public:
    CancellationObservingIndexTranslator(
        std::string key,
        std::unique_ptr<index::IndexBase> index,
        folly::CancellationToken* observed_token,
        std::promise<void>* warmup_started)
        : key_(std::move(key)),
          index_(std::move(index)),
          observed_token_(observed_token),
          warmup_started_(warmup_started),
          meta_(cachinglayer::StorageType::MEMORY,
                cachinglayer::CellIdMappingMode::IDENTICAL,
                cachinglayer::CellDataType::OTHER,
                CacheWarmupPolicy::CacheWarmupPolicy_Async,
                false) {
    }

    size_t
    num_cells() const override {
        return 1;
    }

    cachinglayer::cid_t
    cell_id_of(cachinglayer::uid_t uid) const override {
        return uid;
    }

    std::pair<cachinglayer::ResourceUsage, cachinglayer::ResourceUsage>
    estimated_byte_size_of_cell(cachinglayer::cid_t) const override {
        return {{0, 0}, {0, 0}};
    }

    int64_t
    cells_storage_bytes(
        const std::vector<cachinglayer::cid_t>&) const override {
        return 0;
    }

    const std::string&
    key() const override {
        return key_;
    }

    cachinglayer::Meta*
    meta() override {
        return &meta_;
    }

    std::vector<
        std::pair<cachinglayer::cid_t, std::unique_ptr<index::IndexBase>>>
    get_cells(milvus::OpContext* ctx,
              const std::vector<cachinglayer::cid_t>& cids) override {
        AssertInfo(ctx != nullptr, "warmup context must not be null");
        *observed_token_ = ctx->cancellation_token;
        warmup_started_->set_value();
        std::vector<
            std::pair<cachinglayer::cid_t, std::unique_ptr<index::IndexBase>>>
            result;
        result.reserve(cids.size());
        for (auto cid : cids) {
            AssertInfo(cid == 0, "test index translator has one cell");
            result.emplace_back(cid, std::move(index_));
        }
        return result;
    }

 private:
    std::string key_;
    std::unique_ptr<index::IndexBase> index_;
    folly::CancellationToken* observed_token_;
    std::promise<void>* warmup_started_;
    cachinglayer::Meta meta_;
};

index::CacheIndexBasePtr
CreateCancellationObservingCacheIndex(std::string key,
                                      std::unique_ptr<index::IndexBase> index,
                                      folly::CancellationToken* observed_token,
                                      std::promise<void>* warmup_started,
                                      cachinglayer::internal::DList* dlist) {
    auto translator = std::make_unique<CancellationObservingIndexTranslator>(
        std::move(key), std::move(index), observed_token, warmup_started);
    return std::make_shared<cachinglayer::CacheSlot<index::IndexBase>>(
        std::move(translator),
        dlist,
        false,
        false,
        true,
        std::chrono::seconds(5),
        std::chrono::seconds(5));
}

}  // namespace

const int64_t ROW_COUNT = 2 * 1000;
const int64_t BIAS = 1000;

using Param = std::string;
class SealedTest : public ::testing::TestWithParam<Param> {
 public:
    void
    SetUp() override {
    }
};

TEST(Sealed, without_predicate) {
    auto schema = std::make_shared<Schema>();
    auto dim = 4;
    auto topK = 5;
    auto metric_type = knowhere::metric::L2;
    auto fake_id = schema->AddDebugField(
        "fakevec", DataType::VECTOR_FLOAT, dim, metric_type);
    schema->AddDebugField("age", DataType::FLOAT);
    auto i64_fid = schema->AddDebugField("counter", DataType::INT64);
    schema->set_primary_field_id(i64_fid);

    auto N = ROW_COUNT;

    auto dataset = DataGen(schema, N);
    auto vec_col = dataset.get_col<float>(fake_id);
    for (int64_t i = 0; i < (ROW_COUNT / 2) * dim; ++i) {
        vec_col.push_back(0);
    }
    auto query_ptr = vec_col.data() + BIAS * dim;
    auto segment = CreateGrowingSegment(schema, empty_index_meta);
    segment->PreInsert(N);
    segment->Insert(0,
                    N,
                    dataset.row_ids_.data(),
                    dataset.timestamps_.data(),
                    dataset.raw_);

    ScopedSchemaHandle handle(*schema);
    auto plan_str =
        handle.ParseSearch("", "fakevec", 5, "L2", "{\"nprobe\": 10}", 3);
    auto plan =
        CreateSearchPlanByExpr(schema, plan_str.data(), plan_str.size());
    auto num_queries = 5;
    auto ph_group_raw =
        CreatePlaceholderGroupFromBlob(num_queries, dim, query_ptr);
    auto ph_group =
        ParsePlaceholderGroup(plan.get(), ph_group_raw.SerializeAsString());
    Timestamp timestamp = 1000000;

    std::vector<const PlaceholderGroup*> ph_group_arr = {ph_group.get()};

    auto sr = segment->Search(plan.get(), ph_group.get(), timestamp);
    auto pre_result = SearchResultToJson(*sr);
    milvus::index::CreateIndexInfo create_index_info;
    create_index_info.field_type = DataType::VECTOR_FLOAT;
    create_index_info.metric_type = knowhere::metric::L2;
    create_index_info.index_type = knowhere::IndexEnum::INDEX_FAISS_IVFFLAT;
    create_index_info.index_engine_version =
        knowhere::Version::GetCurrentVersion().VersionNumber();

    auto indexing = milvus::index::IndexFactory::GetInstance().CreateIndex(
        create_index_info, milvus::storage::FileManagerContext());

    auto build_conf =
        knowhere::Json{{knowhere::meta::METRIC_TYPE, knowhere::metric::L2},
                       {knowhere::meta::DIM, std::to_string(dim)},
                       {knowhere::indexparam::NLIST, "100"}};

    auto search_conf = knowhere::Json{{knowhere::indexparam::NPROBE, 10}};

    auto database =
        knowhere::GenDataSet(N, dim, vec_col.data() + (ROW_COUNT / 2) * dim);
    indexing->BuildWithDataset(database, build_conf);

    auto vec_index = dynamic_cast<milvus::index::VectorIndex*>(indexing.get());
    EXPECT_EQ(vec_index->Count(), N);
    EXPECT_EQ(vec_index->GetDim(), dim);
    auto query_dataset = knowhere::GenDataSet(num_queries, dim, query_ptr);

    milvus::SearchInfo searchInfo;
    searchInfo.topk_ = topK;
    searchInfo.metric_type_ = knowhere::metric::L2;
    searchInfo.search_params_ = search_conf;
    SearchResult result;
    vec_index->Query(query_dataset, searchInfo, nullptr, nullptr, result);
    auto ref_result = SearchResultToJson(result);

    LoadIndexInfo load_info;
    load_info.field_id = fake_id.get();
    load_info.index_params = GenIndexParams(indexing.get());
    load_info.cache_index = CreateTestCacheIndex("test", std::move(indexing));
    load_info.index_params["metric_type"] = "L2";

    // load index for vec field, load raw data for scalar field
    auto sealed_segment = CreateSealedWithFieldDataLoaded(schema, dataset);
    sealed_segment->DropFieldData(fake_id);
    sealed_segment->LoadIndex(load_info);

    sr = sealed_segment->Search(plan.get(), ph_group.get(), timestamp);

    auto post_result = SearchResultToJson(*sr);
    std::cout << "ref_result" << std::endl;
    std::cout << ref_result.dump(1) << std::endl;
    std::cout << "post_result" << std::endl;
    std::cout << post_result.dump(1);
    // ASSERT_EQ(ref_result.dump(1), post_result.dump(1));

    sr = sealed_segment->Search(plan.get(), ph_group.get(), 0);
    EXPECT_EQ(sr->get_total_result_count(), 0);
    sr = sealed_segment->Search(plan.get(),
                                ph_group.get(),
                                timestamp,
                                folly::CancellationToken(),
                                0,
                                100);
    EXPECT_EQ(sr->get_total_result_count(), 0);
}

TEST(Sealed, without_search_ef_less_than_limit) {
    auto schema = std::make_shared<Schema>();
    auto dim = 4;
    auto metric_type = knowhere::metric::L2;
    auto fake_id = schema->AddDebugField(
        "fakevec", DataType::VECTOR_FLOAT, dim, metric_type);
    schema->AddDebugField("age", DataType::FLOAT);
    auto i64_fid = schema->AddDebugField("counter", DataType::INT64);
    schema->set_primary_field_id(i64_fid);

    auto N = ROW_COUNT;

    auto dataset = DataGen(schema, N);
    auto vec_col = dataset.get_col<float>(fake_id);
    auto query_ptr = vec_col.data() + BIAS * dim;

    ScopedSchemaHandle handle(*schema);
    auto plan_str =
        handle.ParseSearch("", "fakevec", 100, "L2", "{\"ef\": 10}", 3);
    auto plan =
        CreateSearchPlanByExpr(schema, plan_str.data(), plan_str.size());
    auto num_queries = 5;
    auto ph_group_raw =
        CreatePlaceholderGroupFromBlob(num_queries, dim, query_ptr);
    auto ph_group =
        ParsePlaceholderGroup(plan.get(), ph_group_raw.SerializeAsString());
    Timestamp timestamp = 1000000;

    milvus::index::CreateIndexInfo create_index_info;
    create_index_info.field_type = DataType::VECTOR_FLOAT;
    create_index_info.metric_type = knowhere::metric::L2;
    create_index_info.index_type = knowhere::IndexEnum::INDEX_HNSW;
    create_index_info.index_engine_version =
        knowhere::Version::GetCurrentVersion().VersionNumber();

    auto indexing = milvus::index::IndexFactory::GetInstance().CreateIndex(
        create_index_info, milvus::storage::FileManagerContext());

    auto build_conf =
        knowhere::Json{{knowhere::meta::METRIC_TYPE, knowhere::metric::L2},
                       {knowhere::indexparam::M, "16"},
                       {knowhere::indexparam::EF, "10"}};

    auto database = knowhere::GenDataSet(N, dim, vec_col.data());
    indexing->BuildWithDataset(database, build_conf);

    LoadIndexInfo load_info;
    load_info.field_id = fake_id.get();
    load_info.index_params = GenIndexParams(indexing.get());
    load_info.cache_index = CreateTestCacheIndex("test", std::move(indexing));
    load_info.index_params["metric_type"] = "L2";

    // load index for vec field, load raw data for scalar field
    auto sealed_segment = CreateSealedWithFieldDataLoaded(schema, dataset);
    sealed_segment->DropFieldData(fake_id);
    sealed_segment->LoadIndex(load_info);

    // Test that search fails when ef parameter is less than top-k
    // HNSW index requires ef to be larger than k for proper search
    bool exception_thrown = false;
    try {
        auto sr = sealed_segment->Search(plan.get(), ph_group.get(), timestamp);
        FAIL() << "Expected exception for invalid ef parameter";
    } catch (const std::exception& e) {
        exception_thrown = true;
        std::string error_msg = e.what();
        ASSERT_TRUE(error_msg.find("ef(10) should be larger than k(100)") !=
                    std::string::npos)
            << "Unexpected error message: " << error_msg;
    }
    ASSERT_TRUE(exception_thrown) << "Expected exception was not thrown";
}

TEST(Sealed, with_predicate) {
    auto schema = std::make_shared<Schema>();
    auto dim = 4;
    auto topK = 5;
    auto metric_type = knowhere::metric::L2;
    auto fake_id = schema->AddDebugField(
        "fakevec", DataType::VECTOR_FLOAT, dim, metric_type);
    auto i64_fid = schema->AddDebugField("counter", DataType::INT64);
    schema->set_primary_field_id(i64_fid);

    auto N = ROW_COUNT;

    auto dataset = DataGen(schema, N);
    auto vec_col = dataset.get_col<float>(fake_id);
    auto query_ptr = vec_col.data() + BIAS * dim;
    auto segment = CreateGrowingSegment(schema, empty_index_meta);
    segment->PreInsert(N);
    segment->Insert(0,
                    N,
                    dataset.row_ids_.data(),
                    dataset.timestamps_.data(),
                    dataset.raw_);

    ScopedSchemaHandle handle(*schema);
    // counter >= 1000 AND counter < 1005
    auto plan_str = handle.ParseSearch("counter >= 1000 && counter < 1005",
                                       "fakevec",
                                       5,
                                       "L2",
                                       "{\"nprobe\": 10}",
                                       6);
    auto plan =
        CreateSearchPlanByExpr(schema, plan_str.data(), plan_str.size());
    auto num_queries = 5;
    auto ph_group_raw =
        CreatePlaceholderGroupFromBlob(num_queries, dim, query_ptr);
    auto ph_group =
        ParsePlaceholderGroup(plan.get(), ph_group_raw.SerializeAsString());
    Timestamp timestamp = 1000000;

    std::vector<const PlaceholderGroup*> ph_group_arr = {ph_group.get()};

    auto sr = segment->Search(plan.get(), ph_group.get(), timestamp);
    milvus::index::CreateIndexInfo create_index_info;
    create_index_info.field_type = DataType::VECTOR_FLOAT;
    create_index_info.metric_type = knowhere::metric::L2;
    create_index_info.index_type = knowhere::IndexEnum::INDEX_FAISS_IVFFLAT;
    create_index_info.index_engine_version =
        knowhere::Version::GetCurrentVersion().VersionNumber();
    auto indexing = milvus::index::IndexFactory::GetInstance().CreateIndex(
        create_index_info, milvus::storage::FileManagerContext());

    auto build_conf =
        knowhere::Json{{knowhere::meta::METRIC_TYPE, knowhere::metric::L2},
                       {knowhere::meta::DIM, std::to_string(dim)},
                       {knowhere::indexparam::NLIST, "100"}};

    auto database = knowhere::GenDataSet(N, dim, vec_col.data());
    indexing->BuildWithDataset(database, build_conf);

    auto vec_index = dynamic_cast<index::VectorIndex*>(indexing.get());
    EXPECT_EQ(vec_index->Count(), N);
    EXPECT_EQ(vec_index->GetDim(), dim);

    auto query_dataset = knowhere::GenDataSet(num_queries, dim, query_ptr);

    auto search_conf =
        knowhere::Json{{knowhere::meta::METRIC_TYPE, knowhere::metric::L2},
                       {knowhere::indexparam::NPROBE, 10}};
    milvus::SearchInfo searchInfo;
    searchInfo.topk_ = topK;
    searchInfo.metric_type_ = knowhere::metric::L2;
    searchInfo.search_params_ = search_conf;
    SearchResult result;
    vec_index->Query(query_dataset, searchInfo, nullptr, nullptr, result);

    LoadIndexInfo load_info;
    load_info.field_id = fake_id.get();
    load_info.index_params = GenIndexParams(indexing.get());
    load_info.cache_index = CreateTestCacheIndex("test", std::move(indexing));
    load_info.index_params["metric_type"] = "L2";

    // load index for vec field, load raw data for scalar field
    auto sealed_segment = CreateSealedWithFieldDataLoaded(schema, dataset);
    sealed_segment->DropFieldData(fake_id);
    sealed_segment->LoadIndex(load_info);

    sr = sealed_segment->Search(plan.get(), ph_group.get(), timestamp);

    for (int i = 0; i < num_queries; ++i) {
        auto offset = i * topK;
        ASSERT_EQ(sr->seg_offsets_[offset], BIAS + i);
        ASSERT_EQ(sr->distances_[offset], 0.0);
    }
}

TEST(Sealed, with_predicate_filter_all) {
    auto schema = std::make_shared<Schema>();
    auto dim = 4;
    // auto metric_type = MetricType::METRIC_L2;
    auto metric_type = knowhere::metric::L2;
    auto fake_id = schema->AddDebugField(
        "fakevec", DataType::VECTOR_FLOAT, dim, metric_type);
    auto i64_fid = schema->AddDebugField("counter", DataType::INT64);
    schema->set_primary_field_id(i64_fid);

    auto N = ROW_COUNT;

    auto dataset = DataGen(schema, N);
    auto vec_col = dataset.get_col<float>(fake_id);
    auto query_ptr = vec_col.data() + BIAS * dim;
    ScopedSchemaHandle handle(*schema);
    // counter >= 4200 AND counter < 4199 (impossible range, filters all)
    auto plan_str = handle.ParseSearch("counter >= 4200 && counter < 4199",
                                       "fakevec",
                                       5,
                                       "L2",
                                       "{\"nprobe\": 10}",
                                       6);
    auto plan =
        CreateSearchPlanByExpr(schema, plan_str.data(), plan_str.size());
    auto num_queries = 5;
    auto ph_group_raw =
        CreatePlaceholderGroupFromBlob(num_queries, dim, query_ptr);
    auto ph_group =
        ParsePlaceholderGroup(plan.get(), ph_group_raw.SerializeAsString());
    Timestamp timestamp = 1000000;

    std::vector<const PlaceholderGroup*> ph_group_arr = {ph_group.get()};

    milvus::index::CreateIndexInfo create_index_info;
    create_index_info.field_type = DataType::VECTOR_FLOAT;
    create_index_info.metric_type = knowhere::metric::L2;
    create_index_info.index_type = knowhere::IndexEnum::INDEX_FAISS_IVFFLAT;
    create_index_info.index_engine_version =
        knowhere::Version::GetCurrentVersion().VersionNumber();
    auto ivf_indexing = milvus::index::IndexFactory::GetInstance().CreateIndex(
        create_index_info, milvus::storage::FileManagerContext());

    auto ivf_build_conf =
        knowhere::Json{{knowhere::meta::DIM, std::to_string(dim)},
                       {knowhere::indexparam::NLIST, "100"},
                       {knowhere::meta::METRIC_TYPE, knowhere::metric::L2}};

    auto database = knowhere::GenDataSet(N, dim, vec_col.data());
    ivf_indexing->BuildWithDataset(database, ivf_build_conf);

    auto ivf_vec_index = dynamic_cast<index::VectorIndex*>(ivf_indexing.get());
    EXPECT_EQ(ivf_vec_index->Count(), N);
    EXPECT_EQ(ivf_vec_index->GetDim(), dim);

    LoadIndexInfo load_info;
    load_info.field_id = fake_id.get();
    load_info.index_params = GenIndexParams(ivf_indexing.get());
    load_info.cache_index =
        CreateTestCacheIndex("test", std::move(ivf_indexing));
    load_info.index_params["metric_type"] = "L2";

    // load index for vec field, load raw data for scalar field
    auto ivf_sealed_segment = CreateSealedWithFieldDataLoaded(schema, dataset);
    ivf_sealed_segment->DropFieldData(fake_id);
    ivf_sealed_segment->LoadIndex(load_info);

    auto sr = ivf_sealed_segment->Search(plan.get(), ph_group.get(), timestamp);
    EXPECT_EQ(sr->unity_topK_, 0);
    EXPECT_EQ(sr->get_total_result_count(), 0);

    auto hnsw_conf =
        knowhere::Json{{knowhere::meta::DIM, std::to_string(dim)},
                       {knowhere::indexparam::HNSW_M, "16"},
                       {knowhere::indexparam::EFCONSTRUCTION, "200"},
                       {knowhere::indexparam::EF, "200"},
                       {knowhere::meta::METRIC_TYPE, knowhere::metric::L2}};

    create_index_info.field_type = DataType::VECTOR_FLOAT;
    create_index_info.metric_type = knowhere::metric::L2;
    create_index_info.index_type = knowhere::IndexEnum::INDEX_HNSW;
    create_index_info.index_engine_version =
        knowhere::Version::GetCurrentVersion().VersionNumber();
    auto hnsw_indexing = milvus::index::IndexFactory::GetInstance().CreateIndex(
        create_index_info, milvus::storage::FileManagerContext());
    hnsw_indexing->BuildWithDataset(database, hnsw_conf);

    auto hnsw_vec_index =
        dynamic_cast<index::VectorIndex*>(hnsw_indexing.get());
    EXPECT_EQ(hnsw_vec_index->Count(), N);
    EXPECT_EQ(hnsw_vec_index->GetDim(), dim);

    LoadIndexInfo hnsw_load_info;
    hnsw_load_info.field_id = fake_id.get();
    hnsw_load_info.index_params = GenIndexParams(hnsw_indexing.get());
    hnsw_load_info.cache_index =
        CreateTestCacheIndex("test", std::move(hnsw_indexing));
    hnsw_load_info.index_params["metric_type"] = "L2";

    // load index for vec field, load raw data for scalar field
    auto hnsw_sealed_segment = CreateSealedWithFieldDataLoaded(schema, dataset);
    hnsw_sealed_segment->DropFieldData(fake_id);
    hnsw_sealed_segment->LoadIndex(hnsw_load_info);

    auto sr2 =
        hnsw_sealed_segment->Search(plan.get(), ph_group.get(), timestamp);
    EXPECT_EQ(sr2->unity_topK_, 0);
    EXPECT_EQ(sr2->get_total_result_count(), 0);
}

TEST(Sealed, SegmentInterfaceIsIndexRefineEnabledPropagatesOpContext) {
    auto dim = 4;
    auto N = ROW_COUNT;
    auto metric_type = knowhere::metric::L2;
    auto schema = std::make_shared<Schema>();
    auto fake_id = schema->AddDebugField(
        "fakevec", DataType::VECTOR_FLOAT, dim, metric_type);
    auto i64_fid = schema->AddDebugField("counter", DataType::INT64);
    schema->set_primary_field_id(i64_fid);

    auto dataset = DataGen(schema, N);
    auto vec_col = dataset.get_col<float>(fake_id);

    auto create_index_info = milvus::index::CreateIndexInfo{};
    create_index_info.field_type = DataType::VECTOR_FLOAT;
    create_index_info.metric_type = knowhere::metric::L2;
    create_index_info.index_type = knowhere::IndexEnum::INDEX_HNSW;
    create_index_info.index_engine_version =
        knowhere::Version::GetCurrentVersion().VersionNumber();
    auto indexing = milvus::index::IndexFactory::GetInstance().CreateIndex(
        create_index_info, milvus::storage::FileManagerContext());
    auto build_conf =
        knowhere::Json{{knowhere::meta::DIM, std::to_string(dim)},
                       {knowhere::indexparam::HNSW_M, "16"},
                       {knowhere::indexparam::EFCONSTRUCTION, "200"},
                       {knowhere::indexparam::EF, "200"},
                       {knowhere::meta::METRIC_TYPE, knowhere::metric::L2}};
    auto database = knowhere::GenDataSet(N, dim, vec_col.data());
    indexing->BuildWithDataset(database, build_conf);
    auto expected_refine_enabled =
        dynamic_cast<index::VectorIndex*>(indexing.get())
            ->IsIndexRefineEnabled();

    milvus::OpContext* observed_ctx = nullptr;
    LoadIndexInfo load_info;
    load_info.field_id = fake_id.get();
    load_info.index_params = GenIndexParams(indexing.get());
    load_info.cache_index =
        CreateTestCacheIndex("test", std::move(indexing), &observed_ctx);
    load_info.index_params["metric_type"] = "L2";

    auto sealed_segment = CreateSealedWithFieldDataLoaded(schema, dataset);
    sealed_segment->DropFieldData(fake_id);
    sealed_segment->LoadIndex(load_info);

    auto* segment_interface =
        static_cast<SegmentInterface*>(sealed_segment.get());
    folly::CancellationSource source;
    milvus::OpContext op_context(source.getToken());

    EXPECT_EQ(segment_interface->IsIndexRefineEnabled(&op_context, fake_id),
              expected_refine_enabled);
    EXPECT_EQ(observed_ctx, &op_context);
}

TEST(Sealed, SegmentInterfaceCalcDistByIDsPropagatesOpContext) {
    auto dim = 4;
    auto N = ROW_COUNT;
    auto metric_type = knowhere::metric::L2;
    auto schema = std::make_shared<Schema>();
    auto fake_id = schema->AddDebugField(
        "fakevec", DataType::VECTOR_FLOAT, dim, metric_type);
    auto i64_fid = schema->AddDebugField("counter", DataType::INT64);
    schema->set_primary_field_id(i64_fid);

    auto dataset = DataGen(schema, N);
    auto vec_col = dataset.get_col<float>(fake_id);

    auto create_index_info = milvus::index::CreateIndexInfo{};
    create_index_info.field_type = DataType::VECTOR_FLOAT;
    create_index_info.metric_type = knowhere::metric::L2;
    create_index_info.index_type = knowhere::IndexEnum::INDEX_HNSW;
    create_index_info.index_engine_version =
        knowhere::Version::GetCurrentVersion().VersionNumber();
    auto indexing = milvus::index::IndexFactory::GetInstance().CreateIndex(
        create_index_info, milvus::storage::FileManagerContext());
    auto build_conf =
        knowhere::Json{{knowhere::meta::DIM, std::to_string(dim)},
                       {knowhere::indexparam::HNSW_M, "16"},
                       {knowhere::indexparam::EFCONSTRUCTION, "200"},
                       {knowhere::indexparam::EF, "200"},
                       {knowhere::meta::METRIC_TYPE, knowhere::metric::L2}};
    auto database = knowhere::GenDataSet(N, dim, vec_col.data());
    indexing->BuildWithDataset(database, build_conf);

    milvus::OpContext* observed_ctx = nullptr;
    LoadIndexInfo load_info;
    load_info.field_id = fake_id.get();
    load_info.index_params = GenIndexParams(indexing.get());
    load_info.cache_index =
        CreateTestCacheIndex("test", std::move(indexing), &observed_ctx);
    load_info.index_params["metric_type"] = "L2";

    auto sealed_segment = CreateSealedWithFieldDataLoaded(schema, dataset);
    sealed_segment->DropFieldData(fake_id);
    sealed_segment->LoadIndex(load_info);

    auto* segment_interface =
        static_cast<SegmentInterface*>(sealed_segment.get());
    auto query_dataset = knowhere::GenDataSet(1, dim, vec_col.data());
    auto seg_offsets = std::array<int64_t, 1>{0};
    auto distances = std::array<float, 1>{0.0F};
    folly::CancellationSource source;
    milvus::OpContext op_context(source.getToken());

    EXPECT_TRUE(segment_interface->CalcDistByIDs(&op_context,
                                                 fake_id,
                                                 query_dataset,
                                                 seg_offsets.data(),
                                                 seg_offsets.size(),
                                                 false,
                                                 distances.data()));
    EXPECT_EQ(observed_ctx, &op_context);
}

TEST(Sealed, LoadFieldData) {
    auto dim = 4;
    auto N = ROW_COUNT;
    auto metric_type = knowhere::metric::L2;
    auto schema = std::make_shared<Schema>();
    auto fakevec_id = schema->AddDebugField(
        "fakevec", DataType::VECTOR_FLOAT, dim, metric_type);
    auto counter_id = schema->AddDebugField("counter", DataType::INT64);
    auto double_id = schema->AddDebugField("double", DataType::DOUBLE);
    schema->AddDebugField("nothing", DataType::INT32);
    auto str_id = schema->AddDebugField("str", DataType::VARCHAR);
    schema->AddDebugField("int8", DataType::INT8);
    schema->AddDebugField("int16", DataType::INT16);
    schema->AddDebugField("float", DataType::FLOAT);
    schema->AddDebugField("json", DataType::JSON);
    schema->AddDebugField("geometry", DataType::GEOMETRY);
    schema->AddDebugField("array", DataType::ARRAY, DataType::INT64);
    schema->set_primary_field_id(counter_id);
    auto int8_nullable_id =
        schema->AddDebugField("int8_null", DataType::INT8, true);
    auto int16_nullable_id =
        schema->AddDebugField("int16_null", DataType::INT16, true);
    auto int32_nullable_id =
        schema->AddDebugField("int32_null", DataType::INT32, true);
    auto int64_nullable_id =
        schema->AddDebugField("int64_null", DataType::INT64, true);
    auto double_nullable_id =
        schema->AddDebugField("double_null", DataType::DOUBLE, true);
    auto str_nullable_id =
        schema->AddDebugField("str_null", DataType::VARCHAR, true);
    schema->AddDebugField("float_null", DataType::FLOAT, true);

    auto dataset = DataGen(schema, N);

    auto fakevec = dataset.get_col<float>(fakevec_id);

    auto indexing = GenVecIndexing(
        N, dim, fakevec.data(), knowhere::IndexEnum::INDEX_FAISS_IVFFLAT);
    //
    auto segment = CreateSealedSegment(schema);
    Timestamp timestamp = 1000000;
    ScopedSchemaHandle handle(*schema);
    // double >= -1 AND double < 1
    auto plan_str = handle.ParseSearch("double >= -1 && double < 1",
                                       "fakevec",
                                       5,
                                       "L2",
                                       "{\"nprobe\": 10}",
                                       3);
    auto plan =
        CreateSearchPlanByExpr(schema, plan_str.data(), plan_str.size());
    auto num_queries = 5;
    auto ph_group_raw = CreatePlaceholderGroup(num_queries, dim, 1024);
    auto ph_group =
        ParsePlaceholderGroup(plan.get(), ph_group_raw.SerializeAsString());

    ASSERT_ANY_THROW(segment->Search(plan.get(), ph_group.get(), timestamp));

    segment = CreateSealedWithFieldDataLoaded(schema, dataset);
    segment->Search(plan.get(), ph_group.get(), timestamp);

    segment->DropFieldData(fakevec_id);
    ASSERT_ANY_THROW(segment->Search(plan.get(), ph_group.get(), timestamp));

    LoadIndexInfo vec_info;
    vec_info.field_id = fakevec_id.get();
    vec_info.index_params = GenIndexParams(indexing.get());
    vec_info.cache_index = CreateTestCacheIndex("test", std::move(indexing));
    vec_info.index_params["metric_type"] = knowhere::metric::L2;
    segment->LoadIndex(vec_info);

    ASSERT_EQ(segment->num_chunk(fakevec_id), 1);
    ASSERT_EQ(segment->PinIndex(nullptr, double_id).size(), 0);
    ASSERT_EQ(segment->PinIndex(nullptr, str_id).size(), 0);
    auto chunk_span1 = segment->chunk_data<int64_t>(nullptr, counter_id, 0);
    auto chunk_span2 = segment->chunk_data<double>(nullptr, double_id, 0);
    auto chunk_span3 =
        segment->get_batch_views<std::string_view>(nullptr, str_id, 0, 0, N);
    auto chunk_span4 =
        segment->chunk_data<int8_t>(nullptr, int8_nullable_id, 0);
    auto chunk_span5 =
        segment->chunk_data<int16_t>(nullptr, int16_nullable_id, 0);
    auto chunk_span6 =
        segment->chunk_data<int32_t>(nullptr, int32_nullable_id, 0);
    auto chunk_span7 =
        segment->chunk_data<int64_t>(nullptr, int64_nullable_id, 0);
    auto chunk_span8 =
        segment->chunk_data<double>(nullptr, double_nullable_id, 0);
    auto chunk_span9 = segment->get_batch_views<std::string_view>(
        nullptr, str_nullable_id, 0, 0, N);

    auto ref1 = dataset.get_col<int64_t>(counter_id);
    auto ref2 = dataset.get_col<double>(double_id);
    auto ref3 = dataset.get_col(str_id)->scalars().string_data().data();
    auto ref4 = dataset.get_col<int8_t>(int8_nullable_id);
    auto ref5 = dataset.get_col<int16_t>(int16_nullable_id);
    auto ref6 = dataset.get_col<int32_t>(int32_nullable_id);
    auto ref7 = dataset.get_col<int64_t>(int64_nullable_id);
    auto ref8 = dataset.get_col<double>(double_nullable_id);
    auto ref9 =
        dataset.get_col(str_nullable_id)->scalars().string_data().data();
    auto valid4 = dataset.get_col_valid(int8_nullable_id);
    auto valid5 = dataset.get_col_valid(int16_nullable_id);
    auto valid6 = dataset.get_col_valid(int32_nullable_id);
    auto valid7 = dataset.get_col_valid(int64_nullable_id);
    auto valid8 = dataset.get_col_valid(double_nullable_id);
    auto valid9 = dataset.get_col_valid(str_nullable_id);
    ASSERT_EQ(chunk_span1.get().valid_data(), nullptr);
    ASSERT_EQ(chunk_span2.get().valid_data(), nullptr);
    ASSERT_EQ(chunk_span3.get().second.size(), 0);
    for (int i = 0; i < N; ++i) {
        if (chunk_span1.get().valid_data() == nullptr ||
            chunk_span1.get().valid_data()[i]) {
            ASSERT_EQ(chunk_span1.get().data()[i], ref1[i]);
        }
        if (chunk_span2.get().valid_data() == nullptr ||
            chunk_span2.get().valid_data()[i]) {
            ASSERT_EQ(chunk_span2.get().data()[i], ref2[i]);
        }
        if (chunk_span3.get().second.size() == 0 ||
            chunk_span3.get().second[i]) {
            ASSERT_EQ(chunk_span3.get().first[i], ref3[i]);
        }
        if (chunk_span4.get().valid_data() == nullptr ||
            chunk_span4.get().valid_data()[i]) {
            ASSERT_EQ(chunk_span4.get().data()[i], ref4[i]);
        }
        if (chunk_span5.get().valid_data() == nullptr ||
            chunk_span5.get().valid_data()[i]) {
            ASSERT_EQ(chunk_span5.get().data()[i], ref5[i]);
        }
        if (chunk_span6.get().valid_data() == nullptr ||
            chunk_span6.get().valid_data()[i]) {
            ASSERT_EQ(chunk_span6.get().data()[i], ref6[i]);
        }
        if (chunk_span7.get().valid_data() == nullptr ||
            chunk_span7.get().valid_data()[i]) {
            ASSERT_EQ(chunk_span7.get().data()[i], ref7[i]);
        }
        if (chunk_span8.get().valid_data() == nullptr ||
            chunk_span8.get().valid_data()[i]) {
            ASSERT_EQ(chunk_span8.get().data()[i], ref8[i]);
        }
        if (chunk_span9.get().second.size() == 0 ||
            chunk_span9.get().second[i]) {
            ASSERT_EQ(chunk_span9.get().first[i], ref9[i]);
        }
        ASSERT_EQ(chunk_span4.get().valid_data()[i], valid4[i]);
        ASSERT_EQ(chunk_span5.get().valid_data()[i], valid5[i]);
        ASSERT_EQ(chunk_span6.get().valid_data()[i], valid6[i]);
        ASSERT_EQ(chunk_span7.get().valid_data()[i], valid7[i]);
        ASSERT_EQ(chunk_span8.get().valid_data()[i], valid8[i]);
        ASSERT_EQ(chunk_span9.get().second[i], valid9[i]);
    }

    auto sr = segment->Search(plan.get(), ph_group.get(), timestamp);
    auto json = SearchResultToJson(*sr);
    std::cout << json.dump(1);
}

TEST(Sealed, ClearData) {
    milvus::exec::expression::FunctionFactory& factory =
        milvus::exec::expression::FunctionFactory::Instance();
    factory.Initialize();
    auto dim = 4;
    auto N = ROW_COUNT;
    auto metric_type = knowhere::metric::L2;
    auto schema = std::make_shared<Schema>();
    auto fakevec_id = schema->AddDebugField(
        "fakevec", DataType::VECTOR_FLOAT, dim, metric_type);
    auto counter_id = schema->AddDebugField("counter", DataType::INT64);
    auto double_id = schema->AddDebugField("double", DataType::DOUBLE);
    schema->AddDebugField("nothing", DataType::INT32);
    auto str_id = schema->AddDebugField("str", DataType::VARCHAR);
    schema->AddDebugField("int8", DataType::INT8);
    schema->AddDebugField("int16", DataType::INT16);
    schema->AddDebugField("float", DataType::FLOAT);
    schema->AddDebugField("json", DataType::JSON);
    schema->AddDebugField("geometry", DataType::GEOMETRY);
    schema->AddDebugField("array", DataType::ARRAY, DataType::INT64);
    schema->set_primary_field_id(counter_id);

    auto dataset = DataGen(schema, N);

    auto fakevec = dataset.get_col<float>(fakevec_id);

    auto indexing = GenVecIndexing(
        N, dim, fakevec.data(), knowhere::IndexEnum::INDEX_FAISS_IVFFLAT);

    auto segment = CreateSealedSegment(schema);
    Timestamp timestamp = 1000000;
    ScopedSchemaHandle handle(*schema);
    // double >= -1 AND double < 1
    auto plan_str = handle.ParseSearch("double >= -1 && double < 1",
                                       "fakevec",
                                       5,
                                       "L2",
                                       "{\"nprobe\": 10}",
                                       3);
    auto plan =
        CreateSearchPlanByExpr(schema, plan_str.data(), plan_str.size());
    auto num_queries = 5;
    auto ph_group_raw = CreatePlaceholderGroup(num_queries, dim, 1024);
    auto ph_group =
        ParsePlaceholderGroup(plan.get(), ph_group_raw.SerializeAsString());

    ASSERT_ANY_THROW(segment->Search(plan.get(), ph_group.get(), timestamp));

    segment = CreateSealedWithFieldDataLoaded(schema, dataset);
    segment->Search(plan.get(), ph_group.get(), timestamp);

    segment->DropFieldData(fakevec_id);
    ASSERT_ANY_THROW(segment->Search(plan.get(), ph_group.get(), timestamp));

    LoadIndexInfo vec_info;
    vec_info.field_id = fakevec_id.get();
    vec_info.index_params = GenIndexParams(indexing.get());
    vec_info.cache_index = CreateTestCacheIndex("test", std::move(indexing));
    vec_info.index_params["metric_type"] = knowhere::metric::L2;
    segment->LoadIndex(vec_info);

    ASSERT_EQ(segment->num_chunk(fakevec_id), 1);
    ASSERT_EQ(segment->PinIndex(nullptr, double_id).size(), 0);
    ASSERT_EQ(segment->PinIndex(nullptr, str_id).size(), 0);
    auto chunk_span1 = segment->chunk_data<int64_t>(nullptr, counter_id, 0);
    auto chunk_span2 = segment->chunk_data<double>(nullptr, double_id, 0);
    auto chunk_span3 =
        segment->get_batch_views<std::string_view>(nullptr, str_id, 0, 0, N);
    auto ref1 = dataset.get_col<int64_t>(counter_id);
    auto ref2 = dataset.get_col<double>(double_id);
    auto ref3 = dataset.get_col(str_id)->scalars().string_data().data();
    ASSERT_EQ(chunk_span3.get().second.size(), 0);
    for (int i = 0; i < N; ++i) {
        ASSERT_EQ(chunk_span1.get()[i], ref1[i]);
        ASSERT_EQ(chunk_span2.get()[i], ref2[i]);
        ASSERT_EQ(chunk_span3.get().first[i], ref3[i]);
    }

    auto sr = segment->Search(plan.get(), ph_group.get(), timestamp);
    auto json = SearchResultToJson(*sr);
    std::cout << json.dump(1);

    auto sealed_segment = (ChunkedSegmentSealedImpl*)segment.get();
    sealed_segment->ClearData();
    ASSERT_EQ(sealed_segment->get_row_count(), 0);
    ASSERT_EQ(sealed_segment->get_real_count(), 0);
    ASSERT_ANY_THROW(segment->Search(plan.get(), ph_group.get(), timestamp));
}

TEST(Sealed, LoadFieldDataMmap) {
    auto dim = 4;
    auto N = ROW_COUNT;
    auto metric_type = knowhere::metric::L2;
    auto schema = std::make_shared<Schema>();
    auto fakevec_id = schema->AddDebugField(
        "fakevec", DataType::VECTOR_FLOAT, dim, metric_type);
    auto counter_id = schema->AddDebugField("counter", DataType::INT64);
    auto double_id = schema->AddDebugField("double", DataType::DOUBLE);
    schema->AddDebugField("nothing", DataType::INT32);
    auto str_id = schema->AddDebugField("str", DataType::VARCHAR);
    schema->AddDebugField("int8", DataType::INT8);
    schema->AddDebugField("int16", DataType::INT16);
    schema->AddDebugField("float", DataType::FLOAT);
    schema->AddDebugField("json", DataType::JSON);
    schema->AddDebugField("geometry", DataType::GEOMETRY);
    schema->AddDebugField("array", DataType::ARRAY, DataType::INT64);
    schema->set_primary_field_id(counter_id);

    auto dataset = DataGen(schema, N);

    auto fakevec = dataset.get_col<float>(fakevec_id);

    auto indexing = GenVecIndexing(
        N, dim, fakevec.data(), knowhere::IndexEnum::INDEX_FAISS_IVFFLAT);

    auto segment = CreateSealedSegment(schema);
    Timestamp timestamp = 1000000;
    ScopedSchemaHandle handle(*schema);
    // double >= -1 AND double < 1
    auto plan_str = handle.ParseSearch("double >= -1 && double < 1",
                                       "fakevec",
                                       5,
                                       "L2",
                                       "{\"nprobe\": 10}",
                                       3);
    auto plan =
        CreateSearchPlanByExpr(schema, plan_str.data(), plan_str.size());
    auto num_queries = 5;
    auto ph_group_raw = CreatePlaceholderGroup(num_queries, dim, 1024);
    auto ph_group =
        ParsePlaceholderGroup(plan.get(), ph_group_raw.SerializeAsString());

    ASSERT_ANY_THROW(segment->Search(plan.get(), ph_group.get(), timestamp));

    segment = CreateSealedWithFieldDataLoaded(schema, dataset, true);
    segment->Search(plan.get(), ph_group.get(), timestamp);

    segment->DropFieldData(fakevec_id);
    ASSERT_ANY_THROW(segment->Search(plan.get(), ph_group.get(), timestamp));

    LoadIndexInfo vec_info;
    vec_info.field_id = fakevec_id.get();
    vec_info.index_params = GenIndexParams(indexing.get());
    vec_info.cache_index = CreateTestCacheIndex("test", std::move(indexing));
    vec_info.index_params["metric_type"] = knowhere::metric::L2;
    segment->LoadIndex(vec_info);

    ASSERT_EQ(segment->num_chunk(fakevec_id), 1);
    ASSERT_EQ(segment->PinIndex(nullptr, double_id).size(), 0);
    ASSERT_EQ(segment->PinIndex(nullptr, str_id).size(), 0);
    auto chunk_span1 = segment->chunk_data<int64_t>(nullptr, counter_id, 0);
    auto chunk_span2 = segment->chunk_data<double>(nullptr, double_id, 0);
    auto chunk_span3 =
        segment->get_batch_views<std::string_view>(nullptr, str_id, 0, 0, N);
    auto ref1 = dataset.get_col<int64_t>(counter_id);
    auto ref2 = dataset.get_col<double>(double_id);
    auto ref3 = dataset.get_col(str_id)->scalars().string_data().data();
    ASSERT_EQ(chunk_span3.get().second.size(), 0);
    for (int i = 0; i < N; ++i) {
        ASSERT_EQ(chunk_span1.get()[i], ref1[i]);
        ASSERT_EQ(chunk_span2.get()[i], ref2[i]);
        ASSERT_EQ(chunk_span3.get().first[i], ref3[i]);
    }

    auto sr = segment->Search(plan.get(), ph_group.get(), timestamp);
    auto json = SearchResultToJson(*sr);
    std::cout << json.dump(1);
}

TEST(Sealed, LoadPkScalarIndex) {
    size_t N = ROW_COUNT;
    auto schema = std::make_shared<Schema>();
    auto pk_id = schema->AddDebugField("counter", DataType::INT64);
    schema->AddDebugField("nothing", DataType::INT32);
    schema->set_primary_field_id(pk_id);

    auto dataset = DataGen(schema, N);
    auto segment = CreateSealedWithFieldDataLoaded(schema, dataset);

    LoadIndexInfo pk_index;
    pk_index.field_id = pk_id.get();
    pk_index.field_type = DataType::INT64;
    pk_index.index_params["index_type"] = "STL_SORT";
    auto pk_data = dataset.get_col<int64_t>(pk_id);
    auto index = GenScalarIndexing<int64_t>(N, pk_data.data());
    pk_index.index_params = GenIndexParams(index.get());
    pk_index.cache_index = CreateTestCacheIndex("test", std::move(index));
    segment->LoadIndex(pk_index);
}

TEST(Sealed, LoadScalarIndex) {
    auto dim = 4;
    size_t N = ROW_COUNT;
    auto metric_type = knowhere::metric::L2;
    auto schema = std::make_shared<Schema>();
    auto fakevec_id = schema->AddDebugField(
        "fakevec", DataType::VECTOR_FLOAT, dim, metric_type);
    auto counter_id = schema->AddDebugField("counter", DataType::INT64);
    auto double_id = schema->AddDebugField("double", DataType::DOUBLE);
    auto nothing_id = schema->AddDebugField("nothing", DataType::INT32);
    schema->set_primary_field_id(counter_id);

    auto dataset = DataGen(schema, N);

    auto fakevec = dataset.get_col<float>(fakevec_id);

    auto indexing = GenVecIndexing(
        N, dim, fakevec.data(), knowhere::IndexEnum::INDEX_FAISS_IVFFLAT);

    Timestamp timestamp = 1000000;
    ScopedSchemaHandle handle(*schema);
    // double >= -1 AND double < 1
    auto plan_str = handle.ParseSearch("double >= -1 && double < 1",
                                       "fakevec",
                                       5,
                                       "L2",
                                       "{\"nprobe\": 10}",
                                       3);
    auto plan =
        CreateSearchPlanByExpr(schema, plan_str.data(), plan_str.size());
    auto num_queries = 5;
    auto ph_group_raw = CreatePlaceholderGroup(num_queries, dim, 1024);
    auto ph_group =
        ParsePlaceholderGroup(plan.get(), ph_group_raw.SerializeAsString());

    auto segment = CreateSealedWithFieldDataLoaded(
        schema,
        dataset,
        false,
        GetExcludedFieldIds(schema, {{0, 1, counter_id.get()}}));

    LoadIndexInfo vec_info;
    vec_info.field_id = fakevec_id.get();
    vec_info.field_type = DataType::VECTOR_FLOAT;
    vec_info.index_params = GenIndexParams(indexing.get());
    vec_info.cache_index = CreateTestCacheIndex("test", std::move(indexing));
    vec_info.index_params["metric_type"] = knowhere::metric::L2;
    segment->LoadIndex(vec_info);

    LoadIndexInfo counter_index;
    counter_index.field_id = counter_id.get();
    counter_index.field_type = DataType::INT64;
    counter_index.index_params["index_type"] = "STL_SORT";
    auto counter_data = dataset.get_col<int64_t>(counter_id);
    auto index = GenScalarIndexing<int64_t>(N, counter_data.data());
    counter_index.index_params = GenIndexParams(index.get());
    counter_index.cache_index = CreateTestCacheIndex("test", std::move(index));
    segment->LoadIndex(counter_index);

    LoadIndexInfo double_index;
    double_index.field_id = double_id.get();
    double_index.field_type = DataType::DOUBLE;
    double_index.index_params["index_type"] = "STL_SORT";
    auto double_data = dataset.get_col<double>(double_id);
    auto temp1 = GenScalarIndexing<double>(N, double_data.data());
    double_index.index_params = GenIndexParams(temp1.get());
    double_index.cache_index = CreateTestCacheIndex("test", std::move(temp1));
    segment->LoadIndex(double_index);

    LoadIndexInfo nothing_index;
    nothing_index.field_id = nothing_id.get();
    nothing_index.field_type = DataType::INT32;
    nothing_index.index_params["index_type"] = "STL_SORT";
    auto nothing_data = dataset.get_col<int32_t>(nothing_id);
    auto temp2 = GenScalarIndexing<int32_t>(N, nothing_data.data());
    nothing_index.index_params = GenIndexParams(temp2.get());
    nothing_index.cache_index = CreateTestCacheIndex("test", std::move(temp2));
    segment->LoadIndex(nothing_index);

    auto sr = segment->Search(plan.get(),
                              ph_group.get(),
                              timestamp,
                              folly::CancellationToken(),
                              0,
                              100000);
    auto json = SearchResultToJson(*sr);
    std::cout << json.dump(1);
}

TEST(Sealed, Delete) {
    auto dim = 4;
    auto N = 10;
    auto metric_type = knowhere::metric::L2;
    auto schema = std::make_shared<Schema>();
    auto fakevec_id = schema->AddDebugField(
        "fakevec", DataType::VECTOR_FLOAT, dim, metric_type);
    auto counter_id = schema->AddDebugField("counter", DataType::INT64);
    schema->AddDebugField("double", DataType::DOUBLE);
    schema->AddDebugField("nothing", DataType::INT32);
    schema->set_primary_field_id(counter_id);

    auto dataset = DataGen(schema, N);

    auto fakevec = dataset.get_col<float>(fakevec_id);

    auto segment = CreateSealedSegment(schema);
    Timestamp timestamp = 1000000;
    ScopedSchemaHandle handle(*schema);
    // double >= -1 AND double < 1
    auto plan_str = handle.ParseSearch("double >= -1 && double < 1",
                                       "fakevec",
                                       5,
                                       "L2",
                                       "{\"nprobe\": 10}",
                                       3);
    auto plan =
        CreateSearchPlanByExpr(schema, plan_str.data(), plan_str.size());
    auto num_queries = 5;
    auto ph_group_raw = CreatePlaceholderGroup(num_queries, dim, 1024);
    auto ph_group =
        ParsePlaceholderGroup(plan.get(), ph_group_raw.SerializeAsString());

    ASSERT_ANY_THROW(segment->Search(plan.get(), ph_group.get(), timestamp));

    segment = CreateSealedWithFieldDataLoaded(schema, dataset);

    int64_t row_count = 5;
    std::vector<idx_t> pks{1, 2, 3, 4, 5};
    auto ids = std::make_unique<IdArray>();
    ids->mutable_int_id()->mutable_data()->Add(pks.begin(), pks.end());
    std::vector<Timestamp> timestamps_del{10, 10, 10, 10, 10};

    LoadDeletedRecordInfo info = {timestamps_del.data(), ids.get(), row_count};
    segment->LoadDeletedRecord(info);

    BitsetType bitset(N, false);
    auto bitset_view = BitsetTypeView(bitset);
    segment->mask_with_delete(bitset_view, 10, 11);
    ASSERT_EQ(bitset.count(), pks.size());

    int64_t new_count = 3;
    std::vector<idx_t> new_pks{6, 7, 8};
    auto new_ids = std::make_unique<IdArray>();
    new_ids->mutable_int_id()->mutable_data()->Add(new_pks.begin(),
                                                   new_pks.end());
    std::vector<idx_t> new_timestamps{10, 10, 10};
    segment->Delete(new_count,
                    new_ids.get(),
                    reinterpret_cast<const Timestamp*>(new_timestamps.data()));
}

TEST(Sealed, OverlapDelete) {
    auto dim = 4;
    auto N = 10;
    auto metric_type = knowhere::metric::L2;
    auto schema = std::make_shared<Schema>();
    auto fakevec_id = schema->AddDebugField(
        "fakevec", DataType::VECTOR_FLOAT, dim, metric_type);
    auto counter_id = schema->AddDebugField("counter", DataType::INT64);
    schema->AddDebugField("double", DataType::DOUBLE);
    schema->AddDebugField("nothing", DataType::INT32);
    schema->set_primary_field_id(counter_id);

    auto dataset = DataGen(schema, N);

    auto fakevec = dataset.get_col<float>(fakevec_id);

    auto segment = CreateSealedSegment(schema);
    Timestamp timestamp = 1000000;
    ScopedSchemaHandle handle(*schema);
    // double >= -1 AND double < 1
    auto plan_str = handle.ParseSearch("double >= -1 && double < 1",
                                       "fakevec",
                                       5,
                                       "L2",
                                       "{\"nprobe\": 10}",
                                       3);
    auto plan =
        CreateSearchPlanByExpr(schema, plan_str.data(), plan_str.size());
    auto num_queries = 5;
    auto ph_group_raw = CreatePlaceholderGroup(num_queries, dim, 1024);
    auto ph_group =
        ParsePlaceholderGroup(plan.get(), ph_group_raw.SerializeAsString());

    ASSERT_ANY_THROW(segment->Search(plan.get(), ph_group.get(), timestamp));

    segment = CreateSealedWithFieldDataLoaded(schema, dataset);

    int64_t row_count = 5;
    std::vector<idx_t> pks{1, 2, 3, 4, 5};
    auto ids = std::make_unique<IdArray>();
    ids->mutable_int_id()->mutable_data()->Add(pks.begin(), pks.end());
    std::vector<Timestamp> timestamps_del{10, 10, 10, 10, 10};

    LoadDeletedRecordInfo info = {timestamps_del.data(), ids.get(), row_count};
    segment->LoadDeletedRecord(info);
    ASSERT_EQ(segment->get_deleted_count(), pks.size())
        << "deleted_count=" << segment->get_deleted_count()
        << " pks_count=" << pks.size() << std::endl;

    // Load overlapping delete records
    row_count += 3;
    pks.insert(pks.end(), {6, 7, 8});
    auto new_ids = std::make_unique<IdArray>();
    new_ids->mutable_int_id()->mutable_data()->Add(pks.begin(), pks.end());
    timestamps_del.insert(timestamps_del.end(), {11, 11, 11});
    LoadDeletedRecordInfo overlap_info = {
        timestamps_del.data(), new_ids.get(), row_count};
    segment->LoadDeletedRecord(overlap_info);
    // NOTE: need to change delete timestamp, so not to hit the cache
    ASSERT_EQ(segment->get_deleted_count(), pks.size())
        << "deleted_count=" << segment->get_deleted_count()
        << " pks_count=" << pks.size() << std::endl;
    BitsetType bitset(N, false);
    auto bitset_view = BitsetTypeView(bitset);
    segment->mask_with_delete(bitset_view, 10, 12);
    ASSERT_EQ(bitset.count(), pks.size())
        << "bitset_count=" << bitset.count() << " pks_count=" << pks.size()
        << std::endl;
}

auto
GenMaxFloatVecs(int N, int dim) {
    std::vector<float> vecs;
    for (int i = 0; i < N; i++) {
        for (int j = 0; j < dim; j++) {
            vecs.push_back(std::numeric_limits<float>::max());
        }
    }
    return vecs;
}

auto
GenRandomFloatVecs(int N, int dim) {
    std::vector<float> vecs;
    srand(time(NULL));
    for (int i = 0; i < N; i++) {
        for (int j = 0; j < dim; j++) {
            vecs.push_back(static_cast<float>(rand()) /
                           static_cast<float>(RAND_MAX));
        }
    }
    return vecs;
}

auto
GenQueryVecs(int N, int dim) {
    std::vector<float> vecs;
    for (int i = 0; i < N; i++) {
        for (int j = 0; j < dim; j++) {
            vecs.push_back(1);
        }
    }
    return vecs;
}

TEST(Sealed, BF) {
    auto schema = std::make_shared<Schema>();
    auto dim = 4;
    auto metric_type = "L2";
    auto fake_id = schema->AddDebugField(
        "fakevec", DataType::VECTOR_FLOAT, dim, metric_type);
    auto i64_fid = schema->AddDebugField("counter", DataType::INT64);
    schema->set_primary_field_id(i64_fid);

    size_t N = 100000;

    auto dataset = DataGen(schema, N);
    std::cout << fake_id.get() << std::endl;
    auto segment = CreateSealedWithFieldDataLoaded(
        schema, dataset, false, {fake_id.get()});

    auto vec_data = GenRandomFloatVecs(N, dim);
    auto field_data = storage::CreateFieldData(
        DataType::VECTOR_FLOAT, DataType::NONE, false, dim);
    field_data->FillFieldData(vec_data.data(), N);
    auto cm = milvus::storage::RemoteChunkManagerSingleton::GetInstance()
                  .GetRemoteChunkManager();
    auto load_info = PrepareSingleFieldInsertBinlog(kCollectionID,
                                                    kPartitionID,
                                                    kSegmentID,
                                                    fake_id.get(),
                                                    {field_data},
                                                    cm);

    segment->LoadFieldData(load_info);

    auto topK = 1;
    ScopedSchemaHandle handle(*schema);
    auto plan_str =
        handle.ParseSearch("", "fakevec", topK, "L2", "{\"nprobe\": 10}");
    auto plan =
        CreateSearchPlanByExpr(schema, plan_str.data(), plan_str.size());

    auto num_queries = 10;
    auto query = GenQueryVecs(num_queries, dim);
    auto ph_group_raw = CreatePlaceholderGroup(num_queries, dim, query);
    auto ph_group =
        ParsePlaceholderGroup(plan.get(), ph_group_raw.SerializeAsString());

    auto result = segment->Search(plan.get(), ph_group.get(), MAX_TIMESTAMP);
    auto ves = SearchResultToVector(*result);
    // first: offset, second: distance
    EXPECT_GE(ves[0].first, 0);
    EXPECT_LE(ves[0].first, N);
    EXPECT_LE(ves[0].second, dim);
}

TEST(Sealed, BF_Overflow) {
    auto schema = std::make_shared<Schema>();
    auto dim = 4;
    auto metric_type = "L2";
    auto fake_id = schema->AddDebugField(
        "fakevec", DataType::VECTOR_FLOAT, dim, metric_type);
    auto i64_fid = schema->AddDebugField("counter", DataType::INT64);
    schema->set_primary_field_id(i64_fid);

    size_t N = 10;

    auto dataset = DataGen(schema, N);
    auto segment = CreateSealedWithFieldDataLoaded(
        schema,
        dataset,
        false,
        GetExcludedFieldIds(schema, {0, 1, i64_fid.get()}));

    auto vec_data = GenMaxFloatVecs(N, dim);
    auto field_data = storage::CreateFieldData(
        DataType::VECTOR_FLOAT, DataType::NONE, false, dim);
    field_data->FillFieldData(vec_data.data(), N);
    auto cm = milvus::storage::RemoteChunkManagerSingleton::GetInstance()
                  .GetRemoteChunkManager();
    auto vec_load_info = PrepareSingleFieldInsertBinlog(kCollectionID,
                                                        kPartitionID,
                                                        kSegmentID,
                                                        fake_id.get(),
                                                        {field_data},
                                                        cm);
    segment->LoadFieldData(vec_load_info);

    auto topK = 1;
    ScopedSchemaHandle handle(*schema);
    auto plan_str =
        handle.ParseSearch("", "fakevec", topK, "L2", "{\"nprobe\": 10}");
    auto plan =
        CreateSearchPlanByExpr(schema, plan_str.data(), plan_str.size());

    auto num_queries = 10;
    auto query = GenQueryVecs(num_queries, dim);
    auto ph_group_raw = CreatePlaceholderGroup(num_queries, dim, query);
    auto ph_group =
        ParsePlaceholderGroup(plan.get(), ph_group_raw.SerializeAsString());

    auto result = segment->Search(plan.get(), ph_group.get(), MAX_TIMESTAMP);
    auto ves = SearchResultToVector(*result);
    for (int i = 0; i < num_queries; ++i) {
        EXPECT_EQ(ves[i].first, -1);
    }
}

TEST(Sealed, DeleteCount) {
    {
        auto schema = std::make_shared<Schema>();
        auto pk = schema->AddDebugField("pk", DataType::INT64);
        schema->set_primary_field_id(pk);
        // empty segment
        size_t N = 10;

        auto dataset = DataGen(schema, N);
        auto segment = CreateSealedWithFieldDataLoaded(schema, dataset);

        int64_t c = 10;
        ASSERT_EQ(segment->get_deleted_count(), 0);

        Timestamp begin_ts = 100;
        auto tss = GenTss(c, begin_ts);
        auto pks = GenPKs(c, N);
        auto status = segment->Delete(c, pks.get(), tss.data());
        ASSERT_TRUE(status.ok());

        ASSERT_EQ(segment->get_deleted_count(), 0);
    }
    {
        auto schema = std::make_shared<Schema>();
        auto pk = schema->AddDebugField("pk", DataType::INT64);
        schema->set_primary_field_id(pk);

        int64_t c = 10;
        auto dataset = DataGen(schema, c);
        auto pks = dataset.get_col<int64_t>(pk);
        auto segment = CreateSealedWithFieldDataLoaded(schema, dataset);

        auto iter = std::max_element(pks.begin(), pks.end());
        auto delete_pks = GenPKs(c, *iter);
        Timestamp begin_ts = 100;
        auto tss = GenTss(c, begin_ts);
        auto status = segment->Delete(c, delete_pks.get(), tss.data());
        ASSERT_TRUE(status.ok());

        // 9 of element should be filtered.
        auto cnt = segment->get_deleted_count();
        ASSERT_EQ(cnt, 1);
    }
}

TEST(Sealed, RealCount) {
    milvus::exec::expression::FunctionFactory& factory =
        milvus::exec::expression::FunctionFactory::Instance();
    factory.Initialize();
    auto schema = std::make_shared<Schema>();
    auto pk = schema->AddDebugField("pk", DataType::INT64);
    schema->set_primary_field_id(pk);
    auto segment = CreateSealedSegment(schema);

    ASSERT_EQ(0, segment->get_real_count());

    int64_t c = 10;
    auto dataset = DataGen(schema, c);
    auto pks = dataset.get_col<int64_t>(pk);
    segment = CreateSealedWithFieldDataLoaded(schema, dataset);

    // no delete.
    ASSERT_EQ(c, segment->get_real_count());

    // delete half.
    auto half = c / 2;
    auto del_ids1 = GenPKs(pks.begin(), pks.begin() + half);
    auto del_tss1 = GenTss(half, c);
    auto status = segment->Delete(half, del_ids1.get(), del_tss1.data());
    ASSERT_TRUE(status.ok());
    ASSERT_EQ(c - half, segment->get_real_count());

    // delete duplicate.
    auto del_tss2 = GenTss(half, c + half);
    status = segment->Delete(half, del_ids1.get(), del_tss2.data());
    ASSERT_TRUE(status.ok());
    ASSERT_EQ(c - half, segment->get_real_count());

    // delete all.
    auto del_ids3 = GenPKs(pks.begin(), pks.end());
    auto del_tss3 = GenTss(c, c + half * 2);
    status = segment->Delete(c, del_ids3.get(), del_tss3.data());
    ASSERT_TRUE(status.ok());
    ASSERT_EQ(0, segment->get_real_count());
}

// Test for DeletedRecord atomic snapshot optimization fast path
// When query_timestamp >= max_delete_timestamp, we can directly use the snapshot bitset
TEST(Sealed, DeleteSnapshotOptimizationFastPath) {
    // Save original value and ensure optimization is enabled
    bool original_value = ENABLE_LATEST_DELETE_SNAPSHOT_OPTIMIZATION.load();
    ENABLE_LATEST_DELETE_SNAPSHOT_OPTIMIZATION.store(true);

    auto schema = std::make_shared<Schema>();
    auto pk = schema->AddDebugField("pk", DataType::INT64);
    schema->set_primary_field_id(pk);

    int64_t N = 100;
    auto dataset = DataGen(schema, N);
    auto pks = dataset.get_col<int64_t>(pk);
    auto segment = CreateSealedWithFieldDataLoaded(schema, dataset);

    // Delete some records with timestamps
    int64_t delete_count = 10;
    auto del_ids = GenPKs(pks.begin(), pks.begin() + delete_count);
    Timestamp delete_ts = 1000;
    auto del_tss = GenTss(delete_count, delete_ts);
    auto status = segment->Delete(delete_count, del_ids.get(), del_tss.data());
    ASSERT_TRUE(status.ok());

    // Query with timestamp >= max_delete_ts should use fast path
    BitsetType bitset1(N, false);
    auto bitset_view1 = BitsetTypeView(bitset1);
    Timestamp query_ts_fast =
        delete_ts + delete_count + 100;  // >= max_delete_ts
    segment->mask_with_delete(bitset_view1, N, query_ts_fast);
    ASSERT_EQ(bitset1.count(), delete_count);

    // Query with timestamp < max_delete_ts should use slow path
    BitsetType bitset2(N, false);
    auto bitset_view2 = BitsetTypeView(bitset2);
    Timestamp query_ts_slow =
        delete_ts + 5;  // < max_delete_ts, covers some but not all
    segment->mask_with_delete(bitset_view2, N, query_ts_slow);
    // Should have fewer deletions visible since query_ts is in the middle
    ASSERT_LT(bitset2.count(), delete_count);
    ASSERT_GT(bitset2.count(), 0);

    // Restore original value
    ENABLE_LATEST_DELETE_SNAPSHOT_OPTIMIZATION.store(original_value);
}

// Test for DeletedRecord atomic snapshot optimization disabled
TEST(Sealed, DeleteSnapshotOptimizationDisabled) {
    // Save original value and disable optimization
    bool original_value = ENABLE_LATEST_DELETE_SNAPSHOT_OPTIMIZATION.load();
    ENABLE_LATEST_DELETE_SNAPSHOT_OPTIMIZATION.store(false);

    auto schema = std::make_shared<Schema>();
    auto pk = schema->AddDebugField("pk", DataType::INT64);
    schema->set_primary_field_id(pk);

    int64_t N = 100;
    auto dataset = DataGen(schema, N);
    auto pks = dataset.get_col<int64_t>(pk);
    auto segment = CreateSealedWithFieldDataLoaded(schema, dataset);

    // Delete some records
    int64_t delete_count = 10;
    auto del_ids = GenPKs(pks.begin(), pks.begin() + delete_count);
    Timestamp delete_ts = 1000;
    auto del_tss = GenTss(delete_count, delete_ts);
    auto status = segment->Delete(delete_count, del_ids.get(), del_tss.data());
    ASSERT_TRUE(status.ok());

    // Query should still work correctly via slow path
    BitsetType bitset(N, false);
    auto bitset_view = BitsetTypeView(bitset);
    Timestamp query_ts = delete_ts + delete_count + 100;
    segment->mask_with_delete(bitset_view, N, query_ts);
    ASSERT_EQ(bitset.count(), delete_count);

    // Restore original value
    ENABLE_LATEST_DELETE_SNAPSHOT_OPTIMIZATION.store(original_value);
}

// Test for multiple sequential deletes updating snapshot correctly
TEST(Sealed, DeleteSnapshotMultipleDeletes) {
    bool original_value = ENABLE_LATEST_DELETE_SNAPSHOT_OPTIMIZATION.load();
    ENABLE_LATEST_DELETE_SNAPSHOT_OPTIMIZATION.store(true);

    auto schema = std::make_shared<Schema>();
    auto pk = schema->AddDebugField("pk", DataType::INT64);
    schema->set_primary_field_id(pk);

    int64_t N = 100;
    auto dataset = DataGen(schema, N);
    auto pks = dataset.get_col<int64_t>(pk);
    auto segment = CreateSealedWithFieldDataLoaded(schema, dataset);

    // First batch of deletes
    int64_t delete_count1 = 5;
    auto del_ids1 = GenPKs(pks.begin(), pks.begin() + delete_count1);
    Timestamp delete_ts1 = 1000;
    auto del_tss1 = GenTss(delete_count1, delete_ts1);
    auto status =
        segment->Delete(delete_count1, del_ids1.get(), del_tss1.data());
    ASSERT_TRUE(status.ok());

    // Query after first delete batch
    BitsetType bitset1(N, false);
    auto bitset_view1 = BitsetTypeView(bitset1);
    Timestamp query_ts1 = delete_ts1 + delete_count1 + 100;
    segment->mask_with_delete(bitset_view1, N, query_ts1);
    ASSERT_EQ(bitset1.count(), delete_count1);

    // Second batch of deletes
    int64_t delete_count2 = 5;
    auto del_ids2 = GenPKs(pks.begin() + delete_count1,
                           pks.begin() + delete_count1 + delete_count2);
    Timestamp delete_ts2 = 2000;
    auto del_tss2 = GenTss(delete_count2, delete_ts2);
    status = segment->Delete(delete_count2, del_ids2.get(), del_tss2.data());
    ASSERT_TRUE(status.ok());

    // Query after second delete batch - should see all deletes
    BitsetType bitset2(N, false);
    auto bitset_view2 = BitsetTypeView(bitset2);
    Timestamp query_ts2 = delete_ts2 + delete_count2 + 100;
    segment->mask_with_delete(bitset_view2, N, query_ts2);
    ASSERT_EQ(bitset2.count(), delete_count1 + delete_count2);

    // Query with timestamp between batches - should only see first batch
    BitsetType bitset3(N, false);
    auto bitset_view3 = BitsetTypeView(bitset3);
    Timestamp query_ts_between =
        delete_ts1 + delete_count1 + 50;  // Between batch1 and batch2
    segment->mask_with_delete(bitset_view3, N, query_ts_between);
    ASSERT_EQ(bitset3.count(), delete_count1);

    ENABLE_LATEST_DELETE_SNAPSHOT_OPTIMIZATION.store(original_value);
}

TEST(Sealed, GetVector) {
    auto dim = 4;
    auto N = ROW_COUNT;
    auto metric_type = knowhere::metric::L2;
    auto schema = std::make_shared<Schema>();
    auto fakevec_id = schema->AddDebugField(
        "fakevec", DataType::VECTOR_FLOAT, dim, metric_type);
    auto counter_id = schema->AddDebugField("counter", DataType::INT64);
    schema->AddDebugField("int8", DataType::INT8);
    schema->AddDebugField("int16", DataType::INT16);
    schema->AddDebugField("float", DataType::FLOAT);
    schema->set_primary_field_id(counter_id);

    auto dataset = DataGen(schema, N);

    auto fakevec = dataset.get_col<float>(fakevec_id);

    auto indexing = GenVecIndexing(
        N, dim, fakevec.data(), knowhere::IndexEnum::INDEX_FAISS_IVFFLAT);

    auto segment_sealed = CreateSealedSegment(schema);

    LoadIndexInfo vec_info;
    vec_info.field_id = fakevec_id.get();
    vec_info.index_params = GenIndexParams(indexing.get());
    vec_info.cache_index = CreateTestCacheIndex("test", std::move(indexing));
    vec_info.index_params["metric_type"] = knowhere::metric::L2;
    segment_sealed->LoadIndex(vec_info);

    auto segment =
        dynamic_cast<ChunkedSegmentSealedImpl*>(segment_sealed.get());

    auto has = segment->HasRawData(vec_info.field_id);
    EXPECT_TRUE(has);

    auto ids_ds = GenRandomIds(N);
    auto result = segment->get_vector(nullptr, fakevec_id, ids_ds->GetIds(), N);

    auto vector = result.get()->mutable_vectors()->float_vector().data();
    EXPECT_TRUE(vector.size() == fakevec.size());
    for (size_t i = 0; i < N; ++i) {
        auto id = ids_ds->GetIds()[i];
        for (size_t j = 0; j < dim; ++j) {
            EXPECT_TRUE(vector[i * dim + j] == fakevec[id * dim + j]);
        }
    }
}

TEST(Sealed, LoadArrayFieldData) {
    auto dim = 4;
    auto N = 10;
    auto metric_type = knowhere::metric::L2;
    auto schema = std::make_shared<Schema>();
    auto fakevec_id = schema->AddDebugField(
        "fakevec", DataType::VECTOR_FLOAT, dim, metric_type);
    auto counter_id = schema->AddDebugField("counter", DataType::INT64);
    auto array_id =
        schema->AddDebugField("array", DataType::ARRAY, DataType::INT64);
    schema->set_primary_field_id(counter_id);

    auto dataset = DataGen(schema, N);
    auto fakevec = dataset.get_col<float>(fakevec_id);
    auto segment = CreateSealedSegment(schema);

    ScopedSchemaHandle handle(*schema);
    // array contains 1
    auto plan_str = handle.ParseSearch(
        "array_contains(array, 1)", "fakevec", 5, "L2", "{\"nprobe\": 10}", 3);
    auto plan =
        CreateSearchPlanByExpr(schema, plan_str.data(), plan_str.size());
    auto num_queries = 5;
    auto ph_group_raw = CreatePlaceholderGroup(num_queries, dim, 1024);
    auto ph_group =
        ParsePlaceholderGroup(plan.get(), ph_group_raw.SerializeAsString());

    segment = CreateSealedWithFieldDataLoaded(schema, dataset);
    segment->Search(plan.get(), ph_group.get(), MAX_TIMESTAMP);

    auto ids_ds = GenRandomIds(N);
    auto s = dynamic_cast<ChunkedSegmentSealedImpl*>(segment.get());
    auto int64_result =
        s->bulk_subscript(nullptr, array_id, ids_ds->GetIds(), N);
    auto result_count = int64_result->scalars().array_data().data().size();
    ASSERT_EQ(result_count, N);
}

TEST(Sealed, LoadArrayFieldDataWhenIndexHasRawData) {
    auto dim = 4;
    auto N = 10;
    auto metric_type = knowhere::metric::L2;
    auto schema = std::make_shared<Schema>();
    schema->AddDebugField("fakevec", DataType::VECTOR_FLOAT, dim, metric_type);
    auto counter_id = schema->AddDebugField("counter", DataType::INT64);
    auto array_id =
        schema->AddDebugField("array", DataType::ARRAY, DataType::INT64);
    schema->set_primary_field_id(counter_id);

    auto dataset = DataGen(schema, N);
    auto segment = CreateSealedWithFieldDataLoaded(schema, dataset);

    auto counter_data = dataset.get_col<int64_t>(counter_id);
    auto indexing = GenScalarIndexing<int64_t>(N, counter_data.data());

    LoadIndexInfo array_index;
    array_index.field_id = array_id.get();
    array_index.field_type = DataType::ARRAY;
    array_index.element_type = DataType::INT64;
    array_index.index_params = GenIndexParams(indexing.get());
    LoadResourceRequest request{};
    request.has_raw_data = true;
    array_index.load_resource_request = request;
    array_index.cache_index =
        CreateTestCacheIndex("array_raw_index", std::move(indexing));
    segment->LoadIndex(array_index);

    auto ids_ds = GenRandomIds(N);
    auto s = dynamic_cast<ChunkedSegmentSealedImpl*>(segment.get());
    auto array_result =
        s->bulk_subscript(nullptr, array_id, ids_ds->GetIds(), N);
    ASSERT_EQ(array_result->type(), proto::schema::DataType::Array);
    ASSERT_EQ(array_result->scalars().array_data().data_size(), N);
}

TEST(Sealed, LoadArrayFieldDataWithMMap) {
    auto dim = 4;
    auto N = ROW_COUNT;
    auto metric_type = knowhere::metric::L2;
    auto schema = std::make_shared<Schema>();
    auto fakevec_id = schema->AddDebugField(
        "fakevec", DataType::VECTOR_FLOAT, dim, metric_type);
    auto counter_id = schema->AddDebugField("counter", DataType::INT64);
    schema->AddDebugField("array", DataType::ARRAY, DataType::INT64);
    schema->set_primary_field_id(counter_id);

    auto dataset = DataGen(schema, N);
    auto fakevec = dataset.get_col<float>(fakevec_id);
    auto segment = CreateSealedSegment(schema);

    ScopedSchemaHandle handle(*schema);
    // array contains 1
    auto plan_str = handle.ParseSearch(
        "array_contains(array, 1)", "fakevec", 5, "L2", "{\"nprobe\": 10}", 3);
    auto plan =
        CreateSearchPlanByExpr(schema, plan_str.data(), plan_str.size());
    auto num_queries = 5;
    auto ph_group_raw = CreatePlaceholderGroup(num_queries, dim, 1024);
    auto ph_group =
        ParsePlaceholderGroup(plan.get(), ph_group_raw.SerializeAsString());

    segment = CreateSealedWithFieldDataLoaded(schema, dataset, true);
    segment->Search(plan.get(), ph_group.get(), MAX_TIMESTAMP);
}

TEST(Sealed, SkipIndexSkipUnaryRange) {
    auto schema = std::make_shared<Schema>();
    auto dim = 4;
    auto metrics_type = "L2";
    schema->AddDebugField("fakeVec", DataType::VECTOR_FLOAT, dim, metrics_type);
    auto pk_fid = schema->AddDebugField("pk", DataType::INT64);
    auto i32_fid = schema->AddDebugField("int32_field", DataType::INT32);
    auto i16_fid = schema->AddDebugField("int16_field", DataType::INT16);
    auto i8_fid = schema->AddDebugField("int8_field", DataType::INT8);
    auto float_fid = schema->AddDebugField("float_field", DataType::FLOAT);
    auto double_fid = schema->AddDebugField("double_field", DataType::DOUBLE);
    size_t N = 10;
    auto dataset = DataGen(schema, N);
    auto segment = CreateSealedSegment(schema);
    auto cm = milvus::storage::RemoteChunkManagerSingleton::GetInstance()
                  .GetRemoteChunkManager();
    std::cout << "pk_fid:" << pk_fid.get() << std::endl;

    //test for int64
    std::vector<int64_t> pks = {1, 2, 3, 4, 5, 6, 7, 8, 9, 10};
    auto pk_field_data =
        storage::CreateFieldData(DataType::INT64, DataType::NONE, false, 1, 10);
    pk_field_data->FillFieldData(pks.data(), N);
    auto load_info = PrepareSingleFieldInsertBinlog(kCollectionID,
                                                    kPartitionID,
                                                    kSegmentID,
                                                    pk_fid.get(),
                                                    {pk_field_data},
                                                    cm);
    segment->LoadFieldData(load_info);
    auto& skip_index = segment->GetSkipIndex();
    bool equal_5_skip =
        skip_index.CanSkipUnaryRange<int64_t>(pk_fid, 0, OpType::Equal, 5);
    bool equal_12_skip =
        skip_index.CanSkipUnaryRange<int64_t>(pk_fid, 0, OpType::Equal, 12);
    bool equal_10_skip =
        skip_index.CanSkipUnaryRange<int64_t>(pk_fid, 0, OpType::Equal, 10);
    ASSERT_FALSE(equal_5_skip);
    ASSERT_TRUE(equal_12_skip);
    ASSERT_FALSE(equal_10_skip);
    bool less_than_1_skip =
        skip_index.CanSkipUnaryRange<int64_t>(pk_fid, 0, OpType::LessThan, 1);
    bool less_than_5_skip =
        skip_index.CanSkipUnaryRange<int64_t>(pk_fid, 0, OpType::LessThan, 5);
    ASSERT_TRUE(less_than_1_skip);
    ASSERT_FALSE(less_than_5_skip);
    bool less_equal_than_1_skip =
        skip_index.CanSkipUnaryRange<int64_t>(pk_fid, 0, OpType::LessEqual, 1);
    bool less_equal_than_15_skip =
        skip_index.CanSkipUnaryRange<int64_t>(pk_fid, 0, OpType::LessThan, 15);
    ASSERT_FALSE(less_equal_than_1_skip);
    ASSERT_FALSE(less_equal_than_15_skip);
    bool greater_than_10_skip = skip_index.CanSkipUnaryRange<int64_t>(
        pk_fid, 0, OpType::GreaterThan, 10);
    bool greater_than_5_skip = skip_index.CanSkipUnaryRange<int64_t>(
        pk_fid, 0, OpType::GreaterThan, 5);
    ASSERT_TRUE(greater_than_10_skip);
    ASSERT_FALSE(greater_than_5_skip);
    bool greater_equal_than_10_skip = skip_index.CanSkipUnaryRange<int64_t>(
        pk_fid, 0, OpType::GreaterEqual, 10);
    bool greater_equal_than_5_skip = skip_index.CanSkipUnaryRange<int64_t>(
        pk_fid, 0, OpType::GreaterEqual, 5);
    ASSERT_FALSE(greater_equal_than_10_skip);
    ASSERT_FALSE(greater_equal_than_5_skip);

    //test for int32
    std::vector<int32_t> int32s = {2, 2, 3, 4, 5, 6, 7, 8, 9, 12};
    auto int32_field_data =
        storage::CreateFieldData(DataType::INT32, DataType::NONE, false, 1, 10);
    int32_field_data->FillFieldData(int32s.data(), N);
    load_info = PrepareSingleFieldInsertBinlog(kCollectionID,
                                               kPartitionID,
                                               kSegmentID,
                                               i32_fid.get(),
                                               {int32_field_data},
                                               cm);
    segment->LoadFieldData(load_info);
    less_than_1_skip =
        skip_index.CanSkipUnaryRange<int32_t>(i32_fid, 0, OpType::LessThan, 1);
    ASSERT_TRUE(less_than_1_skip);

    //test for int16
    std::vector<int16_t> int16s = {2, 2, 3, 4, 5, 6, 7, 8, 9, 12};
    auto int16_field_data =
        storage::CreateFieldData(DataType::INT16, DataType::NONE, false, 1, 10);
    int16_field_data->FillFieldData(int16s.data(), N);
    load_info = PrepareSingleFieldInsertBinlog(kCollectionID,
                                               kPartitionID,
                                               kSegmentID,
                                               i16_fid.get(),
                                               {int16_field_data},
                                               cm);
    segment->LoadFieldData(load_info);
    bool less_than_12_skip =
        skip_index.CanSkipUnaryRange<int16_t>(i16_fid, 0, OpType::LessThan, 12);
    ASSERT_FALSE(less_than_12_skip);

    //test for int8
    std::vector<int8_t> int8s = {2, 2, 3, 4, 5, 6, 7, 8, 9, 12};
    auto int8_field_data =
        storage::CreateFieldData(DataType::INT8, DataType::NONE, false, 1, 10);
    int8_field_data->FillFieldData(int8s.data(), N);
    load_info = PrepareSingleFieldInsertBinlog(kCollectionID,
                                               kPartitionID,
                                               kSegmentID,
                                               i8_fid.get(),
                                               {int8_field_data},
                                               cm);
    segment->LoadFieldData(load_info);
    bool greater_than_12_skip = skip_index.CanSkipUnaryRange<int8_t>(
        i8_fid, 0, OpType::GreaterThan, 12);
    ASSERT_TRUE(greater_than_12_skip);

    // test for float
    std::vector<float> floats = {
        1.0, 2.0, 3.0, 4.0, 5.0, 6.0, 7.0, 8.0, 9.0, 10.0};
    auto float_field_data =
        storage::CreateFieldData(DataType::FLOAT, DataType::NONE, false, 1, 10);
    float_field_data->FillFieldData(floats.data(), N);
    load_info = PrepareSingleFieldInsertBinlog(kCollectionID,
                                               kPartitionID,
                                               kSegmentID,
                                               float_fid.get(),
                                               {float_field_data},
                                               cm);
    segment->LoadFieldData(load_info);
    greater_than_10_skip = skip_index.CanSkipUnaryRange<float>(
        float_fid, 0, OpType::GreaterThan, 10.0);
    ASSERT_TRUE(greater_than_10_skip);

    // test for double
    std::vector<double> doubles = {
        1.0, 2.0, 3.0, 4.0, 5.0, 6.0, 7.0, 8.0, 9.0, 10.0};
    auto double_field_data = storage::CreateFieldData(
        DataType::DOUBLE, DataType::NONE, false, 1, 10);
    double_field_data->FillFieldData(doubles.data(), N);
    load_info = PrepareSingleFieldInsertBinlog(kCollectionID,
                                               kPartitionID,
                                               kSegmentID,
                                               double_fid.get(),
                                               {double_field_data},
                                               cm);
    segment->LoadFieldData(load_info);
    greater_than_10_skip = skip_index.CanSkipUnaryRange<double>(
        double_fid, 0, OpType::GreaterThan, 10.0);
    ASSERT_TRUE(greater_than_10_skip);
}

TEST(Sealed, SkipIndexSkipBinaryRange) {
    auto schema = std::make_shared<Schema>();
    auto dim = 4;
    auto metrics_type = "L2";
    schema->AddDebugField("fakeVec", DataType::VECTOR_FLOAT, dim, metrics_type);
    auto pk_fid = schema->AddDebugField("pk", DataType::INT64);
    size_t N = 10;
    auto dataset = DataGen(schema, N);
    auto segment = CreateSealedSegment(schema);
    auto cm = milvus::storage::RemoteChunkManagerSingleton::GetInstance()
                  .GetRemoteChunkManager();
    std::cout << "pk_fid:" << pk_fid.get() << std::endl;

    //test for int64
    std::vector<int64_t> pks = {1, 2, 3, 4, 5, 6, 7, 8, 9, 10};
    auto pk_field_data =
        storage::CreateFieldData(DataType::INT64, DataType::NONE, false, 1, 10);
    pk_field_data->FillFieldData(pks.data(), N);
    auto load_info = PrepareSingleFieldInsertBinlog(kCollectionID,
                                                    kPartitionID,
                                                    kSegmentID,
                                                    pk_fid.get(),
                                                    {pk_field_data},
                                                    cm);
    segment->LoadFieldData(load_info);
    auto& skip_index = segment->GetSkipIndex();
    ASSERT_FALSE(
        skip_index.CanSkipBinaryRange<int64_t>(pk_fid, 0, -3, 1, true, true));
    ASSERT_TRUE(
        skip_index.CanSkipBinaryRange<int64_t>(pk_fid, 0, -3, 1, true, false));

    ASSERT_FALSE(
        skip_index.CanSkipBinaryRange<int64_t>(pk_fid, 0, 7, 9, true, true));
    ASSERT_FALSE(
        skip_index.CanSkipBinaryRange<int64_t>(pk_fid, 0, 8, 12, true, false));

    ASSERT_TRUE(
        skip_index.CanSkipBinaryRange<int64_t>(pk_fid, 0, 10, 12, false, true));
    ASSERT_FALSE(
        skip_index.CanSkipBinaryRange<int64_t>(pk_fid, 0, 10, 12, true, true));
}

TEST(Sealed, SkipIndexSkipUnaryRangeNullable) {
    auto schema = std::make_shared<Schema>();
    auto dim = 4;
    auto metrics_type = "L2";
    schema->AddDebugField("fakeVec", DataType::VECTOR_FLOAT, dim, metrics_type);
    auto i64_fid = schema->AddDebugField("int64_field", DataType::INT64, true);

    auto dataset = DataGen(schema, 5);
    auto segment = CreateSealedSegment(schema);
    auto cm = milvus::storage::RemoteChunkManagerSingleton::GetInstance()
                  .GetRemoteChunkManager();

    //test for int64
    std::vector<int64_t> int64s = {1, 2, 3, 4, 5};
    std::array<uint8_t, 1> valid_data = {0x03};
    auto int64s_field_data =
        storage::CreateFieldData(DataType::INT64, DataType::NONE, true, 1, 5);

    int64s_field_data->FillFieldData(int64s.data(), valid_data.data(), 5, 0);
    auto load_info = PrepareSingleFieldInsertBinlog(kCollectionID,
                                                    kPartitionID,
                                                    kSegmentID,
                                                    i64_fid.get(),
                                                    {int64s_field_data},
                                                    cm);
    segment->LoadFieldData(load_info);
    auto& skip_index = segment->GetSkipIndex();
    bool equal_5_skip =
        skip_index.CanSkipUnaryRange<int64_t>(i64_fid, 0, OpType::Equal, 5);
    bool equal_4_skip =
        skip_index.CanSkipUnaryRange<int64_t>(i64_fid, 0, OpType::Equal, 4);
    bool equal_2_skip =
        skip_index.CanSkipUnaryRange<int64_t>(i64_fid, 0, OpType::Equal, 2);
    bool equal_1_skip =
        skip_index.CanSkipUnaryRange<int64_t>(i64_fid, 0, OpType::Equal, 1);
    ASSERT_TRUE(equal_5_skip);
    ASSERT_TRUE(equal_4_skip);
    ASSERT_FALSE(equal_2_skip);
    ASSERT_FALSE(equal_1_skip);
    bool less_than_1_skip =
        skip_index.CanSkipUnaryRange<int64_t>(i64_fid, 0, OpType::LessThan, 1);
    bool less_than_5_skip =
        skip_index.CanSkipUnaryRange<int64_t>(i64_fid, 0, OpType::LessThan, 5);
    ASSERT_TRUE(less_than_1_skip);
    ASSERT_FALSE(less_than_5_skip);
    bool less_equal_than_1_skip =
        skip_index.CanSkipUnaryRange<int64_t>(i64_fid, 0, OpType::LessEqual, 1);
    bool less_equal_than_15_skip =
        skip_index.CanSkipUnaryRange<int64_t>(i64_fid, 0, OpType::LessThan, 15);
    ASSERT_FALSE(less_equal_than_1_skip);
    ASSERT_FALSE(less_equal_than_15_skip);
    bool greater_than_10_skip = skip_index.CanSkipUnaryRange<int64_t>(
        i64_fid, 0, OpType::GreaterThan, 10);
    bool greater_than_5_skip = skip_index.CanSkipUnaryRange<int64_t>(
        i64_fid, 0, OpType::GreaterThan, 5);
    bool greater_than_2_skip = skip_index.CanSkipUnaryRange<int64_t>(
        i64_fid, 0, OpType::GreaterThan, 2);
    bool greater_than_1_skip = skip_index.CanSkipUnaryRange<int64_t>(
        i64_fid, 0, OpType::GreaterThan, 1);
    ASSERT_TRUE(greater_than_10_skip);
    ASSERT_TRUE(greater_than_5_skip);
    ASSERT_TRUE(greater_than_2_skip);
    ASSERT_FALSE(greater_than_1_skip);
    bool greater_equal_than_3_skip = skip_index.CanSkipUnaryRange<int64_t>(
        i64_fid, 0, OpType::GreaterEqual, 3);
    bool greater_equal_than_2_skip = skip_index.CanSkipUnaryRange<int64_t>(
        i64_fid, 0, OpType::GreaterEqual, 2);
    ASSERT_TRUE(greater_equal_than_3_skip);
    ASSERT_FALSE(greater_equal_than_2_skip);
}

TEST(Sealed, SkipIndexSkipBinaryRangeNullable) {
    auto schema = std::make_shared<Schema>();
    auto dim = 4;
    auto metrics_type = "L2";
    schema->AddDebugField("fakeVec", DataType::VECTOR_FLOAT, dim, metrics_type);
    auto i64_fid = schema->AddDebugField("int64_field", DataType::INT64, true);
    auto dataset = DataGen(schema, 5);
    auto segment = CreateSealedSegment(schema);
    auto cm = milvus::storage::RemoteChunkManagerSingleton::GetInstance()
                  .GetRemoteChunkManager();

    //test for int64
    std::vector<int64_t> int64s = {1, 2, 3, 4, 5};
    std::array<uint8_t, 1> valid_data = {0x03};
    auto int64s_field_data =
        storage::CreateFieldData(DataType::INT64, DataType::NONE, true, 1, 5);

    int64s_field_data->FillFieldData(int64s.data(), valid_data.data(), 5, 0);
    auto load_info = PrepareSingleFieldInsertBinlog(kCollectionID,
                                                    kPartitionID,
                                                    kSegmentID,
                                                    i64_fid.get(),
                                                    {int64s_field_data},
                                                    cm);
    segment->LoadFieldData(load_info);
    auto& skip_index = segment->GetSkipIndex();
    ASSERT_FALSE(
        skip_index.CanSkipBinaryRange<int64_t>(i64_fid, 0, -3, 1, true, true));
    ASSERT_TRUE(
        skip_index.CanSkipBinaryRange<int64_t>(i64_fid, 0, -3, 1, true, false));

    ASSERT_FALSE(
        skip_index.CanSkipBinaryRange<int64_t>(i64_fid, 0, 1, 3, true, true));
    ASSERT_FALSE(
        skip_index.CanSkipBinaryRange<int64_t>(i64_fid, 0, 1, 2, true, false));

    ASSERT_TRUE(
        skip_index.CanSkipBinaryRange<int64_t>(i64_fid, 0, 2, 3, false, true));
    ASSERT_FALSE(
        skip_index.CanSkipBinaryRange<int64_t>(i64_fid, 0, 2, 3, true, true));
}

TEST(Sealed, SkipIndexSkipStringRange) {
    auto schema = std::make_shared<Schema>();
    auto dim = 4;
    auto metrics_type = "L2";
    schema->AddDebugField("pk", DataType::INT64);
    auto string_fid = schema->AddDebugField("string_field", DataType::VARCHAR);
    schema->AddDebugField("fakeVec", DataType::VECTOR_FLOAT, dim, metrics_type);
    size_t N = 5;
    auto dataset = DataGen(schema, N);
    auto segment = CreateSealedSegment(schema);

    //test for string
    std::vector<std::string> strings = {"e", "f", "g", "g", "j"};
    auto string_field_data = storage::CreateFieldData(
        DataType::VARCHAR, DataType::NONE, false, 1, N);
    string_field_data->FillFieldData(strings.data(), N);
    auto cm = milvus::storage::RemoteChunkManagerSingleton::GetInstance()
                  .GetRemoteChunkManager();
    auto load_info = PrepareSingleFieldInsertBinlog(kCollectionID,
                                                    kPartitionID,
                                                    kSegmentID,
                                                    string_fid.get(),
                                                    {string_field_data},
                                                    cm);
    segment->LoadFieldData(load_info);
    auto& skip_index = segment->GetSkipIndex();
    ASSERT_TRUE(skip_index.CanSkipUnaryRange<std::string>(
        string_fid, 0, OpType::Equal, "w"));
    ASSERT_FALSE(skip_index.CanSkipUnaryRange<std::string>(
        string_fid, 0, OpType::Equal, "e"));
    ASSERT_FALSE(skip_index.CanSkipUnaryRange<std::string>(
        string_fid, 0, OpType::Equal, "j"));

    ASSERT_TRUE(skip_index.CanSkipUnaryRange<std::string>(
        string_fid, 0, OpType::LessThan, "e"));
    ASSERT_FALSE(skip_index.CanSkipUnaryRange<std::string>(
        string_fid, 0, OpType::LessEqual, "e"));

    ASSERT_TRUE(skip_index.CanSkipUnaryRange<std::string>(
        string_fid, 0, OpType::GreaterThan, "j"));
    ASSERT_FALSE(skip_index.CanSkipUnaryRange<std::string>(
        string_fid, 0, OpType::GreaterEqual, "j"));
    ASSERT_FALSE(skip_index.CanSkipUnaryRange<int64_t>(
        string_fid, 0, OpType::GreaterEqual, 1));

    ASSERT_TRUE(skip_index.CanSkipBinaryRange<std::string>(
        string_fid, 0, "a", "c", true, true));
    ASSERT_TRUE(skip_index.CanSkipBinaryRange<std::string>(
        string_fid, 0, "c", "e", true, false));
    ASSERT_FALSE(skip_index.CanSkipBinaryRange<std::string>(
        string_fid, 0, "c", "e", true, true));
    ASSERT_FALSE(skip_index.CanSkipBinaryRange<std::string>(
        string_fid, 0, "e", "k", false, true));
    ASSERT_FALSE(skip_index.CanSkipBinaryRange<std::string>(
        string_fid, 0, "j", "k", true, true));
    ASSERT_TRUE(skip_index.CanSkipBinaryRange<std::string>(
        string_fid, 0, "j", "k", false, true));
    ASSERT_FALSE(skip_index.CanSkipBinaryRange<int64_t>(
        string_fid, 0, 1, 2, false, true));
}

TEST(Sealed, QueryAllFields) {
    auto schema = std::make_shared<Schema>();
    auto metric_type = knowhere::metric::L2;
    auto dim = 4;
    auto bool_field = schema->AddDebugField("bool", DataType::BOOL);
    auto int8_field = schema->AddDebugField("int8", DataType::INT8);
    auto int16_field = schema->AddDebugField("int16", DataType::INT16);
    auto int32_field = schema->AddDebugField("int32", DataType::INT32);
    auto int64_field = schema->AddDebugField("int64", DataType::INT64);
    auto float_field = schema->AddDebugField("float", DataType::FLOAT);
    auto double_field = schema->AddDebugField("double", DataType::DOUBLE);
    auto varchar_field = schema->AddDebugField("varchar", DataType::VARCHAR);
    auto json_field = schema->AddDebugField("json", DataType::JSON);
    auto geometry_field = schema->AddDebugField("geometry", DataType::GEOMETRY);
    auto int_array_field =
        schema->AddDebugField("int_array", DataType::ARRAY, DataType::INT8);
    auto long_array_field =
        schema->AddDebugField("long_array", DataType::ARRAY, DataType::INT64);
    auto bool_array_field =
        schema->AddDebugField("bool_array", DataType::ARRAY, DataType::BOOL);
    auto string_array_field = schema->AddDebugField(
        "string_array", DataType::ARRAY, DataType::VARCHAR);
    auto double_array_field = schema->AddDebugField(
        "double_array", DataType::ARRAY, DataType::DOUBLE);
    auto float_array_field =
        schema->AddDebugField("float_array", DataType::ARRAY, DataType::FLOAT);
    auto vec = schema->AddDebugField(
        "embeddings", DataType::VECTOR_FLOAT, dim, metric_type);
    auto float16_vec = schema->AddDebugField(
        "float16_vec", DataType::VECTOR_FLOAT16, dim, metric_type);
    auto bfloat16_vec = schema->AddDebugField(
        "bfloat16_vec", DataType::VECTOR_BFLOAT16, dim, metric_type);
    auto int8_vec = schema->AddDebugField(
        "int8_vec", DataType::VECTOR_INT8, dim, metric_type);
    schema->set_primary_field_id(int64_field);

    std::map<std::string, std::string> index_params = {
        {"index_type", "IVF_FLAT"},
        {"metric_type", metric_type},
        {"nlist", "128"}};
    std::map<std::string, std::string> type_params = {
        {"dim", std::to_string(dim)}};
    FieldIndexMeta fieldIndexMeta(
        vec, std::move(index_params), std::move(type_params));
    std::map<FieldId, FieldIndexMeta> filedMap = {{vec, fieldIndexMeta}};
    IndexMetaPtr metaPtr =
        std::make_shared<CollectionIndexMeta>(100000, std::move(filedMap));
    auto segment_sealed = CreateSealedSegment(schema, metaPtr);
    auto segment =
        dynamic_cast<ChunkedSegmentSealedImpl*>(segment_sealed.get());

    int64_t dataset_size = 1000;
    auto dataset = DataGen(schema, dataset_size);
    segment_sealed = CreateSealedWithFieldDataLoaded(schema, dataset);
    segment = dynamic_cast<ChunkedSegmentSealedImpl*>(segment_sealed.get());

    auto bool_values = dataset.get_col<bool>(bool_field);
    auto int8_values = dataset.get_col<int8_t>(int8_field);
    auto int16_values = dataset.get_col<int16_t>(int16_field);
    auto int32_values = dataset.get_col<int32_t>(int32_field);
    auto int64_values = dataset.get_col<int64_t>(int64_field);
    auto float_values = dataset.get_col<float>(float_field);
    auto double_values = dataset.get_col<double>(double_field);
    auto varchar_values = dataset.get_col<std::string>(varchar_field);
    auto json_values = dataset.get_col<std::string>(json_field);
    auto geometry_values = dataset.get_col<std::string>(geometry_field);
    auto int_array_values = dataset.get_col<ScalarFieldProto>(int_array_field);
    auto long_array_values =
        dataset.get_col<ScalarFieldProto>(long_array_field);
    auto bool_array_values =
        dataset.get_col<ScalarFieldProto>(bool_array_field);
    auto string_array_values =
        dataset.get_col<ScalarFieldProto>(string_array_field);
    auto double_array_values =
        dataset.get_col<ScalarFieldProto>(double_array_field);
    auto float_array_values =
        dataset.get_col<ScalarFieldProto>(float_array_field);
    auto vector_values = dataset.get_col<float>(vec);
    auto float16_vector_values = dataset.get_col<uint8_t>(float16_vec);
    auto bfloat16_vector_values = dataset.get_col<uint8_t>(bfloat16_vec);
    auto int8_vector_values = dataset.get_col<int8>(int8_vec);

    auto ids_ds = GenRandomIds(dataset_size);
    auto bool_result = segment->bulk_subscript(
        nullptr, bool_field, ids_ds->GetIds(), dataset_size);
    auto int8_result = segment->bulk_subscript(
        nullptr, int8_field, ids_ds->GetIds(), dataset_size);
    auto int16_result = segment->bulk_subscript(
        nullptr, int16_field, ids_ds->GetIds(), dataset_size);
    auto int32_result = segment->bulk_subscript(
        nullptr, int32_field, ids_ds->GetIds(), dataset_size);
    auto int64_result = segment->bulk_subscript(
        nullptr, int64_field, ids_ds->GetIds(), dataset_size);
    auto float_result = segment->bulk_subscript(
        nullptr, float_field, ids_ds->GetIds(), dataset_size);
    auto double_result = segment->bulk_subscript(
        nullptr, double_field, ids_ds->GetIds(), dataset_size);
    auto varchar_result = segment->bulk_subscript(
        nullptr, varchar_field, ids_ds->GetIds(), dataset_size);
    auto json_result = segment->bulk_subscript(
        nullptr, json_field, ids_ds->GetIds(), dataset_size);
    auto geometry_result = segment->bulk_subscript(
        nullptr, geometry_field, ids_ds->GetIds(), dataset_size);
    auto int_array_result = segment->bulk_subscript(
        nullptr, int_array_field, ids_ds->GetIds(), dataset_size);
    auto long_array_result = segment->bulk_subscript(
        nullptr, long_array_field, ids_ds->GetIds(), dataset_size);
    auto bool_array_result = segment->bulk_subscript(
        nullptr, bool_array_field, ids_ds->GetIds(), dataset_size);
    auto string_array_result = segment->bulk_subscript(
        nullptr, string_array_field, ids_ds->GetIds(), dataset_size);
    auto double_array_result = segment->bulk_subscript(
        nullptr, double_array_field, ids_ds->GetIds(), dataset_size);
    auto float_array_result = segment->bulk_subscript(
        nullptr, float_array_field, ids_ds->GetIds(), dataset_size);
    auto vec_result =
        segment->bulk_subscript(nullptr, vec, ids_ds->GetIds(), dataset_size);
    auto float16_vec_result = segment->bulk_subscript(
        nullptr, float16_vec, ids_ds->GetIds(), dataset_size);
    auto bfloat16_vec_result = segment->bulk_subscript(
        nullptr, bfloat16_vec, ids_ds->GetIds(), dataset_size);
    auto int8_vec_result = segment->bulk_subscript(
        nullptr, int8_vec, ids_ds->GetIds(), dataset_size);

    EXPECT_EQ(bool_result->scalars().bool_data().data_size(), dataset_size);
    EXPECT_EQ(int8_result->scalars().int_data().data_size(), dataset_size);
    EXPECT_EQ(int16_result->scalars().int_data().data_size(), dataset_size);
    EXPECT_EQ(int32_result->scalars().int_data().data_size(), dataset_size);
    EXPECT_EQ(int64_result->scalars().long_data().data_size(), dataset_size);
    EXPECT_EQ(float_result->scalars().float_data().data_size(), dataset_size);
    EXPECT_EQ(double_result->scalars().double_data().data_size(), dataset_size);
    EXPECT_EQ(varchar_result->scalars().string_data().data_size(),
              dataset_size);
    EXPECT_EQ(json_result->scalars().json_data().data_size(), dataset_size);
    EXPECT_EQ(geometry_result->scalars().geometry_data().data_size(),
              dataset_size);
    EXPECT_EQ(vec_result->vectors().float_vector().data_size(),
              dataset_size * dim);
    EXPECT_EQ(float16_vec_result->vectors().float16_vector().size(),
              dataset_size * dim * 2);
    EXPECT_EQ(bfloat16_vec_result->vectors().bfloat16_vector().size(),
              dataset_size * dim * 2);
    EXPECT_EQ(int8_vec_result->vectors().int8_vector().size(),
              dataset_size * dim);
    EXPECT_EQ(int_array_result->scalars().array_data().data_size(),
              dataset_size);
    EXPECT_EQ(long_array_result->scalars().array_data().data_size(),
              dataset_size);
    EXPECT_EQ(bool_array_result->scalars().array_data().data_size(),
              dataset_size);
    EXPECT_EQ(string_array_result->scalars().array_data().data_size(),
              dataset_size);
    EXPECT_EQ(double_array_result->scalars().array_data().data_size(),
              dataset_size);
    EXPECT_EQ(float_array_result->scalars().array_data().data_size(),
              dataset_size);

    EXPECT_EQ(bool_result->valid_data_size(), 0);
    EXPECT_EQ(int8_result->valid_data_size(), 0);
    EXPECT_EQ(int16_result->valid_data_size(), 0);
    EXPECT_EQ(int32_result->valid_data_size(), 0);
    EXPECT_EQ(int64_result->valid_data_size(), 0);
    EXPECT_EQ(float_result->valid_data_size(), 0);
    EXPECT_EQ(double_result->valid_data_size(), 0);
    EXPECT_EQ(varchar_result->valid_data_size(), 0);
    EXPECT_EQ(json_result->valid_data_size(), 0);
    EXPECT_EQ(int_array_result->valid_data_size(), 0);
    EXPECT_EQ(long_array_result->valid_data_size(), 0);
    EXPECT_EQ(bool_array_result->valid_data_size(), 0);
    EXPECT_EQ(string_array_result->valid_data_size(), 0);
    EXPECT_EQ(double_array_result->valid_data_size(), 0);
    EXPECT_EQ(float_array_result->valid_data_size(), 0);
}

TEST(Sealed, QueryAllNullableFields) {
    auto schema = std::make_shared<Schema>();
    auto metric_type = knowhere::metric::L2;
    auto dim = 4;
    auto bool_field = schema->AddDebugField("bool", DataType::BOOL, true);
    auto int8_field = schema->AddDebugField("int8", DataType::INT8, true);
    auto int16_field = schema->AddDebugField("int16", DataType::INT16, true);
    auto int32_field = schema->AddDebugField("int32", DataType::INT32, true);
    auto int64_field = schema->AddDebugField("int64", DataType::INT64, false);
    auto float_field = schema->AddDebugField("float", DataType::FLOAT, true);
    auto double_field = schema->AddDebugField("double", DataType::DOUBLE, true);
    auto varchar_field =
        schema->AddDebugField("varchar", DataType::VARCHAR, true);
    auto json_field = schema->AddDebugField("json", DataType::JSON, true);
    auto geometry_field =
        schema->AddDebugField("geometry", DataType::GEOMETRY, true);
    auto int_array_field = schema->AddDebugField(
        "int_array", DataType::ARRAY, DataType::INT8, true);
    auto long_array_field = schema->AddDebugField(
        "long_array", DataType::ARRAY, DataType::INT64, true);
    auto bool_array_field = schema->AddDebugField(
        "bool_array", DataType::ARRAY, DataType::BOOL, true);
    auto string_array_field = schema->AddDebugField(
        "string_array", DataType::ARRAY, DataType::VARCHAR, true);
    auto double_array_field = schema->AddDebugField(
        "double_array", DataType::ARRAY, DataType::DOUBLE, true);
    auto float_array_field = schema->AddDebugField(
        "float_array", DataType::ARRAY, DataType::FLOAT, true);
    auto vec = schema->AddDebugField(
        "embeddings", DataType::VECTOR_FLOAT, dim, metric_type);
    schema->set_primary_field_id(int64_field);

    std::map<std::string, std::string> index_params = {
        {"index_type", "IVF_FLAT"},
        {"metric_type", metric_type},
        {"nlist", "128"}};
    std::map<std::string, std::string> type_params = {
        {"dim", std::to_string(dim)}};
    FieldIndexMeta fieldIndexMeta(
        vec, std::move(index_params), std::move(type_params));
    std::map<FieldId, FieldIndexMeta> filedMap = {{vec, fieldIndexMeta}};
    IndexMetaPtr metaPtr =
        std::make_shared<CollectionIndexMeta>(100000, std::move(filedMap));
    auto segment_sealed = CreateSealedSegment(schema, metaPtr);
    auto segment =
        dynamic_cast<ChunkedSegmentSealedImpl*>(segment_sealed.get());

    int64_t dataset_size = 1000;
    auto dataset = DataGen(schema, dataset_size);
    segment_sealed = CreateSealedWithFieldDataLoaded(schema, dataset);
    segment = dynamic_cast<ChunkedSegmentSealedImpl*>(segment_sealed.get());

    auto bool_values = dataset.get_col<bool>(bool_field);
    auto int8_values = dataset.get_col<int8_t>(int8_field);
    auto int16_values = dataset.get_col<int16_t>(int16_field);
    auto int32_values = dataset.get_col<int32_t>(int32_field);
    auto int64_values = dataset.get_col<int64_t>(int64_field);
    auto float_values = dataset.get_col<float>(float_field);
    auto double_values = dataset.get_col<double>(double_field);
    auto varchar_values = dataset.get_col<std::string>(varchar_field);
    auto json_values = dataset.get_col<std::string>(json_field);
    auto geometry_values = dataset.get_col<std::string>(geometry_field);
    auto int_array_values = dataset.get_col<ScalarFieldProto>(int_array_field);
    auto long_array_values =
        dataset.get_col<ScalarFieldProto>(long_array_field);
    auto bool_array_values =
        dataset.get_col<ScalarFieldProto>(bool_array_field);
    auto string_array_values =
        dataset.get_col<ScalarFieldProto>(string_array_field);
    auto double_array_values =
        dataset.get_col<ScalarFieldProto>(double_array_field);
    auto float_array_values =
        dataset.get_col<ScalarFieldProto>(float_array_field);
    auto vector_values = dataset.get_col<float>(vec);

    auto bool_valid_values = dataset.get_col_valid(bool_field);
    auto int8_valid_values = dataset.get_col_valid(int8_field);
    auto int16_valid_values = dataset.get_col_valid(int16_field);
    auto int32_valid_values = dataset.get_col_valid(int32_field);
    auto float_valid_values = dataset.get_col_valid(float_field);
    auto double_valid_values = dataset.get_col_valid(double_field);
    auto varchar_valid_values = dataset.get_col_valid(varchar_field);
    auto json_valid_values = dataset.get_col_valid(json_field);
    auto geometry_valid_values = dataset.get_col_valid(geometry_field);
    auto int_array_valid_values = dataset.get_col_valid(int_array_field);
    auto long_array_valid_values = dataset.get_col_valid(long_array_field);
    auto bool_array_valid_values = dataset.get_col_valid(bool_array_field);
    auto string_array_valid_values = dataset.get_col_valid(string_array_field);
    auto double_array_valid_values = dataset.get_col_valid(double_array_field);
    auto float_array_valid_values = dataset.get_col_valid(float_array_field);

    auto ids_ds = GenRandomIds(dataset_size);
    auto bool_result = segment->bulk_subscript(
        nullptr, bool_field, ids_ds->GetIds(), dataset_size);
    auto int8_result = segment->bulk_subscript(
        nullptr, int8_field, ids_ds->GetIds(), dataset_size);
    auto int16_result = segment->bulk_subscript(
        nullptr, int16_field, ids_ds->GetIds(), dataset_size);
    auto int32_result = segment->bulk_subscript(
        nullptr, int32_field, ids_ds->GetIds(), dataset_size);
    auto int64_result = segment->bulk_subscript(
        nullptr, int64_field, ids_ds->GetIds(), dataset_size);
    auto float_result = segment->bulk_subscript(
        nullptr, float_field, ids_ds->GetIds(), dataset_size);
    auto double_result = segment->bulk_subscript(
        nullptr, double_field, ids_ds->GetIds(), dataset_size);
    auto varchar_result = segment->bulk_subscript(
        nullptr, varchar_field, ids_ds->GetIds(), dataset_size);
    auto json_result = segment->bulk_subscript(
        nullptr, json_field, ids_ds->GetIds(), dataset_size);
    auto geometry_result = segment->bulk_subscript(
        nullptr, geometry_field, ids_ds->GetIds(), dataset_size);
    auto int_array_result = segment->bulk_subscript(
        nullptr, int_array_field, ids_ds->GetIds(), dataset_size);
    auto long_array_result = segment->bulk_subscript(
        nullptr, long_array_field, ids_ds->GetIds(), dataset_size);
    auto bool_array_result = segment->bulk_subscript(
        nullptr, bool_array_field, ids_ds->GetIds(), dataset_size);
    auto string_array_result = segment->bulk_subscript(
        nullptr, string_array_field, ids_ds->GetIds(), dataset_size);
    auto double_array_result = segment->bulk_subscript(
        nullptr, double_array_field, ids_ds->GetIds(), dataset_size);
    auto float_array_result = segment->bulk_subscript(
        nullptr, float_array_field, ids_ds->GetIds(), dataset_size);
    auto vec_result =
        segment->bulk_subscript(nullptr, vec, ids_ds->GetIds(), dataset_size);

    EXPECT_EQ(bool_result->scalars().bool_data().data_size(), dataset_size);
    EXPECT_EQ(int8_result->scalars().int_data().data_size(), dataset_size);
    EXPECT_EQ(int16_result->scalars().int_data().data_size(), dataset_size);
    EXPECT_EQ(int32_result->scalars().int_data().data_size(), dataset_size);
    EXPECT_EQ(int64_result->scalars().long_data().data_size(), dataset_size);
    EXPECT_EQ(float_result->scalars().float_data().data_size(), dataset_size);
    EXPECT_EQ(double_result->scalars().double_data().data_size(), dataset_size);
    EXPECT_EQ(varchar_result->scalars().string_data().data_size(),
              dataset_size);
    EXPECT_EQ(json_result->scalars().json_data().data_size(), dataset_size);
    EXPECT_EQ(geometry_result->scalars().geometry_data().data_size(),
              dataset_size);
    EXPECT_EQ(vec_result->vectors().float_vector().data_size(),
              dataset_size * dim);
    EXPECT_EQ(int_array_result->scalars().array_data().data_size(),
              dataset_size);
    EXPECT_EQ(long_array_result->scalars().array_data().data_size(),
              dataset_size);
    EXPECT_EQ(bool_array_result->scalars().array_data().data_size(),
              dataset_size);
    EXPECT_EQ(string_array_result->scalars().array_data().data_size(),
              dataset_size);
    EXPECT_EQ(double_array_result->scalars().array_data().data_size(),
              dataset_size);
    EXPECT_EQ(float_array_result->scalars().array_data().data_size(),
              dataset_size);

    EXPECT_EQ(bool_result->valid_data_size(), dataset_size);
    EXPECT_EQ(int8_result->valid_data_size(), dataset_size);
    EXPECT_EQ(int16_result->valid_data_size(), dataset_size);
    EXPECT_EQ(int32_result->valid_data_size(), dataset_size);
    EXPECT_EQ(float_result->valid_data_size(), dataset_size);
    EXPECT_EQ(double_result->valid_data_size(), dataset_size);
    EXPECT_EQ(varchar_result->valid_data_size(), dataset_size);
    EXPECT_EQ(json_result->valid_data_size(), dataset_size);
    EXPECT_EQ(geometry_result->valid_data_size(), dataset_size);
    EXPECT_EQ(int_array_result->valid_data_size(), dataset_size);
    EXPECT_EQ(long_array_result->valid_data_size(), dataset_size);
    EXPECT_EQ(bool_array_result->valid_data_size(), dataset_size);
    EXPECT_EQ(string_array_result->valid_data_size(), dataset_size);
    EXPECT_EQ(double_array_result->valid_data_size(), dataset_size);
    EXPECT_EQ(float_array_result->valid_data_size(), dataset_size);
}

using VectorArrayTestParam =
    std::tuple<DataType, std::string, int, std::string>;

void
VerifyVectorResults(const VectorFieldProto& result_vec,
                    const VectorFieldProto& expected_vec,
                    DataType element_type) {
    switch (element_type) {
        case DataType::VECTOR_FLOAT: {
            auto result_data = result_vec.float_vector().data();
            auto expected_data = expected_vec.float_vector().data();
            EXPECT_EQ(result_data.size(), expected_data.size());
            for (int64_t i = 0; i < result_data.size(); ++i) {
                EXPECT_NEAR(result_data[i], expected_data[i], 1e-6f);
            }
            break;
        }
        case DataType::VECTOR_BINARY: {
            auto result_data = result_vec.binary_vector();
            auto expected_data = expected_vec.binary_vector();
            EXPECT_EQ(result_data, expected_data);
            break;
        }
        case DataType::VECTOR_FLOAT16: {
            auto result_data = result_vec.float16_vector();
            auto expected_data = expected_vec.float16_vector();
            EXPECT_EQ(result_data, expected_data);
            break;
        }
        case DataType::VECTOR_BFLOAT16: {
            auto result_data = result_vec.bfloat16_vector();
            auto expected_data = expected_vec.bfloat16_vector();
            EXPECT_EQ(result_data, expected_data);
            break;
        }
        case DataType::VECTOR_INT8: {
            auto result_data = result_vec.int8_vector();
            auto expected_data = expected_vec.int8_vector();
            EXPECT_EQ(result_data, expected_data);
            break;
        }
        default:
            break;
    }
}

class SealedVectorArrayTest
    : public ::testing::TestWithParam<VectorArrayTestParam> {
 protected:
    DataType element_type;
    std::string metric_type;
    int dim;
    std::string test_name;

    void
    SetUp() override {
        auto param = GetParam();
        element_type = std::get<0>(param);
        metric_type = std::get<1>(param);
        dim = std::get<2>(param);
        test_name = std::get<3>(param);

        // Ensure dim is valid for binary vectors
        if (element_type == DataType::VECTOR_BINARY) {
            ASSERT_EQ(dim % 8, 0) << "Binary vector dim must be multiple of 8";
        }
    }
};

TEST_P(SealedVectorArrayTest, QueryVectorArrayAllFields) {
    auto schema = std::make_shared<Schema>();
    auto int64_field = schema->AddDebugField("int64", DataType::INT64);
    auto array_vec = schema->AddDebugVectorArrayField(
        "array_vec", element_type, dim, metric_type);
    schema->set_primary_field_id(int64_field);

    std::map<FieldId, FieldIndexMeta> filedMap{};
    IndexMetaPtr metaPtr =
        std::make_shared<CollectionIndexMeta>(100000, std::move(filedMap));

    int64_t dataset_size = 1000;
    auto dataset = DataGen(schema, dataset_size);
    auto segment_sealed = CreateSealedWithFieldDataLoaded(schema, dataset);
    auto segment =
        dynamic_cast<ChunkedSegmentSealedImpl*>(segment_sealed.get());

    auto int64_values = dataset.get_col<int64_t>(int64_field);
    auto array_vec_values = dataset.get_col<VectorFieldProto>(array_vec);

    auto ids_ds = GenRandomIds(dataset_size);
    auto int64_result = segment->bulk_subscript(
        nullptr, int64_field, ids_ds->GetIds(), dataset_size);
    auto array_vector_result = segment->bulk_subscript(
        nullptr, array_vec, ids_ds->GetIds(), dataset_size);

    EXPECT_EQ(int64_result->scalars().long_data().data_size(), dataset_size);
    EXPECT_EQ(array_vector_result->vectors().vector_array().data_size(),
              dataset_size);

    for (int64_t i = 0; i < dataset_size; ++i) {
        auto result_vec =
            array_vector_result->vectors().vector_array().data()[i];
        auto expected_vec = array_vec_values[ids_ds->GetIds()[i]];
        VerifyVectorResults(result_vec, expected_vec, element_type);
    }

    EXPECT_EQ(int64_result->valid_data_size(), 0);
    EXPECT_EQ(array_vector_result->valid_data_size(), 0);
}

TEST_P(SealedVectorArrayTest, SearchVectorArray) {
    int64_t collection_id = 1;
    int64_t partition_id = 2;
    int64_t segment_id = 3;
    int64_t index_build_id = 4000;
    int64_t index_version = 4000;

    auto schema = std::make_shared<Schema>();
    auto int64_field = schema->AddDebugField("int64", DataType::INT64);
    auto array_vec = schema->AddDebugVectorArrayField(
        "array_vec", element_type, dim, metric_type);
    schema->set_primary_field_id(int64_field);

    auto field_meta = milvus::segcore::gen_field_meta(collection_id,
                                                      partition_id,
                                                      segment_id,
                                                      array_vec.get(),
                                                      DataType::VECTOR_ARRAY,
                                                      element_type,
                                                      false);
    auto index_meta = gen_index_meta(
        segment_id, array_vec.get(), index_build_id, index_version);

    std::map<FieldId, FieldIndexMeta> filedMap{};
    IndexMetaPtr metaPtr =
        std::make_shared<CollectionIndexMeta>(100000, std::move(filedMap));

    int64_t dataset_size = 100;
    auto emb_list_len = 2;
    auto dataset = DataGen(schema, dataset_size, 42, 0, 1, emb_list_len);

    // create field data
    std::string root_path = TestLocalPath;
    auto storage_config = gen_local_storage_config(root_path);
    auto cm = CreateChunkManager(storage_config);
    auto fs = milvus::storage::InitArrowFileSystem(storage_config);
    auto vec_array_col = dataset.get_col<VectorFieldProto>(array_vec);
    std::vector<milvus::VectorArray> vector_arrays;
    for (auto& v : vec_array_col) {
        vector_arrays.push_back(milvus::VectorArray(v));
    }
    auto field_data = storage::CreateFieldData(
        DataType::VECTOR_ARRAY, element_type, false, dim);
    field_data->FillFieldData(vector_arrays.data(), vector_arrays.size());

    // create sealed segment
    auto segment = CreateSealedSegment(schema);
    auto field_data_info = PrepareSingleFieldInsertBinlog(collection_id,
                                                          partition_id,
                                                          segment_id,
                                                          array_vec.get(),
                                                          {field_data},
                                                          cm);
    segment->LoadFieldData(field_data_info);

    // serialize bin logs
    auto payload_reader =
        std::make_shared<milvus::storage::PayloadReader>(field_data);
    storage::InsertData insert_data(payload_reader);
    insert_data.SetFieldDataMeta(field_meta);
    insert_data.SetTimestamps(0, 100);

    auto serialized_bytes = insert_data.Serialize(storage::Remote);

    auto get_binlog_path = [=](int64_t log_id) {
        return fmt::format("{}{}/{}/{}/{}/{}",
                           TestLocalPath,
                           collection_id,
                           partition_id,
                           segment_id,
                           array_vec.get(),
                           log_id);
    };

    auto log_path = get_binlog_path(0);

    auto cm_w = ChunkManagerWrapper(cm);
    cm_w.Write(log_path, serialized_bytes.data(), serialized_bytes.size());

    storage::FileManagerContext ctx(field_meta, index_meta, cm, fs);
    std::vector<std::string> index_files;

    // create index
    milvus::index::CreateIndexInfo create_index_info;
    create_index_info.field_type = DataType::VECTOR_ARRAY;
    create_index_info.metric_type = metric_type;
    create_index_info.index_type = knowhere::IndexEnum::INDEX_HNSW;
    create_index_info.index_engine_version =
        knowhere::Version::GetCurrentVersion().VersionNumber();

    auto emb_list_hnsw_index =
        milvus::index::IndexFactory::GetInstance().CreateIndex(
            create_index_info,
            storage::FileManagerContext(field_meta, index_meta, cm, fs));

    // build index
    Config config;
    config[milvus::index::INDEX_TYPE] = knowhere::IndexEnum::INDEX_HNSW;
    config[INSERT_FILES_KEY] = std::vector<std::string>{log_path};
    config[knowhere::meta::METRIC_TYPE] = create_index_info.metric_type;
    config[knowhere::indexparam::M] = "16";
    config[knowhere::indexparam::EF] = "10";
    config[DIM_KEY] = dim;
    emb_list_hnsw_index->Build(config);

    auto vec_index =
        dynamic_cast<milvus::index::VectorIndex*>(emb_list_hnsw_index.get());
    EXPECT_EQ(vec_index->Count(), dataset_size * emb_list_len);
    EXPECT_EQ(vec_index->GetDim(), dim);

    // search
    auto vec_num = 10;

    // Generate query vectors based on element type
    std::vector<uint8_t> query_vec_bin;
    std::vector<float> query_vec_f32;
    knowhere::DataSetPtr query_dataset;
    if (element_type == DataType::VECTOR_BINARY) {
        auto byte_dim = (dim + 7) / 8;
        auto total_bytes = vec_num * byte_dim;
        query_vec_bin.resize(total_bytes);
        for (size_t i = 0; i < total_bytes; ++i) {
            query_vec_bin[i] = rand() % 256;
        }
        query_dataset =
            knowhere::GenDataSet(vec_num, dim, query_vec_bin.data());
    } else {
        // For float-like types (FLOAT, FLOAT16, BFLOAT16, INT8)
        query_vec_f32 = generate_float_vector(vec_num, dim);
        query_dataset =
            knowhere::GenDataSet(vec_num, dim, query_vec_f32.data());
    }
    std::vector<size_t> query_vec_offsets;
    query_vec_offsets.push_back(0);
    query_vec_offsets.push_back(3);
    query_vec_offsets.push_back(10);
    query_dataset->Set(knowhere::meta::EMB_LIST_OFFSET,
                       const_cast<const size_t*>(query_vec_offsets.data()));

    auto search_conf = knowhere::Json{{knowhere::indexparam::NPROBE, 10}};
    milvus::SearchInfo searchInfo;
    searchInfo.topk_ = 5;
    searchInfo.metric_type_ = metric_type;
    searchInfo.search_params_ = search_conf;
    SearchResult result;
    vec_index->Query(query_dataset, searchInfo, nullptr, nullptr, result);
    auto ref_result = SearchResultToJson(result);
    std::cout << ref_result.dump(1) << std::endl;
    EXPECT_EQ(result.total_nq_, 2);
    EXPECT_EQ(result.distances_.size(), 2 * searchInfo.topk_);

    // create sealed segment
    auto sealed_segment = CreateSealedWithFieldDataLoaded(schema, dataset);

    ScopedSchemaHandle handle(*schema);

    // brute force search
    {
        auto plan_str = handle.ParseSearch(
            "", "array_vec", 5, metric_type, "{\"nprobe\": 10}", 3);
        auto plan =
            CreateSearchPlanByExpr(schema, plan_str.data(), plan_str.size());

        // Create placeholder based on element type
        milvus::proto::common::PlaceholderGroup ph_group_raw;
        if (element_type == DataType::VECTOR_BINARY) {
            auto byte_dim = (dim + 7) / 8;
            auto total_bytes = vec_num * byte_dim;
            std::vector<uint8_t> query_vec(total_bytes);
            for (size_t i = 0; i < total_bytes; ++i) {
                query_vec[i] = rand() % 256;
            }
            ph_group_raw = CreatePlaceholderGroupFromBlob<EmbListBinaryVector>(
                vec_num, dim, query_vec.data(), query_vec_offsets);
        } else if (element_type == DataType::VECTOR_FLOAT16) {
            std::vector<float> float_vec = generate_float_vector(vec_num, dim);
            std::vector<float16> query_vec(vec_num * dim);
            for (size_t i = 0; i < vec_num * dim; ++i) {
                query_vec[i] = float16(float_vec[i]);
            }
            ph_group_raw = CreatePlaceholderGroupFromBlob<EmbListFloat16Vector>(
                vec_num, dim, query_vec.data(), query_vec_offsets);
        } else if (element_type == DataType::VECTOR_BFLOAT16) {
            std::vector<float> float_vec = generate_float_vector(vec_num, dim);
            std::vector<bfloat16> query_vec(vec_num * dim);
            for (size_t i = 0; i < vec_num * dim; ++i) {
                query_vec[i] = bfloat16(float_vec[i]);
            }
            ph_group_raw =
                CreatePlaceholderGroupFromBlob<EmbListBFloat16Vector>(
                    vec_num, dim, query_vec.data(), query_vec_offsets);
        } else if (element_type == DataType::VECTOR_INT8) {
            std::vector<int8_t> query_vec(vec_num * dim);
            for (size_t i = 0; i < vec_num * dim; ++i) {
                query_vec[i] = static_cast<int8_t>(rand() % 256 - 128);
            }
            ph_group_raw = CreatePlaceholderGroupFromBlob<EmbListInt8Vector>(
                vec_num, dim, query_vec.data(), query_vec_offsets);
        } else {
            std::vector<float> query_vec = generate_float_vector(vec_num, dim);
            ph_group_raw = CreatePlaceholderGroupFromBlob<EmbListFloatVector>(
                vec_num, dim, query_vec.data(), query_vec_offsets);
        }

        auto ph_group =
            ParsePlaceholderGroup(plan.get(), ph_group_raw.SerializeAsString());
        Timestamp timestamp = 1000000;
        std::vector<const PlaceholderGroup*> ph_group_arr = {ph_group.get()};

        auto sr = sealed_segment->Search(plan.get(), ph_group.get(), timestamp);
        auto sr_parsed = SearchResultToJson(*sr);
        std::cout << sr_parsed.dump(1) << std::endl;
    }

    // search with index
    {
        LoadIndexInfo load_info;
        load_info.field_id = array_vec.get();
        load_info.field_type = DataType::VECTOR_ARRAY;
        load_info.element_type = element_type;
        load_info.index_params = GenIndexParams(emb_list_hnsw_index.get());
        load_info.cache_index =
            CreateTestCacheIndex("test", std::move(emb_list_hnsw_index));
        load_info.index_params["metric_type"] = metric_type;

        sealed_segment->DropFieldData(array_vec);
        sealed_segment->LoadIndex(load_info);

        auto plan_str = handle.ParseSearch(
            "", "array_vec", 5, metric_type, "{\"nprobe\": 10}", 3);
        auto plan =
            CreateSearchPlanByExpr(schema, plan_str.data(), plan_str.size());

        // Create placeholder based on element type
        milvus::proto::common::PlaceholderGroup ph_group_raw;
        if (element_type == DataType::VECTOR_BINARY) {
            auto byte_dim = (dim + 7) / 8;
            auto total_bytes = vec_num * byte_dim;
            std::vector<uint8_t> query_vec(total_bytes);
            for (size_t i = 0; i < total_bytes; ++i) {
                query_vec[i] = rand() % 256;
            }
            ph_group_raw = CreatePlaceholderGroupFromBlob<EmbListBinaryVector>(
                vec_num, dim, query_vec.data(), query_vec_offsets);
        } else if (element_type == DataType::VECTOR_FLOAT16) {
            std::vector<float> float_vec = generate_float_vector(vec_num, dim);
            std::vector<float16> query_vec(vec_num * dim);
            for (size_t i = 0; i < vec_num * dim; ++i) {
                query_vec[i] = float16(float_vec[i]);
            }
            ph_group_raw = CreatePlaceholderGroupFromBlob<EmbListFloat16Vector>(
                vec_num, dim, query_vec.data(), query_vec_offsets);
        } else if (element_type == DataType::VECTOR_BFLOAT16) {
            std::vector<float> float_vec = generate_float_vector(vec_num, dim);
            std::vector<bfloat16> query_vec(vec_num * dim);
            for (size_t i = 0; i < vec_num * dim; ++i) {
                query_vec[i] = bfloat16(float_vec[i]);
            }
            ph_group_raw =
                CreatePlaceholderGroupFromBlob<EmbListBFloat16Vector>(
                    vec_num, dim, query_vec.data(), query_vec_offsets);
        } else if (element_type == DataType::VECTOR_INT8) {
            std::vector<int8_t> query_vec(vec_num * dim);
            for (size_t i = 0; i < vec_num * dim; ++i) {
                query_vec[i] = static_cast<int8_t>(rand() % 256 - 128);
            }
            ph_group_raw = CreatePlaceholderGroupFromBlob<EmbListInt8Vector>(
                vec_num, dim, query_vec.data(), query_vec_offsets);
        } else {
            std::vector<float> query_vec = generate_float_vector(vec_num, dim);
            ph_group_raw = CreatePlaceholderGroupFromBlob<EmbListFloatVector>(
                vec_num, dim, query_vec.data(), query_vec_offsets);
        }

        auto ph_group =
            ParsePlaceholderGroup(plan.get(), ph_group_raw.SerializeAsString());
        Timestamp timestamp = 1000000;
        std::vector<const PlaceholderGroup*> ph_group_arr = {ph_group.get()};

        auto sr = sealed_segment->Search(plan.get(), ph_group.get(), timestamp);
        auto sr_parsed = SearchResultToJson(*sr);
        std::cout << sr_parsed.dump(1) << std::endl;
    }
}

TEST_P(SealedVectorArrayTest, DISABLED_BulkSubscriptVectorArrayFromIndex) {
    int64_t collection_id = 1;
    int64_t partition_id = 2;
    int64_t segment_id = 3;
    int64_t index_build_id = 5000;
    int64_t index_version = 5000;

    auto schema = std::make_shared<Schema>();
    auto int64_field = schema->AddDebugField("int64", DataType::INT64);
    auto array_vec = schema->AddDebugVectorArrayField(
        "array_vec", element_type, dim, metric_type);
    schema->set_primary_field_id(int64_field);

    auto field_meta = milvus::segcore::gen_field_meta(collection_id,
                                                      partition_id,
                                                      segment_id,
                                                      array_vec.get(),
                                                      DataType::VECTOR_ARRAY,
                                                      element_type,
                                                      false);
    auto index_meta = gen_index_meta(
        segment_id, array_vec.get(), index_build_id, index_version);

    int64_t dataset_size = 100;
    auto emb_list_len = 2;
    auto dataset = DataGen(schema, dataset_size, 42, 0, 1, emb_list_len);

    // prepare field data and binlog
    std::string root_path = TestLocalPath;
    auto storage_config = gen_local_storage_config(root_path);
    auto cm = CreateChunkManager(storage_config);
    auto fs = milvus::storage::InitArrowFileSystem(storage_config);
    auto vec_array_col = dataset.get_col<VectorFieldProto>(array_vec);
    std::vector<milvus::VectorArray> vector_arrays;
    for (auto& v : vec_array_col) {
        vector_arrays.push_back(milvus::VectorArray(v));
    }
    auto field_data = storage::CreateFieldData(
        DataType::VECTOR_ARRAY, element_type, false, dim);
    field_data->FillFieldData(vector_arrays.data(), vector_arrays.size());

    auto segment = CreateSealedSegment(schema);
    auto field_data_info = PrepareSingleFieldInsertBinlog(collection_id,
                                                          partition_id,
                                                          segment_id,
                                                          array_vec.get(),
                                                          {field_data},
                                                          cm);
    segment->LoadFieldData(field_data_info);

    auto sealed = dynamic_cast<ChunkedSegmentSealedImpl*>(segment.get());

    // serialize binlog for index building
    auto payload_reader =
        std::make_shared<milvus::storage::PayloadReader>(field_data);
    storage::InsertData insert_data(payload_reader);
    insert_data.SetFieldDataMeta(field_meta);
    insert_data.SetTimestamps(0, 100);
    auto serialized_bytes = insert_data.Serialize(storage::Remote);

    auto log_path = fmt::format("{}{}/{}/{}/{}/{}",
                                TestLocalPath,
                                collection_id,
                                partition_id,
                                segment_id,
                                array_vec.get(),
                                0);
    auto cm_w = ChunkManagerWrapper(cm);
    cm_w.Write(log_path, serialized_bytes.data(), serialized_bytes.size());

    // build HNSW index (FLAT base = has_raw_data = true)
    milvus::index::CreateIndexInfo create_index_info;
    create_index_info.field_type = DataType::VECTOR_ARRAY;
    create_index_info.metric_type = metric_type;
    create_index_info.index_type = knowhere::IndexEnum::INDEX_HNSW;
    create_index_info.index_engine_version =
        knowhere::Version::GetCurrentVersion().VersionNumber();

    auto emb_list_hnsw_index =
        milvus::index::IndexFactory::GetInstance().CreateIndex(
            create_index_info,
            storage::FileManagerContext(field_meta, index_meta, cm, fs));

    Config config;
    config[milvus::index::INDEX_TYPE] = knowhere::IndexEnum::INDEX_HNSW;
    config[INSERT_FILES_KEY] = std::vector<std::string>{log_path};
    config[knowhere::meta::METRIC_TYPE] = create_index_info.metric_type;
    config[knowhere::indexparam::M] = "16";
    config[knowhere::indexparam::EF] = "10";
    config[DIM_KEY] = dim;
    emb_list_hnsw_index->Build(config);

    // drop field data and load index
    LoadIndexInfo load_info;
    load_info.field_id = array_vec.get();
    load_info.field_type = DataType::VECTOR_ARRAY;
    load_info.element_type = element_type;
    load_info.index_params = GenIndexParams(emb_list_hnsw_index.get());
    load_info.cache_index =
        CreateTestCacheIndex("test", std::move(emb_list_hnsw_index));
    load_info.index_params["metric_type"] = metric_type;

    segment->DropFieldData(array_vec);
    segment->LoadIndex(load_info);

    // bulk_subscript should now go through get_emb_list -> GetEmbListByIds
    auto ids_ds = GenRandomIds(dataset_size);
    auto result = sealed->bulk_subscript(
        nullptr, array_vec, ids_ds->GetIds(), dataset_size);

    ASSERT_NE(result, nullptr);
    EXPECT_EQ(result->vectors().vector_array().data_size(), dataset_size);

    // verify each row's embedding list matches original data
    for (int64_t i = 0; i < dataset_size; ++i) {
        auto result_vec = result->vectors().vector_array().data()[i];
        auto expected_vec = vec_array_col[ids_ds->GetIds()[i]];
        VerifyVectorResults(result_vec, expected_vec, element_type);
    }
}

TEST(SealedVectorArrayNullable, BulkSubscriptEmptyThenSingleVectorArrayRows) {
    constexpr int64_t collection_id = 1;
    constexpr int64_t partition_id = 2;
    constexpr int64_t segment_id = 3;
    constexpr int64_t dim = 8;

    auto schema = std::make_shared<Schema>();
    auto int64_field = schema->AddDebugField("int64", DataType::INT64);
    auto array_vec = schema->AddDebugVectorArrayField("array_vec",
                                                      DataType::VECTOR_FLOAT,
                                                      dim,
                                                      knowhere::metric::MAX_SIM,
                                                      true);
    schema->set_primary_field_id(int64_field);

    VectorFieldProto empty_row;
    empty_row.set_dim(dim);
    empty_row.mutable_float_vector();

    VectorFieldProto single_row;
    single_row.set_dim(dim);
    for (int64_t i = 0; i < dim; ++i) {
        single_row.mutable_float_vector()->add_data(static_cast<float>(i + 1));
    }

    std::vector<milvus::VectorArray> vector_arrays;
    vector_arrays.emplace_back(empty_row);
    vector_arrays.emplace_back(single_row);

    constexpr int64_t row_count = 2;
    std::vector<uint8_t> valid_bitmap((row_count + 7) / 8, 0);
    valid_bitmap[0] = 0b00000011;

    auto field_data = storage::CreateFieldData(
        DataType::VECTOR_ARRAY, DataType::VECTOR_FLOAT, true, dim);
    field_data->FillFieldData(
        vector_arrays.data(), valid_bitmap.data(), row_count, 0);

    auto cm = milvus::storage::RemoteChunkManagerSingleton::GetInstance()
                  .GetRemoteChunkManager();
    auto segment = CreateSealedSegment(schema);
    auto field_data_info = PrepareSingleFieldInsertBinlog(collection_id,
                                                          partition_id,
                                                          segment_id,
                                                          array_vec.get(),
                                                          {field_data},
                                                          cm);
    segment->LoadFieldData(field_data_info);

    auto sealed = dynamic_cast<ChunkedSegmentSealedImpl*>(segment.get());
    ASSERT_NE(sealed, nullptr);

    std::vector<int64_t> offsets = {0, 1};
    auto result = sealed->bulk_subscript(
        nullptr, array_vec, offsets.data(), offsets.size());
    ASSERT_NE(result, nullptr);
    ASSERT_EQ(result->valid_data_size(), row_count);
    EXPECT_TRUE(result->valid_data(0));
    EXPECT_TRUE(result->valid_data(1));

    const auto& rows = result->vectors().vector_array().data();
    ASSERT_EQ(rows.size(), row_count);
    EXPECT_EQ(rows.Get(0).float_vector().data_size(), 0);
    ASSERT_EQ(rows.Get(1).float_vector().data_size(), dim);
    EXPECT_FLOAT_EQ(rows.Get(1).float_vector().data(0), 1.0F);
    EXPECT_FLOAT_EQ(rows.Get(1).float_vector().data(dim - 1),
                    static_cast<float>(dim));
}

TEST(SealedVectorArrayNullable,
     DISABLED_BulkSubscriptVectorArrayFromIndexKeepsLogicalRows) {
    int64_t collection_id = 1;
    int64_t partition_id = 2;
    int64_t segment_id = 3;
    int64_t index_build_id = 5100;
    int64_t index_version = 5100;

    constexpr int64_t dim = 8;
    auto schema = std::make_shared<Schema>();
    auto int64_field = schema->AddDebugField("int64", DataType::INT64);
    auto array_vec = schema->AddDebugVectorArrayField("array_vec",
                                                      DataType::VECTOR_FLOAT,
                                                      dim,
                                                      knowhere::metric::MAX_SIM,
                                                      true);
    schema->set_primary_field_id(int64_field);

    auto field_meta = milvus::segcore::gen_field_meta(collection_id,
                                                      partition_id,
                                                      segment_id,
                                                      array_vec.get(),
                                                      DataType::VECTOR_ARRAY,
                                                      DataType::VECTOR_FLOAT,
                                                      true);
    auto index_meta = gen_index_meta(
        segment_id, array_vec.get(), index_build_id, index_version);

    int64_t dataset_size = 12;
    auto emb_list_len = 2;
    auto dataset = DataGen(schema, dataset_size, 42, 0, 1, emb_list_len);
    auto vec_array_col = dataset.get_col<VectorFieldProto>(array_vec);

    std::vector<uint8_t> valid_bitmap((dataset_size + 7) / 8, 0);
    std::vector<milvus::VectorArray> vector_arrays;
    for (int64_t i = 0; i < dataset_size; ++i) {
        if (i % 3 != 0) {
            valid_bitmap[i >> 3] |= (1 << (i & 0x07));
            vector_arrays.emplace_back(vec_array_col[i]);
        }
    }

    std::string root_path = TestLocalPath;
    auto storage_config = gen_local_storage_config(root_path);
    auto cm = CreateChunkManager(storage_config);
    auto fs = milvus::storage::InitArrowFileSystem(storage_config);
    auto field_data = storage::CreateFieldData(
        DataType::VECTOR_ARRAY, DataType::VECTOR_FLOAT, true, dim);
    field_data->FillFieldData(
        vector_arrays.data(), valid_bitmap.data(), dataset_size, 0);

    auto segment = CreateSealedSegment(schema);
    auto field_data_info = PrepareSingleFieldInsertBinlog(collection_id,
                                                          partition_id,
                                                          segment_id,
                                                          array_vec.get(),
                                                          {field_data},
                                                          cm);
    segment->LoadFieldData(field_data_info);
    auto sealed = dynamic_cast<ChunkedSegmentSealedImpl*>(segment.get());

    auto payload_reader =
        std::make_shared<milvus::storage::PayloadReader>(field_data);
    storage::InsertData insert_data(payload_reader);
    insert_data.SetFieldDataMeta(field_meta);
    insert_data.SetTimestamps(0, 100);
    auto serialized_bytes = insert_data.Serialize(storage::Remote);

    auto log_path = fmt::format("{}{}/{}/{}/{}/{}",
                                TestLocalPath,
                                collection_id,
                                partition_id,
                                segment_id,
                                array_vec.get(),
                                0);
    auto cm_w = ChunkManagerWrapper(cm);
    cm_w.Write(log_path, serialized_bytes.data(), serialized_bytes.size());

    milvus::index::CreateIndexInfo create_index_info;
    create_index_info.field_type = DataType::VECTOR_ARRAY;
    create_index_info.metric_type = knowhere::metric::MAX_SIM;
    create_index_info.index_type = knowhere::IndexEnum::INDEX_HNSW;
    create_index_info.index_engine_version =
        knowhere::Version::GetCurrentVersion().VersionNumber();

    auto emb_list_hnsw_index =
        milvus::index::IndexFactory::GetInstance().CreateIndex(
            create_index_info,
            storage::FileManagerContext(field_meta, index_meta, cm, fs));

    Config config;
    config[milvus::index::INDEX_TYPE] = knowhere::IndexEnum::INDEX_HNSW;
    config[INSERT_FILES_KEY] = std::vector<std::string>{log_path};
    config[knowhere::meta::METRIC_TYPE] = create_index_info.metric_type;
    config[knowhere::indexparam::M] = "16";
    config[knowhere::indexparam::EF] = "10";
    config[DIM_KEY] = dim;
    emb_list_hnsw_index->Build(config);

    auto vec_index =
        dynamic_cast<milvus::index::VectorIndex*>(emb_list_hnsw_index.get());
    ASSERT_NE(vec_index, nullptr);
    EXPECT_TRUE(vec_index->HasValidData());
    EXPECT_EQ(vec_index->GetValidCount(), vector_arrays.size());

    LoadIndexInfo load_info;
    load_info.field_id = array_vec.get();
    load_info.field_type = DataType::VECTOR_ARRAY;
    load_info.element_type = DataType::VECTOR_FLOAT;
    load_info.index_params = GenIndexParams(emb_list_hnsw_index.get());
    load_info.cache_index =
        CreateTestCacheIndex("test", std::move(emb_list_hnsw_index));
    load_info.index_params["metric_type"] = knowhere::metric::MAX_SIM;

    segment->DropFieldData(array_vec);
    segment->LoadIndex(load_info);

    std::vector<int64_t> offsets = {0, 1, 3, 4, 8};
    auto result = sealed->bulk_subscript(
        nullptr, array_vec, offsets.data(), offsets.size());

    ASSERT_NE(result, nullptr);
    ASSERT_EQ(result->valid_data_size(), offsets.size());
    ASSERT_EQ(result->vectors().vector_array().data_size(), offsets.size());

    for (size_t i = 0; i < offsets.size(); ++i) {
        auto logical_offset = offsets[i];
        auto expected_valid = logical_offset % 3 != 0;
        EXPECT_EQ(result->valid_data(i), expected_valid);
        const auto& result_vec = result->vectors().vector_array().data(i);
        if (!expected_valid) {
            EXPECT_TRUE(result_vec.has_float_vector());
            EXPECT_EQ(result_vec.float_vector().data_size(), 0);
            continue;
        }
        VerifyVectorResults(
            result_vec, vec_array_col[logical_offset], DataType::VECTOR_FLOAT);
    }
}

TEST(SealedVectorArrayNullable,
     DISABLED_BulkSubscriptAllNullVectorArrayFromIndex) {
    int64_t collection_id = 1;
    int64_t partition_id = 2;
    int64_t segment_id = 3;
    int64_t index_build_id = 5101;
    int64_t index_version = 5101;

    constexpr int64_t dim = 8;
    constexpr int64_t dataset_size = 8;
    auto schema = std::make_shared<Schema>();
    auto int64_field = schema->AddDebugField("int64", DataType::INT64);
    auto array_vec = schema->AddDebugVectorArrayField("array_vec",
                                                      DataType::VECTOR_FLOAT,
                                                      dim,
                                                      knowhere::metric::MAX_SIM,
                                                      true);
    schema->set_primary_field_id(int64_field);

    auto field_meta = milvus::segcore::gen_field_meta(collection_id,
                                                      partition_id,
                                                      segment_id,
                                                      array_vec.get(),
                                                      DataType::VECTOR_ARRAY,
                                                      DataType::VECTOR_FLOAT,
                                                      true);
    auto index_meta = gen_index_meta(
        segment_id, array_vec.get(), index_build_id, index_version);

    std::string root_path = TestLocalPath;
    auto storage_config = gen_local_storage_config(root_path);
    auto cm = CreateChunkManager(storage_config);
    auto fs = milvus::storage::InitArrowFileSystem(storage_config);
    auto field_data = storage::CreateFieldData(
        DataType::VECTOR_ARRAY, DataType::VECTOR_FLOAT, true, dim);
    std::vector<uint8_t> valid_bitmap((dataset_size + 7) / 8, 0);
    std::vector<milvus::VectorArray> vector_arrays;
    field_data->FillFieldData(
        vector_arrays.data(), valid_bitmap.data(), dataset_size, 0);

    auto segment = CreateSealedSegment(schema);
    auto field_data_info = PrepareSingleFieldInsertBinlog(collection_id,
                                                          partition_id,
                                                          segment_id,
                                                          array_vec.get(),
                                                          {field_data},
                                                          cm);
    segment->LoadFieldData(field_data_info);
    auto sealed = dynamic_cast<ChunkedSegmentSealedImpl*>(segment.get());
    ASSERT_NE(sealed, nullptr);

    auto payload_reader =
        std::make_shared<milvus::storage::PayloadReader>(field_data);
    storage::InsertData insert_data(payload_reader);
    insert_data.SetFieldDataMeta(field_meta);
    insert_data.SetTimestamps(0, 100);
    auto serialized_bytes = insert_data.Serialize(storage::Remote);

    auto log_path = fmt::format("{}{}/{}/{}/{}/{}",
                                TestLocalPath,
                                collection_id,
                                partition_id,
                                segment_id,
                                array_vec.get(),
                                0);
    auto cm_w = ChunkManagerWrapper(cm);
    cm_w.Write(log_path, serialized_bytes.data(), serialized_bytes.size());

    milvus::index::CreateIndexInfo create_index_info;
    create_index_info.field_type = DataType::VECTOR_ARRAY;
    create_index_info.metric_type = knowhere::metric::MAX_SIM;
    create_index_info.index_type = knowhere::IndexEnum::INDEX_HNSW;
    create_index_info.index_engine_version =
        knowhere::Version::GetCurrentVersion().VersionNumber();

    auto emb_list_hnsw_index =
        milvus::index::IndexFactory::GetInstance().CreateIndex(
            create_index_info,
            storage::FileManagerContext(field_meta, index_meta, cm, fs));

    Config config;
    config[milvus::index::INDEX_TYPE] = knowhere::IndexEnum::INDEX_HNSW;
    config[INSERT_FILES_KEY] = std::vector<std::string>{log_path};
    config[knowhere::meta::METRIC_TYPE] = create_index_info.metric_type;
    config[knowhere::indexparam::M] = "16";
    config[knowhere::indexparam::EF] = "10";
    config[DIM_KEY] = dim;
    emb_list_hnsw_index->Build(config);

    auto vec_index =
        dynamic_cast<milvus::index::VectorIndex*>(emb_list_hnsw_index.get());
    ASSERT_NE(vec_index, nullptr);
    EXPECT_TRUE(vec_index->HasValidData());
    EXPECT_EQ(vec_index->GetValidCount(), 0);
    EXPECT_TRUE(vec_index->HasRawData());

    LoadIndexInfo load_info;
    load_info.field_id = array_vec.get();
    load_info.field_type = DataType::VECTOR_ARRAY;
    load_info.element_type = DataType::VECTOR_FLOAT;
    load_info.index_params = GenIndexParams(emb_list_hnsw_index.get());
    load_info.cache_index =
        CreateTestCacheIndex("test", std::move(emb_list_hnsw_index));
    load_info.index_params["metric_type"] = knowhere::metric::MAX_SIM;

    segment->DropFieldData(array_vec);
    segment->LoadIndex(load_info);

    std::vector<int64_t> offsets = {0, 1, 7};
    auto result = sealed->bulk_subscript(
        nullptr, array_vec, offsets.data(), offsets.size());

    ASSERT_NE(result, nullptr);
    ASSERT_EQ(result->valid_data_size(), offsets.size());
    ASSERT_EQ(result->vectors().vector_array().data_size(), offsets.size());
    for (int i = 0; i < result->valid_data_size(); ++i) {
        EXPECT_FALSE(result->valid_data(i));
        EXPECT_EQ(
            result->vectors().vector_array().data(i).float_vector().data_size(),
            0);
    }
}

TEST(SealedVectorArrayFallback,
     BulkSubscriptVectorArrayFallsBackToFieldDataWhenIndexHasNoRawData) {
    constexpr int64_t collection_id = 1;
    constexpr int64_t partition_id = 2;
    constexpr int64_t segment_id = 3;
    constexpr int64_t index_build_id = 7000;
    constexpr int64_t index_version = 7000;
    constexpr int64_t dataset_size = 100;
    constexpr int64_t emb_list_len = 2;
    constexpr int64_t dim = 4;
    const auto element_type = DataType::VECTOR_FLOAT;
    const auto metric_type = knowhere::metric::MAX_SIM;

    auto schema = std::make_shared<Schema>();
    auto int64_field = schema->AddDebugField("int64", DataType::INT64);
    auto array_vec = schema->AddDebugVectorArrayField(
        "array_vec", element_type, dim, metric_type);
    schema->set_primary_field_id(int64_field);

    auto field_meta = milvus::segcore::gen_field_meta(collection_id,
                                                      partition_id,
                                                      segment_id,
                                                      array_vec.get(),
                                                      DataType::VECTOR_ARRAY,
                                                      element_type,
                                                      false);
    auto index_meta = gen_index_meta(
        segment_id, array_vec.get(), index_build_id, index_version);

    auto dataset = DataGen(schema, dataset_size, 42, 0, 1, emb_list_len);

    std::string root_path = TestLocalPath;
    auto storage_config = gen_local_storage_config(root_path);
    auto cm = CreateChunkManager(storage_config);
    auto fs = milvus::storage::InitArrowFileSystem(storage_config);
    auto vec_array_col = dataset.get_col<VectorFieldProto>(array_vec);
    std::vector<milvus::VectorArray> vector_arrays;
    for (auto& v : vec_array_col) {
        vector_arrays.push_back(milvus::VectorArray(v));
    }
    auto field_data = storage::CreateFieldData(
        DataType::VECTOR_ARRAY, element_type, false, dim);
    field_data->FillFieldData(vector_arrays.data(), vector_arrays.size());

    auto segment = CreateSealedSegment(schema);
    auto field_data_info = PrepareSingleFieldInsertBinlog(collection_id,
                                                          partition_id,
                                                          segment_id,
                                                          array_vec.get(),
                                                          {field_data},
                                                          cm);
    segment->LoadFieldData(field_data_info);

    auto sealed = dynamic_cast<ChunkedSegmentSealedImpl*>(segment.get());

    auto payload_reader =
        std::make_shared<milvus::storage::PayloadReader>(field_data);
    storage::InsertData insert_data(payload_reader);
    insert_data.SetFieldDataMeta(field_meta);
    insert_data.SetTimestamps(0, 100);
    auto serialized_bytes = insert_data.Serialize(storage::Remote);

    auto log_path = fmt::format("{}{}/{}/{}/{}/{}",
                                TestLocalPath,
                                collection_id,
                                partition_id,
                                segment_id,
                                array_vec.get(),
                                0);
    auto cm_w = ChunkManagerWrapper(cm);
    cm_w.Write(log_path, serialized_bytes.data(), serialized_bytes.size());

    milvus::index::CreateIndexInfo create_index_info;
    create_index_info.field_type = DataType::VECTOR_ARRAY;
    create_index_info.metric_type = metric_type;
    create_index_info.index_type = knowhere::IndexEnum::INDEX_HNSW_SQ;
    create_index_info.index_engine_version =
        knowhere::Version::GetCurrentVersion().VersionNumber();

    auto emb_list_hnsw_sq_index =
        milvus::index::IndexFactory::GetInstance().CreateIndex(
            create_index_info,
            storage::FileManagerContext(field_meta, index_meta, cm, fs));

    Config config;
    config[milvus::index::INDEX_TYPE] = knowhere::IndexEnum::INDEX_HNSW_SQ;
    config[INSERT_FILES_KEY] = std::vector<std::string>{log_path};
    config[knowhere::meta::METRIC_TYPE] = create_index_info.metric_type;
    config[knowhere::indexparam::M] = "16";
    config[knowhere::indexparam::EF] = "10";
    config[knowhere::indexparam::SQ_TYPE] = "SQ8";
    config[DIM_KEY] = dim;
    emb_list_hnsw_sq_index->Build(config);

    auto index_params = GenIndexParams(emb_list_hnsw_sq_index.get());
    index_params["metric_type"] = metric_type;
    auto request = milvus::index::IndexFactory::GetInstance().IndexLoadResource(
        DataType::VECTOR_ARRAY,
        element_type,
        create_index_info.index_engine_version,
        0,
        index_params,
        false,
        dataset_size * emb_list_len,
        dim);
    ASSERT_FALSE(request.has_raw_data);

    LoadIndexInfo load_info;
    load_info.field_id = array_vec.get();
    load_info.field_type = DataType::VECTOR_ARRAY;
    load_info.element_type = element_type;
    load_info.index_params = index_params;
    load_info.cache_index =
        CreateTestCacheIndex("test", std::move(emb_list_hnsw_sq_index));

    ASSERT_TRUE(segment->HasFieldData(array_vec));
    segment->LoadIndex(load_info);
    ASSERT_TRUE(segment->HasFieldData(array_vec));

    auto ids_ds = GenRandomIds(dataset_size);
    auto result = sealed->bulk_subscript(
        nullptr, array_vec, ids_ds->GetIds(), dataset_size);

    ASSERT_NE(result, nullptr);
    EXPECT_EQ(result->vectors().vector_array().data_size(), dataset_size);

    for (int64_t i = 0; i < dataset_size; ++i) {
        auto result_vec = result->vectors().vector_array().data()[i];
        auto expected_vec = vec_array_col[ids_ds->GetIds()[i]];
        VerifyVectorResults(result_vec, expected_vec, element_type);
    }
}

#ifdef BUILD_DISK_ANN
TEST_P(SealedVectorArrayTest, DISABLED_BulkSubscriptVectorArrayFromDiskIndex) {
    // DiskANN only registers fp32, fp16, bf16
    if (element_type == DataType::VECTOR_BINARY ||
        element_type == DataType::VECTOR_INT8) {
        GTEST_SKIP() << "DiskANN does not support this element type";
    }

    int64_t collection_id = 1;
    int64_t partition_id = 2;
    int64_t segment_id = 3;
    int64_t index_build_id = 6000;
    int64_t index_version = 6000;

    auto schema = std::make_shared<Schema>();
    auto int64_field = schema->AddDebugField("int64", DataType::INT64);
    auto array_vec = schema->AddDebugVectorArrayField(
        "array_vec", element_type, dim, metric_type);
    schema->set_primary_field_id(int64_field);

    auto field_meta = milvus::segcore::gen_field_meta(collection_id,
                                                      partition_id,
                                                      segment_id,
                                                      array_vec.get(),
                                                      DataType::VECTOR_ARRAY,
                                                      element_type,
                                                      false);
    auto index_meta = gen_index_meta(
        segment_id, array_vec.get(), index_build_id, index_version);

    int64_t dataset_size = 100;
    auto emb_list_len = 2;
    auto dataset = DataGen(schema, dataset_size, 42, 0, 1, emb_list_len);

    // prepare field data and binlog
    std::string root_path = TestLocalPath;
    auto storage_config = gen_local_storage_config(root_path);
    auto cm = CreateChunkManager(storage_config);
    auto fs = milvus::storage::InitArrowFileSystem(storage_config);
    auto vec_array_col = dataset.get_col<VectorFieldProto>(array_vec);
    std::vector<milvus::VectorArray> vector_arrays;
    for (auto& v : vec_array_col) {
        vector_arrays.push_back(milvus::VectorArray(v));
    }
    auto field_data = storage::CreateFieldData(
        DataType::VECTOR_ARRAY, element_type, false, dim);
    field_data->FillFieldData(vector_arrays.data(), vector_arrays.size());

    auto segment = CreateSealedSegment(schema);
    auto field_data_info = PrepareSingleFieldInsertBinlog(collection_id,
                                                          partition_id,
                                                          segment_id,
                                                          array_vec.get(),
                                                          {field_data},
                                                          cm);
    segment->LoadFieldData(field_data_info);

    auto sealed = dynamic_cast<ChunkedSegmentSealedImpl*>(segment.get());

    // serialize binlog for index building
    auto payload_reader =
        std::make_shared<milvus::storage::PayloadReader>(field_data);
    storage::InsertData insert_data(payload_reader);
    insert_data.SetFieldDataMeta(field_meta);
    insert_data.SetTimestamps(0, 100);
    auto serialized_bytes = insert_data.Serialize(storage::Remote);

    auto log_path = fmt::format("{}{}/{}/{}/{}/{}",
                                TestLocalPath,
                                collection_id,
                                partition_id,
                                segment_id,
                                array_vec.get(),
                                0);
    auto cm_w = ChunkManagerWrapper(cm);
    cm_w.Write(log_path, serialized_bytes.data(), serialized_bytes.size());

    storage::FileManagerContext file_mgr_ctx(field_meta, index_meta, cm, fs);

    // build DiskANN index
    milvus::index::CreateIndexInfo create_index_info;
    create_index_info.field_type = DataType::VECTOR_ARRAY;
    create_index_info.metric_type = metric_type;
    create_index_info.index_type = knowhere::IndexEnum::INDEX_DISKANN;
    create_index_info.index_engine_version =
        knowhere::Version::GetCurrentVersion().VersionNumber();

    auto disk_index = milvus::index::IndexFactory::GetInstance().CreateIndex(
        create_index_info, file_mgr_ctx);

    Config config;
    config[milvus::index::INDEX_TYPE] = knowhere::IndexEnum::INDEX_DISKANN;
    config[INSERT_FILES_KEY] = std::vector<std::string>{log_path};
    config[knowhere::meta::METRIC_TYPE] = metric_type;
    config[DIM_KEY] = dim;
    config[milvus::index::DISK_ANN_MAX_DEGREE] = std::to_string(24);
    config[milvus::index::DISK_ANN_SEARCH_LIST_SIZE] = std::to_string(56);
    config[milvus::index::DISK_ANN_PQ_CODE_BUDGET] = std::to_string(0.001);
    config[milvus::index::DISK_ANN_BUILD_DRAM_BUDGET] = std::to_string(2);
    config[milvus::index::DISK_ANN_BUILD_THREAD_NUM] = std::to_string(2);
    disk_index->Build(config);

    // Upload to serialize, then reload
    auto upload_result = disk_index->Upload();
    auto index_files = upload_result->GetIndexFiles();
    disk_index.reset();

    auto loaded_index = milvus::index::IndexFactory::GetInstance().CreateIndex(
        create_index_info, file_mgr_ctx);
    auto vec_index =
        dynamic_cast<milvus::index::VectorIndex*>(loaded_index.get());
    auto load_conf = generate_load_conf(knowhere::IndexEnum::INDEX_DISKANN,
                                        metric_type,
                                        dataset_size * emb_list_len);
    load_conf["index_files"] = index_files;
    load_conf[milvus::LOAD_PRIORITY] =
        milvus::proto::common::LoadPriority::HIGH;
    vec_index->Load(milvus::tracer::TraceContext{}, load_conf);

    // drop field data and load index into segment
    LoadIndexInfo load_info;
    load_info.field_id = array_vec.get();
    load_info.field_type = DataType::VECTOR_ARRAY;
    load_info.element_type = element_type;
    load_info.index_params = GenIndexParams(vec_index);
    load_info.cache_index =
        CreateTestCacheIndex("test", std::move(loaded_index));
    load_info.index_params["metric_type"] = metric_type;

    segment->DropFieldData(array_vec);
    segment->LoadIndex(load_info);

    // bulk_subscript should go through get_emb_list -> GetEmbListByIds on disk index
    auto ids_ds = GenRandomIds(dataset_size);
    auto result = sealed->bulk_subscript(
        nullptr, array_vec, ids_ds->GetIds(), dataset_size);

    ASSERT_NE(result, nullptr);
    EXPECT_EQ(result->vectors().vector_array().data_size(), dataset_size);

    for (int64_t i = 0; i < dataset_size; ++i) {
        auto result_vec = result->vectors().vector_array().data()[i];
        auto expected_vec = vec_array_col[ids_ds->GetIds()[i]];
        VerifyVectorResults(result_vec, expected_vec, element_type);
    }
}
#endif  // BUILD_DISK_ANN

INSTANTIATE_TEST_SUITE_P(
    VectorArrayTypes,
    SealedVectorArrayTest,
    ::testing::Values(
        std::make_tuple(DataType::VECTOR_FLOAT, "MAX_SIM", 4, "float_max_sim"),
        std::make_tuple(DataType::VECTOR_FLOAT, "MAX_SIM_L2", 4, "float_l2"),
        std::make_tuple(
            DataType::VECTOR_FLOAT16, "MAX_SIM", 4, "float16_max_sim"),
        std::make_tuple(
            DataType::VECTOR_FLOAT16, "MAX_SIM_L2", 4, "float16_l2"),
        std::make_tuple(
            DataType::VECTOR_BFLOAT16, "MAX_SIM", 4, "bfloat16_max_sim"),
        std::make_tuple(
            DataType::VECTOR_BFLOAT16, "MAX_SIM_L2", 4, "bfloat16_l2"),
        std::make_tuple(DataType::VECTOR_INT8, "MAX_SIM", 4, "int8_max_sim"),
        std::make_tuple(DataType::VECTOR_INT8, "MAX_SIM_L2", 4, "int8_l2"),
        std::make_tuple(
            DataType::VECTOR_BINARY, "MAX_SIM_HAMMING", 32, "binary_hamming"),
        std::make_tuple(
            DataType::VECTOR_BINARY, "MAX_SIM_JACCARD", 32, "binary_jaccard")),
    [](const ::testing::TestParamInfo<VectorArrayTestParam>& info) {
        return std::get<3>(info.param);
    });

// ==================== DropFieldData Guard Tests ====================

TEST(SealedDropFieldData, SkipPKField) {
    // DropFieldData should skip dropping PK field data
    auto schema = std::make_shared<Schema>();
    auto dim = 4;
    auto metric_type = knowhere::metric::L2;
    auto vec_id = schema->AddDebugField(
        "fakevec", DataType::VECTOR_FLOAT, dim, metric_type);
    auto pk_id = schema->AddDebugField("pk", DataType::INT64);
    schema->set_primary_field_id(pk_id);

    auto N = 100;
    auto dataset = DataGen(schema, N);

    auto segment = CreateSealedWithFieldDataLoaded(schema, dataset);

    // PK field data should be present
    EXPECT_TRUE(segment->HasFieldData(pk_id));

    // Try to drop PK field - should be skipped
    segment->DropFieldData(pk_id);

    // PK field data should still be present (protected)
    EXPECT_TRUE(segment->HasFieldData(pk_id));
}

TEST(SealedDropFieldData, DropNonPKField) {
    // DropFieldData should successfully drop non-PK, non-column-group fields
    auto schema = std::make_shared<Schema>();
    auto dim = 4;
    auto metric_type = knowhere::metric::L2;
    auto vec_id = schema->AddDebugField(
        "fakevec", DataType::VECTOR_FLOAT, dim, metric_type);
    auto pk_id = schema->AddDebugField("pk", DataType::INT64);
    schema->set_primary_field_id(pk_id);

    auto N = 100;
    auto dataset = DataGen(schema, N);

    auto segment = CreateSealedWithFieldDataLoaded(schema, dataset);

    // Vec field data should be present
    EXPECT_TRUE(segment->HasFieldData(vec_id));

    // Drop vec field - should succeed (not PK, not column group)
    segment->DropFieldData(vec_id);

    // Vec field data should be gone
    EXPECT_FALSE(segment->HasFieldData(vec_id));
}

TEST(SealedDropFieldData, PKFieldStillDropsBinlogIndex) {
    // When dropping PK field, field data is preserved but binlog index
    // (if present) should still be cleared.
    auto schema = std::make_shared<Schema>();
    auto dim = 4;
    auto metric_type = knowhere::metric::L2;
    auto vec_id = schema->AddDebugField(
        "fakevec", DataType::VECTOR_FLOAT, dim, metric_type);
    auto pk_id = schema->AddDebugField("pk", DataType::INT64);
    schema->set_primary_field_id(pk_id);

    auto N = 100;
    auto dataset = DataGen(schema, N);

    auto segment = CreateSealedWithFieldDataLoaded(schema, dataset);

    // PK field data should be present
    EXPECT_TRUE(segment->HasFieldData(pk_id));

    // Drop PK field - field data should remain
    segment->DropFieldData(pk_id);
    EXPECT_TRUE(segment->HasFieldData(pk_id));

    // Drop again - should be a no-op (idempotent)
    segment->DropFieldData(pk_id);
    EXPECT_TRUE(segment->HasFieldData(pk_id));
}

TEST(SealedSegmentReopen, SchemaOnlyReopenPublishesDefaultFilledState) {
    auto old_schema = std::make_shared<Schema>();
    old_schema->AddDebugField("pk", DataType::INT64);
    old_schema->set_primary_field_id(FieldId(100));
    old_schema->set_schema_version(100);

    auto new_schema = std::make_shared<Schema>();
    new_schema->AddDebugField("pk", DataType::INT64);
    new_schema->AddDebugField("new_field", DataType::INT64);
    new_schema->set_primary_field_id(FieldId(100));
    new_schema->set_schema_version(200);

    auto dataset = DataGen(old_schema, 1);
    auto segment = CreateSealedWithFieldDataLoaded(old_schema, dataset);
    auto* sealed = dynamic_cast<ChunkedSegmentSealedImpl*>(segment.get());
    ASSERT_NE(sealed, nullptr);

    EXPECT_NO_THROW(sealed->Reopen(new_schema));
    EXPECT_TRUE(sealed->HasFieldData(FieldId(101)));

    auto snapshot = sealed->TestGetLoadInfoSnapshot();
    EXPECT_TRUE(snapshot->IsFieldFilledWithDefault(FieldId(101)));
}

TEST(SealedSegmentLoad, LoadPublishesDefaultFilledStateAfterCompact) {
    auto schema = std::make_shared<Schema>();
    auto pk_id = schema->AddDebugField("pk", DataType::INT64);
    auto new_field = schema->AddDebugField("new_field", DataType::INT64);
    schema->set_primary_field_id(pk_id);

    auto segment = CreateSealedSegment(schema);
    auto* sealed = dynamic_cast<ChunkedSegmentSealedImpl*>(segment.get());
    ASSERT_NE(sealed, nullptr);

    const size_t N = 1;
    std::vector<int64_t> pks = {1};
    auto pk_field_data =
        storage::CreateFieldData(DataType::INT64, DataType::NONE, false, 1, N);
    pk_field_data->FillFieldData(pks.data(), N);
    auto cm = milvus::storage::RemoteChunkManagerSingleton::GetInstance()
                  .GetRemoteChunkManager();
    auto pk_load_info = PrepareSingleFieldInsertBinlog(kCollectionID,
                                                       kPartitionID,
                                                       kSegmentID,
                                                       pk_id.get(),
                                                       {pk_field_data},
                                                       cm);

    proto::segcore::SegmentLoadInfo proto;
    proto.set_segmentid(kSegmentID);
    proto.set_partitionid(kPartitionID);
    proto.set_collectionid(kCollectionID);
    proto.set_num_of_rows(N);
    for (const auto& [field_id, field_info] : pk_load_info.field_infos) {
        auto* binlog = proto.add_binlog_paths();
        binlog->set_fieldid(field_id);
        for (int i = 0; i < field_info.insert_files.size(); ++i) {
            auto* log = binlog->add_binlogs();
            log->set_log_path(field_info.insert_files[i]);
            log->set_entries_num(field_info.entries_nums[i]);
            log->set_log_size(field_info.memory_sizes[i]);
        }
    }

    sealed->SetLoadInfo(proto);

    milvus::tracer::TraceContext trace_ctx;
    milvus::OpContext op_ctx;
    ASSERT_NO_THROW(sealed->Load(trace_ctx, &op_ctx));
    EXPECT_TRUE(sealed->HasFieldData(new_field));
    EXPECT_TRUE(
        sealed->TestGetSegmentLoadInfo()->IsFieldFilledWithDefault(new_field));
}

TEST(SealedSegmentReopen, SchemaAwareReopenDiscardsOlderSchema) {
    auto old_schema = std::make_shared<Schema>();
    auto pk_id = old_schema->AddDebugField("pk", DataType::INT64);
    old_schema->set_primary_field_id(pk_id);
    old_schema->set_schema_version(100);

    auto new_schema = std::make_shared<Schema>();
    auto new_pk_id = new_schema->AddDebugField("pk", DataType::INT64);
    auto new_field = new_schema->AddDebugField("new_field", DataType::INT64);
    new_schema->set_primary_field_id(new_pk_id);
    new_schema->set_schema_version(200);

    auto stale_schema = std::make_shared<Schema>();
    auto stale_pk_id = stale_schema->AddDebugField("pk", DataType::INT64);
    stale_schema->set_primary_field_id(stale_pk_id);
    stale_schema->set_schema_version(150);

    auto dataset = DataGen(old_schema, 1);
    auto segment = CreateSealedWithFieldDataLoaded(old_schema, dataset);
    auto* sealed = dynamic_cast<ChunkedSegmentSealedImpl*>(segment.get());
    ASSERT_NE(sealed, nullptr);

    EXPECT_NO_THROW(sealed->Reopen(new_schema));
    ASSERT_TRUE(sealed->TestGetLoadInfoSnapshot()->HasFieldInSchema(new_field));

    proto::segcore::SegmentLoadInfo stale_proto =
        sealed->TestGetSegmentLoadInfo()->GetProto();
    milvus::OpContext op_ctx;
    EXPECT_NO_THROW(sealed->Reopen(&op_ctx, stale_proto, stale_schema));

    auto snapshot = sealed->TestGetLoadInfoSnapshot();
    EXPECT_TRUE(snapshot->HasFieldInSchema(new_field));
    EXPECT_TRUE(sealed->HasFieldData(new_field));
}

TEST(SealedSegmentReopen, DropFieldReopenPublishesSchemaAndLoadInfoTogether) {
    auto old_schema = std::make_shared<Schema>();
    auto pk_id = old_schema->AddDebugField("pk", DataType::INT64);
    auto dropped_id = old_schema->AddDebugField("dropped", DataType::INT64);
    old_schema->set_primary_field_id(pk_id);
    old_schema->set_schema_version(100);

    auto dataset = DataGen(old_schema, 4);
    auto segment = CreateSealedWithFieldDataLoaded(old_schema, dataset);
    auto* sealed = dynamic_cast<ChunkedSegmentSealedImpl*>(segment.get());
    ASSERT_NE(sealed, nullptr);

    ASSERT_TRUE(sealed->FieldAccessible(dropped_id));
    ASSERT_TRUE(
        sealed->TestGetLoadInfoSnapshot()->HasFieldInSchema(dropped_id));

    auto new_schema = std::make_shared<Schema>();
    auto new_pk_id = new_schema->AddDebugField("pk", DataType::INT64);
    new_schema->set_primary_field_id(new_pk_id);
    new_schema->set_schema_version(200);

    EXPECT_NO_THROW(sealed->Reopen(new_schema));

    auto schema_snapshot = sealed->TestGetSchemaSnapshot();
    auto load_info_snapshot = sealed->TestGetLoadInfoSnapshot();
    EXPECT_EQ(schema_snapshot->get_schema_version(), 200);
    EXPECT_EQ(schema_snapshot->get_fields().count(dropped_id), 0);
    EXPECT_FALSE(load_info_snapshot->HasFieldInSchema(dropped_id));
    EXPECT_FALSE(sealed->HasFieldData(dropped_id));
    EXPECT_FALSE(sealed->HasIndex(dropped_id));
    EXPECT_FALSE(sealed->FieldAccessible(dropped_id));
}

TEST(SealedSegmentReopen, TextIndexCancellationGuardDoesNotLeakPendingState) {
    std::map<std::string, std::string> analyzer_params;
    auto schema = std::make_shared<Schema>();
    auto pk_id = schema->AddDebugField("pk", DataType::INT64);
    auto text_fid = schema->AddDebugVarcharField(FieldName("text_field"),
                                                 DataType::VARCHAR,
                                                 /*max_length=*/65535,
                                                 /*nullable=*/false,
                                                 /*enable_match=*/true,
                                                 /*enable_analyzer=*/true,
                                                 analyzer_params,
                                                 std::nullopt);
    schema->set_primary_field_id(pk_id);

    auto segment = CreateSealedSegment(schema);
    auto* sealed = dynamic_cast<ChunkedSegmentSealedImpl*>(segment.get());
    ASSERT_NE(sealed, nullptr);

    EXPECT_FALSE(sealed->TestHasPendingTextIndex(text_fid));
    EXPECT_FALSE(
        sealed->TestGetLoadInfoSnapshot()->HasTextIndexCreated(text_fid));

    sealed->TestRegisterPendingTextIndex(text_fid);

    EXPECT_FALSE(sealed->TestHasPendingTextIndex(text_fid));
    EXPECT_FALSE(
        sealed->TestGetLoadInfoSnapshot()->HasTextIndexCreated(text_fid));
}

TEST(SealedSegmentReopen, RuntimeTextIndexMarkerPublishesToLoadInfoSnapshot) {
    std::map<std::string, std::string> analyzer_params;
    auto schema = std::make_shared<Schema>();
    auto pk_id = schema->AddDebugField("pk", DataType::INT64);
    auto text_fid = schema->AddDebugVarcharField(FieldName("text_field"),
                                                 DataType::VARCHAR,
                                                 /*max_length=*/65535,
                                                 /*nullable=*/false,
                                                 /*enable_match=*/true,
                                                 /*enable_analyzer=*/true,
                                                 analyzer_params,
                                                 std::nullopt);
    schema->set_primary_field_id(pk_id);

    auto segment = CreateSealedSegment(schema);
    auto* sealed = dynamic_cast<ChunkedSegmentSealedImpl*>(segment.get());
    ASSERT_NE(sealed, nullptr);

    proto::segcore::SegmentLoadInfo proto;
    proto.set_segmentid(60001);
    proto.set_num_of_rows(0);
    sealed->SetLoadInfo(proto);

    ASSERT_FALSE(
        sealed->TestGetLoadInfoSnapshot()->HasTextIndexCreated(text_fid));
    sealed->TestRecordTextIndexCreated(text_fid);
    EXPECT_TRUE(
        sealed->TestGetLoadInfoSnapshot()->HasTextIndexCreated(text_fid));
}

TEST(SealedSegmentReopen, SchemaOnlyReopenPreservesCreatedTextIndexState) {
    std::map<std::string, std::string> analyzer_params;
    auto old_schema = std::make_shared<Schema>();
    auto pk_id = old_schema->AddDebugField("pk", DataType::INT64);
    auto text_fid = old_schema->AddDebugVarcharField(FieldName("text_field"),
                                                     DataType::VARCHAR,
                                                     /*max_length=*/65535,
                                                     /*nullable=*/false,
                                                     /*enable_match=*/true,
                                                     /*enable_analyzer=*/true,
                                                     analyzer_params,
                                                     std::nullopt);
    auto new_field = FieldId(text_fid.get() + 1);
    old_schema->set_primary_field_id(pk_id);
    old_schema->set_schema_version(100);

    auto new_schema = std::make_shared<Schema>();
    auto new_pk_id = new_schema->AddDebugField("pk", DataType::INT64);
    auto new_text_fid =
        new_schema->AddDebugVarcharField(FieldName("text_field"),
                                         DataType::VARCHAR,
                                         /*max_length=*/65535,
                                         /*nullable=*/false,
                                         /*enable_match=*/true,
                                         /*enable_analyzer=*/true,
                                         analyzer_params,
                                         std::nullopt);
    ASSERT_EQ(new_schema->AddDebugField("new_field", DataType::INT64),
              new_field);
    new_schema->set_primary_field_id(new_pk_id);
    new_schema->set_schema_version(200);
    ASSERT_EQ(text_fid, new_text_fid);

    auto dataset = DataGen(old_schema, 1);
    auto segment = CreateSealedWithFieldDataLoaded(old_schema, dataset);
    auto* sealed = dynamic_cast<ChunkedSegmentSealedImpl*>(segment.get());
    ASSERT_NE(sealed, nullptr);

    sealed->TestRecordTextIndexCreated(text_fid);
    ASSERT_TRUE(
        sealed->TestGetLoadInfoSnapshot()->HasTextIndexCreated(text_fid));

    EXPECT_NO_THROW(sealed->Reopen(new_schema));
    auto snapshot = sealed->TestGetLoadInfoSnapshot();
    EXPECT_TRUE(snapshot->HasTextIndexCreated(text_fid));
    EXPECT_TRUE(snapshot->HasFieldInSchema(new_field));
}

// Reproducer for issue #49076.
//
// Reopen on a sealed segment does:
//     SegmentLoadInfo new_seg_load_info(new_load_info, schema_);
//     segment_load_info_ = new_seg_load_info;          // <-- Bug 1
//     diff = current.ComputeDiff(new_seg_load_info);
//     ApplyLoadDiff(op_ctx, new_seg_load_info, diff);  // <-- Bug 2
//
// `created_text_indexes_` is runtime-only (not proto-backed), so the RHS is
// always empty; the copy-assign wipes the member (Bug 1). And any Create
// done during this Reopen writes to the local new_seg_load_info, lost on
// return (Bug 2). A later Reopen's ComputeDiff then re-schedules
// CreateTextIndex for the same field at the same deterministic path,
// racing with the still-live holder and causing
// TantivyError: FileDoesNotExist("meta.json").
TEST(SealedSegmentReopen, TextIndexCreatedWipedByReopen) {
    auto schema = std::make_shared<Schema>();
    schema->AddDebugField("pk", DataType::INT64);  // 100
    schema->AddDebugField(
        "vec", DataType::VECTOR_FLOAT, 128, knowhere::metric::L2);  // 101
    std::map<std::string, std::string> analyzer_params;
    schema->AddDebugVarcharField(FieldName("text_field"),
                                 DataType::VARCHAR,
                                 /*max_length=*/65535,
                                 /*nullable=*/false,
                                 /*enable_match=*/true,
                                 /*enable_analyzer=*/true,
                                 analyzer_params,
                                 std::nullopt);  // 102
    schema->set_primary_field_id(FieldId(100));

    const FieldId text_fid(102);

    auto segment = CreateSealedSegment(schema);
    auto* sealed = dynamic_cast<ChunkedSegmentSealedImpl*>(segment.get());
    ASSERT_NE(sealed, nullptr);

    proto::segcore::SegmentLoadInfo proto;
    proto.set_segmentid(49076);
    proto.set_num_of_rows(0);
    sealed->SetLoadInfo(proto);

    // Simulate the state a prior Load would leave behind after running
    // CreateTextIndex for the enable_match field.
    sealed->TestRecordTextIndexCreated(text_fid);
    ASSERT_TRUE(
        sealed->TestGetLoadInfoSnapshot()->HasTextIndexCreated(text_fid));

    // Reopen must preserve the runtime-only created_text_indexes_. Before
    // the fix this fails (member wiped by `segment_load_info_ =
    // new_seg_load_info`); after the fix it passes. The wipe happens before
    // ApplyLoadDiff, so even though ApplyLoadDiff throws here (num_rows_ is
    // unset because no real data was loaded), the state under test is
    // already settled.
    milvus::OpContext op_ctx;
    try {
        sealed->Reopen(&op_ctx, proto);
    } catch (...) {
    }

    EXPECT_TRUE(
        sealed->TestGetLoadInfoSnapshot()->HasTextIndexCreated(text_fid));
}

TEST(SealedSegmentCowState, StagedTextIndexIsInvisibleBeforePublish) {
    auto schema = std::make_shared<Schema>();
    auto pk = schema->AddDebugField("pk", DataType::INT64);
    std::map<std::string, std::string> analyzer_params;
    auto text = schema->AddDebugVarcharField(FieldName("text"),
                                             DataType::VARCHAR,
                                             /*max_length=*/65535,
                                             /*nullable=*/false,
                                             /*enable_match=*/true,
                                             /*enable_analyzer=*/true,
                                             analyzer_params,
                                             std::nullopt);
    schema->set_primary_field_id(pk);

    auto dataset = DataGen(schema, 4);
    auto segment = CreateSealedWithFieldDataLoaded(schema, dataset);
    auto* sealed = dynamic_cast<ChunkedSegmentSealedImpl*>(segment.get());
    ASSERT_NE(sealed, nullptr);

    auto published_before = sealed->TestGetPublishedStateSnapshot();
    ASSERT_EQ(published_before->runtime->text_indexes.count(text), 0);

    auto staged_runtime = sealed->TestCloneMutableRuntimeResourceState();
    sealed->TestCreateTextIndexWithSchema(
        text, schema, nullptr, false, staged_runtime.get());

    EXPECT_EQ(staged_runtime->text_indexes.count(text), 1);
    EXPECT_EQ(published_before->runtime->text_indexes.count(text), 0);
    EXPECT_ANY_THROW(sealed->GetTextIndex(nullptr, text));
    EXPECT_FALSE(sealed->TestGetLoadInfoSnapshot()->HasTextIndexCreated(text));
}

TEST(SealedSegmentCowState, PublishedTextIndexFollowsSnapshotLifetime) {
    auto old_schema = std::make_shared<Schema>();
    auto pk = old_schema->AddDebugField("pk", DataType::INT64);
    std::map<std::string, std::string> analyzer_params;
    auto text = old_schema->AddDebugVarcharField(FieldName("text"),
                                                 DataType::VARCHAR,
                                                 /*max_length=*/65535,
                                                 /*nullable=*/false,
                                                 /*enable_match=*/true,
                                                 /*enable_analyzer=*/true,
                                                 analyzer_params,
                                                 std::nullopt);
    old_schema->set_primary_field_id(pk);
    old_schema->set_schema_version(100);

    auto dataset = DataGen(old_schema, 4);
    auto segment = CreateSealedWithFieldDataLoaded(old_schema, dataset);
    auto* sealed = dynamic_cast<ChunkedSegmentSealedImpl*>(segment.get());
    ASSERT_NE(sealed, nullptr);

    auto before_create = sealed->TestGetPublishedStateSnapshot();
    sealed->CreateTextIndex(text);
    auto published_with_index = sealed->TestGetPublishedStateSnapshot();
    ASSERT_EQ(before_create->runtime->text_indexes.count(text), 0);
    ASSERT_EQ(published_with_index->runtime->text_indexes.count(text), 1);
    EXPECT_TRUE(published_with_index->load_info->HasTextIndexCreated(text));

    auto old_pin = sealed->GetTextIndex(nullptr, text);
    ASSERT_NE(old_pin.get(), nullptr);

    auto new_schema = std::make_shared<Schema>();
    auto new_pk = new_schema->AddDebugField("pk", DataType::INT64);
    new_schema->set_primary_field_id(new_pk);
    new_schema->set_schema_version(200);
    sealed->Reopen(new_schema);

    auto published_after_drop = sealed->TestGetPublishedStateSnapshot();
    EXPECT_EQ(published_after_drop->runtime->text_indexes.count(text), 0);
    EXPECT_ANY_THROW(sealed->GetTextIndex(nullptr, text));
    EXPECT_NE(old_pin.get(), nullptr);
    EXPECT_EQ(published_with_index->runtime->text_indexes.count(text), 1);
}

TEST(SealedSegmentCowState,
     SystemReadyMustNotBeVisibleBeforeTimestampIsPublished) {
    auto schema = std::make_shared<Schema>();
    auto pk = schema->AddDebugField("pk", DataType::INT64);
    schema->AddDebugField(
        "vec", DataType::VECTOR_FLOAT, 16, knowhere::metric::L2);
    schema->set_primary_field_id(pk);

    auto segment = CreateSealedSegment(schema);
    auto* sealed = dynamic_cast<ChunkedSegmentSealedImpl*>(segment.get());
    ASSERT_NE(sealed, nullptr);

    auto before = sealed->TestGetPublishedStateSnapshot();
    ASSERT_FALSE(before->system_field_ready);
    ASSERT_TRUE(before->runtime == nullptr ||
                before->runtime->fields.count(TimestampFieldID) == 0);

    sealed->TestPublishSystemFieldState();

    auto after = sealed->TestGetPublishedStateSnapshot();
    EXPECT_FALSE(after->system_field_ready)
        << "system readiness must be derived from published timestamp state";
    EXPECT_TRUE(after->runtime == nullptr ||
                after->runtime->fields.count(TimestampFieldID) == 0);
}

TEST(SealedSegmentCowState, ReopenNextStateMustNotInheritStaleFieldReadyBit) {
    auto schema = std::make_shared<Schema>();
    auto pk = schema->AddDebugField("pk", DataType::INT64);
    auto payload = schema->AddDebugField("payload", DataType::INT64);
    schema->set_primary_field_id(pk);
    schema->set_schema_version(100);

    auto dataset = DataGen(schema, 1);
    auto segment = CreateSealedWithFieldDataLoaded(schema, dataset);
    auto* sealed = dynamic_cast<ChunkedSegmentSealedImpl*>(segment.get());
    ASSERT_NE(sealed, nullptr);

    auto current = sealed->TestGetPublishedStateSnapshot();
    ASSERT_TRUE(GetFieldBit(current->field_data_ready_bitset, payload));

    auto next_runtime = sealed->TestCloneMutableRuntimeResourceState();
    next_runtime->fields.erase(payload);
    next_runtime->array_offsets_map.erase(payload);

    ChunkedSegmentSealedImpl::StateDelta delta;
    delta.schema = current->schema;
    delta.load_info = current->load_info;
    delta.runtime =
        sealed->TestFreezeRuntimeResourceState(std::move(next_runtime));
    delta.commit_ts = current->commit_ts;

    auto next = sealed->TestBuildNextPublishedState(current, delta);
    EXPECT_FALSE(GetFieldBit(next->field_data_ready_bitset, payload))
        << "field readiness must be rebuilt from published runtime";
}

TEST(SealedSegmentCowState, ReopenNextStateMustNotInheritStaleVectorIndexBit) {
    auto schema = std::make_shared<Schema>();
    auto pk = schema->AddDebugField("pk", DataType::INT64);
    auto vec = schema->AddDebugField(
        "vec", DataType::VECTOR_FLOAT, 16, knowhere::metric::L2);
    schema->set_primary_field_id(pk);
    schema->set_schema_version(100);

    auto segment = CreateSealedSegment(schema);
    auto* sealed = dynamic_cast<ChunkedSegmentSealedImpl*>(segment.get());
    ASSERT_NE(sealed, nullptr);

    sealed->TestPublishVectorIndexFacts(vec,
                                        /*ready=*/true,
                                        /*binlog_ready=*/false,
                                        /*has_raw_data=*/false);

    auto current = sealed->TestGetPublishedStateSnapshot();
    ASSERT_TRUE(GetFieldBit(current->index_ready_bitset, vec));
    ASSERT_FALSE(GetFieldBit(current->binlog_index_bitset, vec));
    ASSERT_TRUE(current->published_index_has_raw_data.count(vec));
    ASSERT_FALSE(current->published_index_has_raw_data.at(vec));

    ChunkedSegmentSealedImpl::StateDelta delta;
    delta.schema = current->schema;
    delta.load_info = current->load_info;
    delta.runtime = current->runtime;
    delta.commit_ts = current->commit_ts;
    delta.published_index_ready_bitset =
        BitsetType(current->published_index_ready_bitset.size());
    delta.published_binlog_index_ready_bitset =
        BitsetType(current->published_binlog_index_ready_bitset.size());
    delta.published_index_has_raw_data = std::unordered_map<FieldId, bool>{};

    auto next = sealed->TestBuildNextPublishedState(current, delta);
    EXPECT_FALSE(GetFieldBit(next->index_ready_bitset, vec))
        << "vector index readiness must be rebuilt from published vector facts";
    EXPECT_FALSE(GetFieldBit(next->binlog_index_bitset, vec));
    EXPECT_TRUE(next->index_has_raw_data.find(vec) ==
                next->index_has_raw_data.end());
}

TEST(SealedSegmentCowState, DropShrunkHighestFieldStateDoesNotReadPastBitset) {
    auto old_schema = std::make_shared<Schema>();
    auto old_pk = old_schema->AddDebugField("pk", DataType::INT64);
    auto dropped_vec = old_schema->AddDebugField(
        "vec", DataType::VECTOR_FLOAT, 4, knowhere::metric::L2);
    old_schema->set_primary_field_id(old_pk);
    old_schema->set_schema_version(100);

    auto new_schema = std::make_shared<Schema>();
    auto new_pk = new_schema->AddDebugField("pk", DataType::INT64);
    new_schema->set_primary_field_id(new_pk);
    new_schema->set_schema_version(200);

    auto segment = CreateSealedSegment(old_schema);
    auto* sealed = dynamic_cast<ChunkedSegmentSealedImpl*>(segment.get());
    ASSERT_NE(sealed, nullptr);

    sealed->TestPublishVectorIndexFacts(dropped_vec,
                                        /*ready=*/true,
                                        /*binlog_ready=*/false,
                                        /*has_raw_data=*/true);
    auto current = sealed->TestGetPublishedStateSnapshot();
    ASSERT_EQ(current->index_ready_bitset.size(), 2);
    ASSERT_TRUE(current->published_index_has_raw_data.count(dropped_vec));

    ChunkedSegmentSealedImpl::StateDelta delta;
    delta.schema = new_schema;
    delta.load_info = std::make_shared<const SegmentLoadInfo>(
        current->load_info->GetProto(), new_schema);
    delta.runtime = current->runtime;
    delta.commit_ts = current->commit_ts;

    auto staged_drop_field =
        sealed->TestBuildNextPublishedState(current, delta);
    ASSERT_EQ(staged_drop_field->index_ready_bitset.size(), 1);
    ASSERT_TRUE(
        staged_drop_field->published_index_has_raw_data.count(dropped_vec));
    ASSERT_NO_THROW(
        sealed->TestDropFieldFromState(*staged_drop_field, dropped_vec));
    EXPECT_FALSE(
        staged_drop_field->published_index_has_raw_data.count(dropped_vec));
    EXPECT_TRUE(staged_drop_field->index_has_raw_data.find(dropped_vec) ==
                staged_drop_field->index_has_raw_data.end());

    auto staged_drop_index =
        sealed->TestBuildNextPublishedState(current, delta);
    ASSERT_EQ(staged_drop_index->index_ready_bitset.size(), 1);
    ASSERT_TRUE(
        staged_drop_index->published_index_has_raw_data.count(dropped_vec));
    ASSERT_NO_THROW(
        sealed->TestDropIndexFromState(*staged_drop_index, dropped_vec));
    EXPECT_FALSE(
        staged_drop_index->published_index_has_raw_data.count(dropped_vec));
    EXPECT_TRUE(staged_drop_index->index_has_raw_data.find(dropped_vec) ==
                staged_drop_index->index_has_raw_data.end());
}

TEST(SealedSegmentCowState, StagedVectorIndexLoadUsesResizedNewSchemaBitset) {
    auto old_schema = std::make_shared<Schema>();
    auto old_pk = old_schema->AddDebugField("pk", DataType::INT64);
    old_schema->set_primary_field_id(old_pk);
    old_schema->set_schema_version(100);

    auto new_schema = std::make_shared<Schema>();
    auto new_pk = new_schema->AddDebugField("pk", DataType::INT64);
    auto new_vec = new_schema->AddDebugField(
        "vec", DataType::VECTOR_FLOAT, 4, knowhere::metric::L2);
    new_schema->set_primary_field_id(new_pk);
    new_schema->set_schema_version(200);

    auto segment = CreateSealedSegment(old_schema);
    auto* sealed = dynamic_cast<ChunkedSegmentSealedImpl*>(segment.get());
    ASSERT_NE(sealed, nullptr);

    auto current = sealed->TestGetPublishedStateSnapshot();
    ASSERT_EQ(current->index_ready_bitset.size(), 1);
    ASSERT_FALSE(current->schema->get_fields().count(new_vec));

    auto runtime = sealed->TestCloneMutableRuntimeResourceState();
    auto staged_load_info = std::make_shared<const SegmentLoadInfo>(
        current->load_info->GetProto(), new_schema);
    ChunkedSegmentSealedImpl::StateDelta initial_delta;
    initial_delta.schema = new_schema;
    initial_delta.load_info = staged_load_info;
    initial_delta.runtime = sealed->TestFreezeRuntimeResourceState(runtime);
    initial_delta.commit_ts = current->commit_ts;
    auto staged = sealed->TestBuildNextPublishedState(current, initial_delta);

    ASSERT_EQ(staged->index_ready_bitset.size(), 2);
    ASSERT_FALSE(GetFieldBit(staged->index_ready_bitset, new_vec));

    milvus::index::CreateIndexInfo create_index_info;
    create_index_info.field_type = DataType::VECTOR_FLOAT;
    create_index_info.metric_type = knowhere::metric::L2;
    create_index_info.index_type = knowhere::IndexEnum::INDEX_FAISS_IVFFLAT;
    create_index_info.index_engine_version =
        knowhere::Version::GetCurrentVersion().VersionNumber();
    auto indexing = milvus::index::IndexFactory::GetInstance().CreateIndex(
        create_index_info, milvus::storage::FileManagerContext());

    LoadIndexInfo load_info;
    load_info.field_id = new_vec.get();
    load_info.index_params = GenIndexParams(indexing.get());
    load_info.index_params["metric_type"] = knowhere::metric::L2;
    load_info.cache_index =
        CreateTestCacheIndex("staged-vector", std::move(indexing));
    load_info.load_resource_request = LoadResourceRequest{0, 0, 0, 0, true};

    ChunkedSegmentSealedImpl::StateDelta final_delta;
    final_delta.schema = new_schema;
    final_delta.load_info = staged_load_info;
    final_delta.commit_ts = current->commit_ts;

    ASSERT_NO_THROW(sealed->TestStageLoadIndexThenPublish(
        load_info,
        false,
        new_schema,
        runtime,
        staged.get(),
        current,
        final_delta,
        [&] {
            auto published_after_staged_load =
                sealed->TestGetPublishedStateSnapshot();
            EXPECT_FALSE(
                published_after_staged_load->schema->get_fields().count(
                    new_vec));
            EXPECT_EQ(published_after_staged_load->index_ready_bitset.size(),
                      1);
            EXPECT_FALSE(sealed->TestVectorIndexReady(new_vec));

            EXPECT_TRUE(GetFieldBit(staged->index_ready_bitset, new_vec));
            EXPECT_TRUE(staged->published_index_has_raw_data.at(new_vec));
        }));
    auto final_state = sealed->TestGetPublishedStateSnapshot();
    EXPECT_TRUE(GetFieldBit(final_state->index_ready_bitset, new_vec));
    EXPECT_TRUE(final_state->index_has_raw_data.at(new_vec));
    EXPECT_TRUE(sealed->TestVectorIndexReady(new_vec));
}

TEST(SealedSegmentCowState,
     WarmupResolverHandlesAddedVectorAbsentFromPublishedBitset) {
    auto old_schema = CreateWarmupPolicySchema(/*include_vector=*/false);
    auto new_schema = CreateWarmupPolicySchema(/*include_vector=*/true);
    const FieldId vec(kWarmupVectorFieldId);

    auto segment = CreateSealedSegment(old_schema);
    auto* sealed = dynamic_cast<ChunkedSegmentSealedImpl*>(segment.get());
    ASSERT_NE(sealed, nullptr);

    auto current = sealed->TestGetPublishedStateSnapshot();
    ASSERT_NE(current, nullptr);
    ASSERT_LE(current->binlog_index_bitset.size(), 1);
    ASSERT_FALSE(current->schema->get_fields().count(vec));

    SegmentLoadInfo staged_load_info(current->load_info->GetProto(),
                                     new_schema);
    auto [has_field_warmup, field_warmup_policy] =
        new_schema->WarmupPolicy(vec, /*is_vector=*/true, /*is_index=*/false);
    ASSERT_TRUE(has_field_warmup);
    ASSERT_EQ(field_warmup_policy, "sync");

    std::string resolved;
    ASSERT_NO_THROW(resolved = sealed->TestResolveFieldDataWarmupPolicy(
                        vec, staged_load_info, new_schema));
    EXPECT_EQ(resolved, "disable");
}

TEST(SealedSegmentCowState,
     StagedStorageV2ColumnGroupUsesVectorIndexWarmupForNoIndexVector) {
    auto schema = CreateWarmupPolicySchema(/*include_vector=*/true);
    const FieldId vec(kWarmupVectorFieldId);

    auto segment = CreateSealedSegment(schema);
    auto* sealed = dynamic_cast<ChunkedSegmentSealedImpl*>(segment.get());
    ASSERT_NE(sealed, nullptr);

    proto::segcore::SegmentLoadInfo load_proto;
    load_proto.set_segmentid(1006);
    load_proto.set_num_of_rows(1);
    SegmentLoadInfo segment_load_info(load_proto, schema);

    ASSERT_EQ(sealed->TestResolveFieldDataGroupWarmupPolicy(
                  {vec}, segment_load_info, schema),
              "disable");

    auto column_groups = MakeWarmupTestColumnGroups();
    auto reader = std::make_shared<WarmupTestReader>(column_groups);
    auto column = sealed->TestStageLoadColumnGroupWithReader(
        column_groups,
        std::make_shared<milvus_storage::api::Properties>(),
        /*index=*/0,
        {vec},
        segment_load_info,
        schema,
        std::move(reader),
        /*eager_load=*/true);

    auto proxy_column = std::dynamic_pointer_cast<ProxyChunkColumn>(column);
    ASSERT_NE(proxy_column, nullptr);
    EXPECT_EQ(proxy_column->TestCacheWarmupPolicy(),
              CacheWarmupPolicy::CacheWarmupPolicy_Disable);
}

TEST(SealedSegmentCowState, StagedVectorIndexSkipsInterimIndexGeneration) {
    auto schema = std::make_shared<Schema>();
    auto pk = schema->AddDebugField("pk", DataType::INT64);
    auto vec = schema->AddDebugField(
        "vec", DataType::VECTOR_FLOAT, 4, knowhere::metric::L2);
    schema->set_primary_field_id(pk);

    std::map<std::string, std::string> index_params = {
        {"index_type", knowhere::IndexEnum::INDEX_FAISS_IVFFLAT},
        {"metric_type", knowhere::metric::L2},
        {"nlist", "64"}};
    std::map<std::string, std::string> type_params = {{"dim", "4"}};
    FieldIndexMeta field_index_meta(
        vec, std::move(index_params), std::move(type_params));
    std::map<FieldId, FieldIndexMeta> field_indexes = {{vec, field_index_meta}};
    IndexMetaPtr index_meta =
        std::make_shared<CollectionIndexMeta>(100000, std::move(field_indexes));

    auto config = SegcoreConfig::default_config();
    config.set_enable_interim_segment_index(true);
    config.set_chunk_rows(1024);
    config.set_dense_vector_intermin_index_type(
        knowhere::IndexEnum::INDEX_FAISS_IVFFLAT_CC);
    config.set_nlist(16);
    config.set_nprobe(16);

    auto dataset = DataGen(schema, 1000);
    auto data_segment = CreateSealedWithFieldDataLoaded(schema, dataset);
    auto* data_sealed =
        dynamic_cast<ChunkedSegmentSealedImpl*>(data_segment.get());
    ASSERT_NE(data_sealed, nullptr);
    auto data_runtime = data_sealed->TestCloneMutableRuntimeResourceState();
    ASSERT_TRUE(data_runtime->fields.count(vec));
    auto loaded_column = data_runtime->fields.at(vec);
    auto row_count = dataset.raw_->num_rows();

    auto unstaged_segment = CreateSealedSegment(schema, index_meta, 0, config);
    auto* unstaged =
        dynamic_cast<ChunkedSegmentSealedImpl*>(unstaged_segment.get());
    ASSERT_NE(unstaged, nullptr);
    ASSERT_TRUE(
        unstaged->TestGenerateInterimIndex(vec, row_count, loaded_column));

    auto segment = CreateSealedSegment(schema, index_meta, 0, config);
    auto* sealed = dynamic_cast<ChunkedSegmentSealedImpl*>(segment.get());
    ASSERT_NE(sealed, nullptr);

    auto current = sealed->TestGetPublishedStateSnapshot();
    auto runtime = sealed->TestCloneMutableRuntimeResourceState();
    auto staged_load_info = std::make_shared<const SegmentLoadInfo>(
        current->load_info->GetProto(), schema);
    ChunkedSegmentSealedImpl::StateDelta initial_delta;
    initial_delta.schema = schema;
    initial_delta.load_info = staged_load_info;
    initial_delta.runtime = sealed->TestFreezeRuntimeResourceState(runtime);
    initial_delta.commit_ts = current->commit_ts;
    auto staged = sealed->TestBuildNextPublishedState(current, initial_delta);

    milvus::index::CreateIndexInfo create_index_info;
    create_index_info.field_type = DataType::VECTOR_FLOAT;
    create_index_info.metric_type = knowhere::metric::L2;
    create_index_info.index_type = knowhere::IndexEnum::INDEX_FAISS_IVFFLAT;
    create_index_info.index_engine_version =
        knowhere::Version::GetCurrentVersion().VersionNumber();
    auto indexing = milvus::index::IndexFactory::GetInstance().CreateIndex(
        create_index_info, milvus::storage::FileManagerContext());

    LoadIndexInfo load_info;
    load_info.field_id = vec.get();
    load_info.index_params = GenIndexParams(indexing.get());
    load_info.index_params["metric_type"] = knowhere::metric::L2;
    load_info.cache_index =
        CreateTestCacheIndex("staged-real-vector", std::move(indexing));
    load_info.load_resource_request = LoadResourceRequest{0, 0, 0, 0, true};

    ChunkedSegmentSealedImpl::StateDelta final_delta;
    final_delta.schema = schema;
    final_delta.load_info = staged_load_info;
    final_delta.commit_ts = current->commit_ts;

    milvus::OpContext op_ctx;
    ASSERT_NO_THROW(sealed->TestStageLoadIndexGenerateInterimThenPublish(
        load_info,
        false,
        schema,
        runtime,
        staged.get(),
        current,
        final_delta,
        loaded_column,
        row_count,
        &op_ctx,
        [&](bool generated_interim) {
            EXPECT_FALSE(generated_interim);
            EXPECT_TRUE(GetFieldBit(staged->index_ready_bitset, vec));
            EXPECT_FALSE(GetFieldBit(staged->binlog_index_bitset, vec));
            EXPECT_TRUE(staged->published_index_has_raw_data.at(vec));
        }));

    auto final_state = sealed->TestGetPublishedStateSnapshot();
    EXPECT_TRUE(GetFieldBit(final_state->index_ready_bitset, vec));
    EXPECT_FALSE(GetFieldBit(final_state->binlog_index_bitset, vec));
    EXPECT_TRUE(final_state->index_has_raw_data.at(vec));
    EXPECT_TRUE(sealed->TestVectorIndexReady(vec));
}

TEST(SealedSegmentCowState,
     StagedVectorIndexRetiresPublishedWarmupOnlyAfterPublish) {
    auto schema = std::make_shared<Schema>();
    auto pk = schema->AddDebugField("pk", DataType::INT64);
    auto vec = schema->AddDebugField(
        "vec", DataType::VECTOR_FLOAT, 4, knowhere::metric::L2);
    schema->set_primary_field_id(pk);

    auto warmup_limit = cachinglayer::ResourceUsage{1024 * 1024, 0};
    auto warmup_dlist = std::make_shared<cachinglayer::internal::DList>(
        false,
        warmup_limit,
        warmup_limit,
        warmup_limit,
        cachinglayer::EvictionConfig{});
    auto segment = CreateSealedSegment(schema);
    auto* sealed = dynamic_cast<ChunkedSegmentSealedImpl*>(segment.get());
    ASSERT_NE(sealed, nullptr);

    auto create_index = [] {
        milvus::index::CreateIndexInfo create_index_info;
        create_index_info.field_type = DataType::VECTOR_FLOAT;
        create_index_info.metric_type = knowhere::metric::L2;
        create_index_info.index_type = knowhere::IndexEnum::INDEX_FAISS_IVFFLAT;
        create_index_info.index_engine_version =
            knowhere::Version::GetCurrentVersion().VersionNumber();
        return milvus::index::IndexFactory::GetInstance().CreateIndex(
            create_index_info, milvus::storage::FileManagerContext());
    };

    folly::CancellationToken published_warmup_token;
    std::promise<void> warmup_started;
    auto warmup_started_future = warmup_started.get_future();
    auto published_index_impl = create_index();
    auto published_index_params = GenIndexParams(published_index_impl.get());
    auto published_cache_index =
        CreateCancellationObservingCacheIndex("published-vector",
                                              std::move(published_index_impl),
                                              &published_warmup_token,
                                              &warmup_started,
                                              warmup_dlist.get());
    auto warmup_pool = std::make_shared<folly::CPUThreadPoolExecutor>(1);
    published_cache_index->Warmup(nullptr, warmup_pool);
    ASSERT_EQ(warmup_started_future.wait_for(std::chrono::seconds(5)),
              std::future_status::ready);
    ASSERT_FALSE(published_warmup_token.isCancellationRequested());

    LoadIndexInfo published_index;
    published_index.field_id = vec.get();
    published_index.index_params = std::move(published_index_params);
    published_index.index_params["metric_type"] = knowhere::metric::L2;
    published_index.cache_index = published_cache_index;
    published_index.load_resource_request =
        LoadResourceRequest{0, 0, 0, 0, true};
    sealed->LoadIndex(published_index);

    auto stage_replacement = [&](bool fail_before_publish) {
        auto current = sealed->TestGetPublishedStateSnapshot();
        auto runtime = sealed->TestCloneMutableRuntimeResourceState();
        auto staged = sealed->TestBuildNextPublishedState(
            current,
            {current->schema,
             current->load_info,
             sealed->TestFreezeRuntimeResourceState(runtime),
             current->commit_ts});

        auto replacement_impl = create_index();
        LoadIndexInfo replacement;
        replacement.field_id = vec.get();
        replacement.index_params = GenIndexParams(replacement_impl.get());
        replacement.index_params["metric_type"] = knowhere::metric::L2;
        replacement.cache_index = CreateTestCacheIndex(
            fail_before_publish ? "rolled-back-vector" : "replacement-vector",
            std::move(replacement_impl));
        replacement.load_resource_request =
            LoadResourceRequest{0, 0, 0, 0, true};

        ChunkedSegmentSealedImpl::StateDelta final_delta;
        final_delta.schema = current->schema;
        final_delta.load_info = current->load_info;
        final_delta.commit_ts = current->commit_ts;

        auto stage_and_publish = [&] {
            sealed->TestStageLoadIndexThenPublish(
                replacement,
                true,
                current->schema,
                runtime,
                staged.get(),
                current,
                final_delta,
                [&] {
                    auto published = sealed->TestGetPublishedStateSnapshot();
                    ASSERT_EQ(
                        published->runtime->vector_indexings.at(vec)->indexing_,
                        published_cache_index);
                    EXPECT_FALSE(
                        published_warmup_token.isCancellationRequested());
                    if (fail_before_publish) {
                        throw std::runtime_error("abort staged publish");
                    }
                });
        };

        if (fail_before_publish) {
            EXPECT_THROW(stage_and_publish(), std::runtime_error);
        } else {
            EXPECT_NO_THROW(stage_and_publish());
        }
    };

    stage_replacement(true);
    EXPECT_FALSE(published_warmup_token.isCancellationRequested());
    auto after_rollback = sealed->TestGetPublishedStateSnapshot();
    ASSERT_EQ(after_rollback->runtime->vector_indexings.at(vec)->indexing_,
              published_cache_index);

    stage_replacement(false);
    EXPECT_TRUE(published_warmup_token.isCancellationRequested());
    auto after_publish = sealed->TestGetPublishedStateSnapshot();
    ASSERT_NE(after_publish->runtime->vector_indexings.at(vec)->indexing_,
              published_cache_index);
}

TEST(SealedSegmentCowState, ReplaceScalarIndexStagesRuntimeUntilFinalPublish) {
    auto schema = std::make_shared<Schema>();
    auto pk = schema->AddDebugField("pk", DataType::INT64);
    auto payload = schema->AddDebugField("payload", DataType::INT64);
    schema->set_primary_field_id(pk);

    auto dataset = DataGen(schema, 4);
    auto segment = CreateSealedWithFieldDataLoaded(schema, dataset);
    auto* sealed = dynamic_cast<ChunkedSegmentSealedImpl*>(segment.get());
    ASSERT_NE(sealed, nullptr);

    LoadIndexInfo first_index;
    first_index.field_id = payload.get();
    first_index.field_type = DataType::INT64;
    first_index.index_params["index_type"] = "STL_SORT";
    auto payload_data = dataset.get_col<int64_t>(payload);
    auto first_impl = GenScalarIndexing<int64_t>(dataset.raw_->num_rows(),
                                                 payload_data.data());
    first_index.index_params = GenIndexParams(first_impl.get());
    first_index.cache_index =
        CreateTestCacheIndex("first", std::move(first_impl));
    auto first_cache_index = first_index.cache_index;
    sealed->LoadIndex(first_index);

    auto before = sealed->TestGetPublishedStateSnapshot();
    ASSERT_NE(before, nullptr);
    ASSERT_NE(before->runtime, nullptr);
    ASSERT_TRUE(before->runtime->scalar_indexings.count(payload) > 0);
    auto published_before_replace =
        before->runtime->scalar_indexings.at(payload);
    ASSERT_EQ(published_before_replace, first_cache_index);

    auto runtime = sealed->TestCloneMutableRuntimeResourceState();
    ASSERT_TRUE(runtime->scalar_indexings.count(payload) > 0);
    auto stale_runtime_index = runtime->scalar_indexings.at(payload);
    ASSERT_EQ(stale_runtime_index, first_cache_index);

    LoadIndexInfo replacement_index;
    replacement_index.field_id = payload.get();
    replacement_index.field_type = DataType::INT64;
    replacement_index.index_params["index_type"] = "STL_SORT";
    auto replacement_impl = GenScalarIndexing<int64_t>(dataset.raw_->num_rows(),
                                                       payload_data.data());
    replacement_index.index_params = GenIndexParams(replacement_impl.get());
    replacement_index.cache_index =
        CreateTestCacheIndex("replacement", std::move(replacement_impl));
    auto replacement_cache_index = replacement_index.cache_index;

    sealed->TestLoadIndex(replacement_index, true, runtime.get());

    ASSERT_TRUE(runtime->scalar_indexings.count(payload) > 0);
    EXPECT_EQ(runtime->scalar_indexings.size(), 1);
    auto updated_runtime_index = runtime->scalar_indexings.at(payload);
    EXPECT_EQ(updated_runtime_index, replacement_cache_index);
    EXPECT_NE(updated_runtime_index, stale_runtime_index);

    auto staged_only = sealed->TestGetPublishedStateSnapshot();
    ASSERT_NE(staged_only, nullptr);
    ASSERT_NE(staged_only->runtime, nullptr);
    ASSERT_TRUE(staged_only->runtime->scalar_indexings.count(payload) > 0);
    EXPECT_EQ(staged_only->runtime->scalar_indexings.at(payload),
              first_cache_index);

    ChunkedSegmentSealedImpl::StateDelta delta;
    delta.schema = staged_only->schema;
    delta.load_info = staged_only->load_info;
    delta.runtime = sealed->TestFreezeRuntimeResourceState(std::move(runtime));
    delta.commit_ts = staged_only->commit_ts;

    auto next = sealed->TestBuildNextPublishedState(staged_only, delta);
    ASSERT_NE(next, nullptr);
    ASSERT_NE(next->runtime, nullptr);
    EXPECT_EQ(next->runtime->scalar_indexings.size(), 1);
    ASSERT_TRUE(next->runtime->scalar_indexings.count(payload) > 0);
    auto published_runtime_index = next->runtime->scalar_indexings.at(payload);
    EXPECT_EQ(published_runtime_index, replacement_cache_index);
    EXPECT_NE(published_runtime_index, stale_runtime_index);
    EXPECT_EQ(next->runtime->ngram_fields.count(payload), 0);
}

TEST(SealedSegmentCowState, ReplacePkStateIsInvisibleUntilFinalPublish) {
    auto schema = std::make_shared<Schema>();
    auto pk = schema->AddDebugField("pk", DataType::INT64);
    schema->set_primary_field_id(pk);

    auto old_dataset = DataGen(schema, 4);
    auto segment = CreateSealedWithFieldDataLoaded(schema, old_dataset);
    auto* sealed = dynamic_cast<ChunkedSegmentSealedImpl*>(segment.get());
    ASSERT_NE(sealed, nullptr);

    auto replacement_dataset = DataGen(schema,
                                       4,
                                       /*seed=*/314159,
                                       /*ts_offset=*/0,
                                       /*repeat_count=*/1,
                                       /*array_len=*/10,
                                       /*group_count=*/1,
                                       /*random_pk=*/true);
    auto replacement_segment =
        CreateSealedWithFieldDataLoaded(schema, replacement_dataset);
    auto* replacement_sealed =
        dynamic_cast<ChunkedSegmentSealedImpl*>(replacement_segment.get());
    ASSERT_NE(replacement_sealed, nullptr);

    auto old_pks = old_dataset.get_col<int64_t>(pk);
    auto replacement_pks = replacement_dataset.get_col<int64_t>(pk);
    auto old_only =
        std::find_if(old_pks.begin(), old_pks.end(), [&](int64_t value) {
            return std::find(replacement_pks.begin(),
                             replacement_pks.end(),
                             value) == replacement_pks.end();
        });
    auto replacement_only = std::find_if(
        replacement_pks.begin(), replacement_pks.end(), [&](int64_t value) {
            return std::find(old_pks.begin(), old_pks.end(), value) ==
                   old_pks.end();
        });
    ASSERT_NE(old_only, old_pks.end());
    ASSERT_NE(replacement_only, replacement_pks.end());
    PkType old_pk = *old_only;
    PkType replacement_pk = *replacement_only;

    auto current = sealed->TestGetPublishedStateSnapshot();
    ASSERT_NE(current->runtime->pk_index_slot, nullptr);
    ASSERT_TRUE(sealed->Contain(old_pk));
    ASSERT_FALSE(sealed->Contain(replacement_pk));
    auto old_pk_index_slot = current->runtime->pk_index_slot;

    auto replacement_state =
        replacement_sealed->TestGetPublishedStateSnapshot();
    auto replacement_column = replacement_state->runtime->fields.at(pk);

    auto stage_replacement = [&](bool fail_before_publish) {
        auto runtime = sealed->TestCloneMutableRuntimeResourceState();
        ChunkedSegmentSealedImpl::StateDelta initial_delta;
        initial_delta.schema = current->schema;
        initial_delta.load_info = current->load_info;
        initial_delta.runtime = sealed->TestFreezeRuntimeResourceState(runtime);
        initial_delta.commit_ts = current->commit_ts;
        auto staged =
            sealed->TestBuildNextPublishedState(current, initial_delta);

        ChunkedSegmentSealedImpl::StateDelta final_delta;
        final_delta.schema = current->schema;
        final_delta.load_info = current->load_info;
        final_delta.commit_ts = current->commit_ts;

        auto stage_and_publish = [&] {
            sealed->TestStageLoadFieldDataThenPublish(
                pk,
                replacement_column,
                replacement_dataset.raw_->num_rows(),
                DataType::INT64,
                schema,
                runtime,
                staged.get(),
                current,
                final_delta,
                [&] {
                    ASSERT_NE(runtime->pk_index_slot, nullptr);
                    EXPECT_NE(runtime->pk_index_slot,
                              current->runtime->pk_index_slot);

                    auto still_published =
                        sealed->TestGetPublishedStateSnapshot();
                    EXPECT_EQ(still_published, current);
                    EXPECT_TRUE(sealed->Contain(old_pk));
                    EXPECT_FALSE(sealed->Contain(replacement_pk));
                    if (fail_before_publish) {
                        throw std::runtime_error("abort staged PK publish");
                    }
                });
        };

        if (fail_before_publish) {
            EXPECT_THROW(stage_and_publish(), std::runtime_error);
        } else {
            EXPECT_NO_THROW(stage_and_publish());
        }
    };

    stage_replacement(true);
    EXPECT_EQ(sealed->TestGetPublishedStateSnapshot(), current);
    EXPECT_TRUE(sealed->Contain(old_pk));
    EXPECT_FALSE(sealed->Contain(replacement_pk));

    stage_replacement(false);
    auto published = sealed->TestGetPublishedStateSnapshot();
    ASSERT_NE(published, current);
    EXPECT_NE(published->runtime->pk_index_slot,
              current->runtime->pk_index_slot);
    EXPECT_FALSE(sealed->Contain(old_pk));
    EXPECT_TRUE(sealed->Contain(replacement_pk));

    EXPECT_EQ(current->runtime->pk_index_slot, old_pk_index_slot);
}

TEST(SealedSegmentCowState, ClearPublishedStateDropsRuntimeSnapshot) {
    auto schema = std::make_shared<Schema>();
    auto pk = schema->AddDebugField("pk", DataType::INT64);
    auto payload = schema->AddDebugField("payload", DataType::INT64);
    schema->set_primary_field_id(pk);

    auto dataset = DataGen(schema, 1);
    auto segment = CreateSealedWithFieldDataLoaded(schema, dataset);
    auto* sealed = dynamic_cast<ChunkedSegmentSealedImpl*>(segment.get());
    ASSERT_NE(sealed, nullptr);

    auto before = sealed->TestGetPublishedStateSnapshot();
    ASSERT_NE(before, nullptr);
    ASSERT_NE(before->runtime, nullptr);
    ASSERT_TRUE(before->runtime->fields.count(payload) > 0);

    sealed->TestClearPublishedState();

    auto after = sealed->TestGetPublishedStateSnapshot();
    ASSERT_NE(after, nullptr);
    EXPECT_EQ(after->runtime, nullptr);
    EXPECT_FALSE(GetFieldBit(after->field_data_ready_bitset, payload));
}

TEST(SealedSegmentCowState, JsonIndexStagesAndFollowsSnapshotLifetime) {
    auto schema = std::make_shared<Schema>();
    auto pk = schema->AddDebugField("pk", DataType::INT64);
    auto json = schema->AddDebugField("payload", DataType::JSON);
    schema->set_primary_field_id(pk);

    auto segment = CreateSealedSegment(schema);
    auto* sealed = dynamic_cast<ChunkedSegmentSealedImpl*>(segment.get());
    ASSERT_NE(sealed, nullptr);

    std::array<int64_t, 4> values = {1, 2, 3, 4};
    auto indexing = GenScalarIndexing<int64_t>(values.size(), values.data());
    LoadIndexInfo load_info;
    load_info.field_id = json.get();
    load_info.field_type = DataType::JSON;
    load_info.index_params = {
        {index::INDEX_TYPE, index::INVERTED_INDEX_TYPE},
        {JSON_PATH, "a"},
        {JSON_CAST_TYPE, "DOUBLE"},
    };
    load_info.cache_index =
        CreateTestCacheIndex("json-runtime", std::move(indexing));
    auto loaded_index = load_info.cache_index;

    auto current = sealed->TestGetPublishedStateSnapshot();
    ASSERT_TRUE(current->runtime->json_indices.empty());
    auto runtime = sealed->TestCloneMutableRuntimeResourceState();
    ChunkedSegmentSealedImpl::StateDelta initial_delta;
    initial_delta.schema = current->schema;
    initial_delta.load_info = current->load_info;
    initial_delta.runtime = sealed->TestFreezeRuntimeResourceState(runtime);
    initial_delta.commit_ts = current->commit_ts;
    auto staged = sealed->TestBuildNextPublishedState(current, initial_delta);

    ChunkedSegmentSealedImpl::StateDelta final_delta;
    final_delta.schema = current->schema;
    final_delta.load_info = current->load_info;
    final_delta.commit_ts = current->commit_ts;
    sealed->TestStageLoadIndexThenPublish(
        load_info,
        false,
        schema,
        runtime,
        staged.get(),
        current,
        final_delta,
        [&] {
            EXPECT_EQ(runtime->json_indices.size(), 1);
            EXPECT_TRUE(current->runtime->json_indices.empty());
            EXPECT_FALSE(sealed->HasJsonIndex(json));
        });

    auto published = sealed->TestGetPublishedStateSnapshot();
    ASSERT_EQ(published->runtime->json_indices.size(), 1);
    EXPECT_EQ(published->runtime->json_indices.front().index, loaded_index);
    EXPECT_TRUE(sealed->HasJsonIndex(json));

    sealed->DropJSONIndex(json, "a");
    auto after_drop = sealed->TestGetPublishedStateSnapshot();
    EXPECT_TRUE(after_drop->runtime->json_indices.empty());
    EXPECT_FALSE(sealed->HasJsonIndex(json));

    ASSERT_EQ(published->runtime->json_indices.size(), 1);
    EXPECT_EQ(published->runtime->json_indices.front().index, loaded_index);
}

TEST(SealedSegmentCowState, JsonIndexReplaceScalarWithNgramErasesScalarPath) {
    auto schema = std::make_shared<Schema>();
    auto pk = schema->AddDebugField("pk", DataType::INT64);
    auto json = schema->AddDebugField("payload", DataType::JSON);
    schema->set_primary_field_id(pk);

    auto segment = CreateSealedSegment(schema);
    auto* sealed = dynamic_cast<ChunkedSegmentSealedImpl*>(segment.get());
    ASSERT_NE(sealed, nullptr);

    std::array<int64_t, 4> values = {1, 2, 3, 4};
    auto make_load_info = [&](std::string key,
                              std::string path,
                              bool is_ngram) {
        LoadIndexInfo info;
        info.field_id = json.get();
        info.field_type = DataType::JSON;
        info.index_params = {
            {index::INDEX_TYPE,
             is_ngram ? index::NGRAM_INDEX_TYPE : index::INVERTED_INDEX_TYPE},
            {JSON_PATH, std::move(path)},
        };
        if (!is_ngram) {
            info.index_params[JSON_CAST_TYPE] = "DOUBLE";
        }
        auto indexing =
            GenScalarIndexing<int64_t>(values.size(), values.data());
        info.cache_index =
            CreateTestCacheIndex(std::move(key), std::move(indexing));
        return info;
    };

    auto sibling = make_load_info("json-scalar-b", "b", false);
    auto sibling_index = sibling.cache_index;
    sealed->LoadIndex(sibling);

    auto original = make_load_info("json-scalar-a", "a", false);
    auto original_index = original.cache_index;
    sealed->LoadIndex(original);

    auto current = sealed->TestGetPublishedStateSnapshot();
    ASSERT_EQ(current->runtime->json_indices.size(), 2);
    EXPECT_TRUE(current->runtime->ngram_indexings.empty());
    EXPECT_FALSE(GetFieldBit(current->index_ready_bitset, json));

    auto runtime = sealed->TestCloneMutableRuntimeResourceState();
    ChunkedSegmentSealedImpl::StateDelta initial_delta;
    initial_delta.schema = current->schema;
    initial_delta.load_info = current->load_info;
    initial_delta.runtime = sealed->TestFreezeRuntimeResourceState(runtime);
    initial_delta.commit_ts = current->commit_ts;
    auto staged = sealed->TestBuildNextPublishedState(current, initial_delta);

    auto replacement = make_load_info("json-ngram-a", "a", true);
    auto replacement_index = replacement.cache_index;
    ChunkedSegmentSealedImpl::StateDelta final_delta;
    final_delta.schema = current->schema;
    final_delta.load_info = current->load_info;
    final_delta.commit_ts = current->commit_ts;
    sealed->TestStageLoadIndexThenPublish(
        replacement,
        true,
        schema,
        runtime,
        staged.get(),
        current,
        final_delta,
        [&] {
            ASSERT_EQ(runtime->json_indices.size(), 1);
            EXPECT_EQ(runtime->json_indices.front().nested_path, "b");
            EXPECT_EQ(runtime->json_indices.front().index, sibling_index);
            ASSERT_EQ(runtime->ngram_indexings.count(json), 1);
            ASSERT_EQ(runtime->ngram_indexings.at(json).count("a"), 1);
            EXPECT_EQ(runtime->ngram_indexings.at(json).at("a"),
                      replacement_index);
            EXPECT_TRUE(GetFieldBit(staged->index_ready_bitset, json));

            ASSERT_EQ(current->runtime->json_indices.size(), 2);
            EXPECT_TRUE(current->runtime->ngram_indexings.empty());
        });

    auto published = sealed->TestGetPublishedStateSnapshot();
    ASSERT_EQ(published->runtime->json_indices.size(), 1);
    EXPECT_EQ(published->runtime->json_indices.front().nested_path, "b");
    EXPECT_EQ(published->runtime->json_indices.front().index, sibling_index);
    ASSERT_EQ(published->runtime->ngram_indexings.count(json), 1);
    EXPECT_EQ(published->runtime->ngram_indexings.at(json).at("a"),
              replacement_index);
    EXPECT_TRUE(GetFieldBit(published->index_ready_bitset, json));

    ASSERT_EQ(current->runtime->json_indices.size(), 2);
    auto old_path = std::find_if(
        current->runtime->json_indices.begin(),
        current->runtime->json_indices.end(),
        [](const auto& index) { return index.nested_path == "a"; });
    ASSERT_NE(old_path, current->runtime->json_indices.end());
    EXPECT_EQ(old_path->index, original_index);
}

TEST(SealedSegmentCowState, JsonIndexReplaceNgramWithScalarErasesNgramPath) {
    auto schema = std::make_shared<Schema>();
    auto pk = schema->AddDebugField("pk", DataType::INT64);
    auto json = schema->AddDebugField("payload", DataType::JSON);
    schema->set_primary_field_id(pk);

    auto segment = CreateSealedSegment(schema);
    auto* sealed = dynamic_cast<ChunkedSegmentSealedImpl*>(segment.get());
    ASSERT_NE(sealed, nullptr);

    std::array<int64_t, 4> values = {1, 2, 3, 4};
    auto make_load_info = [&](std::string key,
                              std::string path,
                              bool is_ngram) {
        LoadIndexInfo info;
        info.field_id = json.get();
        info.field_type = DataType::JSON;
        info.index_params = {
            {index::INDEX_TYPE,
             is_ngram ? index::NGRAM_INDEX_TYPE : index::INVERTED_INDEX_TYPE},
            {JSON_PATH, std::move(path)},
        };
        if (!is_ngram) {
            info.index_params[JSON_CAST_TYPE] = "DOUBLE";
        }
        auto indexing =
            GenScalarIndexing<int64_t>(values.size(), values.data());
        info.cache_index =
            CreateTestCacheIndex(std::move(key), std::move(indexing));
        return info;
    };

    auto sibling = make_load_info("json-scalar-b", "b", false);
    auto sibling_index = sibling.cache_index;
    sealed->LoadIndex(sibling);

    auto original = make_load_info("json-ngram-a", "a", true);
    auto original_index = original.cache_index;
    sealed->LoadIndex(original);

    auto current = sealed->TestGetPublishedStateSnapshot();
    ASSERT_EQ(current->runtime->json_indices.size(), 1);
    ASSERT_EQ(current->runtime->ngram_indexings.count(json), 1);
    EXPECT_EQ(current->runtime->ngram_indexings.at(json).at("a"),
              original_index);
    EXPECT_TRUE(GetFieldBit(current->index_ready_bitset, json));

    auto runtime = sealed->TestCloneMutableRuntimeResourceState();
    ChunkedSegmentSealedImpl::StateDelta initial_delta;
    initial_delta.schema = current->schema;
    initial_delta.load_info = current->load_info;
    initial_delta.runtime = sealed->TestFreezeRuntimeResourceState(runtime);
    initial_delta.commit_ts = current->commit_ts;
    auto staged = sealed->TestBuildNextPublishedState(current, initial_delta);

    auto replacement = make_load_info("json-scalar-a", "a", false);
    auto replacement_index = replacement.cache_index;
    ChunkedSegmentSealedImpl::StateDelta final_delta;
    final_delta.schema = current->schema;
    final_delta.load_info = current->load_info;
    final_delta.commit_ts = current->commit_ts;
    sealed->TestStageLoadIndexThenPublish(
        replacement,
        true,
        schema,
        runtime,
        staged.get(),
        current,
        final_delta,
        [&] {
            EXPECT_TRUE(runtime->ngram_indexings.empty());
            ASSERT_EQ(runtime->json_indices.size(), 2);
            auto replacement_path = std::find_if(
                runtime->json_indices.begin(),
                runtime->json_indices.end(),
                [](const auto& index) { return index.nested_path == "a"; });
            ASSERT_NE(replacement_path, runtime->json_indices.end());
            EXPECT_EQ(replacement_path->index, replacement_index);
            EXPECT_FALSE(GetFieldBit(staged->index_ready_bitset, json));

            ASSERT_EQ(current->runtime->ngram_indexings.count(json), 1);
            EXPECT_EQ(current->runtime->ngram_indexings.at(json).at("a"),
                      original_index);
        });

    auto published = sealed->TestGetPublishedStateSnapshot();
    EXPECT_TRUE(published->runtime->ngram_indexings.empty());
    ASSERT_EQ(published->runtime->json_indices.size(), 2);
    auto sibling_path = std::find_if(
        published->runtime->json_indices.begin(),
        published->runtime->json_indices.end(),
        [](const auto& index) { return index.nested_path == "b"; });
    ASSERT_NE(sibling_path, published->runtime->json_indices.end());
    EXPECT_EQ(sibling_path->index, sibling_index);
    auto replacement_path = std::find_if(
        published->runtime->json_indices.begin(),
        published->runtime->json_indices.end(),
        [](const auto& index) { return index.nested_path == "a"; });
    ASSERT_NE(replacement_path, published->runtime->json_indices.end());
    EXPECT_EQ(replacement_path->index, replacement_index);
    EXPECT_FALSE(GetFieldBit(published->index_ready_bitset, json));

    ASSERT_EQ(current->runtime->ngram_indexings.count(json), 1);
    EXPECT_EQ(current->runtime->ngram_indexings.at(json).at("a"),
              original_index);
}

TEST(SealedSegmentCowState, JsonStatsLivesInRuntimeSnapshot) {
    auto schema = std::make_shared<Schema>();
    auto pk = schema->AddDebugField("pk", DataType::INT64);
    auto json = schema->AddDebugField("payload", DataType::JSON);
    schema->set_primary_field_id(pk);

    auto segment = CreateSealedSegment(schema);
    auto* sealed = dynamic_cast<ChunkedSegmentSealedImpl*>(segment.get());
    ASSERT_NE(sealed, nullptr);

    auto current = sealed->TestGetPublishedStateSnapshot();
    auto load_runtime = sealed->TestCloneMutableRuntimeResourceState();
    load_runtime->json_stats[json] = nullptr;

    ChunkedSegmentSealedImpl::StateDelta load_delta;
    load_delta.schema = current->schema;
    load_delta.load_info = current->load_info;
    load_delta.runtime =
        sealed->TestFreezeRuntimeResourceState(std::move(load_runtime));
    load_delta.commit_ts = current->commit_ts;

    auto after_load = sealed->TestBuildNextPublishedState(current, load_delta);
    ASSERT_NE(after_load, nullptr);
    ASSERT_NE(after_load->runtime, nullptr);
    ASSERT_EQ(after_load->runtime->json_stats.count(json), 1);

    auto drop_runtime =
        std::make_shared<ChunkedSegmentSealedImpl::RuntimeResourceState>(
            *after_load->runtime);
    ASSERT_EQ(drop_runtime->json_stats.count(json), 1);
    drop_runtime->json_stats.erase(json);

    ChunkedSegmentSealedImpl::StateDelta delta;
    delta.schema = after_load->schema;
    delta.load_info = after_load->load_info;
    delta.runtime =
        sealed->TestFreezeRuntimeResourceState(std::move(drop_runtime));
    delta.commit_ts = after_load->commit_ts;

    auto next = sealed->TestBuildNextPublishedState(after_load, delta);
    ASSERT_NE(next, nullptr);
    ASSERT_NE(next->runtime, nullptr);
    EXPECT_EQ(next->runtime->json_stats.count(json), 0);
}

TEST(SealedSegmentCowState, ExternalSyntheticFieldsUseStagedLoadInfoRows) {
    auto schema = std::make_shared<Schema>();
    auto pk = schema->AddDebugField("pk", DataType::INT64);
    schema->set_primary_field_id(pk);
    schema->set_external_source("s3://bucket/data");
    schema->set_external_spec(R"({"format":"parquet"})");

    auto segment = CreateSealedSegment(schema);
    auto* sealed = dynamic_cast<ChunkedSegmentSealedImpl*>(segment.get());
    ASSERT_NE(sealed, nullptr);

    proto::segcore::SegmentLoadInfo published_proto;
    published_proto.set_segmentid(1001);
    published_proto.set_num_of_rows(3);
    sealed->SetLoadInfo(published_proto);

    proto::segcore::SegmentLoadInfo staged_proto;
    staged_proto.set_segmentid(1001);
    staged_proto.set_num_of_rows(7);
    SegmentLoadInfo staged_load_info(staged_proto, schema);

    auto runtime = sealed->TestCloneMutableRuntimeResourceState();
    sealed->TestSynthesizeExternalSystemFields(
        staged_load_info, schema, runtime.get());

    ASSERT_TRUE(runtime->fields.count(pk) > 0);
    ASSERT_NE(runtime->timestamps, nullptr);
    EXPECT_EQ(runtime->timestamps->size(), 7);
    EXPECT_EQ((*runtime->timestamps)[0], 0);
    EXPECT_EQ((*runtime->timestamps)[6], 0);

    auto published = sealed->TestGetPublishedStateSnapshot();
    ASSERT_NE(published, nullptr);
    ASSERT_NE(published->load_info, nullptr);
    EXPECT_EQ(published->load_info->GetNumOfRows(), 3);
}

TEST(SealedSegmentCowState, ExternalZeroRowsClearsStagedPkState) {
    auto schema = std::make_shared<Schema>();
    auto pk = schema->AddDebugField("pk", DataType::INT64);
    schema->set_primary_field_id(pk);
    schema->set_external_source("s3://bucket/data");
    schema->set_external_spec(R"({"format":"parquet"})");

    auto segment = CreateSealedSegment(schema);
    auto* sealed = dynamic_cast<ChunkedSegmentSealedImpl*>(segment.get());
    ASSERT_NE(sealed, nullptr);

    auto runtime = sealed->TestCloneMutableRuntimeResourceState();
    proto::segcore::SegmentLoadInfo nonempty_proto;
    nonempty_proto.set_segmentid(1005);
    nonempty_proto.set_num_of_rows(7);
    SegmentLoadInfo nonempty_load_info(nonempty_proto, schema);
    sealed->TestSynthesizeExternalSystemFields(
        nonempty_load_info, schema, runtime.get());
    ASSERT_TRUE(runtime->fields.count(pk) > 0);
    ASSERT_NE(runtime->virtual_pk2offset, nullptr);

    proto::segcore::SegmentLoadInfo empty_proto;
    empty_proto.set_segmentid(1005);
    empty_proto.set_num_of_rows(0);
    SegmentLoadInfo empty_load_info(empty_proto, schema);
    sealed->TestSynthesizeExternalSystemFields(
        empty_load_info, schema, runtime.get());

    EXPECT_EQ(runtime->fields.count(pk), 0);
    EXPECT_EQ(runtime->pk_index_slot, nullptr);
    ASSERT_NE(runtime->virtual_pk2offset, nullptr);
    EXPECT_FALSE(runtime->virtual_pk2offset->contain(PkType{int64_t{0}}));
    ASSERT_NE(runtime->timestamps, nullptr);
    EXPECT_TRUE(runtime->timestamps->empty());
}

TEST(SealedSegmentCowState,
     PublishRuntimeStateRecomputesExternalSystemFieldReady) {
    auto schema = std::make_shared<Schema>();
    auto pk = schema->AddDebugField("pk", DataType::INT64);
    schema->set_primary_field_id(pk);
    schema->set_external_source("s3://bucket/data");
    schema->set_external_spec(R"({"format":"parquet"})");

    auto segment = CreateSealedSegment(schema);
    auto* sealed = dynamic_cast<ChunkedSegmentSealedImpl*>(segment.get());
    ASSERT_NE(sealed, nullptr);

    proto::segcore::SegmentLoadInfo published_proto;
    published_proto.set_segmentid(1002);
    published_proto.set_num_of_rows(7);
    sealed->SetLoadInfo(published_proto);

    auto runtime = sealed->TestCloneMutableRuntimeResourceState();
    SegmentLoadInfo staged_load_info(published_proto, schema);
    sealed->TestSynthesizeExternalSystemFields(
        staged_load_info, schema, runtime.get());

    auto current = sealed->TestGetPublishedStateSnapshot();
    auto next = sealed->TestBuildNextPublishedState(
        current,
        {current->schema,
         current->load_info,
         sealed->TestFreezeRuntimeResourceState(std::move(runtime)),
         current->commit_ts});

    ASSERT_NE(next, nullptr);
    EXPECT_TRUE(next->system_field_ready);
    ASSERT_NE(next->runtime, nullptr);
    ASSERT_NE(next->runtime->timestamps, nullptr);
    EXPECT_EQ(next->runtime->timestamps->size(), 7);
    EXPECT_EQ((*next->runtime->timestamps)[0], 0);
}

TEST(SealedSegmentCowState,
     ExternalPublicLoadColumnGroupsUsesStagedRuntimeForSyntheticFields) {
    auto schema = std::make_shared<Schema>();
    auto pk = schema->AddDebugField("pk", DataType::INT64);
    schema->set_primary_field_id(pk);
    schema->set_external_source("s3://bucket/data");
    schema->set_external_spec(R"({"format":"parquet"})");

    auto segment = CreateSealedSegment(schema);
    auto* sealed = dynamic_cast<ChunkedSegmentSealedImpl*>(segment.get());
    ASSERT_NE(sealed, nullptr);

    proto::segcore::SegmentLoadInfo published_proto;
    published_proto.set_segmentid(1003);
    published_proto.set_num_of_rows(3);
    sealed->SetLoadInfo(published_proto);

    proto::segcore::SegmentLoadInfo staged_proto;
    staged_proto.set_segmentid(1003);
    staged_proto.set_num_of_rows(7);
    SegmentLoadInfo staged_load_info(staged_proto, schema);

    auto runtime = sealed->TestCloneMutableRuntimeResourceState();
    sealed->TestSynthesizeExternalSystemFields(
        staged_load_info, schema, runtime.get());

    auto current = sealed->TestGetPublishedStateSnapshot();
    auto next = sealed->TestBuildNextPublishedState(
        current,
        {current->schema,
         current->load_info,
         sealed->TestFreezeRuntimeResourceState(std::move(runtime)),
         current->commit_ts});

    ASSERT_NE(next, nullptr);
    ASSERT_NE(next->runtime, nullptr);
    auto it = next->runtime->fields.find(pk);
    ASSERT_TRUE(it != next->runtime->fields.end());
    ASSERT_NE(it->second, nullptr);
    EXPECT_TRUE(next->system_field_ready);
}

TEST(SealedSegmentCowState, ExternalSynthesizeRequiresRuntime) {
    auto schema = std::make_shared<Schema>();
    auto pk = schema->AddDebugField("pk", DataType::INT64);
    schema->set_primary_field_id(pk);
    schema->set_external_source("s3://bucket/data");
    schema->set_external_spec(R"({"format":"parquet"})");

    auto segment = CreateSealedSegment(schema);
    auto* sealed = dynamic_cast<ChunkedSegmentSealedImpl*>(segment.get());
    ASSERT_NE(sealed, nullptr);

    proto::segcore::SegmentLoadInfo published_proto;
    published_proto.set_segmentid(1004);
    published_proto.set_num_of_rows(3);
    sealed->SetLoadInfo(published_proto);

    SegmentLoadInfo staged_load_info(published_proto, schema);
    EXPECT_ANY_THROW(sealed->TestSynthesizeExternalSystemFields(
        staged_load_info, schema, nullptr));
}

TEST(SealedSegmentCowState, ExternalWrapperSynthesizeRequiresRuntime) {
    auto schema = std::make_shared<Schema>();
    auto pk = schema->AddDebugField("pk", DataType::INT64);
    schema->set_primary_field_id(pk);
    schema->set_external_source("s3://bucket/data");
    schema->set_external_spec(R"({"format":"parquet"})");

    auto segment = CreateSealedSegment(schema);
    auto* sealed = dynamic_cast<ChunkedSegmentSealedImpl*>(segment.get());
    ASSERT_NE(sealed, nullptr);

    proto::segcore::SegmentLoadInfo published_proto;
    published_proto.set_segmentid(1005);
    published_proto.set_num_of_rows(3);
    sealed->SetLoadInfo(published_proto);

    EXPECT_ANY_THROW(sealed->TestSynthesizeExternalSystemFields(nullptr));
}