DBImpl.cpp 70.6 KB
Newer Older
1
// Copyright (C) 2019-2020 Zilliz. All rights reserved.
J
jinhai 已提交
2
//
3 4
// Licensed under the Apache License, Version 2.0 (the "License"); you may not use this file except in compliance
// with the License. You may obtain a copy of the License at
J
jinhai 已提交
5
//
6 7 8 9 10
// 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.
J
jinhai 已提交
11

S
starlord 已提交
12
#include "db/DBImpl.h"
Z
Zhiru Zhu 已提交
13 14

#include <assert.h>
15
#include <fiu-local.h>
Z
Zhiru Zhu 已提交
16 17 18 19 20

#include <algorithm>
#include <boost/filesystem.hpp>
#include <chrono>
#include <cstring>
21
#include <functional>
Z
Zhiru Zhu 已提交
22
#include <iostream>
23
#include <limits>
Z
Zhiru Zhu 已提交
24 25 26 27
#include <set>
#include <thread>
#include <utility>

S
starlord 已提交
28
#include "Utils.h"
S
starlord 已提交
29 30
#include "cache/CpuCacheMgr.h"
#include "cache/GpuCacheMgr.h"
31
#include "db/IDGenerator.h"
S
starlord 已提交
32
#include "engine/EngineFactory.h"
S
starlord 已提交
33
#include "insert/MemMenagerFactory.h"
S
starlord 已提交
34
#include "meta/MetaConsts.h"
S
starlord 已提交
35 36
#include "meta/MetaFactory.h"
#include "meta/SqliteMetaImpl.h"
G
groot 已提交
37
#include "metrics/Metrics.h"
S
starlord 已提交
38
#include "scheduler/SchedInst.h"
Y
Yu Kun 已提交
39
#include "scheduler/job/BuildIndexJob.h"
S
starlord 已提交
40 41
#include "scheduler/job/DeleteJob.h"
#include "scheduler/job/SearchJob.h"
42 43 44
#include "segment/SegmentReader.h"
#include "segment/SegmentWriter.h"
#include "utils/Exception.h"
S
starlord 已提交
45
#include "utils/Log.h"
G
groot 已提交
46
#include "utils/StringHelpFunctions.h"
S
starlord 已提交
47
#include "utils/TimeRecorder.h"
48 49
#include "utils/ValidationUtil.h"
#include "wal/WalDefinations.h"
X
Xu Peng 已提交
50

J
jinhai 已提交
51
namespace milvus {
X
Xu Peng 已提交
52
namespace engine {
X
Xu Peng 已提交
53

G
groot 已提交
54 55
namespace {

J
jinhai 已提交
56 57 58
constexpr uint64_t METRIC_ACTION_INTERVAL = 1;
constexpr uint64_t COMPACT_ACTION_INTERVAL = 1;
constexpr uint64_t INDEX_ACTION_INTERVAL = 1;
G
groot 已提交
59

G
groot 已提交
60
static const Status SHUTDOWN_ERROR = Status(DB_ERROR, "Milvus server is shutdown!");
G
groot 已提交
61

S
starlord 已提交
62
}  // namespace
G
groot 已提交
63

Y
Yu Kun 已提交
64
DBImpl::DBImpl(const DBOptions& options)
65
    : options_(options), initialized_(false), merge_thread_pool_(1, 1), index_thread_pool_(1, 1) {
S
starlord 已提交
66
    meta_ptr_ = MetaFactory::Build(options.meta_, options.mode_);
Z
zhiru 已提交
67
    mem_mgr_ = MemManagerFactory::Build(meta_ptr_, options_);
68 69 70 71 72 73 74 75 76 77

    if (options_.wal_enable_) {
        wal::MXLogConfiguration mxlog_config;
        mxlog_config.recovery_error_ignore = options_.recovery_error_ignore_;
        // 2 buffers in the WAL
        mxlog_config.buffer_size = options_.buffer_size_ / 2;
        mxlog_config.mxlog_path = options_.mxlog_path_;
        wal_mgr_ = std::make_shared<wal::WalManager>(mxlog_config);
    }

78 79 80
    SetIdentity("DBImpl");
    AddCacheInsertDataListener();

S
starlord 已提交
81 82 83 84
    Start();
}

DBImpl::~DBImpl() {
85
    RemoveCacheInsertDataListener();
S
starlord 已提交
86 87 88
    Stop();
}

S
starlord 已提交
89
///////////////////////////////////////////////////////////////////////////////////////////////////////////////////
S
starlord 已提交
90
// external api
S
starlord 已提交
91
///////////////////////////////////////////////////////////////////////////////////////////////////////////////////
S
starlord 已提交
92 93
Status
DBImpl::Start() {
94
    if (initialized_.load(std::memory_order_acquire)) {
S
starlord 已提交
95 96 97
        return Status::OK();
    }

S
Shouyu Luo 已提交
98
    // ENGINE_LOG_TRACE << "DB service start";
99
    initialized_.store(true, std::memory_order_release);
S
starlord 已提交
100

101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136
    // wal
    if (options_.wal_enable_) {
        auto error_code = DB_ERROR;
        if (wal_mgr_ != nullptr) {
            error_code = wal_mgr_->Init(meta_ptr_);
        }
        if (error_code != WAL_SUCCESS) {
            throw Exception(error_code, "Wal init error!");
        }

        // recovery
        while (1) {
            wal::MXLogRecord record;
            auto error_code = wal_mgr_->GetNextRecovery(record);
            if (error_code != WAL_SUCCESS) {
                throw Exception(error_code, "Wal recovery error!");
            }
            if (record.type == wal::MXLogType::None) {
                break;
            }

            ExecWalRecord(record);
        }

        // for distribute version, some nodes are read only
        if (options_.mode_ != DBOptions::MODE::CLUSTER_READONLY) {
            // background thread
            bg_wal_thread_ = std::thread(&DBImpl::BackgroundWalTask, this);
        }

    } else {
        // for distribute version, some nodes are read only
        if (options_.mode_ != DBOptions::MODE::CLUSTER_READONLY) {
            // ENGINE_LOG_TRACE << "StartTimerTasks";
            bg_timer_thread_ = std::thread(&DBImpl::BackgroundTimerTask, this);
        }
Z
update  
zhiru 已提交
137
    }
S
starlord 已提交
138

S
starlord 已提交
139 140 141
    return Status::OK();
}

S
starlord 已提交
142 143
Status
DBImpl::Stop() {
144
    if (!initialized_.load(std::memory_order_acquire)) {
S
starlord 已提交
145 146
        return Status::OK();
    }
147

148
    initialized_.store(false, std::memory_order_release);
S
starlord 已提交
149

150 151 152
    if (options_.mode_ != DBOptions::MODE::CLUSTER_READONLY) {
        if (options_.wal_enable_) {
            // wait flush merge/buildindex finish
153
            bg_task_swn_.Notify();
154
            bg_wal_thread_.join();
S
starlord 已提交
155

156 157 158 159 160 161 162
        } else {
            // flush all
            wal::MXLogRecord record;
            record.type = wal::MXLogType::Flush;
            ExecWalRecord(record);

            // wait merge/buildindex finish
163
            bg_task_swn_.Notify();
164 165
            bg_timer_thread_.join();
        }
S
starlord 已提交
166

167
        meta_ptr_->CleanUpShadowFiles();
S
starlord 已提交
168 169
    }

S
Shouyu Luo 已提交
170
    // ENGINE_LOG_TRACE << "DB service stop";
S
starlord 已提交
171
    return Status::OK();
X
Xu Peng 已提交
172 173
}

S
starlord 已提交
174 175
Status
DBImpl::DropAll() {
S
starlord 已提交
176 177 178
    return meta_ptr_->DropAll();
}

S
starlord 已提交
179
Status
Y
Yu Kun 已提交
180
DBImpl::CreateTable(meta::TableSchema& table_schema) {
181
    if (!initialized_.load(std::memory_order_acquire)) {
G
groot 已提交
182
        return SHUTDOWN_ERROR;
S
starlord 已提交
183 184
    }

185
    meta::TableSchema temp_schema = table_schema;
S
starlord 已提交
186
    temp_schema.index_file_size_ *= ONE_MB;  // store as MB
187 188 189 190
    if (options_.wal_enable_) {
        temp_schema.flush_lsn_ = wal_mgr_->CreateTable(table_schema.table_id_);
    }

191
    return meta_ptr_->CreateTable(temp_schema);
192 193
}

S
starlord 已提交
194
Status
195
DBImpl::DropTable(const std::string& table_id) {
196
    if (!initialized_.load(std::memory_order_acquire)) {
G
groot 已提交
197
        return SHUTDOWN_ERROR;
S
starlord 已提交
198 199
    }

200 201 202 203 204
    if (options_.wal_enable_) {
        wal_mgr_->DropTable(table_id);
    }

    return DropTableRecursively(table_id);
G
groot 已提交
205 206
}

S
starlord 已提交
207
Status
Y
Yu Kun 已提交
208
DBImpl::DescribeTable(meta::TableSchema& table_schema) {
209
    if (!initialized_.load(std::memory_order_acquire)) {
G
groot 已提交
210
        return SHUTDOWN_ERROR;
S
starlord 已提交
211 212
    }

S
starlord 已提交
213
    auto stat = meta_ptr_->DescribeTable(table_schema);
S
starlord 已提交
214
    table_schema.index_file_size_ /= ONE_MB;  // return as MB
S
starlord 已提交
215
    return stat;
216 217
}

S
starlord 已提交
218
Status
Y
Yu Kun 已提交
219
DBImpl::HasTable(const std::string& table_id, bool& has_or_not) {
220
    if (!initialized_.load(std::memory_order_acquire)) {
G
groot 已提交
221
        return SHUTDOWN_ERROR;
S
starlord 已提交
222 223
    }

G
groot 已提交
224
    return meta_ptr_->HasTable(table_id, has_or_not);
225 226
}

227 228 229 230 231 232 233 234 235 236 237 238 239 240 241 242 243 244 245 246 247 248 249
Status
DBImpl::HasNativeTable(const std::string& table_id, bool& has_or_not_) {
    if (!initialized_.load(std::memory_order_acquire)) {
        return SHUTDOWN_ERROR;
    }

    engine::meta::TableSchema table_schema;
    table_schema.table_id_ = table_id;
    auto status = DescribeTable(table_schema);
    if (!status.ok()) {
        has_or_not_ = false;
        return status;
    } else {
        if (!table_schema.owner_table_.empty()) {
            has_or_not_ = false;
            return Status(DB_NOT_FOUND, "");
        }

        has_or_not_ = true;
        return Status::OK();
    }
}

S
starlord 已提交
250
Status
Y
Yu Kun 已提交
251
DBImpl::AllTables(std::vector<meta::TableSchema>& table_schema_array) {
252
    if (!initialized_.load(std::memory_order_acquire)) {
G
groot 已提交
253
        return SHUTDOWN_ERROR;
S
starlord 已提交
254 255
    }

256 257 258 259 260 261 262 263 264 265 266 267
    std::vector<meta::TableSchema> all_tables;
    auto status = meta_ptr_->AllTables(all_tables);

    // only return real tables, dont return partition tables
    table_schema_array.clear();
    for (auto& schema : all_tables) {
        if (schema.owner_table_.empty()) {
            table_schema_array.push_back(schema);
        }
    }

    return status;
G
groot 已提交
268 269
}

270 271 272 273 274 275 276 277 278 279 280 281 282 283 284 285 286 287 288 289 290 291 292 293 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
Status
DBImpl::GetTableInfo(const std::string& table_id, TableInfo& table_info) {
    if (!initialized_.load(std::memory_order_acquire)) {
        return SHUTDOWN_ERROR;
    }

    // step1: get all partition ids
    std::vector<std::pair<std::string, std::string>> name2tag = {{table_id, milvus::engine::DEFAULT_PARTITON_TAG}};
    std::vector<meta::TableSchema> partition_array;
    auto status = meta_ptr_->ShowPartitions(table_id, partition_array);
    for (auto& schema : partition_array) {
        name2tag.push_back(std::make_pair(schema.table_id_, schema.partition_tag_));
    }

    // step2: get native table info
    std::vector<int> file_types{meta::TableFileSchema::FILE_TYPE::RAW, meta::TableFileSchema::FILE_TYPE::TO_INDEX,
                                meta::TableFileSchema::FILE_TYPE::INDEX};

    static std::map<int32_t, std::string> index_type_name = {
        {(int32_t)engine::EngineType::FAISS_IDMAP, "IDMAP"},
        {(int32_t)engine::EngineType::FAISS_IVFFLAT, "IVFFLAT"},
        {(int32_t)engine::EngineType::FAISS_IVFSQ8, "IVFSQ8"},
        {(int32_t)engine::EngineType::NSG_MIX, "NSG"},
        {(int32_t)engine::EngineType::FAISS_IVFSQ8H, "IVFSQ8H"},
        {(int32_t)engine::EngineType::FAISS_PQ, "PQ"},
        {(int32_t)engine::EngineType::SPTAG_KDT, "KDT"},
        {(int32_t)engine::EngineType::SPTAG_BKT, "BKT"},
        {(int32_t)engine::EngineType::FAISS_BIN_IDMAP, "IDMAP"},
        {(int32_t)engine::EngineType::FAISS_BIN_IVFFLAT, "IVFFLAT"},
    };

    for (auto& name_tag : name2tag) {
        meta::TableFilesSchema table_files;
        status = meta_ptr_->FilesByType(name_tag.first, file_types, table_files);
        if (!status.ok()) {
            std::string err_msg = "Failed to get table info: " + status.ToString();
            ENGINE_LOG_ERROR << err_msg;
            return Status(DB_ERROR, err_msg);
        }

        std::vector<SegmentStat> segments_stat;
        for (auto& file : table_files) {
            SegmentStat seg_stat;
            seg_stat.name_ = file.segment_id_;
            seg_stat.row_count_ = (int64_t)file.row_count_;
            seg_stat.index_name_ = index_type_name[file.engine_type_];
            seg_stat.data_size_ = (int64_t)file.file_size_;
            segments_stat.emplace_back(seg_stat);
        }

        PartitionStat partition_stat;
        if (name_tag.first == table_id) {
            partition_stat.tag_ = milvus::engine::DEFAULT_PARTITON_TAG;
        } else {
            partition_stat.tag_ = name_tag.second;
        }

        partition_stat.segments_stat_.swap(segments_stat);
        table_info.partitions_stat_.emplace_back(partition_stat);
    }

    return Status::OK();
}

S
starlord 已提交
334
Status
Y
Yu Kun 已提交
335
DBImpl::PreloadTable(const std::string& table_id) {
336
    if (!initialized_.load(std::memory_order_acquire)) {
G
groot 已提交
337
        return SHUTDOWN_ERROR;
S
starlord 已提交
338 339
    }

340
    // step 1: get all table files from parent table
341
    std::vector<size_t> ids;
G
groot 已提交
342
    meta::TableFilesSchema files_array;
343
    auto status = GetFilesToSearch(table_id, ids, files_array);
Y
Yu Kun 已提交
344 345 346
    if (!status.ok()) {
        return status;
    }
Y
Yu Kun 已提交
347

348
    // step 2: get files from partition tables
G
groot 已提交
349 350 351
    std::vector<meta::TableSchema> partition_array;
    status = meta_ptr_->ShowPartitions(table_id, partition_array);
    for (auto& schema : partition_array) {
352
        status = GetFilesToSearch(schema.table_id_, ids, files_array);
G
groot 已提交
353 354
    }

Y
Yu Kun 已提交
355 356
    int64_t size = 0;
    int64_t cache_total = cache::CpuCacheMgr::GetInstance()->CacheCapacity();
Y
Yu Kun 已提交
357 358
    int64_t cache_usage = cache::CpuCacheMgr::GetInstance()->CacheUsage();
    int64_t available_size = cache_total - cache_usage;
Y
Yu Kun 已提交
359

360 361 362 363
    // step 3: load file one by one
    ENGINE_LOG_DEBUG << "Begin pre-load table:" + table_id + ", totally " << files_array.size()
                     << " files need to be pre-loaded";
    TimeRecorderAuto rc("Pre-load table:" + table_id);
G
groot 已提交
364
    for (auto& file : files_array) {
365 366 367 368
        EngineType engine_type;
        if (file.file_type_ == meta::TableFileSchema::FILE_TYPE::RAW ||
            file.file_type_ == meta::TableFileSchema::FILE_TYPE::TO_INDEX ||
            file.file_type_ == meta::TableFileSchema::FILE_TYPE::BACKUP) {
369 370
            engine_type =
                utils::IsBinaryMetricType(file.metric_type_) ? EngineType::FAISS_BIN_IDMAP : EngineType::FAISS_IDMAP;
371 372 373
        } else {
            engine_type = (EngineType)file.engine_type_;
        }
374 375 376 377

        auto json = milvus::json::parse(file.index_params_);
        ExecutionEnginePtr engine =
            EngineFactory::Build(file.dimension_, file.location_, engine_type, (MetricType)file.metric_type_, json);
S
shengjh 已提交
378
        fiu_do_on("DBImpl.PreloadTable.null_engine", engine = nullptr);
G
groot 已提交
379 380 381 382
        if (engine == nullptr) {
            ENGINE_LOG_ERROR << "Invalid engine type";
            return Status(DB_ERROR, "Invalid engine type");
        }
Y
Yu Kun 已提交
383

S
shengjh 已提交
384
        fiu_do_on("DBImpl.PreloadTable.exceed_cache", size = available_size + 1);
385 386 387 388 389 390 391 392 393 394 395

        try {
            fiu_do_on("DBImpl.PreloadTable.engine_throw_exception", throw std::exception());
            std::string msg = "Pre-loaded file: " + file.file_id_ + " size: " + std::to_string(file.file_size_);
            TimeRecorderAuto rc_1(msg);
            engine->Load(true);

            size += engine->Size();
            if (size > available_size) {
                ENGINE_LOG_DEBUG << "Pre-load cancelled since cache is almost full";
                return Status(SERVER_CACHE_FULL, "Cache is full");
Y
Yu Kun 已提交
396
            }
397 398 399 400
        } catch (std::exception& ex) {
            std::string msg = "Pre-load table encounter exception: " + std::string(ex.what());
            ENGINE_LOG_ERROR << msg;
            return Status(DB_ERROR, msg);
Y
Yu Kun 已提交
401 402
        }
    }
G
groot 已提交
403

Y
Yu Kun 已提交
404
    return Status::OK();
Y
Yu Kun 已提交
405 406
}

S
starlord 已提交
407
Status
Y
Yu Kun 已提交
408
DBImpl::UpdateTableFlag(const std::string& table_id, int64_t flag) {
409
    if (!initialized_.load(std::memory_order_acquire)) {
G
groot 已提交
410
        return SHUTDOWN_ERROR;
S
starlord 已提交
411 412
    }

S
starlord 已提交
413 414 415
    return meta_ptr_->UpdateTableFlag(table_id, flag);
}

S
starlord 已提交
416
Status
Y
Yu Kun 已提交
417
DBImpl::GetTableRowCount(const std::string& table_id, uint64_t& row_count) {
418
    if (!initialized_.load(std::memory_order_acquire)) {
G
groot 已提交
419 420 421 422 423 424 425 426 427
        return SHUTDOWN_ERROR;
    }

    return GetTableRowCountRecursively(table_id, row_count);
}

Status
DBImpl::CreatePartition(const std::string& table_id, const std::string& partition_name,
                        const std::string& partition_tag) {
428
    if (!initialized_.load(std::memory_order_acquire)) {
G
groot 已提交
429 430 431
        return SHUTDOWN_ERROR;
    }

432 433 434
    uint64_t lsn = 0;
    meta_ptr_->GetTableFlushLSN(table_id, lsn);
    return meta_ptr_->CreatePartition(table_id, partition_name, partition_tag, lsn);
G
groot 已提交
435 436 437 438
}

Status
DBImpl::DropPartition(const std::string& partition_name) {
439
    if (!initialized_.load(std::memory_order_acquire)) {
G
groot 已提交
440
        return SHUTDOWN_ERROR;
S
starlord 已提交
441 442
    }

443 444 445 446 447 448
    mem_mgr_->EraseMemVector(partition_name);                // not allow insert
    auto status = meta_ptr_->DropPartition(partition_name);  // soft delete table
    if (!status.ok()) {
        ENGINE_LOG_ERROR << status.message();
        return status;
    }
G
groot 已提交
449 450 451 452 453 454 455 456

    // scheduler will determine when to delete table files
    auto nres = scheduler::ResMgrInst::GetInstance()->GetNumOfComputeResource();
    scheduler::DeleteJobPtr job = std::make_shared<scheduler::DeleteJob>(partition_name, meta_ptr_, nres);
    scheduler::JobMgrInst::GetInstance()->Put(job);
    job->WaitAndDelete();

    return Status::OK();
G
groot 已提交
457 458
}

S
starlord 已提交
459
Status
G
groot 已提交
460
DBImpl::DropPartitionByTag(const std::string& table_id, const std::string& partition_tag) {
461
    if (!initialized_.load(std::memory_order_acquire)) {
G
groot 已提交
462 463 464 465 466
        return SHUTDOWN_ERROR;
    }

    std::string partition_name;
    auto status = meta_ptr_->GetPartitionName(table_id, partition_tag, partition_name);
467 468 469 470 471
    if (!status.ok()) {
        ENGINE_LOG_ERROR << status.message();
        return status;
    }

G
groot 已提交
472 473 474 475
    return DropPartition(partition_name);
}

Status
G
groot 已提交
476
DBImpl::ShowPartitions(const std::string& table_id, std::vector<meta::TableSchema>& partition_schema_array) {
477
    if (!initialized_.load(std::memory_order_acquire)) {
G
groot 已提交
478 479 480
        return SHUTDOWN_ERROR;
    }

G
groot 已提交
481
    return meta_ptr_->ShowPartitions(table_id, partition_schema_array);
G
groot 已提交
482 483 484
}

Status
G
groot 已提交
485
DBImpl::InsertVectors(const std::string& table_id, const std::string& partition_tag, VectorsData& vectors) {
S
starlord 已提交
486
    //    ENGINE_LOG_DEBUG << "Insert " << n << " vectors to cache";
487
    if (!initialized_.load(std::memory_order_acquire)) {
G
groot 已提交
488
        return SHUTDOWN_ERROR;
S
starlord 已提交
489
    }
Y
yu yunfeng 已提交
490

491 492 493
    // insert vectors into target table
    // (zhiru): generate ids
    if (vectors.id_array_.empty()) {
J
Jin Hai 已提交
494 495 496 497 498
        SafeIDGenerator& id_generator = SafeIDGenerator::GetInstance();
        Status status = id_generator.GetNextIDNumbers(vectors.vector_count_, vectors.id_array_);
        if (!status.ok()) {
            return status;
        }
499 500
    }

501
    Status status;
502 503 504
    if (options_.wal_enable_) {
        std::string target_table_name;
        status = GetPartitionByTag(table_id, partition_tag, target_table_name);
G
groot 已提交
505 506 507
        if (!status.ok()) {
            return status;
        }
508 509 510 511 512 513

        if (!vectors.float_data_.empty()) {
            wal_mgr_->Insert(table_id, partition_tag, vectors.id_array_, vectors.float_data_);
        } else if (!vectors.binary_data_.empty()) {
            wal_mgr_->Insert(table_id, partition_tag, vectors.id_array_, vectors.binary_data_);
        }
514
        bg_task_swn_.Notify();
515 516 517 518 519 520 521 522 523 524 525 526 527 528 529 530 531 532 533 534 535

    } else {
        wal::MXLogRecord record;
        record.lsn = 0;  // need to get from meta ?
        record.table_id = table_id;
        record.partition_tag = partition_tag;
        record.ids = vectors.id_array_.data();
        record.length = vectors.vector_count_;
        if (vectors.binary_data_.empty()) {
            record.type = wal::MXLogType::InsertVector;
            record.data = vectors.float_data_.data();
            record.data_size = vectors.float_data_.size() * sizeof(float);
        } else {
            record.type = wal::MXLogType::InsertBinary;
            record.ids = vectors.id_array_.data();
            record.length = vectors.vector_count_;
            record.data = vectors.binary_data_.data();
            record.data_size = vectors.binary_data_.size() * sizeof(uint8_t);
        }

        status = ExecWalRecord(record);
G
groot 已提交
536 537
    }

538 539 540 541 542 543 544 545 546 547 548 549 550 551 552 553 554 555 556
    return status;
}

Status
DBImpl::DeleteVector(const std::string& table_id, IDNumber vector_id) {
    IDNumbers ids;
    ids.push_back(vector_id);
    return DeleteVectors(table_id, ids);
}

Status
DBImpl::DeleteVectors(const std::string& table_id, IDNumbers vector_ids) {
    if (!initialized_.load(std::memory_order_acquire)) {
        return SHUTDOWN_ERROR;
    }

    Status status;
    if (options_.wal_enable_) {
        wal_mgr_->DeleteById(table_id, vector_ids);
557
        bg_task_swn_.Notify();
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

    } else {
        wal::MXLogRecord record;
        record.lsn = 0;  // need to get from meta ?
        record.type = wal::MXLogType::Delete;
        record.table_id = table_id;
        record.ids = vector_ids.data();
        record.length = vector_ids.size();

        status = ExecWalRecord(record);
    }

    return status;
}

Status
DBImpl::Flush(const std::string& table_id) {
    if (!initialized_.load(std::memory_order_acquire)) {
        return SHUTDOWN_ERROR;
    }

    Status status;
    bool has_table;
    status = HasTable(table_id, has_table);
    if (!status.ok()) {
        return status;
    }
    if (!has_table) {
        ENGINE_LOG_ERROR << "Table to flush does not exist: " << table_id;
        return Status(DB_NOT_FOUND, "Table to flush does not exist");
    }

    ENGINE_LOG_DEBUG << "Begin flush table: " << table_id;

    if (options_.wal_enable_) {
        ENGINE_LOG_DEBUG << "WAL flush";
        auto lsn = wal_mgr_->Flush(table_id);
        ENGINE_LOG_DEBUG << "wal_mgr_->Flush";
        if (lsn != 0) {
597
            bg_task_swn_.Notify();
598 599 600 601 602 603 604 605 606 607 608 609 610 611 612 613 614 615 616 617 618 619 620 621 622 623 624 625 626 627
            flush_task_swn_.Wait();
            ENGINE_LOG_DEBUG << "flush_task_swn_.Wait()";
        }

    } else {
        ENGINE_LOG_DEBUG << "MemTable flush";
        wal::MXLogRecord record;
        record.type = wal::MXLogType::Flush;
        record.table_id = table_id;
        status = ExecWalRecord(record);
    }

    ENGINE_LOG_DEBUG << "End flush table: " << table_id;

    return status;
}

Status
DBImpl::Flush() {
    if (!initialized_.load(std::memory_order_acquire)) {
        return SHUTDOWN_ERROR;
    }

    ENGINE_LOG_DEBUG << "Begin flush all tables";

    Status status;
    if (options_.wal_enable_) {
        ENGINE_LOG_DEBUG << "WAL flush";
        auto lsn = wal_mgr_->Flush();
        if (lsn != 0) {
628
            bg_task_swn_.Notify();
629 630 631 632 633 634 635 636 637 638 639 640 641 642 643 644 645 646 647 648 649 650 651 652 653 654 655 656 657 658 659 660 661 662 663 664 665
            flush_task_swn_.Wait();
        }
    } else {
        ENGINE_LOG_DEBUG << "MemTable flush";
        wal::MXLogRecord record;
        record.type = wal::MXLogType::Flush;
        status = ExecWalRecord(record);
    }

    ENGINE_LOG_DEBUG << "End flush all tables";

    return status;
}

Status
DBImpl::Compact(const std::string& table_id) {
    if (!initialized_.load(std::memory_order_acquire)) {
        return SHUTDOWN_ERROR;
    }

    engine::meta::TableSchema table_schema;
    table_schema.table_id_ = table_id;
    auto status = DescribeTable(table_schema);
    if (!status.ok()) {
        if (status.code() == DB_NOT_FOUND) {
            ENGINE_LOG_ERROR << "Table to compact does not exist: " << table_id;
            return Status(DB_NOT_FOUND, "Table to compact does not exist");
        } else {
            return status;
        }
    } else {
        if (!table_schema.owner_table_.empty()) {
            ENGINE_LOG_ERROR << "Table to compact does not exist: " << table_id;
            return Status(DB_NOT_FOUND, "Table to compact does not exist");
        }
    }

Z
Zhiru Zhu 已提交
666
    ENGINE_LOG_DEBUG << "Before compacting, wait for build index thread to finish...";
667

Z
update  
Zhiru Zhu 已提交
668
    // WaitBuildIndexFinish();
669

Z
update  
Zhiru Zhu 已提交
670
    const std::lock_guard<std::mutex> index_lock(build_index_mutex_);
Z
Zhiru Zhu 已提交
671
    const std::lock_guard<std::mutex> merge_lock(flush_merge_compact_mutex_);
Z
Zhiru Zhu 已提交
672

Z
Zhiru Zhu 已提交
673
    ENGINE_LOG_DEBUG << "Compacting table: " << table_id;
Z
Zhiru Zhu 已提交
674

675
    // Get files to compact from meta.
Z
Zhiru Zhu 已提交
676 677
    std::vector<int> file_types{meta::TableFileSchema::FILE_TYPE::RAW, meta::TableFileSchema::FILE_TYPE::TO_INDEX,
                                meta::TableFileSchema::FILE_TYPE::BACKUP};
678 679 680 681 682 683 684 685 686 687 688
    meta::TableFilesSchema files_to_compact;
    status = meta_ptr_->FilesByType(table_id, file_types, files_to_compact);
    if (!status.ok()) {
        std::string err_msg = "Failed to get files to compact: " + status.message();
        ENGINE_LOG_ERROR << err_msg;
        return Status(DB_ERROR, err_msg);
    }

    ENGINE_LOG_DEBUG << "Found " << files_to_compact.size() << " segment to compact";

    OngoingFileChecker::GetInstance().MarkOngoingFiles(files_to_compact);
Z
Zhiru Zhu 已提交
689 690

    Status compact_status;
Z
Zhiru Zhu 已提交
691
    for (auto iter = files_to_compact.begin(); iter != files_to_compact.end();) {
G
groot 已提交
692 693 694
        meta::TableFileSchema file = *iter;
        iter = files_to_compact.erase(iter);

Z
Zhiru Zhu 已提交
695 696 697
        // Check if the segment needs compacting
        std::string segment_dir;
        utils::GetParentPath(file.location_, segment_dir);
698

Z
Zhiru Zhu 已提交
699
        segment::SegmentReader segment_reader(segment_dir);
Z
Zhiru Zhu 已提交
700 701
        size_t deleted_docs_size;
        status = segment_reader.ReadDeletedDocsSize(deleted_docs_size);
Z
Zhiru Zhu 已提交
702
        if (!status.ok()) {
G
groot 已提交
703 704
            OngoingFileChecker::GetInstance().UnmarkOngoingFile(file);
            continue;  // skip this file and try compact next one
Z
Zhiru Zhu 已提交
705 706
        }

G
groot 已提交
707
        meta::TableFilesSchema files_to_update;
Z
Zhiru Zhu 已提交
708
        if (deleted_docs_size != 0) {
Z
Zhiru Zhu 已提交
709 710 711 712 713
            compact_status = CompactFile(table_id, file, files_to_update);

            if (!compact_status.ok()) {
                ENGINE_LOG_ERROR << "Compact failed for segment " << file.segment_id_ << ": "
                                 << compact_status.message();
G
groot 已提交
714 715
                OngoingFileChecker::GetInstance().UnmarkOngoingFile(file);
                continue;  // skip this file and try compact next one
Z
Zhiru Zhu 已提交
716 717
            }
        } else {
G
groot 已提交
718
            OngoingFileChecker::GetInstance().UnmarkOngoingFile(file);
G
typo  
groot 已提交
719
            ENGINE_LOG_DEBUG << "Segment " << file.segment_id_ << " has no deleted data. No need to compact";
G
groot 已提交
720
            continue;  // skip this file and try compact next one
721
        }
Z
Zhiru Zhu 已提交
722

G
groot 已提交
723 724
        ENGINE_LOG_DEBUG << "Updating meta after compaction...";
        status = meta_ptr_->UpdateTableFiles(files_to_update);
G
groot 已提交
725
        OngoingFileChecker::GetInstance().UnmarkOngoingFile(file);
G
groot 已提交
726 727 728 729
        if (!status.ok()) {
            compact_status = status;
            break;  // meta error, could not go on
        }
Z
Zhiru Zhu 已提交
730 731
    }

732 733
    OngoingFileChecker::GetInstance().UnmarkOngoingFiles(files_to_compact);

G
groot 已提交
734 735 736
    if (compact_status.ok()) {
        ENGINE_LOG_DEBUG << "Finished compacting table: " << table_id;
    }
737

G
groot 已提交
738
    return compact_status;
739 740 741
}

Status
Z
Zhiru Zhu 已提交
742 743
DBImpl::CompactFile(const std::string& table_id, const meta::TableFileSchema& file,
                    meta::TableFilesSchema& files_to_update) {
744 745 746 747 748 749 750 751 752 753 754 755 756 757 758 759 760 761 762 763 764 765 766 767 768 769 770 771 772 773 774 775 776 777 778 779 780 781 782 783 784 785
    ENGINE_LOG_DEBUG << "Compacting segment " << file.segment_id_ << " for table: " << table_id;

    // Create new table file
    meta::TableFileSchema compacted_file;
    compacted_file.table_id_ = table_id;
    // compacted_file.date_ = date;
    compacted_file.file_type_ = meta::TableFileSchema::NEW_MERGE;  // TODO: use NEW_MERGE for now
    Status status = meta_ptr_->CreateTableFile(compacted_file);

    if (!status.ok()) {
        ENGINE_LOG_ERROR << "Failed to create table file: " << status.message();
        return status;
    }

    // Compact (merge) file to the newly created table file

    std::string new_segment_dir;
    utils::GetParentPath(compacted_file.location_, new_segment_dir);
    auto segment_writer_ptr = std::make_shared<segment::SegmentWriter>(new_segment_dir);

    std::string segment_dir_to_merge;
    utils::GetParentPath(file.location_, segment_dir_to_merge);

    ENGINE_LOG_DEBUG << "Compacting begin...";
    segment_writer_ptr->Merge(segment_dir_to_merge, compacted_file.file_id_);

    // Serialize
    ENGINE_LOG_DEBUG << "Serializing compacted segment...";
    status = segment_writer_ptr->Serialize();
    if (!status.ok()) {
        ENGINE_LOG_ERROR << "Failed to serialize compacted segment: " << status.message();
        compacted_file.file_type_ = meta::TableFileSchema::TO_DELETE;
        auto mark_status = meta_ptr_->UpdateTableFile(compacted_file);
        if (mark_status.ok()) {
            ENGINE_LOG_DEBUG << "Mark file: " << compacted_file.file_id_ << " to to_delete";
        }
        return status;
    }

    // Update table files state
    // if index type isn't IDMAP, set file type to TO_INDEX if file size exceed index_file_size
    // else set file type to RAW, no need to build index
786
    if (!utils::IsRawIndexType(compacted_file.engine_type_)) {
787 788 789 790 791 792 793 794 795 796 797 798 799 800
        compacted_file.file_type_ = (segment_writer_ptr->Size() >= compacted_file.index_file_size_)
                                        ? meta::TableFileSchema::TO_INDEX
                                        : meta::TableFileSchema::RAW;
    } else {
        compacted_file.file_type_ = meta::TableFileSchema::RAW;
    }
    compacted_file.file_size_ = segment_writer_ptr->Size();
    compacted_file.row_count_ = segment_writer_ptr->VectorCount();

    if (compacted_file.row_count_ == 0) {
        ENGINE_LOG_DEBUG << "Compacted segment is empty. Mark it as TO_DELETE";
        compacted_file.file_type_ = meta::TableFileSchema::TO_DELETE;
    }

Z
Zhiru Zhu 已提交
801
    files_to_update.emplace_back(compacted_file);
Z
Zhiru Zhu 已提交
802

Z
Zhiru Zhu 已提交
803 804 805 806 807 808 809 810 811
    // Set all files in segment to TO_DELETE
    auto& segment_id = file.segment_id_;
    meta::TableFilesSchema segment_files;
    status = meta_ptr_->GetTableFilesBySegmentId(segment_id, segment_files);
    if (!status.ok()) {
        return status;
    }
    for (auto& f : segment_files) {
        f.file_type_ = meta::TableFileSchema::FILE_TYPE::TO_DELETE;
Z
Zhiru Zhu 已提交
812 813
        files_to_update.emplace_back(f);
    }
814 815

    ENGINE_LOG_DEBUG << "Compacted segment " << compacted_file.segment_id_ << " from "
Z
Zhiru Zhu 已提交
816 817
                     << std::to_string(file.file_size_) << " bytes to " << std::to_string(compacted_file.file_size_)
                     << " bytes";
818 819 820 821 822 823 824 825 826 827 828 829 830 831 832 833 834 835 836 837 838 839 840 841 842 843 844 845 846 847 848 849 850 851 852 853 854 855 856 857 858 859 860 861 862 863 864 865 866 867 868 869 870 871 872 873 874 875 876 877 878 879 880 881 882 883 884 885 886 887 888 889 890 891 892 893 894 895 896 897 898 899 900 901 902 903 904 905 906 907 908 909 910 911 912 913 914 915 916 917 918 919 920 921 922 923 924 925 926 927 928 929 930 931 932 933 934 935 936 937 938 939 940 941 942 943 944 945 946 947 948 949 950

    if (options_.insert_cache_immediately_) {
        segment_writer_ptr->Cache();
    }

    return status;
}

Status
DBImpl::GetVectorByID(const std::string& table_id, const IDNumber& vector_id, VectorsData& vector) {
    if (!initialized_.load(std::memory_order_acquire)) {
        return SHUTDOWN_ERROR;
    }

    bool has_table;
    auto status = HasTable(table_id, has_table);
    if (!has_table) {
        ENGINE_LOG_ERROR << "Table " << table_id << " does not exist: ";
        return Status(DB_NOT_FOUND, "Table does not exist");
    }
    if (!status.ok()) {
        return status;
    }

    meta::TableFilesSchema files_to_query;

    std::vector<int> file_types{meta::TableFileSchema::FILE_TYPE::RAW, meta::TableFileSchema::FILE_TYPE::TO_INDEX,
                                meta::TableFileSchema::FILE_TYPE::BACKUP};
    meta::TableFilesSchema table_files;
    status = meta_ptr_->FilesByType(table_id, file_types, files_to_query);
    if (!status.ok()) {
        std::string err_msg = "Failed to get files for GetVectorByID: " + status.message();
        ENGINE_LOG_ERROR << err_msg;
        return status;
    }

    std::vector<meta::TableSchema> partition_array;
    status = meta_ptr_->ShowPartitions(table_id, partition_array);
    for (auto& schema : partition_array) {
        meta::TableFilesSchema files;
        status = meta_ptr_->FilesByType(schema.table_id_, file_types, files);
        if (!status.ok()) {
            std::string err_msg = "Failed to get files for GetVectorByID: " + status.message();
            ENGINE_LOG_ERROR << err_msg;
            return status;
        }
        files_to_query.insert(files_to_query.end(), std::make_move_iterator(files.begin()),
                              std::make_move_iterator(files.end()));
    }

    if (files_to_query.empty()) {
        ENGINE_LOG_DEBUG << "No files to get vector by id from";
        return Status::OK();
    }

    cache::CpuCacheMgr::GetInstance()->PrintInfo();
    OngoingFileChecker::GetInstance().MarkOngoingFiles(files_to_query);

    status = GetVectorByIdHelper(table_id, vector_id, vector, files_to_query);

    OngoingFileChecker::GetInstance().UnmarkOngoingFiles(files_to_query);
    cache::CpuCacheMgr::GetInstance()->PrintInfo();

    return status;
}

Status
DBImpl::GetVectorIDs(const std::string& table_id, const std::string& segment_id, IDNumbers& vector_ids) {
    if (!initialized_.load(std::memory_order_acquire)) {
        return SHUTDOWN_ERROR;
    }

    // step 1: check table existence
    bool has_table;
    auto status = HasTable(table_id, has_table);
    if (!has_table) {
        ENGINE_LOG_ERROR << "Table " << table_id << " does not exist: ";
        return Status(DB_NOT_FOUND, "Table does not exist");
    }
    if (!status.ok()) {
        return status;
    }

    //  step 2: find segment
    meta::TableFilesSchema table_files;
    status = meta_ptr_->GetTableFilesBySegmentId(segment_id, table_files);
    if (!status.ok()) {
        return status;
    }

    if (table_files.empty()) {
        return Status(DB_NOT_FOUND, "Segment does not exist");
    }

    // check the segment is belong to this table
    if (table_files[0].table_id_ != table_id) {
        // the segment could be in a partition under this table
        meta::TableSchema table_schema;
        table_schema.table_id_ = table_files[0].table_id_;
        status = DescribeTable(table_schema);
        if (table_schema.owner_table_ != table_id) {
            return Status(DB_NOT_FOUND, "Segment does not belong to this table");
        }
    }

    // step 3: load segment ids and delete offset
    std::string segment_dir;
    engine::utils::GetParentPath(table_files[0].location_, segment_dir);
    segment::SegmentReader segment_reader(segment_dir);

    std::vector<segment::doc_id_t> uids;
    status = segment_reader.LoadUids(uids);
    if (!status.ok()) {
        return status;
    }

    segment::DeletedDocsPtr deleted_docs_ptr;
    status = segment_reader.LoadDeletedDocs(deleted_docs_ptr);
    if (!status.ok()) {
        return status;
    }

    // step 4: construct id array
    // avoid duplicate offset and erase from max offset to min offset
    auto& deleted_offset = deleted_docs_ptr->GetDeletedDocs();
    std::set<segment::offset_t, std::greater<segment::offset_t>> ordered_offset;
    for (segment::offset_t offset : deleted_offset) {
        ordered_offset.insert(offset);
    }
    for (segment::offset_t offset : ordered_offset) {
        uids.erase(uids.begin() + offset);
    }
    vector_ids.swap(uids);
S
starlord 已提交
951

G
groot 已提交
952
    return status;
X
Xu Peng 已提交
953 954
}

955 956 957 958 959 960 961 962 963 964 965 966 967 968 969 970 971 972 973 974 975 976 977 978 979 980 981 982 983 984 985 986 987 988 989 990 991 992
Status
DBImpl::GetVectorByIdHelper(const std::string& table_id, IDNumber vector_id, VectorsData& vector,
                            const meta::TableFilesSchema& files) {
    ENGINE_LOG_DEBUG << "Getting vector by id in " << files.size() << " files";

    for (auto& file : files) {
        // Load bloom filter
        std::string segment_dir;
        engine::utils::GetParentPath(file.location_, segment_dir);
        segment::SegmentReader segment_reader(segment_dir);
        segment::IdBloomFilterPtr id_bloom_filter_ptr;
        segment_reader.LoadBloomFilter(id_bloom_filter_ptr);

        // Check if the id is present in bloom filter.
        if (id_bloom_filter_ptr->Check(vector_id)) {
            // Load uids and check if the id is indeed present. If yes, find its offset.
            std::vector<int64_t> offsets;
            std::vector<segment::doc_id_t> uids;
            auto status = segment_reader.LoadUids(uids);
            if (!status.ok()) {
                return status;
            }

            auto found = std::find(uids.begin(), uids.end(), vector_id);
            if (found != uids.end()) {
                auto offset = std::distance(uids.begin(), found);

                // Check whether the id has been deleted
                segment::DeletedDocsPtr deleted_docs_ptr;
                status = segment_reader.LoadDeletedDocs(deleted_docs_ptr);
                if (!status.ok()) {
                    return status;
                }
                auto& deleted_docs = deleted_docs_ptr->GetDeletedDocs();

                auto deleted = std::find(deleted_docs.begin(), deleted_docs.end(), offset);
                if (deleted == deleted_docs.end()) {
                    // Load raw vector
993
                    bool is_binary = utils::IsBinaryMetricType(file.metric_type_);
994 995 996 997 998 999 1000 1001 1002 1003 1004 1005 1006 1007 1008 1009 1010 1011 1012 1013 1014 1015 1016 1017 1018 1019 1020
                    size_t single_vector_bytes = is_binary ? file.dimension_ / 8 : file.dimension_ * sizeof(float);
                    std::vector<uint8_t> raw_vector;
                    status = segment_reader.LoadVectors(offset * single_vector_bytes, single_vector_bytes, raw_vector);
                    if (!status.ok()) {
                        return status;
                    }

                    vector.vector_count_ = 1;
                    if (is_binary) {
                        vector.binary_data_ = std::move(raw_vector);
                    } else {
                        std::vector<float> float_vector;
                        float_vector.resize(file.dimension_);
                        memcpy(float_vector.data(), raw_vector.data(), single_vector_bytes);
                        vector.float_data_ = std::move(float_vector);
                    }
                    return Status::OK();
                }
            }
        } else {
            continue;
        }
    }

    return Status::OK();
}

S
starlord 已提交
1021
Status
Y
Yu Kun 已提交
1022
DBImpl::CreateIndex(const std::string& table_id, const TableIndex& index) {
1023
    if (!initialized_.load(std::memory_order_acquire)) {
G
groot 已提交
1024 1025 1026
        return SHUTDOWN_ERROR;
    }

G
groot 已提交
1027
    // serialize memory data
1028 1029 1030
    //    std::set<std::string> sync_table_ids;
    //    auto status = SyncMemData(sync_table_ids);
    auto status = Flush();
G
groot 已提交
1031

S
starlord 已提交
1032 1033 1034
    {
        std::unique_lock<std::mutex> lock(build_index_mutex_);

S
starlord 已提交
1035
        // step 1: check index difference
S
starlord 已提交
1036
        TableIndex old_index;
G
groot 已提交
1037
        status = DescribeIndex(table_id, old_index);
S
starlord 已提交
1038
        if (!status.ok()) {
S
starlord 已提交
1039 1040 1041 1042
            ENGINE_LOG_ERROR << "Failed to get table index info for table: " << table_id;
            return status;
        }

S
starlord 已提交
1043
        // step 2: update index info
S
starlord 已提交
1044
        TableIndex new_index = index;
S
starlord 已提交
1045
        new_index.metric_type_ = old_index.metric_type_;  // dont change metric type, it was defined by CreateTable
S
starlord 已提交
1046
        if (!utils::IsSameIndex(old_index, new_index)) {
G
groot 已提交
1047
            status = UpdateTableIndexRecursively(table_id, new_index);
S
starlord 已提交
1048 1049 1050 1051 1052 1053
            if (!status.ok()) {
                return status;
            }
        }
    }

S
starlord 已提交
1054 1055
    // step 3: let merge file thread finish
    // to avoid duplicate data bug
1056 1057
    WaitMergeFileFinish();

S
starlord 已提交
1058
    // step 4: wait and build index
1059
    status = index_failed_checker_.CleanFailedIndexFileOfTable(table_id);
1060
    status = WaitTableIndexRecursively(table_id, index);
S
starlord 已提交
1061

G
groot 已提交
1062
    return status;
S
starlord 已提交
1063 1064
}

S
starlord 已提交
1065
Status
Y
Yu Kun 已提交
1066
DBImpl::DescribeIndex(const std::string& table_id, TableIndex& index) {
1067
    if (!initialized_.load(std::memory_order_acquire)) {
G
groot 已提交
1068 1069 1070
        return SHUTDOWN_ERROR;
    }

S
starlord 已提交
1071 1072 1073
    return meta_ptr_->DescribeTableIndex(table_id, index);
}

S
starlord 已提交
1074
Status
Y
Yu Kun 已提交
1075
DBImpl::DropIndex(const std::string& table_id) {
1076
    if (!initialized_.load(std::memory_order_acquire)) {
G
groot 已提交
1077 1078 1079
        return SHUTDOWN_ERROR;
    }

S
starlord 已提交
1080
    ENGINE_LOG_DEBUG << "Drop index for table: " << table_id;
G
groot 已提交
1081
    return DropTableIndexRecursively(table_id);
S
starlord 已提交
1082 1083
}

S
starlord 已提交
1084
Status
1085
DBImpl::QueryByID(const std::shared_ptr<server::Context>& context, const std::string& table_id,
1086 1087
                  const std::vector<std::string>& partition_tags, uint64_t k, const milvus::json& extra_params,
                  IDNumber vector_id, ResultIds& result_ids, ResultDistances& result_distances) {
1088
    if (!initialized_.load(std::memory_order_acquire)) {
G
groot 已提交
1089
        return SHUTDOWN_ERROR;
S
starlord 已提交
1090 1091
    }

1092 1093 1094
    VectorsData vectors_data = VectorsData();
    vectors_data.id_array_.emplace_back(vector_id);
    vectors_data.vector_count_ = 1;
1095 1096
    Status result =
        Query(context, table_id, partition_tags, k, extra_params, vectors_data, result_ids, result_distances);
Y
yu yunfeng 已提交
1097
    return result;
X
Xu Peng 已提交
1098 1099
}

S
starlord 已提交
1100
Status
Z
Zhiru Zhu 已提交
1101
DBImpl::Query(const std::shared_ptr<server::Context>& context, const std::string& table_id,
1102 1103
              const std::vector<std::string>& partition_tags, uint64_t k, const milvus::json& extra_params,
              const VectorsData& vectors, ResultIds& result_ids, ResultDistances& result_distances) {
Z
Zhiru Zhu 已提交
1104 1105
    auto query_ctx = context->Child("Query");

1106
    if (!initialized_.load(std::memory_order_acquire)) {
G
groot 已提交
1107
        return SHUTDOWN_ERROR;
S
starlord 已提交
1108 1109
    }

G
groot 已提交
1110
    Status status;
1111
    std::vector<size_t> ids;
G
groot 已提交
1112
    meta::TableFilesSchema files_array;
1113

G
groot 已提交
1114 1115 1116
    if (partition_tags.empty()) {
        // no partition tag specified, means search in whole table
        // get all table files from parent table
1117
        status = GetFilesToSearch(table_id, ids, files_array);
G
groot 已提交
1118 1119 1120 1121
        if (!status.ok()) {
            return status;
        }

G
groot 已提交
1122 1123 1124
        std::vector<meta::TableSchema> partition_array;
        status = meta_ptr_->ShowPartitions(table_id, partition_array);
        for (auto& schema : partition_array) {
1125 1126 1127 1128 1129
            status = GetFilesToSearch(schema.table_id_, ids, files_array);
        }

        if (files_array.empty()) {
            return Status::OK();
G
groot 已提交
1130 1131 1132 1133 1134 1135 1136
        }
    } else {
        // get files from specified partitions
        std::set<std::string> partition_name_array;
        GetPartitionsByTags(table_id, partition_tags, partition_name_array);

        for (auto& partition_name : partition_name_array) {
1137 1138 1139 1140 1141
            status = GetFilesToSearch(partition_name, ids, files_array);
        }

        if (files_array.empty()) {
            return Status::OK();
1142 1143 1144
        }
    }

S
starlord 已提交
1145
    cache::CpuCacheMgr::GetInstance()->PrintInfo();  // print cache info before query
1146
    status = QueryAsync(query_ctx, table_id, files_array, k, extra_params, vectors, result_ids, result_distances);
S
starlord 已提交
1147
    cache::CpuCacheMgr::GetInstance()->PrintInfo();  // print cache info after query
Z
Zhiru Zhu 已提交
1148 1149 1150

    query_ctx->GetTraceContext()->GetSpan()->Finish();

S
starlord 已提交
1151
    return status;
G
groot 已提交
1152
}
X
Xu Peng 已提交
1153

S
starlord 已提交
1154
Status
Z
Zhiru Zhu 已提交
1155
DBImpl::QueryByFileID(const std::shared_ptr<server::Context>& context, const std::string& table_id,
1156 1157
                      const std::vector<std::string>& file_ids, uint64_t k, const milvus::json& extra_params,
                      const VectorsData& vectors, ResultIds& result_ids, ResultDistances& result_distances) {
Z
Zhiru Zhu 已提交
1158 1159
    auto query_ctx = context->Child("Query by file id");

1160
    if (!initialized_.load(std::memory_order_acquire)) {
G
groot 已提交
1161
        return SHUTDOWN_ERROR;
S
starlord 已提交
1162 1163
    }

S
starlord 已提交
1164
    // get specified files
1165
    std::vector<size_t> ids;
Y
Yu Kun 已提交
1166
    for (auto& id : file_ids) {
1167
        meta::TableFileSchema table_file;
1168 1169
        table_file.table_id_ = table_id;
        std::string::size_type sz;
J
jinhai 已提交
1170
        ids.push_back(std::stoul(id, &sz));
1171 1172
    }

G
groot 已提交
1173
    meta::TableFilesSchema files_array;
1174
    auto status = GetFilesToSearch(table_id, ids, files_array);
1175 1176
    if (!status.ok()) {
        return status;
1177 1178
    }

S
shengjh 已提交
1179
    fiu_do_on("DBImpl.QueryByFileID.empty_files_array", files_array.clear());
G
groot 已提交
1180
    if (files_array.empty()) {
S
starlord 已提交
1181
        return Status(DB_ERROR, "Invalid file id");
G
groot 已提交
1182 1183
    }

S
starlord 已提交
1184
    cache::CpuCacheMgr::GetInstance()->PrintInfo();  // print cache info before query
1185
    status = QueryAsync(query_ctx, table_id, files_array, k, extra_params, vectors, result_ids, result_distances);
S
starlord 已提交
1186
    cache::CpuCacheMgr::GetInstance()->PrintInfo();  // print cache info after query
Z
Zhiru Zhu 已提交
1187 1188 1189

    query_ctx->GetTraceContext()->GetSpan()->Finish();

S
starlord 已提交
1190
    return status;
1191 1192
}

S
starlord 已提交
1193
Status
Y
Yu Kun 已提交
1194
DBImpl::Size(uint64_t& result) {
1195
    if (!initialized_.load(std::memory_order_acquire)) {
G
groot 已提交
1196
        return SHUTDOWN_ERROR;
S
starlord 已提交
1197 1198
    }

S
starlord 已提交
1199
    return meta_ptr_->Size(result);
S
starlord 已提交
1200 1201 1202
}

///////////////////////////////////////////////////////////////////////////////////////////////////////////////////
S
starlord 已提交
1203
// internal methods
S
starlord 已提交
1204
///////////////////////////////////////////////////////////////////////////////////////////////////////////////////
S
starlord 已提交
1205
Status
Z
Zhiru Zhu 已提交
1206
DBImpl::QueryAsync(const std::shared_ptr<server::Context>& context, const std::string& table_id,
1207 1208
                   const meta::TableFilesSchema& files, uint64_t k, const milvus::json& extra_params,
                   const VectorsData& vectors, ResultIds& result_ids, ResultDistances& result_distances) {
Z
Zhiru Zhu 已提交
1209 1210
    auto query_async_ctx = context->Child("Query Async");

G
groot 已提交
1211
    server::CollectQueryMetrics metrics(vectors.vector_count_);
Y
Yu Kun 已提交
1212

S
starlord 已提交
1213
    TimeRecorder rc("");
G
groot 已提交
1214

1215
    // step 1: construct search job
1216
    auto status = OngoingFileChecker::GetInstance().MarkOngoingFiles(files);
1217

1218
    ENGINE_LOG_DEBUG << "Engine query begin, index file count: " << files.size();
1219
    scheduler::SearchJobPtr job = std::make_shared<scheduler::SearchJob>(query_async_ctx, k, extra_params, vectors);
Y
Yu Kun 已提交
1220
    for (auto& file : files) {
S
starlord 已提交
1221
        scheduler::TableFileSchemaPtr file_ptr = std::make_shared<meta::TableFileSchema>(file);
W
wxyu 已提交
1222
        job->AddIndexFile(file_ptr);
G
groot 已提交
1223 1224
    }

1225
    // step 2: put search job to scheduler and wait result
S
starlord 已提交
1226
    scheduler::JobMgrInst::GetInstance()->Put(job);
W
wxyu 已提交
1227
    job->WaitResult();
1228

1229
    status = OngoingFileChecker::GetInstance().UnmarkOngoingFiles(files);
W
wxyu 已提交
1230 1231
    if (!job->GetStatus().ok()) {
        return job->GetStatus();
1232
    }
G
groot 已提交
1233

1234
    // step 3: construct results
G
groot 已提交
1235 1236
    result_ids = job->GetResultIds();
    result_distances = job->GetResultDistances();
S
starlord 已提交
1237
    rc.ElapseFromBegin("Engine query totally cost");
G
groot 已提交
1238

Z
Zhiru Zhu 已提交
1239 1240
    query_async_ctx->GetTraceContext()->GetSpan()->Finish();

G
groot 已提交
1241 1242 1243
    return Status::OK();
}

S
starlord 已提交
1244 1245
void
DBImpl::BackgroundTimerTask() {
Y
yu yunfeng 已提交
1246
    server::SystemInfo::GetInstance().Init();
X
Xu Peng 已提交
1247
    while (true) {
1248
        if (!initialized_.load(std::memory_order_acquire)) {
1249 1250
            WaitMergeFileFinish();
            WaitBuildIndexFinish();
S
starlord 已提交
1251 1252

            ENGINE_LOG_DEBUG << "DB background thread exit";
G
groot 已提交
1253 1254
            break;
        }
X
Xu Peng 已提交
1255

1256 1257 1258 1259 1260
        if (options_.auto_flush_interval_ > 0) {
            bg_task_swn_.Wait_For(std::chrono::seconds(options_.auto_flush_interval_));
        } else {
            bg_task_swn_.Wait();
        }
X
Xu Peng 已提交
1261

G
groot 已提交
1262
        StartMetricTask();
1263
        StartMergeTask();
G
groot 已提交
1264 1265
        StartBuildIndexTask();
    }
X
Xu Peng 已提交
1266 1267
}

S
starlord 已提交
1268 1269
void
DBImpl::WaitMergeFileFinish() {
1270 1271 1272
    ENGINE_LOG_DEBUG << "Begin WaitMergeFileFinish";
    std::lock_guard<std::mutex> lck(merge_result_mutex_);
    for (auto& iter : merge_thread_results_) {
1273 1274
        iter.wait();
    }
1275
    ENGINE_LOG_DEBUG << "End WaitMergeFileFinish";
1276 1277
}

S
starlord 已提交
1278 1279
void
DBImpl::WaitBuildIndexFinish() {
1280
    ENGINE_LOG_DEBUG << "Begin WaitBuildIndexFinish";
1281
    std::lock_guard<std::mutex> lck(index_result_mutex_);
Y
Yu Kun 已提交
1282
    for (auto& iter : index_thread_results_) {
1283 1284
        iter.wait();
    }
1285
    ENGINE_LOG_DEBUG << "End WaitBuildIndexFinish";
1286 1287
}

S
starlord 已提交
1288 1289
void
DBImpl::StartMetricTask() {
G
groot 已提交
1290
    static uint64_t metric_clock_tick = 0;
1291
    ++metric_clock_tick;
S
starlord 已提交
1292
    if (metric_clock_tick % METRIC_ACTION_INTERVAL != 0) {
G
groot 已提交
1293 1294 1295 1296 1297 1298
        return;
    }

    server::Metrics::GetInstance().KeepingAliveCounterIncrement(METRIC_ACTION_INTERVAL);
    int64_t cache_usage = cache::CpuCacheMgr::GetInstance()->CacheUsage();
    int64_t cache_total = cache::CpuCacheMgr::GetInstance()->CacheCapacity();
S
shengjh 已提交
1299 1300
    fiu_do_on("DBImpl.StartMetricTask.InvalidTotalCache", cache_total = 0);

J
JinHai-CN 已提交
1301 1302 1303 1304 1305 1306 1307
    if (cache_total > 0) {
        double cache_usage_double = cache_usage;
        server::Metrics::GetInstance().CpuCacheUsageGaugeSet(cache_usage_double * 100 / cache_total);
    } else {
        server::Metrics::GetInstance().CpuCacheUsageGaugeSet(0);
    }

Y
Yu Kun 已提交
1308
    server::Metrics::GetInstance().GpuCacheUsageGaugeSet();
G
groot 已提交
1309 1310 1311 1312 1313 1314 1315 1316
    uint64_t size;
    Size(size);
    server::Metrics::GetInstance().DataFileSizeGaugeSet(size);
    server::Metrics::GetInstance().CPUUsagePercentSet();
    server::Metrics::GetInstance().RAMUsagePercentSet();
    server::Metrics::GetInstance().GPUPercentGaugeSet();
    server::Metrics::GetInstance().GPUMemoryUsageGaugeSet();
    server::Metrics::GetInstance().OctetsSet();
S
starlord 已提交
1317

K
kun yu 已提交
1318
    server::Metrics::GetInstance().CPUCoreUsagePercentSet();
K
kun yu 已提交
1319 1320
    server::Metrics::GetInstance().GPUTemperature();
    server::Metrics::GetInstance().CPUTemperature();
1321
    server::Metrics::GetInstance().PushToGateway();
G
groot 已提交
1322 1323
}

S
starlord 已提交
1324
void
1325
DBImpl::StartMergeTask() {
1326
    static uint64_t compact_clock_tick = 0;
1327
    ++compact_clock_tick;
S
starlord 已提交
1328
    if (compact_clock_tick % COMPACT_ACTION_INTERVAL != 0) {
1329 1330 1331
        return;
    }

1332 1333 1334
    if (!options_.wal_enable_) {
        Flush();
    }
1335

1336 1337
    // ENGINE_LOG_DEBUG << "Begin StartMergeTask";
    // merge task has been finished?
1338
    {
1339 1340
        std::lock_guard<std::mutex> lck(merge_result_mutex_);
        if (!merge_thread_results_.empty()) {
1341
            std::chrono::milliseconds span(10);
1342 1343
            if (merge_thread_results_.back().wait_for(span) == std::future_status::ready) {
                merge_thread_results_.pop_back();
1344
            }
G
groot 已提交
1345 1346
        }
    }
X
Xu Peng 已提交
1347

1348
    // add new merge task
1349
    {
1350 1351 1352
        std::lock_guard<std::mutex> lck(merge_result_mutex_);
        if (merge_thread_results_.empty()) {
            // collect merge files for all tables(if merge_table_ids_ is empty) for two reasons:
1353 1354
            // 1. other tables may still has un-merged files
            // 2. server may be closed unexpected, these un-merge files need to be merged when server restart
1355
            if (merge_table_ids_.empty()) {
1356 1357
                std::vector<meta::TableSchema> table_schema_array;
                meta_ptr_->AllTables(table_schema_array);
G
groot 已提交
1358
                for (auto& schema : table_schema_array) {
1359
                    merge_table_ids_.insert(schema.table_id_);
1360 1361 1362 1363
                }
            }

            // start merge file thread
1364 1365 1366
            merge_thread_results_.push_back(
                merge_thread_pool_.enqueue(&DBImpl::BackgroundMerge, this, merge_table_ids_));
            merge_table_ids_.clear();
1367
        }
G
groot 已提交
1368
    }
1369 1370

    // ENGINE_LOG_DEBUG << "End StartMergeTask";
X
Xu Peng 已提交
1371 1372
}

S
starlord 已提交
1373
Status
1374
DBImpl::MergeFiles(const std::string& table_id, const meta::TableFilesSchema& files) {
Z
Zhiru Zhu 已提交
1375
    // const std::lock_guard<std::mutex> lock(flush_merge_compact_mutex_);
1376

S
starlord 已提交
1377
    ENGINE_LOG_DEBUG << "Merge files for table: " << table_id;
S
starlord 已提交
1378

S
starlord 已提交
1379
    // step 1: create table file
X
Xu Peng 已提交
1380
    meta::TableFileSchema table_file;
G
groot 已提交
1381
    table_file.table_id_ = table_id;
1382
    table_file.file_type_ = meta::TableFileSchema::NEW_MERGE;
G
groot 已提交
1383
    Status status = meta_ptr_->CreateTableFile(table_file);
X
Xu Peng 已提交
1384

1385
    if (!status.ok()) {
S
starlord 已提交
1386
        ENGINE_LOG_ERROR << "Failed to create table: " << status.ToString();
1387 1388 1389
        return status;
    }

S
starlord 已提交
1390
    // step 2: merge files
1391
    /*
G
groot 已提交
1392
    ExecutionEnginePtr index =
Y
Yu Kun 已提交
1393 1394
        EngineFactory::Build(table_file.dimension_, table_file.location_, (EngineType)table_file.engine_type_,
                             (MetricType)table_file.metric_type_, table_file.nlist_);
1395
*/
1396
    meta::TableFilesSchema updated;
1397 1398 1399 1400

    std::string new_segment_dir;
    utils::GetParentPath(table_file.location_, new_segment_dir);
    auto segment_writer_ptr = std::make_shared<segment::SegmentWriter>(new_segment_dir);
1401

Y
Yu Kun 已提交
1402
    for (auto& file : files) {
Y
Yu Kun 已提交
1403
        server::CollectMergeFilesMetrics metrics;
1404 1405 1406
        std::string segment_dir_to_merge;
        utils::GetParentPath(file.location_, segment_dir_to_merge);
        segment_writer_ptr->Merge(segment_dir_to_merge, table_file.file_id_);
1407
        auto file_schema = file;
G
groot 已提交
1408
        file_schema.file_type_ = meta::TableFileSchema::TO_DELETE;
1409
        updated.push_back(file_schema);
1410 1411
        auto size = segment_writer_ptr->Size();
        if (size >= file_schema.index_file_size_) {
S
starlord 已提交
1412
            break;
S
starlord 已提交
1413
        }
1414 1415
    }

S
starlord 已提交
1416
    // step 3: serialize to disk
S
starlord 已提交
1417
    try {
1418
        status = segment_writer_ptr->Serialize();
S
shengjh 已提交
1419 1420
        fiu_do_on("DBImpl.MergeFiles.Serialize_ThrowException", throw std::exception());
        fiu_do_on("DBImpl.MergeFiles.Serialize_ErrorStatus", status = Status(DB_ERROR, ""));
Y
Yu Kun 已提交
1421
    } catch (std::exception& ex) {
S
starlord 已提交
1422
        std::string msg = "Serialize merged index encounter exception: " + std::string(ex.what());
S
starlord 已提交
1423
        ENGINE_LOG_ERROR << msg;
G
groot 已提交
1424 1425
        status = Status(DB_ERROR, msg);
    }
Y
yu yunfeng 已提交
1426

G
groot 已提交
1427
    if (!status.ok()) {
1428 1429
        ENGINE_LOG_ERROR << "Failed to persist merged segment: " << new_segment_dir << ". Error: " << status.message();

G
groot 已提交
1430
        // if failed to serialize merge file to disk
1431
        // typical error: out of disk space, out of memory or permission denied
S
starlord 已提交
1432 1433 1434
        table_file.file_type_ = meta::TableFileSchema::TO_DELETE;
        status = meta_ptr_->UpdateTableFile(table_file);
        ENGINE_LOG_DEBUG << "Failed to update file to index, mark file: " << table_file.file_id_ << " to to_delete";
X
Xu Peng 已提交
1435

G
groot 已提交
1436
        return status;
S
starlord 已提交
1437 1438
    }

S
starlord 已提交
1439
    // step 4: update table files state
1440
    // if index type isn't IDMAP, set file type to TO_INDEX if file size exceed index_file_size
S
starlord 已提交
1441
    // else set file type to RAW, no need to build index
1442
    if (!utils::IsRawIndexType(table_file.engine_type_)) {
1443 1444 1445
        table_file.file_type_ = (segment_writer_ptr->Size() >= table_file.index_file_size_)
                                    ? meta::TableFileSchema::TO_INDEX
                                    : meta::TableFileSchema::RAW;
1446 1447 1448
    } else {
        table_file.file_type_ = meta::TableFileSchema::RAW;
    }
1449 1450
    table_file.file_size_ = segment_writer_ptr->Size();
    table_file.row_count_ = segment_writer_ptr->VectorCount();
X
Xu Peng 已提交
1451
    updated.push_back(table_file);
G
groot 已提交
1452
    status = meta_ptr_->UpdateTableFiles(updated);
1453 1454
    ENGINE_LOG_DEBUG << "New merged segment " << table_file.segment_id_ << " of size " << segment_writer_ptr->Size()
                     << " bytes";
1455

S
starlord 已提交
1456
    if (options_.insert_cache_immediately_) {
1457
        segment_writer_ptr->Cache();
S
starlord 已提交
1458
    }
X
Xu Peng 已提交
1459

1460 1461 1462
    return status;
}

S
starlord 已提交
1463
Status
Y
Yu Kun 已提交
1464
DBImpl::BackgroundMergeFiles(const std::string& table_id) {
Z
Zhiru Zhu 已提交
1465
    const std::lock_guard<std::mutex> lock(flush_merge_compact_mutex_);
1466 1467

    meta::TableFilesSchema raw_files;
G
groot 已提交
1468
    auto status = meta_ptr_->FilesToMerge(table_id, raw_files);
X
Xu Peng 已提交
1469
    if (!status.ok()) {
S
starlord 已提交
1470
        ENGINE_LOG_ERROR << "Failed to get merge files for table: " << table_id;
X
Xu Peng 已提交
1471 1472
        return status;
    }
1473

1474 1475 1476 1477
    if (raw_files.size() < options_.merge_trigger_number_) {
        ENGINE_LOG_TRACE << "Files number not greater equal than merge trigger number, skip merge action";
        return Status::OK();
    }
1478

1479 1480 1481
    status = OngoingFileChecker::GetInstance().MarkOngoingFiles(raw_files);
    MergeFiles(table_id, raw_files);
    status = OngoingFileChecker::GetInstance().UnmarkOngoingFiles(raw_files);
G
groot 已提交
1482

1483 1484
    if (!initialized_.load(std::memory_order_acquire)) {
        ENGINE_LOG_DEBUG << "Server will shutdown, skip merge action for table: " << table_id;
1485
    }
X
Xu Peng 已提交
1486

G
groot 已提交
1487 1488
    return Status::OK();
}
1489

S
starlord 已提交
1490
void
1491 1492
DBImpl::BackgroundMerge(std::set<std::string> table_ids) {
    // ENGINE_LOG_TRACE << " Background merge thread start";
S
starlord 已提交
1493

G
groot 已提交
1494
    Status status;
Y
Yu Kun 已提交
1495
    for (auto& table_id : table_ids) {
G
groot 已提交
1496 1497
        status = BackgroundMergeFiles(table_id);
        if (!status.ok()) {
S
starlord 已提交
1498
            ENGINE_LOG_ERROR << "Merge files for table " << table_id << " failed: " << status.ToString();
G
groot 已提交
1499
        }
S
starlord 已提交
1500

1501
        if (!initialized_.load(std::memory_order_acquire)) {
S
starlord 已提交
1502 1503 1504
            ENGINE_LOG_DEBUG << "Server will shutdown, skip merge action";
            break;
        }
G
groot 已提交
1505
    }
X
Xu Peng 已提交
1506

G
groot 已提交
1507
    meta_ptr_->Archive();
Z
update  
zhiru 已提交
1508

1509
    {
G
groot 已提交
1510
        uint64_t ttl = 10 * meta::SECOND;  // default: file will be hard-deleted few seconds after soft-deleted
1511
        if (options_.mode_ == DBOptions::MODE::CLUSTER_WRITABLE) {
1512
            ttl = meta::HOUR;
1513
        }
G
groot 已提交
1514

1515
        meta_ptr_->CleanUpFilesWithTTL(ttl);
Z
update  
zhiru 已提交
1516
    }
S
starlord 已提交
1517

1518
    // ENGINE_LOG_TRACE << " Background merge thread exit";
G
groot 已提交
1519
}
X
Xu Peng 已提交
1520

S
starlord 已提交
1521 1522
void
DBImpl::StartBuildIndexTask(bool force) {
G
groot 已提交
1523
    static uint64_t index_clock_tick = 0;
1524
    ++index_clock_tick;
S
starlord 已提交
1525
    if (!force && (index_clock_tick % INDEX_ACTION_INTERVAL != 0)) {
G
groot 已提交
1526 1527 1528
        return;
    }

S
starlord 已提交
1529
    // build index has been finished?
1530 1531 1532 1533 1534 1535 1536
    {
        std::lock_guard<std::mutex> lck(index_result_mutex_);
        if (!index_thread_results_.empty()) {
            std::chrono::milliseconds span(10);
            if (index_thread_results_.back().wait_for(span) == std::future_status::ready) {
                index_thread_results_.pop_back();
            }
G
groot 已提交
1537 1538 1539
        }
    }

S
starlord 已提交
1540
    // add new build index task
1541 1542 1543
    {
        std::lock_guard<std::mutex> lck(index_result_mutex_);
        if (index_thread_results_.empty()) {
S
starlord 已提交
1544
            index_thread_results_.push_back(index_thread_pool_.enqueue(&DBImpl::BackgroundBuildIndex, this));
1545
        }
G
groot 已提交
1546
    }
X
Xu Peng 已提交
1547 1548
}

S
starlord 已提交
1549 1550
void
DBImpl::BackgroundBuildIndex() {
P
peng.xu 已提交
1551
    std::unique_lock<std::mutex> lock(build_index_mutex_);
1552
    meta::TableFilesSchema to_index_files;
G
groot 已提交
1553
    meta_ptr_->FilesToIndex(to_index_files);
1554
    Status status = index_failed_checker_.IgnoreFailedIndexFiles(to_index_files);
1555

1556
    if (!to_index_files.empty()) {
G
groot 已提交
1557
        ENGINE_LOG_DEBUG << "Background build index thread begin";
1558
        status = OngoingFileChecker::GetInstance().MarkOngoingFiles(to_index_files);
1559

1560
        // step 2: put build index task to scheduler
G
groot 已提交
1561
        std::vector<std::pair<scheduler::BuildIndexJobPtr, scheduler::TableFileSchemaPtr>> job2file_map;
1562
        for (auto& file : to_index_files) {
G
groot 已提交
1563
            scheduler::BuildIndexJobPtr job = std::make_shared<scheduler::BuildIndexJob>(meta_ptr_, options_);
1564 1565
            scheduler::TableFileSchemaPtr file_ptr = std::make_shared<meta::TableFileSchema>(file);
            job->AddToIndexFiles(file_ptr);
G
groot 已提交
1566
            scheduler::JobMgrInst::GetInstance()->Put(job);
G
groot 已提交
1567
            job2file_map.push_back(std::make_pair(job, file_ptr));
1568
        }
G
groot 已提交
1569

G
groot 已提交
1570
        // step 3: wait build index finished and mark failed files
G
groot 已提交
1571 1572 1573 1574 1575 1576 1577 1578
        for (auto iter = job2file_map.begin(); iter != job2file_map.end(); ++iter) {
            scheduler::BuildIndexJobPtr job = iter->first;
            meta::TableFileSchema& file_schema = *(iter->second.get());
            job->WaitBuildIndexFinish();
            if (!job->GetStatus().ok()) {
                Status status = job->GetStatus();
                ENGINE_LOG_ERROR << "Building index job " << job->id() << " failed: " << status.ToString();

1579
                index_failed_checker_.MarkFailedIndexFile(file_schema, status.message());
G
groot 已提交
1580 1581
            } else {
                ENGINE_LOG_DEBUG << "Building index job " << job->id() << " succeed.";
G
groot 已提交
1582 1583

                index_failed_checker_.MarkSucceedIndexFile(file_schema);
G
groot 已提交
1584
            }
1585
            status = OngoingFileChecker::GetInstance().UnmarkOngoingFile(file_schema);
1586
        }
G
groot 已提交
1587 1588

        ENGINE_LOG_DEBUG << "Background build index thread finished";
Y
Yu Kun 已提交
1589
    }
X
Xu Peng 已提交
1590 1591
}

G
groot 已提交
1592 1593 1594 1595 1596 1597 1598 1599 1600 1601 1602 1603
Status
DBImpl::GetFilesToBuildIndex(const std::string& table_id, const std::vector<int>& file_types,
                             meta::TableFilesSchema& files) {
    files.clear();
    auto status = meta_ptr_->FilesByType(table_id, file_types, files);

    // only build index for files that row count greater than certain threshold
    for (auto it = files.begin(); it != files.end();) {
        if ((*it).file_type_ == static_cast<int>(meta::TableFileSchema::RAW) &&
            (*it).row_count_ < meta::BUILD_INDEX_THRESHOLD) {
            it = files.erase(it);
        } else {
1604
            ++it;
G
groot 已提交
1605 1606 1607 1608 1609 1610
        }
    }

    return Status::OK();
}

G
groot 已提交
1611
Status
1612
DBImpl::GetFilesToSearch(const std::string& table_id, const std::vector<size_t>& file_ids,
G
groot 已提交
1613
                         meta::TableFilesSchema& files) {
1614 1615
    ENGINE_LOG_DEBUG << "Collect files from table: " << table_id;

1616 1617
    meta::TableFilesSchema search_files;
    auto status = meta_ptr_->FilesToSearch(table_id, file_ids, search_files);
G
groot 已提交
1618 1619 1620 1621
    if (!status.ok()) {
        return status;
    }

1622 1623 1624
    for (auto& file : search_files) {
        files.push_back(file);
    }
G
groot 已提交
1625 1626 1627
    return Status::OK();
}

1628 1629 1630 1631 1632 1633 1634 1635 1636 1637 1638 1639 1640 1641 1642 1643 1644 1645 1646 1647 1648 1649 1650 1651 1652 1653 1654
Status
DBImpl::GetPartitionByTag(const std::string& table_id, const std::string& partition_tag, std::string& partition_name) {
    Status status;

    if (partition_tag.empty()) {
        partition_name = table_id;

    } else {
        // trim side-blank of tag, only compare valid characters
        // for example: " ab cd " is treated as "ab cd"
        std::string valid_tag = partition_tag;
        server::StringHelpFunctions::TrimStringBlank(valid_tag);

        if (valid_tag == milvus::engine::DEFAULT_PARTITON_TAG) {
            partition_name = table_id;
            return status;
        }

        status = meta_ptr_->GetPartitionName(table_id, partition_tag, partition_name);
        if (!status.ok()) {
            ENGINE_LOG_ERROR << status.message();
        }
    }

    return status;
}

G
groot 已提交
1655 1656 1657
Status
DBImpl::GetPartitionsByTags(const std::string& table_id, const std::vector<std::string>& partition_tags,
                            std::set<std::string>& partition_name_array) {
G
groot 已提交
1658 1659
    std::vector<meta::TableSchema> partition_array;
    auto status = meta_ptr_->ShowPartitions(table_id, partition_array);
G
groot 已提交
1660 1661

    for (auto& tag : partition_tags) {
1662 1663 1664 1665
        // trim side-blank of tag, only compare valid characters
        // for example: " ab cd " is treated as "ab cd"
        std::string valid_tag = tag;
        server::StringHelpFunctions::TrimStringBlank(valid_tag);
1666 1667 1668 1669 1670 1671

        if (valid_tag == milvus::engine::DEFAULT_PARTITON_TAG) {
            partition_name_array.insert(table_id);
            return status;
        }

G
groot 已提交
1672
        for (auto& schema : partition_array) {
1673
            if (server::StringHelpFunctions::IsRegexMatch(schema.partition_tag_, valid_tag)) {
G
groot 已提交
1674 1675 1676 1677 1678 1679 1680 1681 1682
                partition_name_array.insert(schema.table_id_);
            }
        }
    }

    return Status::OK();
}

Status
1683
DBImpl::DropTableRecursively(const std::string& table_id) {
G
groot 已提交
1684 1685 1686 1687
    // dates partly delete files of the table but currently we don't support
    ENGINE_LOG_DEBUG << "Prepare to delete table " << table_id;

    Status status;
1688 1689
    if (options_.wal_enable_) {
        wal_mgr_->DropTable(table_id);
G
groot 已提交
1690 1691
    }

1692 1693 1694 1695 1696 1697 1698 1699 1700 1701
    status = mem_mgr_->EraseMemVector(table_id);  // not allow insert
    status = meta_ptr_->DropTable(table_id);      // soft delete table
    index_failed_checker_.CleanFailedIndexFileOfTable(table_id);

    // scheduler will determine when to delete table files
    auto nres = scheduler::ResMgrInst::GetInstance()->GetNumOfComputeResource();
    scheduler::DeleteJobPtr job = std::make_shared<scheduler::DeleteJob>(table_id, meta_ptr_, nres);
    scheduler::JobMgrInst::GetInstance()->Put(job);
    job->WaitAndDelete();

G
groot 已提交
1702 1703 1704
    std::vector<meta::TableSchema> partition_array;
    status = meta_ptr_->ShowPartitions(table_id, partition_array);
    for (auto& schema : partition_array) {
1705
        status = DropTableRecursively(schema.table_id_);
S
shengjh 已提交
1706
        fiu_do_on("DBImpl.DropTableRecursively.failed", status = Status(DB_ERROR, ""));
G
groot 已提交
1707 1708 1709 1710 1711 1712 1713 1714 1715 1716 1717 1718 1719
        if (!status.ok()) {
            return status;
        }
    }

    return Status::OK();
}

Status
DBImpl::UpdateTableIndexRecursively(const std::string& table_id, const TableIndex& index) {
    DropIndex(table_id);

    auto status = meta_ptr_->UpdateTableIndex(table_id, index);
S
shengjh 已提交
1720 1721
    fiu_do_on("DBImpl.UpdateTableIndexRecursively.fail_update_table_index",
              status = Status(DB_META_TRANSACTION_FAILED, ""));
G
groot 已提交
1722 1723 1724 1725 1726
    if (!status.ok()) {
        ENGINE_LOG_ERROR << "Failed to update table index info for table: " << table_id;
        return status;
    }

G
groot 已提交
1727 1728 1729
    std::vector<meta::TableSchema> partition_array;
    status = meta_ptr_->ShowPartitions(table_id, partition_array);
    for (auto& schema : partition_array) {
G
groot 已提交
1730 1731 1732 1733 1734 1735 1736 1737 1738 1739
        status = UpdateTableIndexRecursively(schema.table_id_, index);
        if (!status.ok()) {
            return status;
        }
    }

    return Status::OK();
}

Status
1740
DBImpl::WaitTableIndexRecursively(const std::string& table_id, const TableIndex& index) {
G
groot 已提交
1741 1742 1743
    // for IDMAP type, only wait all NEW file converted to RAW file
    // for other type, wait NEW/RAW/NEW_MERGE/NEW_INDEX/TO_INDEX files converted to INDEX files
    std::vector<int> file_types;
1744
    if (utils::IsRawIndexType(index.engine_type_)) {
G
groot 已提交
1745 1746 1747 1748 1749 1750 1751 1752 1753 1754 1755 1756 1757 1758 1759
        file_types = {
            static_cast<int32_t>(meta::TableFileSchema::NEW),
            static_cast<int32_t>(meta::TableFileSchema::NEW_MERGE),
        };
    } else {
        file_types = {
            static_cast<int32_t>(meta::TableFileSchema::RAW),
            static_cast<int32_t>(meta::TableFileSchema::NEW),
            static_cast<int32_t>(meta::TableFileSchema::NEW_MERGE),
            static_cast<int32_t>(meta::TableFileSchema::NEW_INDEX),
            static_cast<int32_t>(meta::TableFileSchema::TO_INDEX),
        };
    }

    // get files to build index
G
groot 已提交
1760 1761
    meta::TableFilesSchema table_files;
    auto status = GetFilesToBuildIndex(table_id, file_types, table_files);
G
groot 已提交
1762 1763
    int times = 1;

G
groot 已提交
1764
    while (!table_files.empty()) {
G
groot 已提交
1765
        ENGINE_LOG_DEBUG << "Non index files detected! Will build index " << times;
1766
        if (!utils::IsRawIndexType(index.engine_type_)) {
G
groot 已提交
1767 1768 1769 1770
            status = meta_ptr_->UpdateTableFilesToIndex(table_id);
        }

        std::this_thread::sleep_for(std::chrono::milliseconds(std::min(10 * 1000, times * 100)));
G
groot 已提交
1771
        GetFilesToBuildIndex(table_id, file_types, table_files);
1772
        ++times;
G
groot 已提交
1773

1774
        index_failed_checker_.IgnoreFailedIndexFiles(table_files);
G
groot 已提交
1775 1776 1777
    }

    // build index for partition
G
groot 已提交
1778 1779 1780
    std::vector<meta::TableSchema> partition_array;
    status = meta_ptr_->ShowPartitions(table_id, partition_array);
    for (auto& schema : partition_array) {
1781 1782
        status = WaitTableIndexRecursively(schema.table_id_, index);
        fiu_do_on("DBImpl.WaitTableIndexRecursively.fail_build_table_Index_for_partition",
S
shengjh 已提交
1783
                  status = Status(DB_ERROR, ""));
G
groot 已提交
1784 1785 1786 1787 1788
        if (!status.ok()) {
            return status;
        }
    }

G
groot 已提交
1789
    // failed to build index for some files, return error
1790 1791
    std::string err_msg;
    index_failed_checker_.GetErrMsgForTable(table_id, err_msg);
1792
    fiu_do_on("DBImpl.WaitTableIndexRecursively.not_empty_err_msg", err_msg.append("fiu"));
1793 1794
    if (!err_msg.empty()) {
        return Status(DB_ERROR, err_msg);
G
groot 已提交
1795 1796
    }

G
groot 已提交
1797 1798 1799 1800 1801 1802
    return Status::OK();
}

Status
DBImpl::DropTableIndexRecursively(const std::string& table_id) {
    ENGINE_LOG_DEBUG << "Drop index for table: " << table_id;
1803
    index_failed_checker_.CleanFailedIndexFileOfTable(table_id);
G
groot 已提交
1804 1805 1806 1807 1808 1809
    auto status = meta_ptr_->DropTableIndex(table_id);
    if (!status.ok()) {
        return status;
    }

    // drop partition index
G
groot 已提交
1810 1811 1812
    std::vector<meta::TableSchema> partition_array;
    status = meta_ptr_->ShowPartitions(table_id, partition_array);
    for (auto& schema : partition_array) {
G
groot 已提交
1813
        status = DropTableIndexRecursively(schema.table_id_);
S
shengjh 已提交
1814 1815
        fiu_do_on("DBImpl.DropTableIndexRecursively.fail_drop_table_Index_for_partition",
                  status = Status(DB_ERROR, ""));
G
groot 已提交
1816 1817 1818 1819 1820 1821 1822 1823 1824 1825 1826 1827 1828 1829 1830 1831 1832
        if (!status.ok()) {
            return status;
        }
    }

    return Status::OK();
}

Status
DBImpl::GetTableRowCountRecursively(const std::string& table_id, uint64_t& row_count) {
    row_count = 0;
    auto status = meta_ptr_->Count(table_id, row_count);
    if (!status.ok()) {
        return status;
    }

    // get partition row count
G
groot 已提交
1833 1834 1835
    std::vector<meta::TableSchema> partition_array;
    status = meta_ptr_->ShowPartitions(table_id, partition_array);
    for (auto& schema : partition_array) {
G
groot 已提交
1836 1837
        uint64_t partition_row_count = 0;
        status = GetTableRowCountRecursively(schema.table_id_, partition_row_count);
S
shengjh 已提交
1838 1839
        fiu_do_on("DBImpl.GetTableRowCountRecursively.fail_get_table_rowcount_for_partition",
                  status = Status(DB_ERROR, ""));
G
groot 已提交
1840 1841 1842 1843 1844 1845 1846 1847 1848 1849
        if (!status.ok()) {
            return status;
        }

        row_count += partition_row_count;
    }

    return Status::OK();
}

1850 1851 1852 1853 1854 1855 1856 1857 1858 1859 1860 1861 1862 1863 1864 1865 1866 1867 1868 1869 1870 1871 1872 1873 1874 1875 1876 1877 1878 1879 1880 1881 1882 1883 1884 1885 1886 1887 1888 1889 1890 1891 1892 1893 1894 1895 1896 1897 1898 1899 1900 1901 1902 1903 1904 1905 1906 1907 1908 1909 1910 1911 1912 1913 1914 1915 1916 1917 1918 1919 1920 1921 1922 1923 1924 1925 1926 1927 1928 1929 1930 1931 1932 1933 1934 1935 1936 1937 1938 1939 1940 1941 1942 1943 1944 1945 1946 1947 1948 1949 1950 1951 1952 1953 1954 1955 1956 1957 1958 1959 1960 1961 1962 1963 1964 1965 1966 1967 1968 1969 1970 1971 1972 1973 1974 1975 1976 1977 1978 1979 1980 1981 1982 1983 1984 1985 1986 1987 1988 1989 1990 1991 1992 1993 1994 1995 1996 1997 1998 1999 2000
Status
DBImpl::ExecWalRecord(const wal::MXLogRecord& record) {
    fiu_return_on("DBImpl.ExexWalRecord.return", Status(););

    auto tables_flushed = [&](const std::set<std::string>& table_ids) -> uint64_t {
        if (table_ids.empty()) {
            return 0;
        }

        uint64_t max_lsn = 0;
        if (options_.wal_enable_) {
            for (auto& table : table_ids) {
                uint64_t lsn = 0;
                meta_ptr_->GetTableFlushLSN(table, lsn);
                wal_mgr_->TableFlushed(table, lsn);
                if (lsn > max_lsn) {
                    max_lsn = lsn;
                }
            }
        }

        std::lock_guard<std::mutex> lck(merge_result_mutex_);
        for (auto& table : table_ids) {
            merge_table_ids_.insert(table);
        }
        return max_lsn;
    };

    Status status;

    switch (record.type) {
        case wal::MXLogType::InsertBinary: {
            std::string target_table_name;
            status = GetPartitionByTag(record.table_id, record.partition_tag, target_table_name);
            if (!status.ok()) {
                return status;
            }

            std::set<std::string> flushed_tables;
            status = mem_mgr_->InsertVectors(target_table_name, record.length, record.ids,
                                             (record.data_size / record.length / sizeof(uint8_t)),
                                             (const u_int8_t*)record.data, record.lsn, flushed_tables);
            // even though !status.ok, run
            tables_flushed(flushed_tables);

            // metrics
            milvus::server::CollectInsertMetrics metrics(record.length, status);
            break;
        }

        case wal::MXLogType::InsertVector: {
            std::string target_table_name;
            status = GetPartitionByTag(record.table_id, record.partition_tag, target_table_name);
            if (!status.ok()) {
                return status;
            }

            std::set<std::string> flushed_tables;
            status = mem_mgr_->InsertVectors(target_table_name, record.length, record.ids,
                                             (record.data_size / record.length / sizeof(float)),
                                             (const float*)record.data, record.lsn, flushed_tables);
            // even though !status.ok, run
            tables_flushed(flushed_tables);

            // metrics
            milvus::server::CollectInsertMetrics metrics(record.length, status);
            break;
        }

        case wal::MXLogType::Delete: {
            std::vector<meta::TableSchema> partition_array;
            status = meta_ptr_->ShowPartitions(record.table_id, partition_array);
            if (!status.ok()) {
                return status;
            }

            std::vector<std::string> table_ids{record.table_id};
            for (auto& partition : partition_array) {
                auto& partition_table_id = partition.table_id_;
                table_ids.emplace_back(partition_table_id);
            }

            if (record.length == 1) {
                for (auto& table_id : table_ids) {
                    status = mem_mgr_->DeleteVector(table_id, *record.ids, record.lsn);
                    if (!status.ok()) {
                        return status;
                    }
                }
            } else {
                for (auto& table_id : table_ids) {
                    status = mem_mgr_->DeleteVectors(table_id, record.length, record.ids, record.lsn);
                    if (!status.ok()) {
                        return status;
                    }
                }
            }
            break;
        }

        case wal::MXLogType::Flush: {
            if (!record.table_id.empty()) {
                // flush one table
                std::vector<meta::TableSchema> partition_array;
                status = meta_ptr_->ShowPartitions(record.table_id, partition_array);
                if (!status.ok()) {
                    return status;
                }

                std::vector<std::string> table_ids{record.table_id};
                for (auto& partition : partition_array) {
                    auto& partition_table_id = partition.table_id_;
                    table_ids.emplace_back(partition_table_id);
                }

                std::set<std::string> flushed_tables;
                for (auto& table_id : table_ids) {
                    const std::lock_guard<std::mutex> lock(flush_merge_compact_mutex_);
                    status = mem_mgr_->Flush(table_id);
                    if (!status.ok()) {
                        break;
                    }
                    flushed_tables.insert(table_id);
                }

                tables_flushed(flushed_tables);

            } else {
                // flush all tables
                std::set<std::string> table_ids;
                {
                    const std::lock_guard<std::mutex> lock(flush_merge_compact_mutex_);
                    status = mem_mgr_->Flush(table_ids);
                }

                uint64_t lsn = tables_flushed(table_ids);
                if (options_.wal_enable_) {
                    wal_mgr_->RemoveOldFiles(lsn);
                }
            }
            break;
        }
    }

    return status;
}

void
DBImpl::BackgroundWalTask() {
    server::SystemInfo::GetInstance().Init();

2001
    std::chrono::system_clock::time_point next_auto_flush_time;
2002
    auto get_next_auto_flush_time = [&]() {
2003
        return std::chrono::system_clock::now() + std::chrono::seconds(options_.auto_flush_interval_);
2004
    };
2005 2006 2007
    if (options_.auto_flush_interval_ > 0) {
        next_auto_flush_time = get_next_auto_flush_time();
    }
2008 2009 2010 2011 2012 2013 2014 2015 2016 2017 2018 2019 2020 2021

    wal::MXLogRecord record;

    auto auto_flush = [&]() {
        record.type = wal::MXLogType::Flush;
        record.table_id.clear();
        ExecWalRecord(record);

        StartMetricTask();
        StartMergeTask();
        StartBuildIndexTask();
    };

    while (true) {
2022 2023 2024 2025 2026
        if (options_.auto_flush_interval_ > 0) {
            if (std::chrono::system_clock::now() >= next_auto_flush_time) {
                auto_flush();
                next_auto_flush_time = get_next_auto_flush_time();
            }
2027 2028 2029 2030 2031 2032 2033 2034 2035 2036 2037 2038 2039 2040 2041
        }

        auto error_code = wal_mgr_->GetNextRecord(record);
        if (error_code != WAL_SUCCESS) {
            ENGINE_LOG_ERROR << "WAL background GetNextRecord error";
            break;
        }

        if (record.type != wal::MXLogType::None) {
            ExecWalRecord(record);
            if (record.type == wal::MXLogType::Flush) {
                // user req flush
                flush_task_swn_.Notify();

                // if user flush all manually, update auto flush also
2042
                if (record.table_id.empty() && options_.auto_flush_interval_ > 0) {
2043 2044 2045 2046 2047 2048 2049 2050 2051 2052 2053 2054 2055
                    next_auto_flush_time = get_next_auto_flush_time();
                }
            }

        } else {
            if (!initialized_.load(std::memory_order_acquire)) {
                auto_flush();
                WaitMergeFileFinish();
                WaitBuildIndexFinish();
                ENGINE_LOG_DEBUG << "WAL background thread exit";
                break;
            }

2056 2057 2058 2059 2060
            if (options_.auto_flush_interval_ > 0) {
                bg_task_swn_.Wait_Until(next_auto_flush_time);
            } else {
                bg_task_swn_.Wait();
            }
2061 2062 2063 2064
        }
    }
}

2065 2066 2067 2068 2069
void
DBImpl::OnCacheInsertDataChanged(bool value) {
    options_.insert_cache_immediately_ = value;
}

S
starlord 已提交
2070 2071
}  // namespace engine
}  // namespace milvus