test_db.cpp 22.0 KB
Newer Older
J
jinhai 已提交
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17
// Licensed to the Apache Software Foundation (ASF) 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.

Y
youny626 已提交
18 19
#include "cache/CpuCacheMgr.h"
#include "db/Constants.h"
G
groot 已提交
20
#include "db/DB.h"
Y
youny626 已提交
21
#include "db/DBFactory.h"
X
Xu Peng 已提交
22
#include "db/DBImpl.h"
S
starlord 已提交
23
#include "db/meta/MetaConsts.h"
Y
youny626 已提交
24
#include "db/utils.h"
25
#include "server/Config.h"
Y
youny626 已提交
26
#include "utils/CommonUtil.h"
G
groot 已提交
27

J
jinhai 已提交
28 29 30
#include <gtest/gtest.h>
#include <boost/filesystem.hpp>
#include <random>
Y
youny626 已提交
31
#include <thread>
32

G
groot 已提交
33 34
namespace {

Y
youny626 已提交
35
static const char* TABLE_NAME = "test_group";
S
starlord 已提交
36 37 38 39 40 41
static constexpr int64_t TABLE_DIM = 256;
static constexpr int64_t VECTOR_COUNT = 25000;
static constexpr int64_t INSERT_LOOP = 1000;
static constexpr int64_t SECONDS_EACH_HOUR = 3600;
static constexpr int64_t DAY_SECONDS = 24 * 60 * 60;

S
starlord 已提交
42
milvus::engine::meta::TableSchema
S
starlord 已提交
43
BuildTableSchema() {
S
starlord 已提交
44
    milvus::engine::meta::TableSchema table_info;
S
starlord 已提交
45 46 47 48
    table_info.dimension_ = TABLE_DIM;
    table_info.table_id_ = TABLE_NAME;
    return table_info;
}
G
groot 已提交
49

S
starlord 已提交
50
void
Y
youny626 已提交
51
BuildVectors(int64_t n, std::vector<float>& vectors) {
S
starlord 已提交
52 53
    vectors.clear();
    vectors.resize(n * TABLE_DIM);
Y
youny626 已提交
54
    float* data = vectors.data();
S
starlord 已提交
55 56 57
    for (int i = 0; i < n; i++) {
        for (int j = 0; j < TABLE_DIM; j++) data[TABLE_DIM * i + j] = drand48();
        data[TABLE_DIM * i] += i / 2000.;
G
groot 已提交
58
    }
S
starlord 已提交
59
}
G
groot 已提交
60

S
starlord 已提交
61 62 63 64 65 66 67 68
std::string
CurrentTmDate(int64_t offset_day = 0) {
    time_t tt;
    time(&tt);
    tt = tt + 8 * SECONDS_EACH_HOUR;
    tt = tt + 24 * SECONDS_EACH_HOUR * offset_day;
    tm t;
    gmtime_r(&tt, &t);
69

Y
youny626 已提交
70 71
    std::string str =
        std::to_string(t.tm_year + 1900) + "-" + std::to_string(t.tm_mon + 1) + "-" + std::to_string(t.tm_mday);
72

S
starlord 已提交
73 74
    return str;
}
75

S
starlord 已提交
76
void
Y
youny626 已提交
77 78
ConvertTimeRangeToDBDates(const std::string& start_value, const std::string& end_value,
                          std::vector<milvus::engine::meta::DateT>& dates) {
S
starlord 已提交
79
    dates.clear();
80

S
starlord 已提交
81 82
    time_t tt_start, tt_end;
    tm tm_start, tm_end;
S
starlord 已提交
83
    if (!milvus::server::CommonUtil::TimeStrToTime(start_value, tt_start, tm_start)) {
S
starlord 已提交
84 85
        return;
    }
86

S
starlord 已提交
87
    if (!milvus::server::CommonUtil::TimeStrToTime(end_value, tt_end, tm_end)) {
S
starlord 已提交
88 89
        return;
    }
90

Y
youny626 已提交
91
    int64_t days = (tt_end > tt_start) ? (tt_end - tt_start) / DAY_SECONDS : (tt_start - tt_end) / DAY_SECONDS;
S
starlord 已提交
92 93 94
    if (days == 0) {
        return;
    }
95

S
starlord 已提交
96 97 98
    for (int64_t i = 0; i < days; i++) {
        time_t tt_day = tt_start + DAY_SECONDS * i;
        tm tm_day;
S
starlord 已提交
99
        milvus::server::CommonUtil::ConvertTime(tt_day, tm_day);
100

Y
youny626 已提交
101
        int64_t date = tm_day.tm_year * 10000 + tm_day.tm_mon * 100 + tm_day.tm_mday;  // according to db logic
S
starlord 已提交
102
        dates.push_back(date);
103
    }
G
groot 已提交
104 105
}

Y
youny626 已提交
106
}  // namespace
S
starlord 已提交
107

X
Xu Peng 已提交
108 109
TEST_F(DBTest, CONFIG_TEST) {
    {
S
starlord 已提交
110
        ASSERT_ANY_THROW(milvus::engine::ArchiveConf conf("wrong"));
111
        /* EXPECT_DEATH(engine::ArchiveConf conf("wrong"), ""); */
X
Xu Peng 已提交
112 113
    }
    {
S
starlord 已提交
114
        milvus::engine::ArchiveConf conf("delete");
X
Xu Peng 已提交
115 116
        ASSERT_EQ(conf.GetType(), "delete");
        auto criterias = conf.GetCriterias();
S
starlord 已提交
117
        ASSERT_EQ(criterias.size(), 0);
X
Xu Peng 已提交
118 119
    }
    {
S
starlord 已提交
120
        milvus::engine::ArchiveConf conf("swap");
X
Xu Peng 已提交
121 122
        ASSERT_EQ(conf.GetType(), "swap");
        auto criterias = conf.GetCriterias();
S
starlord 已提交
123
        ASSERT_EQ(criterias.size(), 0);
X
Xu Peng 已提交
124 125
    }
    {
S
starlord 已提交
126 127 128
        ASSERT_ANY_THROW(milvus::engine::ArchiveConf conf1("swap", "disk:"));
        ASSERT_ANY_THROW(milvus::engine::ArchiveConf conf2("swap", "disk:a"));
        milvus::engine::ArchiveConf conf("swap", "disk:1024");
X
Xu Peng 已提交
129
        auto criterias = conf.GetCriterias();
S
starlord 已提交
130 131
        ASSERT_EQ(criterias.size(), 1);
        ASSERT_EQ(criterias["disk"], 1024);
X
Xu Peng 已提交
132 133
    }
    {
S
starlord 已提交
134 135 136
        ASSERT_ANY_THROW(milvus::engine::ArchiveConf conf1("swap", "days:"));
        ASSERT_ANY_THROW(milvus::engine::ArchiveConf conf2("swap", "days:a"));
        milvus::engine::ArchiveConf conf("swap", "days:100");
X
Xu Peng 已提交
137
        auto criterias = conf.GetCriterias();
S
starlord 已提交
138 139
        ASSERT_EQ(criterias.size(), 1);
        ASSERT_EQ(criterias["days"], 100);
X
Xu Peng 已提交
140 141
    }
    {
S
starlord 已提交
142 143 144
        ASSERT_ANY_THROW(milvus::engine::ArchiveConf conf1("swap", "days:"));
        ASSERT_ANY_THROW(milvus::engine::ArchiveConf conf2("swap", "days:a"));
        milvus::engine::ArchiveConf conf("swap", "days:100;disk:200");
X
Xu Peng 已提交
145
        auto criterias = conf.GetCriterias();
S
starlord 已提交
146 147 148
        ASSERT_EQ(criterias.size(), 2);
        ASSERT_EQ(criterias["days"], 100);
        ASSERT_EQ(criterias["disk"], 200);
X
Xu Peng 已提交
149 150 151
    }
}

152
TEST_F(DBTest, DB_TEST) {
S
starlord 已提交
153
    milvus::engine::meta::TableSchema table_info = BuildTableSchema();
S
starlord 已提交
154
    auto stat = db_->CreateTable(table_info);
G
groot 已提交
155

S
starlord 已提交
156
    milvus::engine::meta::TableSchema table_info_get;
G
groot 已提交
157
    table_info_get.table_id_ = TABLE_NAME;
G
groot 已提交
158
    stat = db_->DescribeTable(table_info_get);
S
starlord 已提交
159
    ASSERT_TRUE(stat.ok());
G
groot 已提交
160
    ASSERT_EQ(table_info_get.dimension_, TABLE_DIM);
G
groot 已提交
161

S
starlord 已提交
162 163
    milvus::engine::IDNumbers vector_ids;
    milvus::engine::IDNumbers target_ids;
X
Xu Peng 已提交
164

G
groot 已提交
165 166 167
    int64_t nb = 50;
    std::vector<float> xb;
    BuildVectors(nb, xb);
G
groot 已提交
168

G
groot 已提交
169 170 171
    int64_t qb = 5;
    std::vector<float> qxb;
    BuildVectors(qb, qxb);
X
Xu Peng 已提交
172

X
Xu Peng 已提交
173
    std::thread search([&]() {
G
groot 已提交
174 175
        milvus::engine::ResultIds result_ids;
        milvus::engine::ResultDistances result_distances;
X
Xu Peng 已提交
176 177 178 179 180
        int k = 10;
        std::this_thread::sleep_for(std::chrono::seconds(2));

        INIT_TIMER;
        std::stringstream ss;
G
groot 已提交
181 182
        uint64_t count = 0;
        uint64_t prev_count = 0;
X
Xu Peng 已提交
183

S
starlord 已提交
184
        for (auto j = 0; j < 10; ++j) {
X
Xu Peng 已提交
185
            ss.str("");
X
Xu Peng 已提交
186
            db_->Size(count);
X
Xu Peng 已提交
187
            prev_count = count;
X
Xu Peng 已提交
188 189

            START_TIMER;
G
groot 已提交
190 191 192

            std::vector<std::string> tags;
            stat = db_->Query(TABLE_NAME, tags, k, qb, 10, qxb.data(), result_ids, result_distances);
S
starlord 已提交
193
            ss << "Search " << j << " With Size " << count / milvus::engine::M << " M";
X
Xu Peng 已提交
194 195
            STOP_TIMER(ss.str());

S
starlord 已提交
196
            ASSERT_TRUE(stat.ok());
G
groot 已提交
197 198
            for (auto i = 0; i < qb; ++i) {
                ASSERT_EQ(result_ids[i*k], target_ids[i]);
X
Xu Peng 已提交
199
                ss.str("");
G
groot 已提交
200 201 202
                ss << "Result [" << i << "]:";
                for (auto t = 0; t < k; t++) {
                    ss << result_ids[i * k + t] << " ";
X
Xu Peng 已提交
203
                }
204
                /* LOG(DEBUG) << ss.str(); */
X
Xu Peng 已提交
205
            }
X
Xu Peng 已提交
206
            ASSERT_TRUE(count >= prev_count);
X
Xu Peng 已提交
207 208 209 210
            std::this_thread::sleep_for(std::chrono::seconds(1));
        }
    });

G
groot 已提交
211
    int loop = INSERT_LOOP;
X
Xu Peng 已提交
212

S
starlord 已提交
213 214
    for (auto i = 0; i < loop; ++i) {
        if (i == 40) {
G
groot 已提交
215
            db_->InsertVectors(TABLE_NAME, "", qb, qxb.data(), target_ids);
216
            ASSERT_EQ(target_ids.size(), qb);
X
Xu Peng 已提交
217
        } else {
G
groot 已提交
218
            db_->InsertVectors(TABLE_NAME, "", nb, xb.data(), vector_ids);
X
Xu Peng 已提交
219
        }
X
Xu Peng 已提交
220
        std::this_thread::sleep_for(std::chrono::microseconds(1));
X
Xu Peng 已提交
221
    }
X
xj.lin 已提交
222

X
Xu Peng 已提交
223
    search.join();
S
starlord 已提交
224 225 226

    uint64_t count;
    stat = db_->GetTableRowCount(TABLE_NAME, count);
S
starlord 已提交
227
    ASSERT_TRUE(stat.ok());
S
starlord 已提交
228 229
    ASSERT_GT(count, 0);
}
X
xj.lin 已提交
230

231
TEST_F(DBTest, SEARCH_TEST) {
S
starlord 已提交
232 233
    std::string config_path(CONFIG_PATH);
    config_path += CONFIG_FILE;
Y
youny626 已提交
234
    milvus::server::Config& config = milvus::server::Config::GetInstance();
S
starlord 已提交
235
    milvus::Status s = config.LoadConfigFile(config_path);
236

S
starlord 已提交
237
    milvus::engine::meta::TableSchema table_info = BuildTableSchema();
S
starlord 已提交
238
    auto stat = db_->CreateTable(table_info);
X
xj.lin 已提交
239

S
starlord 已提交
240
    milvus::engine::meta::TableSchema table_info_get;
G
groot 已提交
241
    table_info_get.table_id_ = TABLE_NAME;
G
groot 已提交
242
    stat = db_->DescribeTable(table_info_get);
S
starlord 已提交
243
    ASSERT_TRUE(stat.ok());
G
groot 已提交
244
    ASSERT_EQ(table_info_get.dimension_, TABLE_DIM);
X
xj.lin 已提交
245 246

    // prepare raw data
G
groot 已提交
247
    size_t nb = VECTOR_COUNT;
X
xj.lin 已提交
248 249
    size_t nq = 10;
    size_t k = 5;
S
starlord 已提交
250 251 252
    std::vector<float> xb(nb * TABLE_DIM);
    std::vector<float> xq(nq * TABLE_DIM);
    std::vector<int64_t> ids(nb);
X
xj.lin 已提交
253 254 255 256

    std::random_device rd;
    std::mt19937 gen(rd());
    std::uniform_real_distribution<> dis_xt(-1.0, 1.0);
S
starlord 已提交
257
    for (size_t i = 0; i < nb * TABLE_DIM; i++) {
X
xj.lin 已提交
258
        xb[i] = dis_xt(gen);
S
starlord 已提交
259
        if (i < nb) {
X
xj.lin 已提交
260 261 262
            ids[i] = i;
        }
    }
S
starlord 已提交
263
    for (size_t i = 0; i < nq * TABLE_DIM; i++) {
X
xj.lin 已提交
264 265 266 267
        xq[i] = dis_xt(gen);
    }

    // result data
Y
youny626 已提交
268
    // std::vector<long> nns_gt(k*nq);
S
starlord 已提交
269
    std::vector<int64_t> nns(k * nq);  // nns = nearst neg search
Y
youny626 已提交
270
    // std::vector<float> dis_gt(k*nq);
S
starlord 已提交
271
    std::vector<float> dis(k * nq);
X
xj.lin 已提交
272 273 274 275

    // insert data
    const int batch_size = 100;
    for (int j = 0; j < nb / batch_size; ++j) {
G
groot 已提交
276
        stat = db_->InsertVectors(TABLE_NAME, "", batch_size, xb.data() + batch_size * j * TABLE_DIM, ids);
Y
youny626 已提交
277 278 279
        if (j == 200) {
            sleep(1);
        }
S
starlord 已提交
280
        ASSERT_TRUE(stat.ok());
X
xj.lin 已提交
281 282
    }

S
starlord 已提交
283
    milvus::engine::TableIndex index;
Y
youny626 已提交
284 285
    index.engine_type_ = (int)milvus::engine::EngineType::FAISS_IDMAP;
    db_->CreateIndex(TABLE_NAME, index);  // wait until build index finish
X
xj.lin 已提交
286

S
starlord 已提交
287
    {
G
groot 已提交
288 289 290 291
        std::vector<std::string> tags;
        milvus::engine::ResultIds result_ids;
        milvus::engine::ResultDistances result_distances;
        stat = db_->Query(TABLE_NAME, tags, k, nq, 10, xq.data(), result_ids, result_distances);
S
starlord 已提交
292
        ASSERT_TRUE(stat.ok());
S
starlord 已提交
293 294
    }

Y
youny626 已提交
295
    {  // search by specify index file
S
starlord 已提交
296
        milvus::engine::meta::DatesT dates;
297
        std::vector<std::string> file_ids = {"1", "2", "3", "4", "5", "6"};
G
groot 已提交
298 299 300
        milvus::engine::ResultIds result_ids;
        milvus::engine::ResultDistances result_distances;
        stat = db_->QueryByFileID(TABLE_NAME, file_ids, k, nq, 10, xq.data(), dates, result_ids, result_distances);
S
starlord 已提交
301
        ASSERT_TRUE(stat.ok());
X
xj.lin 已提交
302
    }
X
xj.lin 已提交
303

Z
update  
zhiru 已提交
304
#ifdef CUSTOMIZATION
Y
youny626 已提交
305
    // test FAISS_IVFSQ8H optimizer
306
    index.engine_type_ = (int)milvus::engine::EngineType::FAISS_IVFSQ8H;
Y
youny626 已提交
307
    db_->CreateIndex(TABLE_NAME, index);  // wait until build index finish
308 309 310
    std::vector<std::string> partition_tag;
    milvus::engine::ResultIds result_ids;
    milvus::engine::ResultDistances result_dists;
311 312

    {
313 314 315
        result_ids.clear();
        result_dists.clear();
        stat = db_->Query(TABLE_NAME, partition_tag, k, nq, 10, xq.data(), result_ids, result_dists);
316 317 318
        ASSERT_TRUE(stat.ok());
    }

Y
Yu Kun 已提交
319
    {
320 321 322
        result_ids.clear();
        result_dists.clear();
        stat = db_->Query(TABLE_NAME, partition_tag, k, 200, 10, xq.data(), result_ids, result_dists);
Y
Yu Kun 已提交
323 324 325
        ASSERT_TRUE(stat.ok());
    }

Y
youny626 已提交
326
    {  // search by specify index file
327 328
        milvus::engine::meta::DatesT dates;
        std::vector<std::string> file_ids = {"1", "2", "3", "4", "5", "6"};
329 330 331
        result_ids.clear();
        result_dists.clear();
        stat = db_->QueryByFileID(TABLE_NAME, file_ids, k, nq, 10, xq.data(), dates, result_ids, result_dists);
332 333
        ASSERT_TRUE(stat.ok());
    }
Y
Yu Kun 已提交
334

X
xiaojun.lin 已提交
335
#endif
S
starlord 已提交
336
}
Y
c  
yu yunfeng 已提交
337

Y
Yu Kun 已提交
338
TEST_F(DBTest, PRELOADTABLE_TEST) {
S
starlord 已提交
339
    milvus::engine::meta::TableSchema table_info = BuildTableSchema();
S
starlord 已提交
340
    auto stat = db_->CreateTable(table_info);
Y
Yu Kun 已提交
341

S
starlord 已提交
342
    milvus::engine::meta::TableSchema table_info_get;
Y
Yu Kun 已提交
343 344
    table_info_get.table_id_ = TABLE_NAME;
    stat = db_->DescribeTable(table_info_get);
S
starlord 已提交
345
    ASSERT_TRUE(stat.ok());
Y
Yu Kun 已提交
346 347
    ASSERT_EQ(table_info_get.dimension_, TABLE_DIM);

348
    int64_t nb = VECTOR_COUNT;
Y
Yu Kun 已提交
349 350 351
    std::vector<float> xb;
    BuildVectors(nb, xb);

Y
Yu Kun 已提交
352
    int loop = 5;
S
starlord 已提交
353
    for (auto i = 0; i < loop; ++i) {
S
starlord 已提交
354
        milvus::engine::IDNumbers vector_ids;
G
groot 已提交
355
        db_->InsertVectors(TABLE_NAME, "", nb, xb.data(), vector_ids);
356
        ASSERT_EQ(vector_ids.size(), nb);
Y
Yu Kun 已提交
357
    }
358

S
starlord 已提交
359
    milvus::engine::TableIndex index;
Y
youny626 已提交
360 361
    index.engine_type_ = (int)milvus::engine::EngineType::FAISS_IDMAP;
    db_->CreateIndex(TABLE_NAME, index);  // wait until build index finish
Y
Yu Kun 已提交
362

S
starlord 已提交
363
    int64_t prev_cache_usage = milvus::cache::CpuCacheMgr::GetInstance()->CacheUsage();
Y
Yu Kun 已提交
364
    stat = db_->PreloadTable(TABLE_NAME);
S
starlord 已提交
365
    ASSERT_TRUE(stat.ok());
S
starlord 已提交
366
    int64_t cur_cache_usage = milvus::cache::CpuCacheMgr::GetInstance()->CacheUsage();
Y
Yu Kun 已提交
367 368 369
    ASSERT_TRUE(prev_cache_usage < cur_cache_usage);
}

S
starlord 已提交
370 371 372
TEST_F(DBTest, SHUTDOWN_TEST) {
    db_->Stop();

S
starlord 已提交
373
    milvus::engine::meta::TableSchema table_info = BuildTableSchema();
S
starlord 已提交
374
    auto stat = db_->CreateTable(table_info);
S
starlord 已提交
375 376 377 378 379 380 381 382 383
    ASSERT_FALSE(stat.ok());

    stat = db_->DescribeTable(table_info);
    ASSERT_FALSE(stat.ok());

    bool has_table = false;
    stat = db_->HasTable(table_info.table_id_, has_table);
    ASSERT_FALSE(stat.ok());

S
starlord 已提交
384
    milvus::engine::IDNumbers ids;
G
groot 已提交
385
    stat = db_->InsertVectors(table_info.table_id_, "", 0, nullptr, ids);
S
starlord 已提交
386 387 388 389 390 391 392 393 394
    ASSERT_FALSE(stat.ok());

    stat = db_->PreloadTable(table_info.table_id_);
    ASSERT_FALSE(stat.ok());

    uint64_t row_count = 0;
    stat = db_->GetTableRowCount(table_info.table_id_, row_count);
    ASSERT_FALSE(stat.ok());

S
starlord 已提交
395
    milvus::engine::TableIndex index;
S
starlord 已提交
396 397 398 399 400 401
    stat = db_->CreateIndex(table_info.table_id_, index);
    ASSERT_FALSE(stat.ok());

    stat = db_->DescribeIndex(table_info.table_id_, index);
    ASSERT_FALSE(stat.ok());

G
groot 已提交
402
    std::vector<std::string> tags;
S
starlord 已提交
403
    milvus::engine::meta::DatesT dates;
G
groot 已提交
404 405 406
    milvus::engine::ResultIds result_ids;
    milvus::engine::ResultDistances result_distances;
    stat = db_->Query(table_info.table_id_, tags, 1, 1, 1, nullptr, dates, result_ids, result_distances);
S
starlord 已提交
407 408
    ASSERT_FALSE(stat.ok());
    std::vector<std::string> file_ids;
G
groot 已提交
409
    stat = db_->QueryByFileID(table_info.table_id_, file_ids, 1, 1, 1, nullptr, dates, result_ids, result_distances);
S
starlord 已提交
410 411
    ASSERT_FALSE(stat.ok());

G
groot 已提交
412
    stat = db_->DropTable(table_info.table_id_, dates);
S
starlord 已提交
413 414 415 416
    ASSERT_FALSE(stat.ok());
}

TEST_F(DBTest, INDEX_TEST) {
S
starlord 已提交
417
    milvus::engine::meta::TableSchema table_info = BuildTableSchema();
S
starlord 已提交
418
    auto stat = db_->CreateTable(table_info);
S
starlord 已提交
419 420 421 422 423

    int64_t nb = VECTOR_COUNT;
    std::vector<float> xb;
    BuildVectors(nb, xb);

S
starlord 已提交
424
    milvus::engine::IDNumbers vector_ids;
G
groot 已提交
425
    db_->InsertVectors(TABLE_NAME, "", nb, xb.data(), vector_ids);
S
starlord 已提交
426 427
    ASSERT_EQ(vector_ids.size(), nb);

S
starlord 已提交
428
    milvus::engine::TableIndex index;
Y
youny626 已提交
429 430
    index.engine_type_ = (int)milvus::engine::EngineType::FAISS_IVFSQ8;
    index.metric_type_ = (int)milvus::engine::MetricType::IP;
S
starlord 已提交
431 432 433
    stat = db_->CreateIndex(table_info.table_id_, index);
    ASSERT_TRUE(stat.ok());

Y
youny626 已提交
434
    index.engine_type_ = (int)milvus::engine::EngineType::FAISS_IVFFLAT;
Y
Yu Kun 已提交
435 436 437 438 439 440 441 442 443
    stat = db_->CreateIndex(table_info.table_id_, index);
    ASSERT_TRUE(stat.ok());

#ifdef CUSTOMIZATION
    index.engine_type_ = (int)milvus::engine::EngineType::FAISS_IVFSQ8H;
    stat = db_->CreateIndex(table_info.table_id_, index);
    ASSERT_TRUE(stat.ok());
#endif

S
starlord 已提交
444
    milvus::engine::TableIndex index_out;
S
starlord 已提交
445
    stat = db_->DescribeIndex(table_info.table_id_, index_out);
S
starlord 已提交
446 447 448
    ASSERT_TRUE(stat.ok());
    ASSERT_EQ(index.engine_type_, index_out.engine_type_);
    ASSERT_EQ(index.nlist_, index_out.nlist_);
S
starlord 已提交
449
    ASSERT_EQ(table_info.metric_type_, index_out.metric_type_);
S
starlord 已提交
450 451 452 453 454

    stat = db_->DropIndex(table_info.table_id_);
    ASSERT_TRUE(stat.ok());
}

G
groot 已提交
455 456 457 458 459 460 461 462 463 464 465 466 467 468 469
TEST_F(DBTest, PARTITION_TEST) {
    milvus::engine::meta::TableSchema table_info = BuildTableSchema();
    auto stat = db_->CreateTable(table_info);
    ASSERT_TRUE(stat.ok());

    // create partition and insert data
    const int64_t PARTITION_COUNT = 5;
    const int64_t INSERT_BATCH = 2000;
    std::string table_name = TABLE_NAME;
    for (int64_t i = 0; i < PARTITION_COUNT; i++) {
        std::string partition_tag = std::to_string(i);
        std::string partition_name = table_name + "_" + partition_tag;
        stat = db_->CreatePartition(table_name, partition_name, partition_tag);
        ASSERT_TRUE(stat.ok());

470 471 472 473 474 475 476
        // not allow nested partition
        stat = db_->CreatePartition(partition_name, "dumy", "dummy");
        ASSERT_FALSE(stat.ok());

        // not allow duplicated partition
        stat = db_->CreatePartition(table_name, partition_name, partition_tag);
        ASSERT_FALSE(stat.ok());
G
groot 已提交
477 478 479 480 481 482 483 484 485 486 487 488 489 490 491 492 493 494 495 496 497 498 499 500 501 502 503 504 505 506 507 508 509 510 511 512 513 514 515 516 517 518 519 520 521 522 523 524 525 526 527 528 529 530 531 532 533 534 535 536 537 538 539 540 541 542 543 544 545 546 547 548 549 550 551 552 553 554 555 556 557 558 559 560 561

        std::vector<float> xb;
        BuildVectors(INSERT_BATCH, xb);

        milvus::engine::IDNumbers vector_ids;
        vector_ids.resize(INSERT_BATCH);
        for (int64_t k = 0; k < INSERT_BATCH; k++) {
            vector_ids[k] = i*INSERT_BATCH + k;
        }

        db_->InsertVectors(table_name, partition_tag, INSERT_BATCH, xb.data(), vector_ids);
        ASSERT_EQ(vector_ids.size(), INSERT_BATCH);
    }

    //duplicated partition is not allowed
    stat = db_->CreatePartition(table_name, "", "0");
    ASSERT_FALSE(stat.ok());

    std::vector<milvus::engine::meta::TableSchema> partiton_schema_array;
    stat = db_->ShowPartitions(table_name, partiton_schema_array);
    ASSERT_TRUE(stat.ok());
    ASSERT_EQ(partiton_schema_array.size(), PARTITION_COUNT);
    for (int64_t i = 0; i < PARTITION_COUNT; i++) {
        ASSERT_EQ(partiton_schema_array[i].table_id_, table_name + "_" + std::to_string(i));
    }

    { // build index
        milvus::engine::TableIndex index;
        index.engine_type_ = (int) milvus::engine::EngineType::FAISS_IVFFLAT;
        index.metric_type_ = (int) milvus::engine::MetricType::L2;
        stat = db_->CreateIndex(table_info.table_id_, index);
        ASSERT_TRUE(stat.ok());

        uint64_t row_count = 0;
        stat = db_->GetTableRowCount(TABLE_NAME, row_count);
        ASSERT_TRUE(stat.ok());
        ASSERT_EQ(row_count, INSERT_BATCH*PARTITION_COUNT);
    }

    { // search
        const int64_t nq = 5;
        const int64_t topk = 10;
        const int64_t nprobe = 10;
        std::vector<float> xq;
        BuildVectors(nq, xq);

        // specify partition tags
        std::vector<std::string> tags = {"0", std::to_string(PARTITION_COUNT - 1)};
        milvus::engine::ResultIds result_ids;
        milvus::engine::ResultDistances result_distances;
        stat = db_->Query(TABLE_NAME, tags, topk, nq, nprobe, xq.data(), result_ids, result_distances);
        ASSERT_TRUE(stat.ok());
        ASSERT_EQ(result_ids.size()/topk, nq);

        // search in whole table
        tags.clear();
        result_ids.clear();
        result_distances.clear();
        stat = db_->Query(TABLE_NAME, tags, topk, nq, nprobe, xq.data(), result_ids, result_distances);
        ASSERT_TRUE(stat.ok());
        ASSERT_EQ(result_ids.size()/topk, nq);

        // search in all partitions(tag regex match)
        tags.push_back("\\d");
        result_ids.clear();
        result_distances.clear();
        stat = db_->Query(TABLE_NAME, tags, topk, nq, nprobe, xq.data(), result_ids, result_distances);
        ASSERT_TRUE(stat.ok());
        ASSERT_EQ(result_ids.size()/topk, nq);
    }

    stat = db_->DropPartition(table_name + "_0");
    ASSERT_TRUE(stat.ok());

    stat = db_->DropPartitionByTag(table_name, "1");
    ASSERT_TRUE(stat.ok());

    stat = db_->DropIndex(table_name);
    ASSERT_TRUE(stat.ok());

    milvus::engine::meta::DatesT dates;
    stat = db_->DropTable(table_name, dates);
    ASSERT_TRUE(stat.ok());
}

G
groot 已提交
562
TEST_F(DBTest2, ARHIVE_DISK_CHECK) {
S
starlord 已提交
563
    milvus::engine::meta::TableSchema table_info = BuildTableSchema();
S
starlord 已提交
564
    auto stat = db_->CreateTable(table_info);
Z
zhiru 已提交
565

S
starlord 已提交
566
    std::vector<milvus::engine::meta::TableSchema> table_schema_array;
G
groot 已提交
567
    stat = db_->AllTables(table_schema_array);
S
starlord 已提交
568
    ASSERT_TRUE(stat.ok());
G
groot 已提交
569
    bool bfound = false;
Y
youny626 已提交
570
    for (auto& schema : table_schema_array) {
S
starlord 已提交
571
        if (schema.table_id_ == TABLE_NAME) {
G
groot 已提交
572 573 574 575 576 577
            bfound = true;
            break;
        }
    }
    ASSERT_TRUE(bfound);

S
starlord 已提交
578
    milvus::engine::meta::TableSchema table_info_get;
G
groot 已提交
579 580
    table_info_get.table_id_ = TABLE_NAME;
    stat = db_->DescribeTable(table_info_get);
S
starlord 已提交
581
    ASSERT_TRUE(stat.ok());
G
groot 已提交
582
    ASSERT_EQ(table_info_get.dimension_, TABLE_DIM);
Z
zhiru 已提交
583

G
groot 已提交
584
    uint64_t size;
Z
zhiru 已提交
585 586
    db_->Size(size);

G
groot 已提交
587 588 589
    int64_t nb = 10;
    std::vector<float> xb;
    BuildVectors(nb, xb);
Z
zhiru 已提交
590

G
groot 已提交
591
    int loop = INSERT_LOOP;
S
starlord 已提交
592
    for (auto i = 0; i < loop; ++i) {
S
starlord 已提交
593
        milvus::engine::IDNumbers vector_ids;
G
groot 已提交
594
        db_->InsertVectors(TABLE_NAME, "", nb, xb.data(), vector_ids);
Z
zhiru 已提交
595 596 597
        std::this_thread::sleep_for(std::chrono::microseconds(1));
    }

G
groot 已提交
598
    std::this_thread::sleep_for(std::chrono::seconds(1));
Z
zhiru 已提交
599 600 601

    db_->Size(size);
    LOG(DEBUG) << "size=" << size;
S
starlord 已提交
602
    ASSERT_LE(size, 1 * milvus::engine::G);
S
starlord 已提交
603
}
Z
zhiru 已提交
604

G
groot 已提交
605
TEST_F(DBTest2, DELETE_TEST) {
S
starlord 已提交
606
    milvus::engine::meta::TableSchema table_info = BuildTableSchema();
S
starlord 已提交
607
    auto stat = db_->CreateTable(table_info);
G
groot 已提交
608

S
starlord 已提交
609
    milvus::engine::meta::TableSchema table_info_get;
G
groot 已提交
610 611
    table_info_get.table_id_ = TABLE_NAME;
    stat = db_->DescribeTable(table_info_get);
S
starlord 已提交
612
    ASSERT_TRUE(stat.ok());
G
groot 已提交
613

S
starlord 已提交
614 615 616
    bool has_table = false;
    db_->HasTable(TABLE_NAME, has_table);
    ASSERT_TRUE(has_table);
G
groot 已提交
617 618 619 620

    uint64_t size;
    db_->Size(size);

621
    int64_t nb = VECTOR_COUNT;
G
groot 已提交
622 623 624
    std::vector<float> xb;
    BuildVectors(nb, xb);

S
starlord 已提交
625
    milvus::engine::IDNumbers vector_ids;
G
groot 已提交
626
    stat = db_->InsertVectors(TABLE_NAME, "", nb, xb.data(), vector_ids);
S
starlord 已提交
627
    milvus::engine::TableIndex index;
628
    stat = db_->CreateIndex(TABLE_NAME, index);
G
groot 已提交
629

S
starlord 已提交
630
    std::vector<milvus::engine::meta::DateT> dates;
G
groot 已提交
631
    stat = db_->DropTable(TABLE_NAME, dates);
G
groot 已提交
632 633
    std::this_thread::sleep_for(std::chrono::seconds(2));
    ASSERT_TRUE(stat.ok());
S
starlord 已提交
634 635 636

    db_->HasTable(TABLE_NAME, has_table);
    ASSERT_FALSE(has_table);
S
starlord 已提交
637
}
638 639

TEST_F(DBTest2, DELETE_BY_RANGE_TEST) {
S
starlord 已提交
640
    milvus::engine::meta::TableSchema table_info = BuildTableSchema();
S
starlord 已提交
641
    auto stat = db_->CreateTable(table_info);
642

S
starlord 已提交
643
    milvus::engine::meta::TableSchema table_info_get;
644 645
    table_info_get.table_id_ = TABLE_NAME;
    stat = db_->DescribeTable(table_info_get);
S
starlord 已提交
646
    ASSERT_TRUE(stat.ok());
647 648 649 650 651 652 653

    bool has_table = false;
    db_->HasTable(TABLE_NAME, has_table);
    ASSERT_TRUE(has_table);

    uint64_t size;
    db_->Size(size);
654
    ASSERT_EQ(size, 0UL);
655

656
    int64_t nb = VECTOR_COUNT;
657 658 659
    std::vector<float> xb;
    BuildVectors(nb, xb);

S
starlord 已提交
660
    milvus::engine::IDNumbers vector_ids;
G
groot 已提交
661
    stat = db_->InsertVectors(TABLE_NAME, "", nb, xb.data(), vector_ids);
S
starlord 已提交
662
    milvus::engine::TableIndex index;
663 664 665 666
    stat = db_->CreateIndex(TABLE_NAME, index);

    db_->Size(size);
    ASSERT_NE(size, 0UL);
667

S
starlord 已提交
668
    std::vector<milvus::engine::meta::DateT> dates;
669 670
    std::string start_value = CurrentTmDate();
    std::string end_value = CurrentTmDate(1);
671 672
    ConvertTimeRangeToDBDates(start_value, end_value, dates);

G
groot 已提交
673
    stat = db_->DropTable(TABLE_NAME, dates);
S
starlord 已提交
674
    ASSERT_TRUE(stat.ok());
675 676 677 678

    uint64_t row_count = 0;
    db_->GetTableRowCount(TABLE_NAME, row_count);
    ASSERT_EQ(row_count, 0UL);
S
starlord 已提交
679
}