test_db.cpp 23.9 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) {
232
    milvus::scheduler::OptimizerInst::GetInstance()->Init();
S
starlord 已提交
233 234
    std::string config_path(CONFIG_PATH);
    config_path += CONFIG_FILE;
Y
youny626 已提交
235
    milvus::server::Config& config = milvus::server::Config::GetInstance();
S
starlord 已提交
236
    milvus::Status s = config.LoadConfigFile(config_path);
237

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

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

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

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

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

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

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

S
starlord 已提交
288
    {
G
groot 已提交
289 290 291 292
        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 已提交
293
        ASSERT_TRUE(stat.ok());
294 295 296 297 298 299 300 301 302 303 304 305 306 307 308 309 310 311 312 313 314 315 316 317 318 319 320 321 322 323 324 325 326 327 328 329 330 331 332 333 334 335
        stat = db_->Query(TABLE_NAME, tags, k, 1100, 10, xq.data(), result_ids, result_distances);
        ASSERT_TRUE(stat.ok());
    }

    index.engine_type_ = (int)milvus::engine::EngineType::FAISS_IVFFLAT;
    db_->CreateIndex(TABLE_NAME, index);  // wait until build index finish

    {
        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);
        ASSERT_TRUE(stat.ok());
        stat = db_->Query(TABLE_NAME, tags, k, 1100, 10, xq.data(), result_ids, result_distances);
        ASSERT_TRUE(stat.ok());
    }

    index.engine_type_ = (int)milvus::engine::EngineType::FAISS_IVFSQ8;
    db_->CreateIndex(TABLE_NAME, index);  // wait until build index finish

    {
        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);
        ASSERT_TRUE(stat.ok());
        stat = db_->Query(TABLE_NAME, tags, k, 1100, 10, xq.data(), result_ids, result_distances);
        ASSERT_TRUE(stat.ok());
    }

#ifdef CUSTOMIZATION
    index.engine_type_ = (int)milvus::engine::EngineType::FAISS_IVFSQ8H;
    db_->CreateIndex(TABLE_NAME, index);  // wait until build index finish

    {
        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);
        ASSERT_TRUE(stat.ok());
        stat = db_->Query(TABLE_NAME, tags, k, 1100, 10, xq.data(), result_ids, result_distances);
        ASSERT_TRUE(stat.ok());
S
starlord 已提交
336
    }
337
#endif
S
starlord 已提交
338

Y
youny626 已提交
339
    {  // search by specify index file
S
starlord 已提交
340
        milvus::engine::meta::DatesT dates;
341
        std::vector<std::string> file_ids = {"1", "2", "3", "4", "5", "6"};
G
groot 已提交
342 343 344
        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 已提交
345
        ASSERT_TRUE(stat.ok());
X
xj.lin 已提交
346
    }
X
xj.lin 已提交
347

Z
update  
zhiru 已提交
348
#ifdef CUSTOMIZATION
Y
youny626 已提交
349
    // test FAISS_IVFSQ8H optimizer
350
    index.engine_type_ = (int)milvus::engine::EngineType::FAISS_IVFSQ8H;
Y
youny626 已提交
351
    db_->CreateIndex(TABLE_NAME, index);  // wait until build index finish
352 353 354
    std::vector<std::string> partition_tag;
    milvus::engine::ResultIds result_ids;
    milvus::engine::ResultDistances result_dists;
355 356

    {
357 358 359
        result_ids.clear();
        result_dists.clear();
        stat = db_->Query(TABLE_NAME, partition_tag, k, nq, 10, xq.data(), result_ids, result_dists);
360 361 362
        ASSERT_TRUE(stat.ok());
    }

Y
Yu Kun 已提交
363
    {
364 365 366
        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 已提交
367 368 369
        ASSERT_TRUE(stat.ok());
    }

Y
youny626 已提交
370
    {  // search by specify index file
371 372
        milvus::engine::meta::DatesT dates;
        std::vector<std::string> file_ids = {"1", "2", "3", "4", "5", "6"};
373 374 375
        result_ids.clear();
        result_dists.clear();
        stat = db_->QueryByFileID(TABLE_NAME, file_ids, k, nq, 10, xq.data(), dates, result_ids, result_dists);
376 377
        ASSERT_TRUE(stat.ok());
    }
Y
Yu Kun 已提交
378

X
xiaojun.lin 已提交
379
#endif
S
starlord 已提交
380
}
Y
c  
yu yunfeng 已提交
381

Y
Yu Kun 已提交
382
TEST_F(DBTest, PRELOADTABLE_TEST) {
S
starlord 已提交
383
    milvus::engine::meta::TableSchema table_info = BuildTableSchema();
S
starlord 已提交
384
    auto stat = db_->CreateTable(table_info);
Y
Yu Kun 已提交
385

S
starlord 已提交
386
    milvus::engine::meta::TableSchema table_info_get;
Y
Yu Kun 已提交
387 388
    table_info_get.table_id_ = TABLE_NAME;
    stat = db_->DescribeTable(table_info_get);
S
starlord 已提交
389
    ASSERT_TRUE(stat.ok());
Y
Yu Kun 已提交
390 391
    ASSERT_EQ(table_info_get.dimension_, TABLE_DIM);

392
    int64_t nb = VECTOR_COUNT;
Y
Yu Kun 已提交
393 394 395
    std::vector<float> xb;
    BuildVectors(nb, xb);

Y
Yu Kun 已提交
396
    int loop = 5;
S
starlord 已提交
397
    for (auto i = 0; i < loop; ++i) {
S
starlord 已提交
398
        milvus::engine::IDNumbers vector_ids;
G
groot 已提交
399
        db_->InsertVectors(TABLE_NAME, "", nb, xb.data(), vector_ids);
400
        ASSERT_EQ(vector_ids.size(), nb);
Y
Yu Kun 已提交
401
    }
402

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

S
starlord 已提交
407
    int64_t prev_cache_usage = milvus::cache::CpuCacheMgr::GetInstance()->CacheUsage();
Y
Yu Kun 已提交
408
    stat = db_->PreloadTable(TABLE_NAME);
S
starlord 已提交
409
    ASSERT_TRUE(stat.ok());
S
starlord 已提交
410
    int64_t cur_cache_usage = milvus::cache::CpuCacheMgr::GetInstance()->CacheUsage();
Y
Yu Kun 已提交
411 412 413
    ASSERT_TRUE(prev_cache_usage < cur_cache_usage);
}

S
starlord 已提交
414 415 416
TEST_F(DBTest, SHUTDOWN_TEST) {
    db_->Stop();

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 424 425 426 427
    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 已提交
428
    milvus::engine::IDNumbers ids;
G
groot 已提交
429
    stat = db_->InsertVectors(table_info.table_id_, "", 0, nullptr, ids);
S
starlord 已提交
430 431 432 433 434 435 436 437 438
    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 已提交
439
    milvus::engine::TableIndex index;
S
starlord 已提交
440 441 442 443 444 445
    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 已提交
446
    std::vector<std::string> tags;
S
starlord 已提交
447
    milvus::engine::meta::DatesT dates;
G
groot 已提交
448 449 450
    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 已提交
451 452
    ASSERT_FALSE(stat.ok());
    std::vector<std::string> file_ids;
G
groot 已提交
453
    stat = db_->QueryByFileID(table_info.table_id_, file_ids, 1, 1, 1, nullptr, dates, result_ids, result_distances);
S
starlord 已提交
454 455
    ASSERT_FALSE(stat.ok());

G
groot 已提交
456
    stat = db_->DropTable(table_info.table_id_, dates);
S
starlord 已提交
457 458 459 460
    ASSERT_FALSE(stat.ok());
}

TEST_F(DBTest, INDEX_TEST) {
S
starlord 已提交
461
    milvus::engine::meta::TableSchema table_info = BuildTableSchema();
S
starlord 已提交
462
    auto stat = db_->CreateTable(table_info);
S
starlord 已提交
463 464 465 466 467

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

S
starlord 已提交
468
    milvus::engine::IDNumbers vector_ids;
G
groot 已提交
469
    db_->InsertVectors(TABLE_NAME, "", nb, xb.data(), vector_ids);
S
starlord 已提交
470 471
    ASSERT_EQ(vector_ids.size(), nb);

S
starlord 已提交
472
    milvus::engine::TableIndex index;
Y
youny626 已提交
473 474
    index.engine_type_ = (int)milvus::engine::EngineType::FAISS_IVFSQ8;
    index.metric_type_ = (int)milvus::engine::MetricType::IP;
S
starlord 已提交
475 476 477
    stat = db_->CreateIndex(table_info.table_id_, index);
    ASSERT_TRUE(stat.ok());

Y
youny626 已提交
478
    index.engine_type_ = (int)milvus::engine::EngineType::FAISS_IVFFLAT;
Y
Yu Kun 已提交
479 480 481 482 483 484 485 486 487
    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 已提交
488
    milvus::engine::TableIndex index_out;
S
starlord 已提交
489
    stat = db_->DescribeIndex(table_info.table_id_, index_out);
S
starlord 已提交
490 491 492
    ASSERT_TRUE(stat.ok());
    ASSERT_EQ(index.engine_type_, index_out.engine_type_);
    ASSERT_EQ(index.nlist_, index_out.nlist_);
S
starlord 已提交
493
    ASSERT_EQ(table_info.metric_type_, index_out.metric_type_);
S
starlord 已提交
494 495 496 497 498

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

G
groot 已提交
499 500 501 502 503 504 505 506 507 508 509 510 511 512 513
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());

514 515 516 517 518 519 520
        // 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 已提交
521 522 523 524 525 526 527 528 529 530 531 532 533 534 535 536 537 538

        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());

G
groot 已提交
539 540
    std::vector<milvus::engine::meta::TableSchema> partition_schema_array;
    stat = db_->ShowPartitions(table_name, partition_schema_array);
G
groot 已提交
541
    ASSERT_TRUE(stat.ok());
G
groot 已提交
542
    ASSERT_EQ(partition_schema_array.size(), PARTITION_COUNT);
G
groot 已提交
543
    for (int64_t i = 0; i < PARTITION_COUNT; i++) {
G
groot 已提交
544
        ASSERT_EQ(partition_schema_array[i].table_id_, table_name + "_" + std::to_string(i));
G
groot 已提交
545 546 547 548 549 550 551 552 553 554 555 556 557 558 559 560 561 562 563 564 565 566 567 568 569 570 571 572 573 574 575 576 577 578 579 580 581 582 583 584 585 586 587 588 589 590 591 592 593 594 595 596 597 598 599 600 601 602 603 604 605
    }

    { // 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 已提交
606
TEST_F(DBTest2, ARHIVE_DISK_CHECK) {
S
starlord 已提交
607
    milvus::engine::meta::TableSchema table_info = BuildTableSchema();
S
starlord 已提交
608
    auto stat = db_->CreateTable(table_info);
Z
zhiru 已提交
609

S
starlord 已提交
610
    std::vector<milvus::engine::meta::TableSchema> table_schema_array;
G
groot 已提交
611
    stat = db_->AllTables(table_schema_array);
S
starlord 已提交
612
    ASSERT_TRUE(stat.ok());
G
groot 已提交
613
    bool bfound = false;
Y
youny626 已提交
614
    for (auto& schema : table_schema_array) {
S
starlord 已提交
615
        if (schema.table_id_ == TABLE_NAME) {
G
groot 已提交
616 617 618 619 620 621
            bfound = true;
            break;
        }
    }
    ASSERT_TRUE(bfound);

S
starlord 已提交
622
    milvus::engine::meta::TableSchema table_info_get;
G
groot 已提交
623 624
    table_info_get.table_id_ = TABLE_NAME;
    stat = db_->DescribeTable(table_info_get);
S
starlord 已提交
625
    ASSERT_TRUE(stat.ok());
G
groot 已提交
626
    ASSERT_EQ(table_info_get.dimension_, TABLE_DIM);
Z
zhiru 已提交
627

G
groot 已提交
628
    uint64_t size;
Z
zhiru 已提交
629 630
    db_->Size(size);

G
groot 已提交
631 632 633
    int64_t nb = 10;
    std::vector<float> xb;
    BuildVectors(nb, xb);
Z
zhiru 已提交
634

G
groot 已提交
635
    int loop = INSERT_LOOP;
S
starlord 已提交
636
    for (auto i = 0; i < loop; ++i) {
S
starlord 已提交
637
        milvus::engine::IDNumbers vector_ids;
G
groot 已提交
638
        db_->InsertVectors(TABLE_NAME, "", nb, xb.data(), vector_ids);
Z
zhiru 已提交
639 640 641
        std::this_thread::sleep_for(std::chrono::microseconds(1));
    }

G
groot 已提交
642
    std::this_thread::sleep_for(std::chrono::seconds(1));
Z
zhiru 已提交
643 644 645

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

G
groot 已提交
649
TEST_F(DBTest2, DELETE_TEST) {
S
starlord 已提交
650
    milvus::engine::meta::TableSchema table_info = BuildTableSchema();
S
starlord 已提交
651
    auto stat = db_->CreateTable(table_info);
G
groot 已提交
652

S
starlord 已提交
653
    milvus::engine::meta::TableSchema table_info_get;
G
groot 已提交
654 655
    table_info_get.table_id_ = TABLE_NAME;
    stat = db_->DescribeTable(table_info_get);
S
starlord 已提交
656
    ASSERT_TRUE(stat.ok());
G
groot 已提交
657

S
starlord 已提交
658 659 660
    bool has_table = false;
    db_->HasTable(TABLE_NAME, has_table);
    ASSERT_TRUE(has_table);
G
groot 已提交
661 662 663 664

    uint64_t size;
    db_->Size(size);

665
    int64_t nb = VECTOR_COUNT;
G
groot 已提交
666 667 668
    std::vector<float> xb;
    BuildVectors(nb, xb);

S
starlord 已提交
669
    milvus::engine::IDNumbers vector_ids;
G
groot 已提交
670
    stat = db_->InsertVectors(TABLE_NAME, "", nb, xb.data(), vector_ids);
S
starlord 已提交
671
    milvus::engine::TableIndex index;
672
    stat = db_->CreateIndex(TABLE_NAME, index);
G
groot 已提交
673

S
starlord 已提交
674
    std::vector<milvus::engine::meta::DateT> dates;
G
groot 已提交
675
    stat = db_->DropTable(TABLE_NAME, dates);
G
groot 已提交
676 677
    std::this_thread::sleep_for(std::chrono::seconds(2));
    ASSERT_TRUE(stat.ok());
S
starlord 已提交
678 679 680

    db_->HasTable(TABLE_NAME, has_table);
    ASSERT_FALSE(has_table);
S
starlord 已提交
681
}
682 683

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

S
starlord 已提交
687
    milvus::engine::meta::TableSchema table_info_get;
688 689
    table_info_get.table_id_ = TABLE_NAME;
    stat = db_->DescribeTable(table_info_get);
S
starlord 已提交
690
    ASSERT_TRUE(stat.ok());
691 692 693 694 695 696 697

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

    uint64_t size;
    db_->Size(size);
698
    ASSERT_EQ(size, 0UL);
699

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

S
starlord 已提交
704
    milvus::engine::IDNumbers vector_ids;
G
groot 已提交
705
    stat = db_->InsertVectors(TABLE_NAME, "", nb, xb.data(), vector_ids);
S
starlord 已提交
706
    milvus::engine::TableIndex index;
707 708 709 710
    stat = db_->CreateIndex(TABLE_NAME, index);

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

S
starlord 已提交
712
    std::vector<milvus::engine::meta::DateT> dates;
G
groot 已提交
713 714
    std::string start_value = CurrentTmDate(-5);
    std::string end_value = CurrentTmDate(5);
715 716
    ConvertTimeRangeToDBDates(start_value, end_value, dates);

G
groot 已提交
717
    stat = db_->DropTable(TABLE_NAME, dates);
S
starlord 已提交
718
    ASSERT_TRUE(stat.ok());
719 720 721 722

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