milvus/internal/core/unittest/test_data_codec.cpp

// Licensed to the LF AI & Data foundation under one
// or more contributor license agreements. See the NOTICE file
// distributed with this work for additional information
// regarding copyright ownership. The ASF licenses this file
// to you 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 "storage/DataCodec.h"
#include "storage/InsertData.h"
#include "storage/IndexData.h"
#include "common/Consts.h"
#include "utils/Json.h"

using namespace milvus;

TEST(storage, InsertDataFloat) {
    std::vector<float> data = {1, 2, 3, 4, 5};
    storage::Payload payload{storage::DataType::FLOAT,
                             reinterpret_cast<const uint8_t*>(data.data()),
                             int(data.size())};
    auto field_data = std::make_shared<storage::FieldData>(payload);

    storage::InsertData insert_data(field_data);
    storage::FieldDataMeta field_data_meta{100, 101, 102, 103};
    insert_data.SetFieldDataMeta(field_data_meta);
    insert_data.SetTimestamps(0, 100);

    auto serialized_bytes = insert_data.Serialize(storage::StorageType::Remote);
    auto new_insert_data = storage::DeserializeFileData(
        reinterpret_cast<const uint8_t*>(serialized_bytes.data()),
        serialized_bytes.size());
    ASSERT_EQ(new_insert_data->GetCodecType(), storage::InsertDataType);
    ASSERT_EQ(new_insert_data->GetTimeRage(),
              std::make_pair(Timestamp(0), Timestamp(100)));
    auto new_payload = new_insert_data->GetPayload();
    ASSERT_EQ(new_payload->data_type, storage::DataType::FLOAT);
    ASSERT_EQ(new_payload->rows, data.size());
    std::vector<float> new_data(data.size());
    memcpy(new_data.data(),
           new_payload->raw_data,
           new_payload->rows * sizeof(float));
    ASSERT_EQ(data, new_data);
}

TEST(storage, InsertDataVectorFloat) {
    std::vector<float> data = {1, 2, 3, 4, 5, 6, 7, 8};
    int DIM = 2;
    storage::Payload payload{storage::DataType::VECTOR_FLOAT,
                             reinterpret_cast<const uint8_t*>(data.data()),
                             int(data.size()) / DIM,
                             DIM};
    auto field_data = std::make_shared<storage::FieldData>(payload);

    storage::InsertData insert_data(field_data);
    storage::FieldDataMeta field_data_meta{100, 101, 102, 103};
    insert_data.SetFieldDataMeta(field_data_meta);
    insert_data.SetTimestamps(0, 100);

    auto serialized_bytes = insert_data.Serialize(storage::StorageType::Remote);
    auto new_insert_data = storage::DeserializeFileData(
        reinterpret_cast<const uint8_t*>(serialized_bytes.data()),
        serialized_bytes.size());
    ASSERT_EQ(new_insert_data->GetCodecType(), storage::InsertDataType);
    ASSERT_EQ(new_insert_data->GetTimeRage(),
              std::make_pair(Timestamp(0), Timestamp(100)));
    auto new_payload = new_insert_data->GetPayload();
    ASSERT_EQ(new_payload->data_type, storage::DataType::VECTOR_FLOAT);
    ASSERT_EQ(new_payload->rows, data.size() / DIM);
    std::vector<float> new_data(data.size());
    memcpy(new_data.data(),
           new_payload->raw_data,
           new_payload->rows * sizeof(float) * DIM);
    ASSERT_EQ(data, new_data);
}

TEST(storage, LocalInsertDataVectorFloat) {
    std::vector<float> data = {1, 2, 3, 4, 5, 6, 7, 8};
    int DIM = 2;
    storage::Payload payload{storage::DataType::VECTOR_FLOAT,
                             reinterpret_cast<const uint8_t*>(data.data()),
                             int(data.size()) / DIM,
                             DIM};
    auto field_data = std::make_shared<storage::FieldData>(payload);

    storage::InsertData insert_data(field_data);
    storage::FieldDataMeta field_data_meta{100, 101, 102, 103};
    insert_data.SetFieldDataMeta(field_data_meta);

    auto serialized_bytes =
        insert_data.Serialize(storage::StorageType::LocalDisk);
    auto new_insert_data = storage::DeserializeLocalInsertFileData(
        reinterpret_cast<const uint8_t*>(serialized_bytes.data()),
        serialized_bytes.size(),
        storage::DataType::VECTOR_FLOAT);
    ASSERT_EQ(new_insert_data->GetCodecType(), storage::InsertDataType);
    auto new_payload = new_insert_data->GetPayload();
    ASSERT_EQ(new_payload->data_type, storage::DataType::VECTOR_FLOAT);
    ASSERT_EQ(new_payload->rows, data.size() / DIM);
    std::vector<float> new_data(data.size());
    memcpy(new_data.data(),
           new_payload->raw_data,
           new_payload->rows * sizeof(float) * DIM);
    ASSERT_EQ(data, new_data);
}

TEST(storage, LocalIndexData) {
    std::vector<uint8_t> data = {1, 2, 3, 4, 5, 6, 7, 8};
    storage::Payload payload{storage::DataType::INT8,
                             reinterpret_cast<const uint8_t*>(data.data()),
                             int(data.size())};
    auto field_data = std::make_shared<storage::FieldData>(payload);
    storage::IndexData indexData_data(field_data);
    auto serialized_bytes =
        indexData_data.Serialize(storage::StorageType::LocalDisk);

    auto new_index_data = storage::DeserializeLocalIndexFileData(
        reinterpret_cast<const uint8_t*>(serialized_bytes.data()),
        serialized_bytes.size());
    ASSERT_EQ(new_index_data->GetCodecType(), storage::IndexDataType);
    auto new_payload = new_index_data->GetPayload();
    ASSERT_EQ(new_payload->data_type, storage::DataType::INT8);
    ASSERT_EQ(new_payload->rows, data.size());
    std::vector<uint8_t> new_data(data.size());
    memcpy(new_data.data(),
           new_payload->raw_data,
           new_payload->rows * sizeof(uint8_t));
    ASSERT_EQ(data, new_data);
}