test_db.cpp 21.5 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

    {
Y
youny626 已提交
310 311
        milvus::engine::QueryResults results;
        stat = db_->Query(TABLE_NAME, k, nq, 10, xq.data(), results);
312 313 314
        ASSERT_TRUE(stat.ok());
    }

Y
Yu Kun 已提交
315
    {
Y
youny626 已提交
316 317
        milvus::engine::QueryResults large_nq_results;
        stat = db_->Query(TABLE_NAME, k, 200, 10, xq.data(), large_nq_results);
Y
Yu Kun 已提交
318 319 320
        ASSERT_TRUE(stat.ok());
    }

Y
youny626 已提交
321
    {  // search by specify index file
322 323
        milvus::engine::meta::DatesT dates;
        std::vector<std::string> file_ids = {"1", "2", "3", "4", "5", "6"};
Y
youny626 已提交
324 325
        milvus::engine::QueryResults results;
        stat = db_->Query(TABLE_NAME, file_ids, k, nq, 10, xq.data(), dates, results);
326 327
        ASSERT_TRUE(stat.ok());
    }
Y
Yu Kun 已提交
328

X
xiaojun.lin 已提交
329
#endif
S
starlord 已提交
330
}
Y
c  
yu yunfeng 已提交
331

Y
Yu Kun 已提交
332
TEST_F(DBTest, PRELOADTABLE_TEST) {
S
starlord 已提交
333
    milvus::engine::meta::TableSchema table_info = BuildTableSchema();
S
starlord 已提交
334
    auto stat = db_->CreateTable(table_info);
Y
Yu Kun 已提交
335

S
starlord 已提交
336
    milvus::engine::meta::TableSchema table_info_get;
Y
Yu Kun 已提交
337 338
    table_info_get.table_id_ = TABLE_NAME;
    stat = db_->DescribeTable(table_info_get);
S
starlord 已提交
339
    ASSERT_TRUE(stat.ok());
Y
Yu Kun 已提交
340 341
    ASSERT_EQ(table_info_get.dimension_, TABLE_DIM);

342
    int64_t nb = VECTOR_COUNT;
Y
Yu Kun 已提交
343 344 345
    std::vector<float> xb;
    BuildVectors(nb, xb);

Y
Yu Kun 已提交
346
    int loop = 5;
S
starlord 已提交
347
    for (auto i = 0; i < loop; ++i) {
S
starlord 已提交
348
        milvus::engine::IDNumbers vector_ids;
G
groot 已提交
349
        db_->InsertVectors(TABLE_NAME, "", nb, xb.data(), vector_ids);
350
        ASSERT_EQ(vector_ids.size(), nb);
Y
Yu Kun 已提交
351
    }
352

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

S
starlord 已提交
357
    int64_t prev_cache_usage = milvus::cache::CpuCacheMgr::GetInstance()->CacheUsage();
Y
Yu Kun 已提交
358
    stat = db_->PreloadTable(TABLE_NAME);
S
starlord 已提交
359
    ASSERT_TRUE(stat.ok());
S
starlord 已提交
360
    int64_t cur_cache_usage = milvus::cache::CpuCacheMgr::GetInstance()->CacheUsage();
Y
Yu Kun 已提交
361 362 363
    ASSERT_TRUE(prev_cache_usage < cur_cache_usage);
}

S
starlord 已提交
364 365 366
TEST_F(DBTest, SHUTDOWN_TEST) {
    db_->Stop();

S
starlord 已提交
367
    milvus::engine::meta::TableSchema table_info = BuildTableSchema();
S
starlord 已提交
368
    auto stat = db_->CreateTable(table_info);
S
starlord 已提交
369 370 371 372 373 374 375 376 377
    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 已提交
378
    milvus::engine::IDNumbers ids;
G
groot 已提交
379
    stat = db_->InsertVectors(table_info.table_id_, "", 0, nullptr, ids);
S
starlord 已提交
380 381 382 383 384 385 386 387 388
    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 已提交
389
    milvus::engine::TableIndex index;
S
starlord 已提交
390 391 392 393 394 395
    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 已提交
396
    std::vector<std::string> tags;
S
starlord 已提交
397
    milvus::engine::meta::DatesT dates;
G
groot 已提交
398 399 400
    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 已提交
401 402
    ASSERT_FALSE(stat.ok());
    std::vector<std::string> file_ids;
G
groot 已提交
403
    stat = db_->QueryByFileID(table_info.table_id_, file_ids, 1, 1, 1, nullptr, dates, result_ids, result_distances);
S
starlord 已提交
404 405
    ASSERT_FALSE(stat.ok());

G
groot 已提交
406
    stat = db_->DropTable(table_info.table_id_, dates);
S
starlord 已提交
407 408 409 410
    ASSERT_FALSE(stat.ok());
}

TEST_F(DBTest, INDEX_TEST) {
S
starlord 已提交
411
    milvus::engine::meta::TableSchema table_info = BuildTableSchema();
S
starlord 已提交
412
    auto stat = db_->CreateTable(table_info);
S
starlord 已提交
413 414 415 416 417

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

S
starlord 已提交
418
    milvus::engine::IDNumbers vector_ids;
G
groot 已提交
419
    db_->InsertVectors(TABLE_NAME, "", nb, xb.data(), vector_ids);
S
starlord 已提交
420 421
    ASSERT_EQ(vector_ids.size(), nb);

S
starlord 已提交
422
    milvus::engine::TableIndex index;
Y
youny626 已提交
423 424
    index.engine_type_ = (int)milvus::engine::EngineType::FAISS_IVFSQ8;
    index.metric_type_ = (int)milvus::engine::MetricType::IP;
S
starlord 已提交
425 426 427
    stat = db_->CreateIndex(table_info.table_id_, index);
    ASSERT_TRUE(stat.ok());

Y
youny626 已提交
428
    index.engine_type_ = (int)milvus::engine::EngineType::FAISS_IVFFLAT;
Y
Yu Kun 已提交
429 430 431 432 433 434 435 436 437
    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 已提交
438
    milvus::engine::TableIndex index_out;
S
starlord 已提交
439
    stat = db_->DescribeIndex(table_info.table_id_, index_out);
S
starlord 已提交
440 441 442
    ASSERT_TRUE(stat.ok());
    ASSERT_EQ(index.engine_type_, index_out.engine_type_);
    ASSERT_EQ(index.nlist_, index_out.nlist_);
S
starlord 已提交
443
    ASSERT_EQ(table_info.metric_type_, index_out.metric_type_);
S
starlord 已提交
444 445 446 447 448

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

G
groot 已提交
449 450 451 452 453 454 455 456 457 458 459 460 461 462 463 464 465 466 467 468 469 470 471 472 473 474 475 476 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
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());


        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 已提交
549
TEST_F(DBTest2, ARHIVE_DISK_CHECK) {
S
starlord 已提交
550
    milvus::engine::meta::TableSchema table_info = BuildTableSchema();
S
starlord 已提交
551
    auto stat = db_->CreateTable(table_info);
Z
zhiru 已提交
552

S
starlord 已提交
553
    std::vector<milvus::engine::meta::TableSchema> table_schema_array;
G
groot 已提交
554
    stat = db_->AllTables(table_schema_array);
S
starlord 已提交
555
    ASSERT_TRUE(stat.ok());
G
groot 已提交
556
    bool bfound = false;
Y
youny626 已提交
557
    for (auto& schema : table_schema_array) {
S
starlord 已提交
558
        if (schema.table_id_ == TABLE_NAME) {
G
groot 已提交
559 560 561 562 563 564
            bfound = true;
            break;
        }
    }
    ASSERT_TRUE(bfound);

S
starlord 已提交
565
    milvus::engine::meta::TableSchema table_info_get;
G
groot 已提交
566 567
    table_info_get.table_id_ = TABLE_NAME;
    stat = db_->DescribeTable(table_info_get);
S
starlord 已提交
568
    ASSERT_TRUE(stat.ok());
G
groot 已提交
569
    ASSERT_EQ(table_info_get.dimension_, TABLE_DIM);
Z
zhiru 已提交
570

G
groot 已提交
571
    uint64_t size;
Z
zhiru 已提交
572 573
    db_->Size(size);

G
groot 已提交
574 575 576
    int64_t nb = 10;
    std::vector<float> xb;
    BuildVectors(nb, xb);
Z
zhiru 已提交
577

G
groot 已提交
578
    int loop = INSERT_LOOP;
S
starlord 已提交
579
    for (auto i = 0; i < loop; ++i) {
S
starlord 已提交
580
        milvus::engine::IDNumbers vector_ids;
G
groot 已提交
581
        db_->InsertVectors(TABLE_NAME, "", nb, xb.data(), vector_ids);
Z
zhiru 已提交
582 583 584
        std::this_thread::sleep_for(std::chrono::microseconds(1));
    }

G
groot 已提交
585
    std::this_thread::sleep_for(std::chrono::seconds(1));
Z
zhiru 已提交
586 587 588

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

G
groot 已提交
592
TEST_F(DBTest2, DELETE_TEST) {
S
starlord 已提交
593
    milvus::engine::meta::TableSchema table_info = BuildTableSchema();
S
starlord 已提交
594
    auto stat = db_->CreateTable(table_info);
G
groot 已提交
595

S
starlord 已提交
596
    milvus::engine::meta::TableSchema table_info_get;
G
groot 已提交
597 598
    table_info_get.table_id_ = TABLE_NAME;
    stat = db_->DescribeTable(table_info_get);
S
starlord 已提交
599
    ASSERT_TRUE(stat.ok());
G
groot 已提交
600

S
starlord 已提交
601 602 603
    bool has_table = false;
    db_->HasTable(TABLE_NAME, has_table);
    ASSERT_TRUE(has_table);
G
groot 已提交
604 605 606 607

    uint64_t size;
    db_->Size(size);

608
    int64_t nb = VECTOR_COUNT;
G
groot 已提交
609 610 611
    std::vector<float> xb;
    BuildVectors(nb, xb);

S
starlord 已提交
612
    milvus::engine::IDNumbers vector_ids;
G
groot 已提交
613
    stat = db_->InsertVectors(TABLE_NAME, "", nb, xb.data(), vector_ids);
S
starlord 已提交
614
    milvus::engine::TableIndex index;
615
    stat = db_->CreateIndex(TABLE_NAME, index);
G
groot 已提交
616

S
starlord 已提交
617
    std::vector<milvus::engine::meta::DateT> dates;
G
groot 已提交
618
    stat = db_->DropTable(TABLE_NAME, dates);
G
groot 已提交
619 620
    std::this_thread::sleep_for(std::chrono::seconds(2));
    ASSERT_TRUE(stat.ok());
S
starlord 已提交
621 622 623

    db_->HasTable(TABLE_NAME, has_table);
    ASSERT_FALSE(has_table);
S
starlord 已提交
624
}
625 626

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

S
starlord 已提交
630
    milvus::engine::meta::TableSchema table_info_get;
631 632
    table_info_get.table_id_ = TABLE_NAME;
    stat = db_->DescribeTable(table_info_get);
S
starlord 已提交
633
    ASSERT_TRUE(stat.ok());
634 635 636 637 638 639 640

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

    uint64_t size;
    db_->Size(size);
641
    ASSERT_EQ(size, 0UL);
642

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

S
starlord 已提交
647
    milvus::engine::IDNumbers vector_ids;
G
groot 已提交
648
    stat = db_->InsertVectors(TABLE_NAME, "", nb, xb.data(), vector_ids);
S
starlord 已提交
649
    milvus::engine::TableIndex index;
650 651 652 653
    stat = db_->CreateIndex(TABLE_NAME, index);

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

S
starlord 已提交
655
    std::vector<milvus::engine::meta::DateT> dates;
656 657
    std::string start_value = CurrentTmDate();
    std::string end_value = CurrentTmDate(1);
658 659
    ConvertTimeRangeToDBDates(start_value, end_value, dates);

G
groot 已提交
660
    stat = db_->DropTable(TABLE_NAME, dates);
S
starlord 已提交
661
    ASSERT_TRUE(stat.ok());
662 663 664 665

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