// 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 <folly/FBVector.h>
#include <gtest/gtest.h>
#include <nlohmann/json.hpp>
#include <stddef.h>
#include <algorithm>
#include <array>
#include <cstdint>
#include <filesystem>
#include <iostream>
#include <map>
#include <memory>
#include <numeric>
#include <optional>
#include <string>
#include <thread>
#include <tuple>
#include <utility>
#include <vector>

#include "NamedType/named_type_impl.hpp"
#include "bitset/bitset.h"
#include "bitset/detail/element_vectorized.h"
#include "cachinglayer/Utils.h"
#include "common/Consts.h"
#include "common/EasyAssert.h"
#include "common/IndexMeta.h"
#include "common/QueryResult.h"
#include "common/Schema.h"
#include "common/Types.h"
#include "common/Utils.h"
#include "common/VectorTrait.h"
#include "common/protobuf_utils.h"
#include "expr/ITypeExpr.h"
#include "filemanager/InputStream.h"
#include "gtest/gtest.h"
#include "knowhere/comp/index_param.h"
#include "knowhere/dataset.h"
#include "knowhere/object.h"
#include "knowhere/sparse_utils.h"
#include "pb/common.pb.h"
#include "pb/schema.pb.h"
#include "pb/segcore.pb.h"
#include "plan/PlanNode.h"
#include "query/ExecPlanNodeVisitor.h"
#include "query/Plan.h"
#include "query/PlanNode.h"
#include "segcore/ConcurrentVector.h"
#include "segcore/InsertRecord.h"
#include "segcore/SegcoreConfig.h"
#include "segcore/SegmentGrowing.h"
#include "segcore/SegmentGrowingImpl.h"
#include "segcore/Utils.h"
#include "test_utils/DataGen.h"
#include "test_utils/ManifestTestUtil.h"
#include "test_utils/storage_test_utils.h"

using namespace milvus::segcore;
using namespace milvus;

namespace {

void
AddStorageV3SystemFields(const SchemaPtr& schema) {
    schema->AddField(
        FieldName("RowID"), RowFieldID, DataType::INT64, false, std::nullopt);
    schema->AddField(FieldName("Timestamp"),
                     TimestampFieldID,
                     DataType::INT64,
                     false,
                     std::nullopt);
}

bool
ManifestHasField(const milvus::test::V3SegmentTestData& test_data,
                 FieldId field_id) {
    auto field_column = std::to_string(field_id.get());
    auto column_groups = test_data.GetColumnGroups();
    for (size_t i = 0; i < column_groups->size(); ++i) {
        const auto& columns = column_groups->at(i)->columns;
        if (std::find(columns.begin(), columns.end(), field_column) !=
            columns.end()) {
            return true;
        }
    }
    return false;
}

}  // namespace

TEST(Growing, DeleteCount) {
    auto schema = std::make_shared<Schema>();
    auto pk = schema->AddDebugField("pk", DataType::INT64);
    schema->set_primary_field_id(pk);
    auto segment = CreateGrowingSegment(schema, empty_index_meta);

    int64_t c = 10;
    auto offset = 0;

    auto dataset = DataGen(schema, c);
    auto pks = dataset.get_col<int64_t>(pk);
    segment->Insert(offset,
                    c,
                    dataset.row_ids_.data(),
                    dataset.timestamps_.data(),
                    dataset.raw_);

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

    auto cnt = segment->get_deleted_count();
    ASSERT_EQ(cnt, c);
}

TEST(Growing, RealCount) {
    auto schema = std::make_shared<Schema>();
    auto pk = schema->AddDebugField("pk", DataType::INT64);
    schema->set_primary_field_id(pk);
    auto segment = CreateGrowingSegment(schema, empty_index_meta);

    int64_t c = 10;
    auto offset = 0;
    auto dataset = DataGen(schema, c);
    auto pks = dataset.get_col<int64_t>(pk);
    segment->Insert(offset,
                    c,
                    dataset.row_ids_.data(),
                    dataset.timestamps_.data(),
                    dataset.raw_);

    // 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_offset2 = segment->get_deleted_count();
    ASSERT_EQ(del_offset2, half);
    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_offset3 = segment->get_deleted_count();
    ASSERT_EQ(del_offset3, half);
    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(Growing, LoadStorageV3ManifestCapsRowsAtCheckpoint) {
    auto schema = std::make_shared<Schema>();
    AddStorageV3SystemFields(schema);
    auto pk = schema->AddDebugField("pk", DataType::INT64);
    constexpr int64_t dim = 4;
    auto vec = schema->AddDebugField(
        "vec", DataType::VECTOR_FLOAT, dim, knowhere::metric::L2);

    std::map<std::string, std::string> analyzer_params;
    auto text = schema->AddDebugVarcharField(FieldName("text"),
                                             DataType::VARCHAR,
                                             65535,
                                             false,
                                             true,
                                             true,
                                             analyzer_params,
                                             std::nullopt);
    schema->set_primary_field_id(pk);

    auto base_path = (std::filesystem::path(TestLocalPath) /
                      "growing_recovery_checkpoint_row_cap")
                         .string();
    std::filesystem::remove_all(base_path);
    milvus::test::V3SegmentTestData test_data(
        schema, 2, 3, dim, TestLocalPath, base_path);
    ASSERT_EQ(test_data.NumColumnGroups(), 2);
    ASSERT_TRUE(ManifestHasField(test_data, RowFieldID));
    ASSERT_TRUE(ManifestHasField(test_data, TimestampFieldID));

    constexpr int64_t checkpoint_rows = 4;
    milvus::proto::segcore::SegmentLoadInfo load_info;
    load_info.set_collectionid(1);
    load_info.set_partitionid(2);
    load_info.set_segmentid(3);
    load_info.set_storageversion(STORAGE_V3);
    load_info.set_num_of_rows(checkpoint_rows);
    load_info.set_manifest_path(test_data.ManifestPathJson());
    load_info.set_insert_channel("by-dev-rootcoord-dml_0_1v0");

    auto segment =
        CreateGrowingSegment(schema, empty_index_meta, load_info.segmentid());
    segment->SetLoadInfo(load_info);
    milvus::tracer::TraceContext trace_ctx;
    segment->Load(trace_ctx, nullptr);
    ASSERT_EQ(segment->get_row_count(), checkpoint_rows);
    EXPECT_EQ(segment->get_real_count(), checkpoint_rows);

    std::vector<int64_t> row_ids = {checkpoint_rows};
    std::vector<Timestamp> timestamps = {100};
    std::vector<int64_t> pks = {100000};
    std::vector<std::string> texts = {"text after recovery checkpoint"};
    std::vector<float> vectors(dim, 1.0F);

    auto insert_data = std::make_unique<InsertRecordProto>();
    insert_data->set_num_rows(1);
    insert_data->mutable_fields_data()->AddAllocated(
        CreateDataArrayFrom(pks.data(), nullptr, 1, (*schema)[pk]).release());
    insert_data->mutable_fields_data()->AddAllocated(
        CreateDataArrayFrom(texts.data(), nullptr, 1, (*schema)[text])
            .release());
    insert_data->mutable_fields_data()->AddAllocated(
        CreateDataArrayFrom(vectors.data(), nullptr, 1, (*schema)[vec])
            .release());

    auto offset = segment->PreInsert(1);
    ASSERT_EQ(offset, checkpoint_rows);
    ASSERT_NO_THROW(segment->Insert(
        offset, 1, row_ids.data(), timestamps.data(), insert_data.get()));
    EXPECT_EQ(segment->get_row_count(), checkpoint_rows + 1);

    std::filesystem::remove_all(base_path);
}

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

    auto base_path = (std::filesystem::path(TestLocalPath) /
                      "growing_recovery_short_required_rows")
                         .string();
    std::filesystem::remove_all(base_path);
    milvus::test::V3SegmentTestData test_data(
        schema, 1, 2, 1, TestLocalPath, base_path);
    ASSERT_TRUE(ManifestHasField(test_data, RowFieldID));
    ASSERT_TRUE(ManifestHasField(test_data, TimestampFieldID));

    milvus::proto::segcore::SegmentLoadInfo load_info;
    load_info.set_collectionid(1);
    load_info.set_partitionid(2);
    load_info.set_segmentid(4);
    load_info.set_storageversion(STORAGE_V3);
    load_info.set_num_of_rows(test_data.TotalRows() + 1);
    load_info.set_manifest_path(test_data.ManifestPathJson());
    load_info.set_insert_channel("by-dev-rootcoord-dml_0_1v0");

    auto segment =
        CreateGrowingSegment(schema, empty_index_meta, load_info.segmentid());
    segment->SetLoadInfo(load_info);
    milvus::tracer::TraceContext trace_ctx;
    ASSERT_ANY_THROW(segment->Load(trace_ctx, nullptr));
    EXPECT_EQ(segment->get_row_count(), 0);

    std::filesystem::remove_all(base_path);
}

TEST(Growing, InsertSkipsMissingFunctionOutputField) {
    auto schema = std::make_shared<Schema>();
    schema->set_schema_version(2);
    auto pk = schema->AddDebugField("pk", DataType::INT64);
    schema->set_primary_field_id(pk);
    auto sparse = FieldId(pk.get() + 1);
    schema->AddField(FieldName("sparse_out"),
                     sparse,
                     DataType::VECTOR_SPARSE_U32_F32,
                     0,
                     std::nullopt,
                     false);
    schema->add_function_output_field_id(sparse);

    auto segment = CreateGrowingSegment(schema, empty_index_meta);

    auto insert_schema = std::make_shared<Schema>();
    insert_schema->set_schema_version(1);
    insert_schema->AddField(
        FieldName("pk"), pk, DataType::INT64, false, std::nullopt);
    insert_schema->set_primary_field_id(pk);

    constexpr int64_t row_count = 5;
    auto dataset = DataGen(insert_schema, row_count);
    segment->PreInsert(row_count);
    ASSERT_NO_THROW(segment->Insert(0,
                                    row_count,
                                    dataset.row_ids_.data(),
                                    dataset.timestamps_.data(),
                                    dataset.raw_));
    EXPECT_FALSE(segment->FieldAccessible(sparse));
}

TEST(Growing, MissingStructArrayOffsetsReturnsEmptyForOldRows) {
    auto old_schema = std::make_shared<Schema>();
    old_schema->set_schema_version(1);
    auto pk = old_schema->AddDebugField("pk", DataType::INT64);
    old_schema->set_primary_field_id(pk);
    auto segment = CreateGrowingSegment(old_schema, empty_index_meta);

    constexpr int64_t row_count = 5;
    auto dataset = DataGen(old_schema, row_count);
    segment->Insert(0,
                    row_count,
                    dataset.row_ids_.data(),
                    dataset.timestamps_.data(),
                    dataset.raw_);

    auto new_schema = std::make_shared<Schema>();
    new_schema->set_schema_version(2);
    new_schema->AddField(
        FieldName("pk"), pk, DataType::INT64, false, std::nullopt);
    new_schema->set_primary_field_id(pk);
    auto label = FieldId(pk.get() + 1);
    new_schema->AddField(FieldName("chunks[label]"),
                         label,
                         DataType::ARRAY,
                         DataType::VARCHAR,
                         true);
    segment->Reopen(new_schema);

    auto offsets = segment->GetArrayOffsets(label);
    ASSERT_NE(offsets, nullptr);
    EXPECT_EQ(offsets->GetRowCount(), row_count);
    EXPECT_EQ(offsets->GetTotalElementCount(), 0);
    for (int64_t i = 0; i <= row_count; ++i) {
        auto [start, end] = offsets->ElementIDRangeOfRow(i);
        EXPECT_EQ(start, 0);
        EXPECT_EQ(end, 0);
    }
}

TEST(Growing, LoadMissingStructArrayOffsetsReturnsEmptyForOldRows) {
    auto schema = std::make_shared<Schema>();
    auto pk = schema->AddDebugField("pk", DataType::INT64);
    schema->set_primary_field_id(pk);
    auto label =
        schema->AddDebugArrayField("chunks[label]", DataType::VARCHAR, true);
    auto score =
        schema->AddDebugArrayField("chunks[score]", DataType::INT32, true);

    constexpr int64_t row_count = 5;
    auto dataset = DataGen(schema, row_count);
    auto config = SegcoreConfig::default_config();
    auto segment = CreateGrowingWithFieldDataLoaded(
        schema,
        empty_index_meta,
        config,
        dataset,
        false,
        std::vector<int64_t>{label.get(), score.get()});

    auto* growing = dynamic_cast<SegmentGrowingImpl*>(segment.get());
    ASSERT_NE(growing, nullptr);
    growing->FillAbsentFields();
    ASSERT_EQ(growing->get_row_count(), row_count);

    auto offsets = growing->GetArrayOffsets(label);
    ASSERT_NE(offsets, nullptr);
    auto score_offsets = growing->GetArrayOffsets(score);
    ASSERT_NE(score_offsets, nullptr);
    EXPECT_EQ(offsets.get(), score_offsets.get());
    EXPECT_EQ(offsets->GetRowCount(), row_count);
    EXPECT_EQ(offsets->GetTotalElementCount(), 0);
    for (int64_t i = 0; i <= row_count; ++i) {
        auto [start, end] = offsets->ElementIDRangeOfRow(i);
        EXPECT_EQ(start, 0);
        EXPECT_EQ(end, 0);
    }
}

class GrowingTest
    : public ::testing::TestWithParam<
          std::tuple</*index type*/ std::string, knowhere::MetricType>> {
 public:
    void
    SetUp() override {
        index_type = std::get<0>(GetParam());
        metric_type = std::get<1>(GetParam());
        if (index_type == knowhere::IndexEnum::INDEX_FAISS_IVFFLAT ||
            index_type == knowhere::IndexEnum::INDEX_FAISS_IVFFLAT_CC) {
            data_type = DataType::VECTOR_FLOAT;
        } else if (index_type ==
                       knowhere::IndexEnum::INDEX_SPARSE_INVERTED_INDEX ||
                   index_type == knowhere::IndexEnum::INDEX_SPARSE_WAND) {
            data_type = DataType::VECTOR_SPARSE_U32_F32;
        } else {
            ASSERT_TRUE(false);
        }
    }
    knowhere::MetricType metric_type;
    std::string index_type;
    DataType data_type;
};

INSTANTIATE_TEST_SUITE_P(
    FloatGrowingTest,
    GrowingTest,
    ::testing::Combine(
        ::testing::Values(knowhere::IndexEnum::INDEX_FAISS_IVFFLAT,
                          knowhere::IndexEnum::INDEX_FAISS_IVFFLAT_CC),
        ::testing::Values(knowhere::metric::L2,
                          knowhere::metric::IP,
                          knowhere::metric::COSINE)));

INSTANTIATE_TEST_SUITE_P(
    SparseFloatGrowingTest,
    GrowingTest,
    ::testing::Combine(
        ::testing::Values(knowhere::IndexEnum::INDEX_SPARSE_INVERTED_INDEX,
                          knowhere::IndexEnum::INDEX_SPARSE_WAND),
        ::testing::Values(knowhere::metric::IP)));

TEST_P(GrowingTest, FillData) {
    auto schema = std::make_shared<Schema>();
    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 timestamptz_field =
        schema->AddDebugField("timestamptz", DataType::TIMESTAMPTZ);
    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", data_type, 128, metric_type);
    schema->set_primary_field_id(int64_field);

    std::map<std::string, std::string> index_params = {
        {"index_type", index_type},
        {"metric_type", metric_type},
        {"nlist", "128"}};
    std::map<std::string, std::string> type_params = {{"dim", "128"}};
    FieldIndexMeta fieldIndexMeta(
        vec, std::move(index_params), std::move(type_params));
    auto config = SegcoreConfig::default_config();
    config.set_chunk_rows(1024);
    config.set_enable_interim_segment_index(true);
    std::map<FieldId, FieldIndexMeta> filedMap = {{vec, fieldIndexMeta}};
    IndexMetaPtr metaPtr =
        std::make_shared<CollectionIndexMeta>(100000, std::move(filedMap));
    auto segment_growing = CreateGrowingSegment(schema, metaPtr, 1, config);
    auto segment = dynamic_cast<SegmentGrowingImpl*>(segment_growing.get());

    int64_t per_batch = 1000;
    int64_t n_batch = 3;
    int64_t dim = 128;
    for (int64_t i = 0; i < n_batch; i++) {
        auto dataset = DataGen(schema, per_batch);

        auto offset = segment->PreInsert(per_batch);
        segment->Insert(offset,
                        per_batch,
                        dataset.row_ids_.data(),
                        dataset.timestamps_.data(),
                        dataset.raw_);
        auto num_inserted = (i + 1) * per_batch;
        auto ids_ds = GenRandomIds(num_inserted);
        auto bool_result = segment->bulk_subscript(
            nullptr, bool_field, ids_ds->GetIds(), num_inserted);
        auto int8_result = segment->bulk_subscript(
            nullptr, int8_field, ids_ds->GetIds(), num_inserted);
        auto int16_result = segment->bulk_subscript(
            nullptr, int16_field, ids_ds->GetIds(), num_inserted);
        auto int32_result = segment->bulk_subscript(
            nullptr, int32_field, ids_ds->GetIds(), num_inserted);
        auto int64_result = segment->bulk_subscript(
            nullptr, int64_field, ids_ds->GetIds(), num_inserted);
        auto float_result = segment->bulk_subscript(
            nullptr, float_field, ids_ds->GetIds(), num_inserted);
        auto double_result = segment->bulk_subscript(
            nullptr, double_field, ids_ds->GetIds(), num_inserted);
        auto timestamptz_result = segment->bulk_subscript(
            nullptr, timestamptz_field, ids_ds->GetIds(), num_inserted);
        auto varchar_result = segment->bulk_subscript(
            nullptr, varchar_field, ids_ds->GetIds(), num_inserted);
        auto json_result = segment->bulk_subscript(
            nullptr, json_field, ids_ds->GetIds(), num_inserted);
        auto geometry_result = segment->bulk_subscript(
            nullptr, geometry_field, ids_ds->GetIds(), num_inserted);
        auto int_array_result = segment->bulk_subscript(
            nullptr, int_array_field, ids_ds->GetIds(), num_inserted);
        auto long_array_result = segment->bulk_subscript(
            nullptr, long_array_field, ids_ds->GetIds(), num_inserted);
        auto bool_array_result = segment->bulk_subscript(
            nullptr, bool_array_field, ids_ds->GetIds(), num_inserted);
        auto string_array_result = segment->bulk_subscript(
            nullptr, string_array_field, ids_ds->GetIds(), num_inserted);
        auto double_array_result = segment->bulk_subscript(
            nullptr, double_array_field, ids_ds->GetIds(), num_inserted);
        auto float_array_result = segment->bulk_subscript(
            nullptr, float_array_field, ids_ds->GetIds(), num_inserted);
        auto vec_result = segment->bulk_subscript(
            nullptr, vec, ids_ds->GetIds(), num_inserted);
        // checking result data
        EXPECT_EQ(bool_result->scalars().bool_data().data_size(), num_inserted);
        EXPECT_EQ(int8_result->scalars().int_data().data_size(), num_inserted);
        EXPECT_EQ(int16_result->scalars().int_data().data_size(), num_inserted);
        EXPECT_EQ(int32_result->scalars().int_data().data_size(), num_inserted);
        EXPECT_EQ(int64_result->scalars().long_data().data_size(),
                  num_inserted);
        EXPECT_EQ(timestamptz_result->scalars().timestamptz_data().data_size(),
                  num_inserted);
        EXPECT_EQ(float_result->scalars().float_data().data_size(),
                  num_inserted);
        EXPECT_EQ(double_result->scalars().double_data().data_size(),
                  num_inserted);
        EXPECT_EQ(varchar_result->scalars().string_data().data_size(),
                  num_inserted);
        EXPECT_EQ(json_result->scalars().json_data().data_size(), num_inserted);
        EXPECT_EQ(geometry_result->scalars().geometry_data().data_size(),
                  num_inserted);
        if (data_type == DataType::VECTOR_FLOAT) {
            EXPECT_EQ(vec_result->vectors().float_vector().data_size(),
                      num_inserted * dim);
        } else if (data_type == DataType::VECTOR_SPARSE_U32_F32) {
            EXPECT_EQ(
                vec_result->vectors().sparse_float_vector().contents_size(),
                num_inserted);
        } else {
            ASSERT_TRUE(false);
        }
        EXPECT_EQ(int_array_result->scalars().array_data().data_size(),
                  num_inserted);
        EXPECT_EQ(long_array_result->scalars().array_data().data_size(),
                  num_inserted);
        EXPECT_EQ(bool_array_result->scalars().array_data().data_size(),
                  num_inserted);
        EXPECT_EQ(string_array_result->scalars().array_data().data_size(),
                  num_inserted);
        EXPECT_EQ(double_array_result->scalars().array_data().data_size(),
                  num_inserted);
        EXPECT_EQ(float_array_result->scalars().array_data().data_size(),
                  num_inserted);

        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(timestamptz_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(Growing, FillNullableData) {
    auto schema = std::make_shared<Schema>();
    auto metric_type = knowhere::metric::L2;
    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);
    auto float_field = schema->AddDebugField("float", DataType::FLOAT, true);
    auto double_field = schema->AddDebugField("double", DataType::DOUBLE, true);
    auto timestamptz_field =
        schema->AddDebugField("timestamptz", DataType::TIMESTAMPTZ, true);
    auto varchar_field =
        schema->AddDebugField("varchar", DataType::VARCHAR, true);
    auto json_field = schema->AddDebugField("json", DataType::JSON, 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, 128, 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", "128"}};
    FieldIndexMeta fieldIndexMeta(
        vec, std::move(index_params), std::move(type_params));
    auto config = SegcoreConfig::default_config();
    config.set_chunk_rows(1024);
    config.set_enable_interim_segment_index(true);
    std::map<FieldId, FieldIndexMeta> filedMap = {{vec, fieldIndexMeta}};
    IndexMetaPtr metaPtr =
        std::make_shared<CollectionIndexMeta>(100000, std::move(filedMap));
    auto segment_growing = CreateGrowingSegment(schema, metaPtr, 1, config);
    auto segment = dynamic_cast<SegmentGrowingImpl*>(segment_growing.get());

    int64_t per_batch = 1000;
    int64_t n_batch = 3;
    int64_t dim = 128;
    for (int64_t i = 0; i < n_batch; i++) {
        auto dataset = DataGen(schema, per_batch);

        auto offset = segment->PreInsert(per_batch);
        segment->Insert(offset,
                        per_batch,
                        dataset.row_ids_.data(),
                        dataset.timestamps_.data(),
                        dataset.raw_);
        auto num_inserted = (i + 1) * per_batch;
        auto ids_ds = GenRandomIds(num_inserted);
        auto bool_result = segment->bulk_subscript(
            nullptr, bool_field, ids_ds->GetIds(), num_inserted);
        auto int8_result = segment->bulk_subscript(
            nullptr, int8_field, ids_ds->GetIds(), num_inserted);
        auto int16_result = segment->bulk_subscript(
            nullptr, int16_field, ids_ds->GetIds(), num_inserted);
        auto int32_result = segment->bulk_subscript(
            nullptr, int32_field, ids_ds->GetIds(), num_inserted);
        auto int64_result = segment->bulk_subscript(
            nullptr, int64_field, ids_ds->GetIds(), num_inserted);
        auto float_result = segment->bulk_subscript(
            nullptr, float_field, ids_ds->GetIds(), num_inserted);
        auto double_result = segment->bulk_subscript(
            nullptr, double_field, ids_ds->GetIds(), num_inserted);
        auto timestamptz_result = segment->bulk_subscript(
            nullptr, timestamptz_field, ids_ds->GetIds(), num_inserted);
        auto varchar_result = segment->bulk_subscript(
            nullptr, varchar_field, ids_ds->GetIds(), num_inserted);
        auto json_result = segment->bulk_subscript(
            nullptr, json_field, ids_ds->GetIds(), num_inserted);
        auto int_array_result = segment->bulk_subscript(
            nullptr, int_array_field, ids_ds->GetIds(), num_inserted);
        auto long_array_result = segment->bulk_subscript(
            nullptr, long_array_field, ids_ds->GetIds(), num_inserted);
        auto bool_array_result = segment->bulk_subscript(
            nullptr, bool_array_field, ids_ds->GetIds(), num_inserted);
        auto string_array_result = segment->bulk_subscript(
            nullptr, string_array_field, ids_ds->GetIds(), num_inserted);
        auto double_array_result = segment->bulk_subscript(
            nullptr, double_array_field, ids_ds->GetIds(), num_inserted);
        auto float_array_result = segment->bulk_subscript(
            nullptr, float_array_field, ids_ds->GetIds(), num_inserted);
        auto vec_result = segment->bulk_subscript(
            nullptr, vec, ids_ds->GetIds(), num_inserted);

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

class GrowingNullableTest : public ::testing::TestWithParam<
                                std::tuple</*data_type*/ DataType,
                                           /*metric_type*/ knowhere::MetricType,
                                           /*index_type*/ std::string,
                                           /*null_percent*/ int,
                                           /*enable_interim_index*/ bool,
                                           /*use_iterator*/ bool>> {
 public:
    void
    SetUp() override {
        std::tie(data_type,
                 metric_type,
                 index_type,
                 null_percent,
                 enable_interim_index,
                 use_iterator) = GetParam();
    }

    DataType data_type;
    knowhere::MetricType metric_type;
    std::string index_type;
    int null_percent;
    bool enable_interim_index;
    bool use_iterator;
};

static std::vector<
    std::tuple<DataType, knowhere::MetricType, std::string, int, bool, bool>>
GenerateGrowingNullableTestParams() {
    std::vector<
        std::
            tuple<DataType, knowhere::MetricType, std::string, int, bool, bool>>
        params;

    // Dense float vectors with IVF_FLAT
    std::vector<std::tuple<DataType, knowhere::MetricType, std::string>>
        base_configs = {
            {DataType::VECTOR_FLOAT,
             knowhere::metric::L2,
             knowhere::IndexEnum::INDEX_FAISS_IVFFLAT},
            {DataType::VECTOR_FLOAT,
             knowhere::metric::IP,
             knowhere::IndexEnum::INDEX_FAISS_IVFFLAT},
            {DataType::VECTOR_FLOAT,
             knowhere::metric::COSINE,
             knowhere::IndexEnum::INDEX_FAISS_IVFFLAT},
            {DataType::VECTOR_SPARSE_U32_F32,
             knowhere::metric::IP,
             knowhere::IndexEnum::INDEX_SPARSE_INVERTED_INDEX},
        };

    std::vector<int> null_percents = {0, 20, 100};

    std::vector<bool> interim_index_configs = {true, false};

    std::vector<bool> iterator_configs = {false, true};

    for (const auto& [dtype, metric, idx_type] : base_configs) {
        for (int null_pct : null_percents) {
            for (bool enable_interim : interim_index_configs) {
                for (bool use_iter : iterator_configs) {
                    // Skip iterator for sparse vectors (not supported)
                    if (use_iter && dtype == DataType::VECTOR_SPARSE_U32_F32) {
                        continue;
                    }
                    params.push_back({dtype,
                                      metric,
                                      idx_type,
                                      null_pct,
                                      enable_interim,
                                      use_iter});
                }
            }
        }
    }
    return params;
}

INSTANTIATE_TEST_SUITE_P(
    NullableVectorParameters,
    GrowingNullableTest,
    ::testing::ValuesIn(GenerateGrowingNullableTestParams()));

TEST_P(GrowingNullableTest, SearchAndQueryNullableVectors) {
    using namespace milvus::query;

    bool nullable = true;

    auto schema = std::make_shared<Schema>();
    auto int64_field = schema->AddDebugField("int64", DataType::INT64);
    int64_t dim = 8;
    auto vec = schema->AddDebugField(
        "embeddings", data_type, dim, metric_type, nullable);
    schema->set_primary_field_id(int64_field);

    std::map<std::string, std::string> index_params;
    std::map<std::string, std::string> type_params;
    if (data_type == DataType::VECTOR_SPARSE_U32_F32) {
        index_params = {{"index_type", index_type},
                        {"metric_type", metric_type}};
        type_params = {};
    } else {
        index_params = {{"index_type", index_type},
                        {"metric_type", metric_type},
                        {"nlist", "128"}};
        type_params = {{"dim", std::to_string(dim)}};
    }
    FieldIndexMeta fieldIndexMeta(
        vec, std::move(index_params), std::move(type_params));
    auto config = SegcoreConfig::default_config();
    config.set_chunk_rows(1024);
    config.set_enable_interim_segment_index(enable_interim_index);
    // Explicitly set interim index type to avoid contamination from other tests
    config.set_dense_vector_intermin_index_type(
        knowhere::IndexEnum::INDEX_FAISS_IVFFLAT_CC);
    std::map<FieldId, FieldIndexMeta> filedMap = {{vec, fieldIndexMeta}};
    IndexMetaPtr metaPtr =
        std::make_shared<CollectionIndexMeta>(100000, std::move(filedMap));
    auto segment_growing = CreateGrowingSegment(schema, metaPtr, 1, config);
    auto segment = dynamic_cast<SegmentGrowingImpl*>(segment_growing.get());

    int64_t batch_size = 2000;
    int64_t num_rounds = 10;
    int64_t topk = 5;
    // Iterator only supports single query
    int64_t num_queries = use_iterator ? 1 : 2;
    Timestamp timestamp = 10000000;

    // Prepare search plan using ScopedSchemaHandle
    milvus::segcore::ScopedSchemaHandle schema_handle(*schema);
    std::vector<char> plan_str;
    if (data_type == DataType::VECTOR_SPARSE_U32_F32) {
        plan_str = schema_handle.ParseSearch(
            "",                               // expression (no filter)
            "embeddings",                     // vector field name
            topk,                             // topk
            metric_type,                      // metric_type
            R"({"drop_ratio_search": 0.1})",  // search_params
            3);                               // round_decimal
    } else if (use_iterator) {
        plan_str = schema_handle.ParseSearchIterator(
            "",                           // expression (no filter)
            "embeddings",                 // vector field name
            topk,                         // topk
            metric_type,                  // metric_type
            R"({"nprobe": 10})",          // search_params
            static_cast<uint32_t>(topk),  // batch_size
            "",                           // token (empty)
            std::nullopt,                 // last_bound (none)
            3);                           // round_decimal
    } else {
        plan_str =
            schema_handle.ParseSearch("",            // expression (no filter)
                                      "embeddings",  // vector field name
                                      topk,          // topk
                                      metric_type,   // metric_type
                                      R"({"nprobe": 10})",  // search_params
                                      3);                   // round_decimal
    }
    auto plan =
        CreateSearchPlanByExpr(schema, plan_str.data(), plan_str.size());

    // Create query vectors
    proto::common::PlaceholderGroup ph_group_raw;
    if (data_type == DataType::VECTOR_SPARSE_U32_F32) {
        ph_group_raw = CreateSparseFloatPlaceholderGroup(num_queries, 42);
    } else {
        auto query_data = generate_float_vector(num_queries, dim);
        ph_group_raw =
            CreatePlaceholderGroupFromBlob(num_queries, dim, query_data.data());
    }

    auto ph_group =
        ParsePlaceholderGroup(plan.get(), ph_group_raw.SerializeAsString());

    // Store all inserted data for verification
    // For nullable vectors, data is stored sparsely (only valid vectors)
    // We need a mapping from logical offset to physical offset
    std::vector<float> all_float_vectors;  // Physical storage (only valid)
    std::vector<knowhere::sparse::SparseRow<float>> all_sparse_vectors;
    std::vector<bool> all_valid_data;  // Logical storage (all rows)
    std::vector<int64_t>
        logical_to_physical;  // Maps logical offset to physical

    // Insert data in multiple rounds and test after each round
    for (int64_t round = 0; round < num_rounds; round++) {
        int64_t total_rows = (round + 1) * batch_size;
        int64_t expected_valid_count =
            total_rows - (total_rows * null_percent / 100);

        auto dataset = DataGen(schema,
                               batch_size,
                               42 + round,
                               0,
                               1,
                               10,
                               1,
                               false,
                               true,
                               false,
                               null_percent);

        // Build logical to physical mapping for this batch
        int64_t base_physical = all_float_vectors.size() / dim;
        if (data_type == DataType::VECTOR_SPARSE_U32_F32) {
            base_physical = all_sparse_vectors.size();
        }

        auto valid_data_from_dataset = dataset.get_col_valid(vec);
        int64_t physical_idx = base_physical;
        for (size_t i = 0; i < valid_data_from_dataset.size(); i++) {
            if (valid_data_from_dataset[i]) {
                logical_to_physical.push_back(physical_idx);
                physical_idx++;
            } else {
                logical_to_physical.push_back(-1);  // null
            }
        }

        // Get original data directly from proto (sparse storage for nullable)
        // Data is stored sparsely - only valid vectors are in the proto
        if (data_type == DataType::VECTOR_FLOAT) {
            auto field_data = dataset.get_col(vec);
            auto& float_data = field_data->vectors().float_vector().data();
            all_float_vectors.insert(
                all_float_vectors.end(), float_data.begin(), float_data.end());
        } else if (data_type == DataType::VECTOR_SPARSE_U32_F32) {
            auto field_data = dataset.get_col(vec);
            auto& sparse_array = field_data->vectors().sparse_float_vector();
            for (int i = 0; i < sparse_array.contents_size(); i++) {
                auto& content = sparse_array.contents(i);
                auto row = CopyAndWrapSparseRow(content.data(), content.size());
                all_sparse_vectors.push_back(std::move(row));
            }
        }
        all_valid_data.insert(all_valid_data.end(),
                              valid_data_from_dataset.begin(),
                              valid_data_from_dataset.end());

        auto offset = segment->PreInsert(batch_size);
        segment->Insert(offset,
                        batch_size,
                        dataset.row_ids_.data(),
                        dataset.timestamps_.data(),
                        dataset.raw_);

        auto& insert_record = segment->get_insert_record();
        ASSERT_TRUE(insert_record.is_valid_data_exist(vec));

        auto valid_data_ptr = insert_record.get_data_base(vec);
        auto valid_data = valid_data_ptr->get_valid_data();

        // Test search
        auto sr =
            segment_growing->Search(plan.get(), ph_group.get(), timestamp);

        ASSERT_EQ(sr->total_nq_, num_queries);
        ASSERT_EQ(sr->unity_topK_, topk);

        if (expected_valid_count == 0) {
            auto total_results = sr->get_total_result_count();
            EXPECT_EQ(total_results, 0)
                << "Round " << round
                << ": 100% null should return 0 results, but got "
                << total_results;
        } else {
            // Verify search results don't contain null vectors
            for (size_t i = 0; i < sr->seg_offsets_.size(); i++) {
                auto seg_offset = sr->seg_offsets_[i];
                if (seg_offset < 0) {
                    continue;
                }
                ASSERT_TRUE(valid_data[seg_offset])
                    << "Round " << round
                    << ": Search returned null vector at offset " << seg_offset;
            }
        }

        auto vec_result = segment->bulk_subscript(
            nullptr, vec, sr->seg_offsets_.data(), sr->seg_offsets_.size());
        ASSERT_TRUE(vec_result != nullptr);

        if (data_type == DataType::VECTOR_FLOAT) {
            auto& float_data = vec_result->vectors().float_vector();
            size_t valid_idx = 0;
            for (size_t i = 0; i < sr->seg_offsets_.size(); i++) {
                auto offset = sr->seg_offsets_[i];
                if (offset < 0) {
                    continue;  // Skip invalid offsets
                }
                auto physical_idx = logical_to_physical[offset];
                for (int d = 0; d < dim; d++) {
                    float expected_val =
                        all_float_vectors[physical_idx * dim + d];
                    float actual_val = float_data.data(valid_idx * dim + d);
                    ASSERT_FLOAT_EQ(expected_val, actual_val)
                        << "Round " << round << ": Mismatch at logical offset "
                        << offset << " dim " << d;
                }
                valid_idx++;
            }
        } else if (data_type == DataType::VECTOR_SPARSE_U32_F32) {
            auto& sparse_data = vec_result->vectors().sparse_float_vector();
            size_t valid_idx = 0;
            for (size_t i = 0; i < sr->seg_offsets_.size(); i++) {
                auto offset = sr->seg_offsets_[i];
                if (offset < 0) {
                    continue;  // Skip invalid offsets
                }
                auto physical_idx = logical_to_physical[offset];
                auto& content = sparse_data.contents(valid_idx);
                auto retrieved_row =
                    CopyAndWrapSparseRow(content.data(), content.size());
                const auto& expected_row = all_sparse_vectors[physical_idx];
                ASSERT_EQ(retrieved_row.size(), expected_row.size())
                    << "Round " << round
                    << ": Sparse vector size mismatch at logical offset "
                    << offset;
                for (size_t j = 0; j < retrieved_row.size(); j++) {
                    ASSERT_EQ(retrieved_row[j].id, expected_row[j].id)
                        << "Round " << round
                        << ": Sparse vector id mismatch at logical offset "
                        << offset << " element " << j;
                    ASSERT_FLOAT_EQ(retrieved_row[j].val, expected_row[j].val)
                        << "Round " << round
                        << ": Sparse vector val mismatch at logical offset "
                        << offset << " element " << j;
                }
                valid_idx++;
            }
        }
    }
}

TEST_P(GrowingTest, FillVectorArrayData) {
    auto schema = std::make_shared<Schema>();
    auto int64_field = schema->AddDebugField("int64", DataType::INT64);
    auto array_float_vector = schema->AddDebugVectorArrayField(
        "array_float_vector", DataType::VECTOR_FLOAT, 128, metric_type);
    schema->set_primary_field_id(int64_field);

    auto config = SegcoreConfig::default_config();
    config.set_chunk_rows(1024);
    config.set_enable_interim_segment_index(true);
    std::map<FieldId, FieldIndexMeta> filedMap = {};
    IndexMetaPtr metaPtr =
        std::make_shared<CollectionIndexMeta>(100000, std::move(filedMap));
    auto segment_growing = CreateGrowingSegment(schema, metaPtr, 1, config);
    auto segment = dynamic_cast<SegmentGrowingImpl*>(segment_growing.get());
    int64_t per_batch = 1000;
    int64_t n_batch = 3;
    for (int64_t i = 0; i < n_batch; i++) {
        auto dataset = DataGen(schema, per_batch);

        auto offset = segment->PreInsert(per_batch);
        segment->Insert(offset,
                        per_batch,
                        dataset.row_ids_.data(),
                        dataset.timestamps_.data(),
                        dataset.raw_);
        auto num_inserted = (i + 1) * per_batch;
        auto ids_ds = GenRandomIds(num_inserted);
        auto int64_result = segment->bulk_subscript(
            nullptr, int64_field, ids_ds->GetIds(), num_inserted);
        auto array_float_vector_result = segment->bulk_subscript(
            nullptr, array_float_vector, ids_ds->GetIds(), num_inserted);

        EXPECT_EQ(int64_result->scalars().long_data().data_size(),
                  num_inserted);
        EXPECT_EQ(
            array_float_vector_result->vectors().vector_array().data_size(),
            num_inserted);

        if (i == 0) {
            // Verify vector array data
            auto verify_float_vectors = [](auto arr1, auto arr2) {
                static constexpr float EPSILON = 1e-6;
                EXPECT_EQ(arr1.size(), arr2.size());
                for (int64_t i = 0; i < arr1.size(); ++i) {
                    EXPECT_NEAR(arr1[i], arr2[i], EPSILON);
                }
            };

            auto array_vec_values =
                dataset.get_col<VectorFieldProto>(array_float_vector);
            for (int64_t i = 0; i < per_batch; ++i) {
                auto arrow_array = array_float_vector_result->vectors()
                                       .vector_array()
                                       .data()[i]
                                       .float_vector()
                                       .data();
                auto expected_array =
                    array_vec_values[ids_ds->GetIds()[i]].float_vector().data();
                verify_float_vectors(arrow_array, expected_array);
            }
        }

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

TEST(GrowingTest, EmptyVectorArrayRowsInitializeElementOneof) {
    auto schema = std::make_shared<Schema>();
    auto array_float_vector =
        schema->AddDebugVectorArrayField("array_float_vector",
                                         DataType::VECTOR_FLOAT,
                                         8,
                                         knowhere::metric::L2,
                                         true);

    auto data =
        CreateEmptyVectorDataArray(2, schema->operator[](array_float_vector));

    ASSERT_EQ(data->vectors().vector_array().data_size(), 2);
    for (const auto& row : data->vectors().vector_array().data()) {
        EXPECT_EQ(row.dim(), 8);
        EXPECT_EQ(row.data_case(),
                  proto::schema::VectorField::DataCase::kFloatVector);
        EXPECT_TRUE(row.float_vector().data().empty());
        EXPECT_NO_THROW(milvus::VectorArray(row));
    }
}

TEST(GrowingTest, QueryNullableVectorArrayUsesPhysicalOffsets) {
    auto schema = std::make_shared<Schema>();
    auto int64_field = schema->AddDebugField("int64", DataType::INT64);
    auto array_float_vector =
        schema->AddDebugVectorArrayField("array_float_vector",
                                         DataType::VECTOR_FLOAT,
                                         4,
                                         knowhere::metric::L2,
                                         true);
    schema->set_primary_field_id(int64_field);

    auto config = SegcoreConfig::default_config();
    config.set_chunk_rows(1024);
    config.set_enable_interim_segment_index(false);
    auto segment_growing =
        CreateGrowingSegment(schema, empty_index_meta, 1, config);
    auto segment = dynamic_cast<SegmentGrowingImpl*>(segment_growing.get());
    ASSERT_NE(segment, nullptr);

    auto insert_record_proto = std::make_unique<InsertRecordProto>();
    insert_record_proto->set_num_rows(2);

    auto pk_data = insert_record_proto->add_fields_data();
    pk_data->set_field_id(int64_field.get());
    pk_data->set_type(proto::schema::DataType::Int64);
    pk_data->mutable_scalars()->mutable_long_data()->add_data(10);
    pk_data->mutable_scalars()->mutable_long_data()->add_data(11);

    auto array_data = insert_record_proto->add_fields_data();
    array_data->set_field_id(array_float_vector.get());
    array_data->set_type(proto::schema::DataType::ArrayOfVector);
    array_data->add_valid_data(false);
    array_data->add_valid_data(true);
    array_data->mutable_vectors()->set_dim(4);
    auto vector_array = array_data->mutable_vectors()->mutable_vector_array();
    vector_array->set_dim(4);
    vector_array->set_element_type(proto::schema::DataType::FloatVector);
    auto valid_row = vector_array->add_data();
    valid_row->set_dim(4);
    for (auto value : {1.0F, 2.0F, 3.0F, 4.0F, 5.0F, 6.0F, 7.0F, 8.0F}) {
        valid_row->mutable_float_vector()->add_data(value);
    }

    std::vector<int64_t> row_ids = {0, 1};
    std::vector<Timestamp> timestamps = {100, 101};
    auto offset = segment->PreInsert(2);
    segment->Insert(offset,
                    2,
                    row_ids.data(),
                    timestamps.data(),
                    insert_record_proto.get());

    std::array<int64_t, 2> offsets = {0, 1};
    std::unique_ptr<DataArray> result;
    EXPECT_NO_THROW(
        result = segment->bulk_subscript(
            nullptr, array_float_vector, offsets.data(), offsets.size()));
    ASSERT_NE(result, nullptr);

    ASSERT_EQ(result->valid_data_size(), 2);
    EXPECT_FALSE(result->valid_data(0));
    EXPECT_TRUE(result->valid_data(1));
    ASSERT_EQ(result->vectors().vector_array().data_size(), 2);
    EXPECT_TRUE(
        result->vectors().vector_array().data(0).float_vector().data().empty());
    const auto& row = result->vectors().vector_array().data(1);
    ASSERT_EQ(row.float_vector().data_size(), 8);
    EXPECT_FLOAT_EQ(row.float_vector().data(0), 1.0F);
    EXPECT_FLOAT_EQ(row.float_vector().data(7), 8.0F);
}

TEST(GrowingTest, QueryNullableVectorArrayStoresRowDenseInputByValidData) {
    auto schema = std::make_shared<Schema>();
    auto int64_field = schema->AddDebugField("int64", DataType::INT64);
    auto array_float_vector =
        schema->AddDebugVectorArrayField("array_float_vector",
                                         DataType::VECTOR_FLOAT,
                                         4,
                                         knowhere::metric::L2,
                                         true);
    schema->set_primary_field_id(int64_field);

    auto config = SegcoreConfig::default_config();
    config.set_chunk_rows(1024);
    config.set_enable_interim_segment_index(false);
    auto segment_growing =
        CreateGrowingSegment(schema, empty_index_meta, 1, config);
    auto segment = dynamic_cast<SegmentGrowingImpl*>(segment_growing.get());
    ASSERT_NE(segment, nullptr);

    auto insert_record_proto = std::make_unique<InsertRecordProto>();
    insert_record_proto->set_num_rows(2);

    auto pk_data = insert_record_proto->add_fields_data();
    pk_data->set_field_id(int64_field.get());
    pk_data->set_type(proto::schema::DataType::Int64);
    pk_data->mutable_scalars()->mutable_long_data()->add_data(10);
    pk_data->mutable_scalars()->mutable_long_data()->add_data(11);

    auto array_data = insert_record_proto->add_fields_data();
    array_data->set_field_id(array_float_vector.get());
    array_data->set_type(proto::schema::DataType::ArrayOfVector);
    array_data->add_valid_data(false);
    array_data->add_valid_data(true);
    array_data->mutable_vectors()->set_dim(4);
    auto vector_array = array_data->mutable_vectors()->mutable_vector_array();
    vector_array->set_dim(4);
    vector_array->set_element_type(proto::schema::DataType::FloatVector);
    auto null_row = vector_array->add_data();
    null_row->set_dim(4);
    null_row->mutable_float_vector();
    auto valid_row = vector_array->add_data();
    valid_row->set_dim(4);
    for (auto value : {1.0F, 2.0F, 3.0F, 4.0F, 5.0F, 6.0F, 7.0F, 8.0F}) {
        valid_row->mutable_float_vector()->add_data(value);
    }

    std::vector<int64_t> row_ids = {0, 1};
    std::vector<Timestamp> timestamps = {100, 101};
    auto offset = segment->PreInsert(2);
    segment->Insert(offset,
                    2,
                    row_ids.data(),
                    timestamps.data(),
                    insert_record_proto.get());

    std::array<int64_t, 2> offsets = {0, 1};
    std::unique_ptr<DataArray> result;
    EXPECT_NO_THROW(
        result = segment->bulk_subscript(
            nullptr, array_float_vector, offsets.data(), offsets.size()));
    ASSERT_NE(result, nullptr);

    ASSERT_EQ(result->valid_data_size(), 2);
    EXPECT_FALSE(result->valid_data(0));
    EXPECT_TRUE(result->valid_data(1));
    ASSERT_EQ(result->vectors().vector_array().data_size(), 2);
    EXPECT_TRUE(
        result->vectors().vector_array().data(0).float_vector().data().empty());
    const auto& row = result->vectors().vector_array().data(1);
    ASSERT_EQ(row.float_vector().data_size(), 8);
    EXPECT_FLOAT_EQ(row.float_vector().data(0), 1.0F);
    EXPECT_FLOAT_EQ(row.float_vector().data(7), 8.0F);
}

TEST(GrowingTest, LoadVectorArrayData) {
    auto schema = std::make_shared<Schema>();
    auto metric_type = knowhere::metric::MAX_SIM;
    auto int64_field = schema->AddDebugField("int64", DataType::INT64);
    auto array_float_vector = schema->AddDebugVectorArrayField(
        "array_vec", DataType::VECTOR_FLOAT, 128, metric_type);
    schema->set_primary_field_id(int64_field);

    auto config = SegcoreConfig::default_config();
    config.set_chunk_rows(1024);
    config.set_enable_interim_segment_index(true);
    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_growing =
        CreateGrowingWithFieldDataLoaded(schema, metaPtr, config, dataset);
    auto segment = segment_growing.get();

    // Verify data
    auto int64_values = dataset.get_col<int64_t>(int64_field);
    auto array_vec_values =
        dataset.get_col<VectorFieldProto>(array_float_vector);

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

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

    auto verify_float_vectors = [](auto arr1, auto arr2) {
        static constexpr float EPSILON = 1e-6;
        EXPECT_EQ(arr1.size(), arr2.size());
        for (int64_t i = 0; i < arr1.size(); ++i) {
            EXPECT_NEAR(arr1[i], arr2[i], EPSILON);
        }
    };

    for (int64_t i = 0; i < dataset_size; ++i) {
        auto arrow_array = array_float_vector_result->vectors()
                               .vector_array()
                               .data()[i]
                               .float_vector()
                               .data();
        auto expected_array =
            array_vec_values[ids_ds->GetIds()[i]].float_vector().data();
        verify_float_vectors(arrow_array, expected_array);
    }
}

TEST(GrowingTest, SearchVectorArray) {
    using namespace milvus::query;

    auto schema = std::make_shared<Schema>();
    auto metric_type = knowhere::metric::MAX_SIM;

    auto dim = 32;

    // Add fields
    auto int64_field = schema->AddDebugField("int64", DataType::INT64);
    auto array_vec = schema->AddDebugVectorArrayField(
        "array_vec", DataType::VECTOR_FLOAT, dim, metric_type);
    schema->set_primary_field_id(int64_field);

    // Configure segment
    auto config = SegcoreConfig::default_config();
    config.set_chunk_rows(1024);
    config.set_enable_interim_segment_index(true);

    std::map<std::string, std::string> index_params = {
        {"index_type", knowhere::IndexEnum::INDEX_HNSW},
        {"metric_type", metric_type},
        {"nlist", "128"}};
    std::map<std::string, std::string> type_params = {
        {"dim", std::to_string(dim)}};
    FieldIndexMeta fieldIndexMeta(
        array_vec, std::move(index_params), std::move(type_params));
    std::map<FieldId, FieldIndexMeta> fieldMap = {{array_vec, fieldIndexMeta}};

    IndexMetaPtr metaPtr =
        std::make_shared<CollectionIndexMeta>(100000, std::move(fieldMap));
    auto segment = CreateGrowingSegment(schema, metaPtr, 1, config);

    // Insert data
    int64_t N = 100;
    uint64_t seed = 42;
    int emb_list_len = 5;  // Each row contains 5 vectors
    auto dataset = DataGen(schema, N, seed, 0, 1, emb_list_len);

    auto offset = 0;
    segment->Insert(offset,
                    N,
                    dataset.row_ids_.data(),
                    dataset.timestamps_.data(),
                    dataset.raw_);

    // Prepare search query
    int vec_num = 10;  // Total number of query vectors
    std::vector<float> query_vec = generate_float_vector(vec_num, dim);

    // Create query dataset with offsets for VectorArray
    std::vector<size_t> query_vec_offsets;
    query_vec_offsets.push_back(0);  // First query has 3 vectors
    query_vec_offsets.push_back(3);
    query_vec_offsets.push_back(10);  // Second query has 7 vectors

    // Create search plan using ScopedSchemaHandle
    milvus::segcore::ScopedSchemaHandle schema_handle(*schema);
    auto plan_str =
        schema_handle.ParseSearch("",           // expression (no filter)
                                  "array_vec",  // vector field name
                                  5,            // topk
                                  "MAX_SIM",    // metric_type
                                  R"({"nprobe": 10})",  // search_params
                                  3);                   // round_decimal
    auto plan =
        CreateSearchPlanByExpr(schema, plan_str.data(), plan_str.size());

    // Use CreatePlaceholderGroupFromBlob for VectorArray
    auto ph_group_raw = CreatePlaceholderGroupFromBlob<EmbListFloatVector>(
        vec_num, dim, query_vec.data(), query_vec_offsets);
    auto ph_group =
        ParsePlaceholderGroup(plan.get(), ph_group_raw.SerializeAsString());

    // Execute search
    Timestamp timestamp = 10000000;
    auto sr = segment->Search(plan.get(), ph_group.get(), timestamp);
    auto sr_parsed = SearchResultToJson(*sr);
    std::cout << sr_parsed.dump(1) << std::endl;
}

TEST(Growing, TestMaskWithTTLField) {
    auto schema = std::make_shared<Schema>();
    auto pk_fid = schema->AddDebugField("pk", DataType::INT64, false);
    auto ttl_fid =
        schema->AddDebugField("ttl_field", DataType::TIMESTAMPTZ, false);
    schema->set_primary_field_id(pk_fid);
    schema->set_ttl_field_id(ttl_fid);

    auto segment = CreateGrowingSegment(schema, empty_index_meta);
    auto segment_impl = dynamic_cast<SegmentGrowingImpl*>(segment.get());
    ASSERT_NE(segment_impl, nullptr);

    int64_t test_data_count = 100;

    uint64_t base_ts = 1000000000ULL << 18;
    std::vector<Timestamp> ts_data(test_data_count);
    for (int i = 0; i < test_data_count; i++) {
        ts_data[i] = base_ts + i;
    }

    std::vector<idx_t> row_ids(test_data_count);
    std::iota(row_ids.begin(), row_ids.end(), 0);

    std::vector<int64_t> pk_data(test_data_count);
    std::iota(pk_data.begin(), pk_data.end(), 0);

    uint64_t base_physical_us = (base_ts >> 18) * 1000;
    std::vector<int64_t> ttl_data(test_data_count);
    for (int i = 0; i < test_data_count; i++) {
        if (i < test_data_count / 2) {
            ttl_data[i] = static_cast<int64_t>(base_physical_us - 10);
        } else {
            ttl_data[i] = static_cast<int64_t>(base_physical_us + 10);
        }
    }

    auto insert_record_proto = std::make_unique<InsertRecordProto>();
    insert_record_proto->set_num_rows(test_data_count);

    {
        auto field_data = insert_record_proto->add_fields_data();
        field_data->set_field_id(pk_fid.get());
        field_data->set_type(proto::schema::DataType::Int64);
        auto* scalars = field_data->mutable_scalars();
        auto* data = scalars->mutable_long_data();
        for (auto v : pk_data) {
            data->add_data(v);
        }
    }

    {
        auto field_data = insert_record_proto->add_fields_data();
        field_data->set_field_id(ttl_fid.get());
        field_data->set_type(proto::schema::DataType::Timestamptz);
        auto* scalars = field_data->mutable_scalars();
        auto* data = scalars->mutable_timestamptz_data();
        for (auto v : ttl_data) {
            data->add_data(v);
        }
    }

    auto offset = segment->PreInsert(test_data_count);
    segment->Insert(offset,
                    test_data_count,
                    row_ids.data(),
                    ts_data.data(),
                    insert_record_proto.get());

    // Test TTL field filtering using CompileExpressions pathway
    Timestamp query_ts = base_ts + test_data_count;
    int64_t active_count = segment->get_active_count(query_ts);

    // Create an expression list with AlwaysTrueExpr - CompileExpressions will
    // automatically add TTL field filtering expression
    auto always_true_expr = std::make_shared<expr::AlwaysTrueExpr>();
    auto plan = std::make_shared<plan::FilterBitsNode>(DEFAULT_PLANNODE_ID,
                                                       always_true_expr);

    // Execute query expression - this will trigger CompileExpressions which
    // automatically adds TTL field filtering expression
    BitsetType bitset =
        query::ExecuteQueryExpr(plan, segment.get(), active_count, query_ts);

    // Note: ExecuteQueryExpr already flips the bitset, so bitset[i] = true
    // means the row matches (not expired), bitset[i] = false means expired
    BitsetTypeView bitset_view(bitset);

    // Verify results:
    // After ExecuteQueryExpr, bitset[i] = true means row matches (not expired)
    // bitset[i] = false means row is filtered out (expired)
    // - i < test_data_count / 2: expired (TTL < current time), should be expired (bitset_view[i] = false)
    // - i >= test_data_count / 2: not expired, should NOT be expired (bitset_view[i] = true)
    int expired_count = 0;
    for (int i = 0; i < test_data_count; i++) {
        if (!bitset_view[i]) {
            expired_count++;
        }
    }

    EXPECT_EQ(expired_count, test_data_count / 2);
    for (int i = 0; i < test_data_count / 2; i++) {
        EXPECT_FALSE(bitset_view[i]) << "Row " << i << " should be expired";
    }
    for (int i = test_data_count / 2; i < test_data_count; i++) {
        EXPECT_TRUE(bitset_view[i]) << "Row " << i << " should not be expired";
    }
}

// Test TTL field filtering with nullable field for Growing segment
TEST(Growing, TestMaskWithNullableTTLField) {
    // Create schema with nullable TTL field
    auto schema = std::make_shared<Schema>();
    auto pk_fid = schema->AddDebugField("pk", DataType::INT64, false);
    auto ttl_fid = schema->AddDebugField(
        "ttl_field", DataType::TIMESTAMPTZ, true);  // nullable
    schema->set_primary_field_id(pk_fid);
    schema->set_ttl_field_id(ttl_fid);

    auto segment = CreateGrowingSegment(schema, empty_index_meta);
    auto segment_impl = dynamic_cast<SegmentGrowingImpl*>(segment.get());
    ASSERT_NE(segment_impl, nullptr);

    int64_t test_data_count = 100;

    // Generate timestamp data
    uint64_t base_ts = 1000000000ULL << 18;
    std::vector<Timestamp> ts_data(test_data_count);
    for (int i = 0; i < test_data_count; i++) {
        ts_data[i] = base_ts + i;
    }

    // Generate row IDs
    std::vector<idx_t> row_ids(test_data_count);
    std::iota(row_ids.begin(), row_ids.end(), 0);

    // Generate PK data
    std::vector<int64_t> pk_data(test_data_count);
    std::iota(pk_data.begin(), pk_data.end(), 0);

    // Generate TTL data with some nulls
    uint64_t base_physical_us = (base_ts >> 18) * 1000;
    std::vector<int64_t> ttl_data(test_data_count);
    std::vector<bool> valid_data(test_data_count);
    for (int i = 0; i < test_data_count; i++) {
        if (i % 4 == 0) {
            // Null value - should not expire
            ttl_data[i] = 0;
            valid_data[i] = false;
        } else if (i % 4 == 1) {
            // Expired
            ttl_data[i] = static_cast<int64_t>(base_physical_us - 10);
            valid_data[i] = true;
        } else {
            // Not expired
            ttl_data[i] = static_cast<int64_t>(base_physical_us + 10);
            valid_data[i] = true;
        }
    }

    // Create insert record proto
    auto insert_record_proto = std::make_unique<InsertRecordProto>();
    insert_record_proto->set_num_rows(test_data_count);

    // Add PK field data
    {
        auto field_data = insert_record_proto->add_fields_data();
        field_data->set_field_id(pk_fid.get());
        field_data->set_type(proto::schema::DataType::Int64);
        auto* scalars = field_data->mutable_scalars();
        auto* data = scalars->mutable_long_data();
        for (auto v : pk_data) {
            data->add_data(v);
        }
    }

    // Add nullable TTL field data
    {
        auto field_data = insert_record_proto->add_fields_data();
        field_data->set_field_id(ttl_fid.get());
        field_data->set_type(proto::schema::DataType::Timestamptz);
        auto* scalars = field_data->mutable_scalars();
        auto* data = scalars->mutable_timestamptz_data();
        for (auto v : ttl_data) {
            data->add_data(v);
        }
        // Add valid_data for nullable field
        // Note: valid_data[i] = false means null, valid_data[i] = true means non-null
        for (size_t i = 0; i < valid_data.size(); ++i) {
            field_data->add_valid_data(valid_data[i]);
        }
        // Verify valid_data was added correctly
        ASSERT_EQ(field_data->valid_data_size(), test_data_count)
            << "valid_data size mismatch: expected " << test_data_count
            << ", got " << field_data->valid_data_size();
    }

    // Insert data
    auto offset = segment->PreInsert(test_data_count);
    segment->Insert(offset,
                    test_data_count,
                    row_ids.data(),
                    ts_data.data(),
                    insert_record_proto.get());

    // Test TTL field filtering using CompileExpressions pathway
    Timestamp query_ts = base_ts + test_data_count;
    int64_t active_count = segment->get_active_count(query_ts);

    // Create an expression list with AlwaysTrueExpr - CompileExpressions will
    // automatically add TTL field filtering expression
    auto always_true_expr = std::make_shared<expr::AlwaysTrueExpr>();
    auto plan = std::make_shared<plan::FilterBitsNode>(DEFAULT_PLANNODE_ID,
                                                       always_true_expr);

    // Execute query expression - this will trigger CompileExpressions which
    // automatically adds TTL field filtering expression
    BitsetType bitset =
        query::ExecuteQueryExpr(plan, segment.get(), active_count, query_ts);

    // Note: ExecuteQueryExpr already flips the bitset, so bitset[i] = true
    // means the row matches (not expired), bitset[i] = false means expired
    BitsetTypeView bitset_view(bitset);

    // Verify results:
    // After ExecuteQueryExpr, bitset[i] = true means row matches (not expired)
    // bitset[i] = false means row is filtered out (expired)
    // - i % 4 == 0: null, should NOT be expired (bitset_view[i] = true)
    // - i % 4 == 1: expired (TTL < current time), should be expired (bitset_view[i] = false)
    // - i % 4 == 2 or 3: not expired, should NOT be expired (bitset_view[i] = true)
    int expired_count = 0;
    for (int i = 0; i < test_data_count; i++) {
        if (i % 4 == 0) {
            // Null value should not be expired (should match)
            EXPECT_TRUE(bitset_view[i])
                << "Row " << i << " (null) should not be expired";
        } else if (i % 4 == 1) {
            // Should be expired (should be filtered out)
            EXPECT_FALSE(bitset_view[i]) << "Row " << i << " should be expired";
            expired_count++;
        } else {
            // Should not be expired (should match)
            EXPECT_TRUE(bitset_view[i])
                << "Row " << i << " should not be expired";
        }
    }

    EXPECT_EQ(expired_count, test_data_count / 4);
}

TEST(Growing, FilterOnlySearchUsesEntityTTLPhysicalTime) {
    using namespace milvus::query;

    auto schema = std::make_shared<Schema>();
    auto dim = 4;
    auto metric_type = knowhere::metric::L2;
    schema->AddDebugField("vec", DataType::VECTOR_FLOAT, dim, metric_type);
    auto pk_fid = schema->AddDebugField("pk", DataType::INT64, false);
    auto ttl_fid =
        schema->AddDebugField("ttl_field", DataType::TIMESTAMPTZ, false);
    schema->set_primary_field_id(pk_fid);
    schema->set_ttl_field_id(ttl_fid);

    auto segment = CreateGrowingSegment(schema, empty_index_meta);
    ASSERT_NE(dynamic_cast<SegmentGrowingImpl*>(segment.get()), nullptr);

    constexpr int64_t test_data_count = 100;
    // 90/10 split: first 10 rows expire before the TTL cutoff,
    // remaining 90 are still alive.  An inverted-bitset bug would
    // report 10 instead of 90.
    constexpr int64_t expired_count = 10;
    constexpr int64_t expected_valid_count =
        test_data_count - expired_count;  // 90
    auto dataset = DataGen(schema, test_data_count);

    constexpr int64_t entity_ttl_physical_time_us = 1770026400000000LL;
    for (auto& field_data : *dataset.raw_->mutable_fields_data()) {
        if (field_data.field_id() != ttl_fid.get()) {
            continue;
        }

        auto* ttl_values =
            field_data.mutable_scalars()->mutable_timestamptz_data();
        ttl_values->clear_data();
        for (int i = 0; i < test_data_count; ++i) {
            ttl_values->add_data(i < expired_count
                                     ? entity_ttl_physical_time_us - 10
                                     : entity_ttl_physical_time_us + 10);
        }
    }

    segment->PreInsert(test_data_count);
    segment->Insert(0,
                    test_data_count,
                    dataset.row_ids_.data(),
                    dataset.timestamps_.data(),
                    dataset.raw_);

    Timestamp query_ts = test_data_count + 10;
    int64_t active_count = segment->get_active_count(query_ts);
    ASSERT_EQ(active_count, test_data_count);

    // Cross-check via ExecuteQueryExpr to confirm the expected count.
    auto always_true_expr = std::make_shared<expr::AlwaysTrueExpr>();
    auto expected_plan = std::make_shared<plan::FilterBitsNode>(
        DEFAULT_PLANNODE_ID, always_true_expr);
    auto expected_bitset = query::ExecuteQueryExpr(expected_plan,
                                                   segment.get(),
                                                   active_count,
                                                   query_ts,
                                                   entity_ttl_physical_time_us);
    BitsetTypeView expected_matches(expected_bitset);
    int64_t cross_check_count = 0;
    for (int i = 0; i < test_data_count; ++i) {
        if (expected_matches[i]) {
            ++cross_check_count;
        }
    }
    ASSERT_EQ(cross_check_count, expected_valid_count);

    milvus::segcore::ScopedSchemaHandle schema_handle(*schema);
    auto plan_str = schema_handle.ParseSearch(
        "pk >= 0", "vec", test_data_count, metric_type, R"({"nprobe": 10})", 3);
    auto plan =
        CreateSearchPlanByExpr(schema, plan_str.data(), plan_str.size());
    ASSERT_NE(plan, nullptr);

    auto query_data = generate_float_vector(1, dim);
    auto ph_group_raw =
        CreatePlaceholderGroupFromBlob(1, dim, query_data.data());
    auto ph_group =
        ParsePlaceholderGroup(plan.get(), ph_group_raw.SerializeAsString());
    ASSERT_NE(ph_group, nullptr);

    auto search_result = segment->Search(plan.get(),
                                         ph_group.get(),
                                         query_ts,
                                         folly::CancellationToken(),
                                         0,
                                         0,
                                         entity_ttl_physical_time_us,
                                         true);

    EXPECT_EQ(search_result->valid_count_, expected_valid_count);
}

// Resource tracking tests for growing segments
TEST(Growing, EmptySegmentResourceEstimation) {
    auto schema = std::make_shared<Schema>();
    auto dim = 128;
    auto metric_type = knowhere::metric::L2;
    schema->AddDebugField("vec", DataType::VECTOR_FLOAT, dim, metric_type);
    auto pk_fid = schema->AddDebugField("pk", DataType::INT64);
    schema->set_primary_field_id(pk_fid);

    auto segment = CreateGrowingSegment(schema, empty_index_meta);
    auto* segment_impl = dynamic_cast<SegmentGrowingImpl*>(segment.get());
    ASSERT_NE(segment_impl, nullptr);

    // Empty segment should have zero resource usage
    auto resource = segment_impl->EstimateSegmentResourceUsage();
    EXPECT_EQ(resource.memory_bytes, 0);
    EXPECT_EQ(resource.file_bytes, 0);
}

TEST(Growing, ResourceEstimationAfterInsert) {
    auto schema = std::make_shared<Schema>();
    auto dim = 128;
    auto metric_type = knowhere::metric::L2;
    schema->AddDebugField("vec", DataType::VECTOR_FLOAT, dim, metric_type);
    auto pk_fid = schema->AddDebugField("pk", DataType::INT64);
    schema->set_primary_field_id(pk_fid);

    auto segment = CreateGrowingSegment(schema, empty_index_meta);
    auto* segment_impl = dynamic_cast<SegmentGrowingImpl*>(segment.get());
    ASSERT_NE(segment_impl, nullptr);

    // Insert some data
    const int64_t N = 1000;
    auto dataset = DataGen(schema, N);
    segment->PreInsert(N);
    segment->Insert(0,
                    N,
                    dataset.row_ids_.data(),
                    dataset.timestamps_.data(),
                    dataset.raw_);

    // After insert, resource usage should be positive
    auto resource = segment_impl->EstimateSegmentResourceUsage();
    EXPECT_GT(resource.memory_bytes, 0);

    // Memory should include at least:
    // - Vector data: N * dim * sizeof(float) = 1000 * 128 * 4 = 512000 bytes
    // - Timestamps: N * sizeof(Timestamp) = 1000 * 8 = 8000 bytes
    // - RowIDs: N * sizeof(int64_t) = 1000 * 8 = 8000 bytes
    // - PK field: N * sizeof(int64_t) = 1000 * 8 = 8000 bytes
    // Plus safety margin of 1.2x
    int64_t expected_min_size = N * dim * sizeof(float) +
                                N * sizeof(Timestamp) + N * sizeof(int64_t) +
                                N * sizeof(int64_t);
    EXPECT_GE(resource.memory_bytes, expected_min_size);
}

TEST(Growing, ResourceIncrementsWithMoreInserts) {
    auto schema = std::make_shared<Schema>();
    auto dim = 128;
    auto metric_type = knowhere::metric::L2;
    schema->AddDebugField("vec", DataType::VECTOR_FLOAT, dim, metric_type);
    auto pk_fid = schema->AddDebugField("pk", DataType::INT64);
    schema->set_primary_field_id(pk_fid);

    auto segment = CreateGrowingSegment(schema, empty_index_meta);
    auto* segment_impl = dynamic_cast<SegmentGrowingImpl*>(segment.get());
    ASSERT_NE(segment_impl, nullptr);

    // First insert
    const int64_t N1 = 500;
    auto dataset1 = DataGen(schema, N1, 42, 0);
    segment->PreInsert(N1);
    segment->Insert(0,
                    N1,
                    dataset1.row_ids_.data(),
                    dataset1.timestamps_.data(),
                    dataset1.raw_);
    auto resource1 = segment_impl->EstimateSegmentResourceUsage();

    // Second insert
    const int64_t N2 = 500;
    auto dataset2 = DataGen(schema, N2, 43, N1);
    segment->PreInsert(N2);
    segment->Insert(N1,
                    N2,
                    dataset2.row_ids_.data(),
                    dataset2.timestamps_.data(),
                    dataset2.raw_);
    auto resource2 = segment_impl->EstimateSegmentResourceUsage();

    // Resource should increase after second insert
    EXPECT_GT(resource2.memory_bytes, resource1.memory_bytes);
}

TEST(Growing, ResourceTrackingAfterDelete) {
    auto schema = std::make_shared<Schema>();
    auto dim = 64;
    auto metric_type = knowhere::metric::L2;
    schema->AddDebugField("vec", DataType::VECTOR_FLOAT, dim, metric_type);
    auto pk_fid = schema->AddDebugField("pk", DataType::INT64);
    schema->set_primary_field_id(pk_fid);

    auto segment = CreateGrowingSegment(schema, empty_index_meta);
    auto* segment_impl = dynamic_cast<SegmentGrowingImpl*>(segment.get());
    ASSERT_NE(segment_impl, nullptr);

    // Insert data first
    const int64_t N = 100;
    auto dataset = DataGen(schema, N);
    segment->PreInsert(N);
    segment->Insert(0,
                    N,
                    dataset.row_ids_.data(),
                    dataset.timestamps_.data(),
                    dataset.raw_);

    auto resource_before_delete = segment_impl->EstimateSegmentResourceUsage();
    EXPECT_GT(resource_before_delete.memory_bytes, 0);

    // Delete some rows
    auto pks = dataset.get_col<int64_t>(pk_fid);
    auto del_pks = GenPKs(pks.begin(), pks.begin() + 5);
    auto del_tss = GenTss(5, N);
    auto status = segment->Delete(5, del_pks.get(), del_tss.data());
    EXPECT_TRUE(status.ok());

    // Resource estimation should still work after delete
    auto resource_after_delete = segment_impl->EstimateSegmentResourceUsage();
    EXPECT_GT(resource_after_delete.memory_bytes, 0);
}

TEST(Growing, ConcurrentInsertResourceTracking) {
    auto schema = std::make_shared<Schema>();
    auto dim = 32;
    auto metric_type = knowhere::metric::L2;
    schema->AddDebugField("vec", DataType::VECTOR_FLOAT, dim, metric_type);
    auto pk_fid = schema->AddDebugField("pk", DataType::INT64);
    schema->set_primary_field_id(pk_fid);

    auto segment = CreateGrowingSegment(schema, empty_index_meta);
    auto* segment_impl = dynamic_cast<SegmentGrowingImpl*>(segment.get());
    ASSERT_NE(segment_impl, nullptr);

    const int num_threads = 4;
    const int64_t rows_per_thread = 100;
    std::vector<std::thread> threads;

    // Reserve space for all rows upfront
    int64_t total_rows = num_threads * rows_per_thread;
    segment->PreInsert(total_rows);

    // Concurrent inserts from multiple threads
    for (int t = 0; t < num_threads; ++t) {
        threads.emplace_back([&, t]() {
            auto dataset =
                DataGen(schema, rows_per_thread, 42 + t, t * rows_per_thread);
            segment->Insert(t * rows_per_thread,
                            rows_per_thread,
                            dataset.row_ids_.data(),
                            dataset.timestamps_.data(),
                            dataset.raw_);
        });
    }

    for (auto& thread : threads) {
        thread.join();
    }

    // Verify total row count
    EXPECT_EQ(segment->get_row_count(), total_rows);

    // Verify resource estimation is consistent and positive
    auto resource = segment_impl->EstimateSegmentResourceUsage();
    EXPECT_GT(resource.memory_bytes, 0);
}

TEST(Growing, MultipleFieldsResourceEstimation) {
    // Create schema with multiple fields
    auto schema = std::make_shared<Schema>();
    auto dim = 64;
    auto metric_type = knowhere::metric::L2;
    schema->AddDebugField("vec", DataType::VECTOR_FLOAT, dim, metric_type);
    auto pk_fid = schema->AddDebugField("pk", DataType::INT64);
    schema->AddDebugField("age", DataType::FLOAT);
    schema->AddDebugField("score", DataType::DOUBLE);
    schema->set_primary_field_id(pk_fid);

    auto segment = CreateGrowingSegment(schema, empty_index_meta);
    auto* segment_impl = dynamic_cast<SegmentGrowingImpl*>(segment.get());
    ASSERT_NE(segment_impl, nullptr);

    // Insert data
    const int64_t N = 500;
    auto dataset = DataGen(schema, N);
    segment->PreInsert(N);
    segment->Insert(0,
                    N,
                    dataset.row_ids_.data(),
                    dataset.timestamps_.data(),
                    dataset.raw_);

    auto resource = segment_impl->EstimateSegmentResourceUsage();

    // Memory should include all fields:
    // - Vector: N * dim * sizeof(float) = 500 * 64 * 4 = 128000 bytes
    // - pk (int64): N * 8 = 4000 bytes
    // - age (float): N * 4 = 2000 bytes
    // - score (double): N * 8 = 4000 bytes
    // - Timestamps: N * 8 = 4000 bytes
    // Plus safety margin
    int64_t min_expected = N * dim * sizeof(float) + N * sizeof(int64_t) +
                           N * sizeof(float) + N * sizeof(double) +
                           N * sizeof(Timestamp);
    EXPECT_GE(resource.memory_bytes, min_expected);
}
