// 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 <gtest/gtest.h>
#include "common/Schema.h"
#include "query/Plan.h"

#include "segcore/reduce_c.h"
#include "segcore/plan_c.h"
#include "segcore/segment_c.h"
#include "test_utils/DataGen.h"
#include "test_utils/c_api_test_utils.h"
#include "test_utils/storage_test_utils.h"
#include "test_utils/cachinglayer_test_utils.h"
#include "test_utils/storage_test_utils.h"

using namespace milvus;
using namespace milvus::query;
using namespace milvus::segcore;
using namespace milvus::storage;
using namespace milvus::tracer;

const char* METRICS_TYPE = "metric_type";

TEST(GroupBY, SealedIndex) {
    using namespace milvus;
    using namespace milvus::query;
    using namespace milvus::segcore;

    //0. prepare schema
    int dim = 64;
    auto schema = std::make_shared<Schema>();
    auto vec_fid = schema->AddDebugField(
        "fakevec", DataType::VECTOR_FLOAT, dim, knowhere::metric::L2);
    auto int8_fid = schema->AddDebugField("int8", DataType::INT8);  // id 101
    auto int16_fid = schema->AddDebugField("int16", DataType::INT16);
    auto int32_fid = schema->AddDebugField("int32", DataType::INT32);
    auto int64_fid = schema->AddDebugField("int64", DataType::INT64);
    auto str_fid = schema->AddDebugField("string1", DataType::VARCHAR);
    auto bool_fid = schema->AddDebugField("bool", DataType::BOOL);
    auto timestamptz_fid =
        schema->AddDebugField("timestamptz", DataType::TIMESTAMPTZ);
    schema->set_primary_field_id(str_fid);
    size_t N = 50;

    //2. load raw data
    auto raw_data = DataGen(schema, N, 42, 0, 8, 10, 1, false, false);
    auto segment = CreateSealedWithFieldDataLoaded(schema, raw_data);

    //3. load index
    auto vector_data = raw_data.get_col<float>(vec_fid);
    auto indexing = GenVecIndexing(
        N, dim, vector_data.data(), knowhere::IndexEnum::INDEX_HNSW);
    LoadIndexInfo load_index_info;
    load_index_info.field_id = vec_fid.get();
    load_index_info.index_params = GenIndexParams(indexing.get());
    load_index_info.cache_index =
        CreateTestCacheIndex("test", std::move(indexing));
    load_index_info.index_params[METRICS_TYPE] = knowhere::metric::L2;
    segment->LoadIndex(load_index_info);
    int topK = 15;
    int group_size = 3;

    //4. search group by int8
    {
        const char* raw_plan = R"(vector_anns: <
                                        field_id: 100
                                        query_info: <
                                          topk: 15
                                          metric_type: "L2"
                                          search_params: "{\"ef\": 10}"
                                          group_by_field_id: 101
                                          group_size: 3
                                        >
                                        placeholder_tag: "$0"

         >)";
        proto::plan::PlanNode plan_node;
        auto ok =
            google::protobuf::TextFormat::ParseFromString(raw_plan, &plan_node);
        auto plan = CreateSearchPlanFromPlanNode(schema, plan_node);
        auto num_queries = 1;
        auto seed = 1024;
        auto ph_group_raw = CreatePlaceholderGroup(num_queries, dim, seed);
        auto ph_group =
            ParsePlaceholderGroup(plan.get(), ph_group_raw.SerializeAsString());
        auto search_result =
            segment->Search(plan.get(), ph_group.get(), MAX_TIMESTAMP);
        CheckGroupBySearchResult(*search_result, topK, num_queries, false);

        auto& group_by_values = search_result->group_by_values_.value();
        ASSERT_EQ(20, group_by_values.size());
        //as the total data is 0,0,....6,6, so there will be 7 buckets with [3,3,3,3,3,3,2] items respectively
        //so there will be 20 items returned

        int size = group_by_values.size();
        std::unordered_map<int8_t, int> i8_map;
        float lastDistance = 0.0;
        for (size_t i = 0; i < size; i++) {
            if (std::holds_alternative<int8_t>(group_by_values[i].value())) {
                int8_t g_val = std::get<int8_t>(group_by_values[i].value());
                i8_map[g_val] += 1;
                ASSERT_TRUE(i8_map[g_val] <= group_size);
                //for every group, the number of hits should not exceed group_size
                auto distance = search_result->distances_.at(i);
                ASSERT_TRUE(
                    lastDistance <=
                    distance);  //distance should be decreased as metrics_type is L2
                lastDistance = distance;
            }
        }
        ASSERT_TRUE(i8_map.size() <= topK);
        ASSERT_TRUE(i8_map.size() == 7);
    }

    //5. search group by int16
    {
        const char* raw_plan = R"(vector_anns: <
                                        field_id: 100
                                        query_info: <
                                          topk: 100
                                          metric_type: "L2"
                                          search_params: "{\"ef\": 10}"
                                          group_by_field_id: 102
                                          group_size: 3
                                        >
                                        placeholder_tag: "$0"

         >)";

        proto::plan::PlanNode plan_node;
        auto ok =
            google::protobuf::TextFormat::ParseFromString(raw_plan, &plan_node);
        auto plan = CreateSearchPlanFromPlanNode(schema, plan_node);
        auto num_queries = 1;
        auto seed = 1024;
        auto ph_group_raw = CreatePlaceholderGroup(num_queries, dim, seed);
        auto ph_group =
            ParsePlaceholderGroup(plan.get(), ph_group_raw.SerializeAsString());
        auto search_result =
            segment->Search(plan.get(), ph_group.get(), MAX_TIMESTAMP);
        CheckGroupBySearchResult(*search_result, topK, num_queries, false);

        auto& group_by_values = search_result->group_by_values_.value();
        int size = group_by_values.size();
        ASSERT_EQ(20, size);
        //as the total data is 0,0,....6,6, so there will be 7 buckets with [3,3,3,3,3,3,2] items respectively
        //so there will be 20 items returned

        std::unordered_map<int16_t, int> i16_map;
        float lastDistance = 0.0;
        for (size_t i = 0; i < size; i++) {
            if (std::holds_alternative<int16_t>(group_by_values[i].value())) {
                int16_t g_val = std::get<int16_t>(group_by_values[i].value());
                i16_map[g_val] += 1;
                ASSERT_TRUE(i16_map[g_val] <= group_size);
                auto distance = search_result->distances_.at(i);
                ASSERT_TRUE(
                    lastDistance <=
                    distance);  //distance should be decreased as metrics_type is L2
                lastDistance = distance;
            }
        }
        ASSERT_TRUE(i16_map.size() == 7);
    }
    //6. search group by int32
    {
        const char* raw_plan = R"(vector_anns: <
                                        field_id: 100
                                        query_info: <
                                          topk: 100
                                          metric_type: "L2"
                                          search_params: "{\"ef\": 10}"
                                          group_by_field_id: 103
                                          group_size: 3
                                        >
                                        placeholder_tag: "$0"

         >)";

        proto::plan::PlanNode plan_node;
        auto ok =
            google::protobuf::TextFormat::ParseFromString(raw_plan, &plan_node);
        auto plan = CreateSearchPlanFromPlanNode(schema, plan_node);
        auto num_queries = 1;
        auto seed = 1024;
        auto ph_group_raw = CreatePlaceholderGroup(num_queries, dim, seed);
        auto ph_group =
            ParsePlaceholderGroup(plan.get(), ph_group_raw.SerializeAsString());
        auto search_result =
            segment->Search(plan.get(), ph_group.get(), MAX_TIMESTAMP);
        CheckGroupBySearchResult(*search_result, topK, num_queries, false);

        auto& group_by_values = search_result->group_by_values_.value();
        int size = group_by_values.size();
        ASSERT_EQ(20, size);
        //as the total data is 0,0,....6,6, so there will be 7 buckets with [3,3,3,3,3,3,2] items respectively
        //so there will be 20 items returned

        std::unordered_map<int32_t, int> i32_map;
        float lastDistance = 0.0;
        for (size_t i = 0; i < size; i++) {
            if (std::holds_alternative<int32_t>(group_by_values[i].value())) {
                int16_t g_val = std::get<int32_t>(group_by_values[i].value());
                i32_map[g_val] += 1;
                ASSERT_TRUE(i32_map[g_val] <= group_size);
                auto distance = search_result->distances_.at(i);
                ASSERT_TRUE(
                    lastDistance <=
                    distance);  //distance should be decreased as metrics_type is L2
                lastDistance = distance;
            }
        }
        ASSERT_TRUE(i32_map.size() == 7);
    }

    //7. search group by int64
    {
        const char* raw_plan = R"(vector_anns: <
                                        field_id: 100
                                        query_info: <
                                          topk: 100
                                          metric_type: "L2"
                                          search_params: "{\"ef\": 10}"
                                          group_by_field_id: 104
                                          group_size: 3
                                        >
                                        placeholder_tag: "$0"

         >)";

        proto::plan::PlanNode plan_node;
        auto ok =
            google::protobuf::TextFormat::ParseFromString(raw_plan, &plan_node);
        auto plan = CreateSearchPlanFromPlanNode(schema, plan_node);
        auto num_queries = 1;
        auto seed = 1024;
        auto ph_group_raw = CreatePlaceholderGroup(num_queries, dim, seed);
        auto ph_group =
            ParsePlaceholderGroup(plan.get(), ph_group_raw.SerializeAsString());
        auto search_result =
            segment->Search(plan.get(), ph_group.get(), MAX_TIMESTAMP);
        CheckGroupBySearchResult(*search_result, topK, num_queries, false);
        auto& group_by_values = search_result->group_by_values_.value();
        int size = group_by_values.size();
        ASSERT_EQ(20, size);
        //as the total data is 0,0,....6,6, so there will be 7 buckets with [3,3,3,3,3,3,2] items respectively
        //so there will be 20 items returned

        std::unordered_map<int64_t, int> i64_map;
        float lastDistance = 0.0;
        for (size_t i = 0; i < size; i++) {
            if (std::holds_alternative<int64_t>(group_by_values[i].value())) {
                int16_t g_val = std::get<int64_t>(group_by_values[i].value());
                i64_map[g_val] += 1;
                ASSERT_TRUE(i64_map[g_val] <= group_size);
                auto distance = search_result->distances_.at(i);
                ASSERT_TRUE(
                    lastDistance <=
                    distance);  //distance should be decreased as metrics_type is L2
                lastDistance = distance;
            }
        }
        ASSERT_TRUE(i64_map.size() == 7);
    }

    //8. search group by string
    {
        const char* raw_plan = R"(vector_anns: <
                                        field_id: 100
                                        query_info: <
                                          topk: 100
                                          metric_type: "L2"
                                          search_params: "{\"ef\": 10}"
                                          group_by_field_id: 105
                                          group_size: 3
                                        >
                                        placeholder_tag: "$0"

         >)";

        proto::plan::PlanNode plan_node;
        auto ok =
            google::protobuf::TextFormat::ParseFromString(raw_plan, &plan_node);
        auto plan = CreateSearchPlanFromPlanNode(schema, plan_node);
        auto num_queries = 1;
        auto seed = 1024;
        auto ph_group_raw = CreatePlaceholderGroup(num_queries, dim, seed);
        auto ph_group =
            ParsePlaceholderGroup(plan.get(), ph_group_raw.SerializeAsString());
        auto search_result =
            segment->Search(plan.get(), ph_group.get(), MAX_TIMESTAMP);
        CheckGroupBySearchResult(*search_result, topK, num_queries, false);
        auto& group_by_values = search_result->group_by_values_.value();
        ASSERT_EQ(20, group_by_values.size());
        int size = group_by_values.size();

        std::unordered_map<std::string, int> strs_map;
        float lastDistance = 0.0;
        for (size_t i = 0; i < size; i++) {
            if (std::holds_alternative<std::string>(
                    group_by_values[i].value())) {
                std::string g_val = std::move(
                    std::get<std::string>(group_by_values[i].value()));
                strs_map[g_val] += 1;
                ASSERT_TRUE(strs_map[g_val] <= group_size);
                auto distance = search_result->distances_.at(i);
                ASSERT_TRUE(
                    lastDistance <=
                    distance);  //distance should be decreased as metrics_type is L2
                lastDistance = distance;
            }
        }
        ASSERT_TRUE(strs_map.size() == 7);
    }

    //9. search group by bool
    {
        const char* raw_plan = R"(vector_anns: <
                                        field_id: 100
                                        query_info: <
                                          topk: 100
                                          metric_type: "L2"
                                          search_params: "{\"ef\": 10}"
                                          group_by_field_id: 106
                                          group_size: 3
                                        >
                                        placeholder_tag: "$0"

         >)";

        proto::plan::PlanNode plan_node;
        auto ok =
            google::protobuf::TextFormat::ParseFromString(raw_plan, &plan_node);
        auto plan = CreateSearchPlanFromPlanNode(schema, plan_node);
        auto num_queries = 1;
        auto seed = 1024;
        auto ph_group_raw = CreatePlaceholderGroup(num_queries, dim, seed);
        auto ph_group =
            ParsePlaceholderGroup(plan.get(), ph_group_raw.SerializeAsString());
        auto search_result =
            segment->Search(plan.get(), ph_group.get(), MAX_TIMESTAMP);
        CheckGroupBySearchResult(*search_result, topK, num_queries, false);

        auto& group_by_values = search_result->group_by_values_.value();
        int size = group_by_values.size();
        ASSERT_EQ(size, 6);
        //as there are only two possible values: true, false
        //for each group, there are at most 3 items, so the final size of group_by_vals is 3 * 2 = 6

        std::unordered_map<bool, int> bools_map;
        float lastDistance = 0.0;
        for (size_t i = 0; i < size; i++) {
            if (std::holds_alternative<bool>(group_by_values[i].value())) {
                bool g_val = std::get<bool>(group_by_values[i].value());
                bools_map[g_val] += 1;
                ASSERT_TRUE(bools_map[g_val] <= group_size);
                auto distance = search_result->distances_.at(i);
                ASSERT_TRUE(
                    lastDistance <=
                    distance);  //distance should be decreased as metrics_type is L2
                lastDistance = distance;
            }
        }
        ASSERT_TRUE(bools_map.size() == 2);  //bool values cannot exceed two
    }

    //10. search group by timestamptz
    {
        const char* raw_plan = R"(vector_anns: <
                                        field_id: 100
                                        query_info: <
                                          topk: 100
                                          metric_type: "L2"
                                          search_params: "{\"ef\": 10}"
                                          group_by_field_id: 107
                                          group_size: 3
                                        >
                                        placeholder_tag: "$0"

         >)";

        proto::plan::PlanNode plan_node;
        auto ok =
            google::protobuf::TextFormat::ParseFromString(raw_plan, &plan_node);
        auto plan = CreateSearchPlanFromPlanNode(schema, plan_node);
        auto num_queries = 1;
        auto seed = 1024;
        auto ph_group_raw = CreatePlaceholderGroup(num_queries, dim, seed);
        auto ph_group =
            ParsePlaceholderGroup(plan.get(), ph_group_raw.SerializeAsString());
        auto search_result =
            segment->Search(plan.get(), ph_group.get(), MAX_TIMESTAMP);
        CheckGroupBySearchResult(*search_result, topK, num_queries, false);
        auto& group_by_values = search_result->group_by_values_.value();
        int size = group_by_values.size();
        ASSERT_EQ(20, size);
        //as the total data is 0,0,....6,6, so there will be 7 buckets with [3,3,3,3,3,3,2] items respectively
        //so there will be 20 items returned

        std::unordered_map<int64_t, int> timestamptz_map;
        float lastDistance = 0.0;
        for (size_t i = 0; i < size; i++) {
            if (std::holds_alternative<int64_t>(group_by_values[i].value())) {
                int16_t g_val = std::get<int64_t>(group_by_values[i].value());
                timestamptz_map[g_val] += 1;
                ASSERT_TRUE(timestamptz_map[g_val] <= group_size);
                auto distance = search_result->distances_.at(i);
                ASSERT_TRUE(lastDistance <= distance);
                lastDistance = distance;
            }
        }
        ASSERT_TRUE(timestamptz_map.size() == 7);
    }
}

TEST(GroupBY, SealedData) {
    using namespace milvus;
    using namespace milvus::query;
    using namespace milvus::segcore;

    //0. prepare schema
    int dim = 64;
    auto schema = std::make_shared<Schema>();
    auto vec_fid = schema->AddDebugField(
        "fakevec", DataType::VECTOR_FLOAT, dim, knowhere::metric::L2);
    auto int8_fid = schema->AddDebugField("int8", DataType::INT8, true);
    auto int16_fid = schema->AddDebugField("int16", DataType::INT16);
    auto int32_fid = schema->AddDebugField("int32", DataType::INT32);
    auto int64_fid = schema->AddDebugField("int64", DataType::INT64);
    auto str_fid = schema->AddDebugField("string1", DataType::VARCHAR);
    auto bool_fid = schema->AddDebugField("bool", DataType::BOOL);
    schema->set_primary_field_id(str_fid);
    size_t N = 100;

    //2. load raw data
    auto raw_data = DataGen(schema, N, 42, 0, 20, 10, 1, false, false);
    auto segment = CreateSealedWithFieldDataLoaded(schema, raw_data);

    int topK = 10;
    int group_size = 5;
    //3. search group by int8
    {
        const char* raw_plan = R"(vector_anns: <
                                        field_id: 100
                                        query_info: <
                                          topk: 10
                                          metric_type: "L2"
                                          search_params: "{\"ef\": 10}"
                                          group_by_field_id: 101,
                                          group_size: 5,
                                          strict_group_size: true,
                                        >
                                        placeholder_tag: "$0"

         >)";
        auto plan_str = translate_text_plan_to_binary_plan(raw_plan);
        auto plan =
            CreateSearchPlanByExpr(schema, plan_str.data(), plan_str.size());
        auto num_queries = 1;
        auto seed = 1024;
        auto ph_group_raw = CreatePlaceholderGroup(num_queries, dim, seed);
        auto ph_group =
            ParsePlaceholderGroup(plan.get(), ph_group_raw.SerializeAsString());
        auto search_result =
            segment->Search(plan.get(), ph_group.get(), MAX_TIMESTAMP);
        CheckGroupBySearchResult(*search_result, topK, num_queries, false);

        auto& group_by_values = search_result->group_by_values_.value();
        int size = group_by_values.size();
        // groups are: (0, 1, 2, 3, 4, null), counts are: (10, 10, 10, 10 ,10, 50)
        ASSERT_EQ(30, size);

        std::unordered_map<int8_t, int> i8_map;
        float lastDistance = 0.0;
        for (size_t i = 0; i < size; i++) {
            if (group_by_values[i].has_value()) {
                int8_t g_val = std::get<int8_t>(group_by_values[i].value());
                i8_map[g_val] += 1;
            } else {
                i8_map[-1] += 1;
            }
            auto distance = search_result->distances_.at(i);
            ASSERT_TRUE(
                lastDistance <=
                distance);  //distance should be decreased as metrics_type is L2
            lastDistance = distance;
        }

        ASSERT_EQ(i8_map.size(), 6);
        for (const auto& it : i8_map) {
            ASSERT_TRUE(it.second == group_size)
                << "unexpected count on group " << it.first;
        }
    }
}

TEST(GroupBY, Reduce) {
    using namespace milvus;
    using namespace milvus::query;
    using namespace milvus::segcore;

    //0. prepare schema
    int dim = 64;
    auto schema = std::make_shared<Schema>();
    auto vec_fid = schema->AddDebugField(
        "fakevec", DataType::VECTOR_FLOAT, dim, knowhere::metric::L2);
    auto int64_fid = schema->AddDebugField("int64", DataType::INT64, false);
    auto int32_fid = schema->AddDebugField("int32", DataType::INT32, false);
    auto int16_fid = schema->AddDebugField("int16", DataType::INT16, false);
    auto int8_fid = schema->AddDebugField("int8", DataType::INT8, false);
    auto bool_fid = schema->AddDebugField("bool", DataType::BOOL, false);
    auto string_fid = schema->AddDebugField("string", DataType::VARCHAR, false);
    auto fp16_fid = schema->AddDebugField(
        "fakevec_fp16", DataType::VECTOR_FLOAT16, dim, knowhere::metric::L2);
    auto bf16_fid = schema->AddDebugField(
        "fakevec_bf16", DataType::VECTOR_BFLOAT16, dim, knowhere::metric::L2);
    schema->set_primary_field_id(int64_fid);

    //1. load raw data
    size_t N = 100;
    uint64_t seed = 512;
    uint64_t ts_offset = 0;
    int repeat_count_1 = 2;
    int repeat_count_2 = 5;
    auto raw_data1 = DataGen(
        schema, N, seed, ts_offset, repeat_count_1, 10, 10, false, false);
    auto raw_data2 = DataGen(
        schema, N, seed, ts_offset, repeat_count_2, 10, 10, false, false);

    auto segment1 = CreateSealedWithFieldDataLoaded(schema, raw_data1);
    auto segment2 = CreateSealedWithFieldDataLoaded(schema, raw_data2);

    //3. load index
    auto vector_data_1 = raw_data1.get_col<float>(vec_fid);
    auto indexing_1 = GenVecIndexing(
        N, dim, vector_data_1.data(), knowhere::IndexEnum::INDEX_HNSW);
    LoadIndexInfo load_index_info_1;
    load_index_info_1.field_id = vec_fid.get();
    load_index_info_1.index_params = GenIndexParams(indexing_1.get());
    load_index_info_1.cache_index =
        CreateTestCacheIndex("test", std::move(indexing_1));
    load_index_info_1.index_params[METRICS_TYPE] = knowhere::metric::L2;
    segment1->LoadIndex(load_index_info_1);

    auto vector_data_2 = raw_data2.get_col<float>(vec_fid);
    auto indexing_2 = GenVecIndexing(
        N, dim, vector_data_2.data(), knowhere::IndexEnum::INDEX_HNSW);
    LoadIndexInfo load_index_info_2;
    load_index_info_2.field_id = vec_fid.get();
    load_index_info_2.index_params = GenIndexParams(indexing_2.get());
    load_index_info_2.cache_index =
        CreateTestCacheIndex("test", std::move(indexing_2));
    load_index_info_2.index_params[METRICS_TYPE] = knowhere::metric::L2;
    segment2->LoadIndex(load_index_info_2);
    CSegmentInterface c_segment_1 = segment1.release();
    CSegmentInterface c_segment_2 = segment2.release();

    //4. search group by respectively
    auto num_queries = 10;
    auto topK = 10;
    int group_size = 3;
    auto slice_nqs = std::vector<int64_t>{num_queries / 2, num_queries / 2};
    auto slice_topKs = std::vector<int64_t>{topK / 2, topK};

    // Lambda function to execute search and reduce with given raw plan
    auto executeGroupBySearchAndReduce = [&](const char* raw_plan) {
        auto plan_str = translate_text_plan_to_binary_plan(raw_plan);
        auto plan =
            CreateSearchPlanByExpr(schema, plan_str.data(), plan_str.size());
        auto ph_group_raw = CreatePlaceholderGroup(num_queries, dim, seed);
        auto ph_group =
            ParsePlaceholderGroup(plan.get(), ph_group_raw.SerializeAsString());
        CPlaceholderGroup c_ph_group = ph_group.release();
        CSearchPlan c_plan = plan.release();

        CSearchResult c_search_res_1;
        CSearchResult c_search_res_2;
        auto status = CSearch(
            c_segment_1, c_plan, c_ph_group, MAX_TIMESTAMP, &c_search_res_1);
        ASSERT_EQ(status.error_code, Success);
        status = CSearch(
            c_segment_2, c_plan, c_ph_group, MAX_TIMESTAMP, &c_search_res_2);
        ASSERT_EQ(status.error_code, Success);
        std::vector<CSearchResult> results;
        results.push_back(c_search_res_1);
        results.push_back(c_search_res_2);

        CSearchResultDataBlobs cSearchResultData;
        status = ReduceSearchResultsAndFillData({},
                                                &cSearchResultData,
                                                c_plan,
                                                results.data(),
                                                results.size(),
                                                slice_nqs.data(),
                                                slice_topKs.data(),
                                                slice_nqs.size());
        CheckSearchResultDuplicate(results, group_size);
        DeleteSearchResult(c_search_res_1);
        DeleteSearchResult(c_search_res_2);
        DeleteSearchResultDataBlobs(cSearchResultData);
        DeleteSearchPlan(c_plan);
        DeletePlaceholderGroup(c_ph_group);
    };

    // Execute the test with the original plan (INT64)
    const char* raw_plan = R"(vector_anns: <
                                        field_id: 100
                                        query_info: <
                                          topk: 10
                                          metric_type: "L2"
                                          search_params: "{\"ef\": 10}"
                                          group_by_field_id: 101
                                          group_size: 3
                                        >
                                        placeholder_tag: "$0"
         >)";
    executeGroupBySearchAndReduce(raw_plan);

    // Test Case: Group by INT32 field
    const char* raw_plan_int32 = R"(vector_anns: <
                                        field_id: 100
                                        query_info: <
                                          topk: 10
                                          metric_type: "L2"
                                          search_params: "{\"ef\": 10}"
                                          group_by_field_id: 102
                                          group_size: 3
                                        >
                                        placeholder_tag: "$0"
         >)";
    executeGroupBySearchAndReduce(raw_plan_int32);

    // Test Case: Group by INT16 field
    const char* raw_plan_int16 = R"(vector_anns: <
                                        field_id: 100
                                        query_info: <
                                          topk: 10
                                          metric_type: "L2"
                                          search_params: "{\"ef\": 10}"
                                          group_by_field_id: 103
                                          group_size: 3
                                        >
                                        placeholder_tag: "$0"
         >)";
    executeGroupBySearchAndReduce(raw_plan_int16);

    // Test Case: Group by INT8 field
    const char* raw_plan_int8 = R"(vector_anns: <
                                        field_id: 100
                                        query_info: <
                                          topk: 10
                                          metric_type: "L2"
                                          search_params: "{\"ef\": 10}"
                                          group_by_field_id: 104
                                          group_size: 3
                                        >
                                        placeholder_tag: "$0"
         >)";
    executeGroupBySearchAndReduce(raw_plan_int8);

    // Test Case: Group by BOOL field
    const char* raw_plan_bool = R"(vector_anns: <
                                        field_id: 100
                                        query_info: <
                                          topk: 10
                                          metric_type: "L2"
                                          search_params: "{\"ef\": 10}"
                                          group_by_field_id: 105
                                          group_size: 3
                                        >
                                        placeholder_tag: "$0"
         >)";
    executeGroupBySearchAndReduce(raw_plan_bool);

    // Test Case: Group by VARCHAR field
    const char* raw_plan_string = R"(vector_anns: <
                                        field_id: 100
                                        query_info: <
                                          topk: 10
                                          metric_type: "L2"
                                          search_params: "{\"ef\": 10}"
                                          group_by_field_id: 106
                                          group_size: 3
                                        >
                                        placeholder_tag: "$0"
         >)";
    executeGroupBySearchAndReduce(raw_plan_string);

    DeleteSegment(c_segment_1);
    DeleteSegment(c_segment_2);
}

TEST(GroupBY, GrowingRawData) {
    //0. set up growing segment
    int dim = 128;
    uint64_t seed = 512;
    auto schema = std::make_shared<Schema>();
    auto metric_type = knowhere::metric::L2;
    auto int64_field_id = schema->AddDebugField("int64", DataType::INT64);
    auto int32_field_id = schema->AddDebugField("int32", DataType::INT32);
    auto vec_field_id = schema->AddDebugField(
        "embeddings", DataType::VECTOR_FLOAT, 128, metric_type);
    schema->set_primary_field_id(int64_field_id);

    auto config = SegcoreConfig::default_config();
    auto original_chunk_rows = config.get_chunk_rows();
    DeferLambda([&]() { config.set_chunk_rows(original_chunk_rows); });
    config.set_chunk_rows(128);
    config.set_enable_interim_segment_index(
        false);  //no growing index, test brute force
    auto segment_growing = CreateGrowingSegment(schema, nullptr, 1, config);
    auto segment_growing_impl =
        dynamic_cast<SegmentGrowingImpl*>(segment_growing.get());

    //1. prepare raw data in growing segment
    int64_t rows_per_batch = 512;
    int n_batch = 3;
    for (int i = 0; i < n_batch; i++) {
        auto data_set =
            DataGen(schema, rows_per_batch, 42, 0, 8, 10, 1, false, false);
        auto offset = segment_growing_impl->PreInsert(rows_per_batch);
        segment_growing_impl->Insert(offset,
                                     rows_per_batch,
                                     data_set.row_ids_.data(),
                                     data_set.timestamps_.data(),
                                     data_set.raw_);
    }

    //2. Search group by
    auto num_queries = 10;
    auto topK = 100;
    int group_size = 1;
    const char* raw_plan = R"(vector_anns: <
                                        field_id: 102
                                        query_info: <
                                          topk: 100
                                          metric_type: "L2"
                                          search_params: "{\"ef\": 10}"
                                          group_by_field_id: 101
                                          group_size: 1
                                        >
                                        placeholder_tag: "$0"

         >)";
    auto plan_str = translate_text_plan_to_binary_plan(raw_plan);
    auto plan =
        CreateSearchPlanByExpr(schema, plan_str.data(), plan_str.size());
    auto ph_group_raw = CreatePlaceholderGroup(num_queries, dim, seed);
    auto ph_group =
        ParsePlaceholderGroup(plan.get(), ph_group_raw.SerializeAsString());
    auto search_result =
        segment_growing_impl->Search(plan.get(), ph_group.get(), MAX_TIMESTAMP);
    CheckGroupBySearchResult(*search_result, topK, num_queries, true);

    auto& group_by_values = search_result->group_by_values_.value();
    int size = group_by_values.size();
    ASSERT_EQ(size, 640);
    //as the number of data is 512 and repeated count is 8, the group number is 64 for every query
    //and the total group number should be 640
    int expected_group_count = 64;
    int idx = 0;
    for (int i = 0; i < num_queries; i++) {
        std::unordered_set<int32_t> i32_set;
        float lastDistance = 0.0;
        for (int j = 0; j < expected_group_count; j++) {
            if (std::holds_alternative<int32_t>(group_by_values[idx].value())) {
                int32_t g_val = std::get<int32_t>(group_by_values[idx].value());
                ASSERT_FALSE(
                    i32_set.count(g_val) >
                    0);  //as the group_size is 1, there should not be any duplication for group_by value
                i32_set.insert(g_val);
                auto distance = search_result->distances_.at(idx);
                ASSERT_TRUE(lastDistance <= distance);
                lastDistance = distance;
            }
            idx++;
        }
    }
}

TEST(GroupBY, GrowingIndex) {
    //0. set up growing segment
    int dim = 128;
    uint64_t seed = 512;
    auto schema = std::make_shared<Schema>();
    auto metric_type = knowhere::metric::L2;
    auto int64_field_id = schema->AddDebugField("int64", DataType::INT64);
    auto int32_field_id = schema->AddDebugField("int32", DataType::INT32);
    auto vec_field_id = schema->AddDebugField(
        "embeddings", DataType::VECTOR_FLOAT, 128, metric_type);
    schema->set_primary_field_id(int64_field_id);

    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_field_id, std::move(index_params), std::move(type_params));
    std::map<FieldId, FieldIndexMeta> fieldMap = {
        {vec_field_id, fieldIndexMeta}};
    IndexMetaPtr metaPtr =
        std::make_shared<CollectionIndexMeta>(10000, std::move(fieldMap));

    auto config = SegcoreConfig::default_config();
    auto original_chunk_rows = config.get_chunk_rows();
    DeferLambda([&]() { config.set_chunk_rows(original_chunk_rows); });
    config.set_chunk_rows(128);
    config.set_enable_interim_segment_index(
        true);  //no growing index, test growing inter index
    config.set_nlist(128);
    auto segment_growing = CreateGrowingSegment(schema, metaPtr, 1, config);
    auto segment_growing_impl =
        dynamic_cast<SegmentGrowingImpl*>(segment_growing.get());

    //1. prepare raw data in growing segment
    int64_t rows_per_batch = 1024;
    int n_batch = 10;
    for (int i = 0; i < n_batch; i++) {
        auto data_set =
            DataGen(schema, rows_per_batch, 42, 0, 8, 10, 10, false, false);
        auto offset = segment_growing_impl->PreInsert(rows_per_batch);
        segment_growing_impl->Insert(offset,
                                     rows_per_batch,
                                     data_set.row_ids_.data(),
                                     data_set.timestamps_.data(),
                                     data_set.raw_);
    }

    //2. Search group by int32
    auto num_queries = 10;
    auto topK = 100;
    int group_size = 3;
    const char* raw_plan = R"(vector_anns: <
                                        field_id: 102
                                        query_info: <
                                          topk: 100
                                          metric_type: "L2"
                                          search_params: "{\"ef\": 10}"
                                          group_by_field_id: 101
                                          group_size: 3
                                          strict_group_size: true
                                        >
                                        placeholder_tag: "$0"

         >)";
    auto plan_str = translate_text_plan_to_binary_plan(raw_plan);
    auto plan =
        CreateSearchPlanByExpr(schema, plan_str.data(), plan_str.size());
    auto ph_group_raw = CreatePlaceholderGroup(num_queries, dim, seed);
    auto ph_group =
        ParsePlaceholderGroup(plan.get(), ph_group_raw.SerializeAsString());
    auto search_result =
        segment_growing_impl->Search(plan.get(), ph_group.get(), MAX_TIMESTAMP);
    CheckGroupBySearchResult(*search_result, topK, num_queries, true);

    auto& group_by_values = search_result->group_by_values_.value();
    auto size = group_by_values.size();
    int expected_group_count = 100;
    ASSERT_EQ(size, expected_group_count * group_size * num_queries);
    int idx = 0;
    for (int i = 0; i < num_queries; i++) {
        std::unordered_map<int32_t, int> i32_map;
        float lastDistance = 0.0;
        for (int j = 0; j < expected_group_count * group_size; j++) {
            if (std::holds_alternative<int32_t>(group_by_values[idx].value())) {
                int32_t g_val = std::get<int32_t>(group_by_values[idx].value());
                i32_map[g_val] += 1;
                ASSERT_TRUE(i32_map[g_val] <= group_size);
                auto distance = search_result->distances_.at(idx);
                ASSERT_TRUE(
                    lastDistance <=
                    distance);  //distance should be decreased as metrics_type is L2
                lastDistance = distance;
            }
            idx++;
        }
        ASSERT_EQ(i32_map.size(), expected_group_count);
        for (const auto& map_pair : i32_map) {
            ASSERT_EQ(group_size, map_pair.second);
        }
    }
}