DBImpl.cpp 104.5 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>
24
#include <mutex>
25
#include <queue>
Z
Zhiru Zhu 已提交
26 27
#include <set>
#include <thread>
28
#include <unordered_map>
Z
Zhiru Zhu 已提交
29 30
#include <utility>

S
starlord 已提交
31
#include "Utils.h"
S
starlord 已提交
32 33
#include "cache/CpuCacheMgr.h"
#include "cache/GpuCacheMgr.h"
34
#include "db/IDGenerator.h"
G
groot 已提交
35
#include "db/merge/MergeManagerFactory.h"
S
starlord 已提交
36
#include "engine/EngineFactory.h"
37
#include "index/knowhere/knowhere/index/vector_index/helpers/BuilderSuspend.h"
38
#include "index/thirdparty/faiss/utils/distances.h"
39
#include "insert/MemManagerFactory.h"
S
starlord 已提交
40
#include "meta/MetaConsts.h"
S
starlord 已提交
41 42
#include "meta/MetaFactory.h"
#include "meta/SqliteMetaImpl.h"
G
groot 已提交
43
#include "metrics/Metrics.h"
G
groot 已提交
44
#include "scheduler/Definition.h"
S
starlord 已提交
45
#include "scheduler/SchedInst.h"
Y
Yu Kun 已提交
46
#include "scheduler/job/BuildIndexJob.h"
S
starlord 已提交
47 48
#include "scheduler/job/DeleteJob.h"
#include "scheduler/job/SearchJob.h"
49 50 51
#include "segment/SegmentReader.h"
#include "segment/SegmentWriter.h"
#include "utils/Exception.h"
S
starlord 已提交
52
#include "utils/Log.h"
G
groot 已提交
53
#include "utils/StringHelpFunctions.h"
S
starlord 已提交
54
#include "utils/TimeRecorder.h"
55 56
#include "utils/ValidationUtil.h"
#include "wal/WalDefinations.h"
X
Xu Peng 已提交
57

58 59
#include "search/TaskInst.h"

J
jinhai 已提交
60
namespace milvus {
X
Xu Peng 已提交
61
namespace engine {
X
Xu Peng 已提交
62

G
groot 已提交
63
namespace {
G
groot 已提交
64 65
constexpr uint64_t BACKGROUND_METRIC_INTERVAL = 1;
constexpr uint64_t BACKGROUND_INDEX_INTERVAL = 1;
G
groot 已提交
66
constexpr uint64_t WAIT_BUILD_INDEX_INTERVAL = 5;
G
groot 已提交
67

68 69 70 71 72 73 74 75
constexpr const char* JSON_ROW_COUNT = "row_count";
constexpr const char* JSON_PARTITIONS = "partitions";
constexpr const char* JSON_PARTITION_TAG = "tag";
constexpr const char* JSON_SEGMENTS = "segments";
constexpr const char* JSON_SEGMENT_NAME = "name";
constexpr const char* JSON_INDEX_NAME = "index_name";
constexpr const char* JSON_DATA_SIZE = "data_size";

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

S
starlord 已提交
78
}  // namespace
G
groot 已提交
79

Y
Yu Kun 已提交
80
DBImpl::DBImpl(const DBOptions& options)
81
    : options_(options), initialized_(false), merge_thread_pool_(1, 1), index_thread_pool_(1, 1) {
S
starlord 已提交
82
    meta_ptr_ = MetaFactory::Build(options.meta_, options.mode_);
Z
zhiru 已提交
83
    mem_mgr_ = MemManagerFactory::Build(meta_ptr_, options_);
G
groot 已提交
84
    merge_mgr_ptr_ = MergeManagerFactory::Build(meta_ptr_, options_);
85 86 87 88 89 90 91 92 93 94

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

95 96
    SetIdentity("DBImpl");
    AddCacheInsertDataListener();
97
    AddUseBlasThresholdListener();
98

S
starlord 已提交
99 100 101 102 103 104 105
    Start();
}

DBImpl::~DBImpl() {
    Stop();
}

S
starlord 已提交
106
///////////////////////////////////////////////////////////////////////////////////////////////////////////////////
S
starlord 已提交
107
// external api
S
starlord 已提交
108
///////////////////////////////////////////////////////////////////////////////////////////////////////////////////
S
starlord 已提交
109 110
Status
DBImpl::Start() {
111
    if (initialized_.load(std::memory_order_acquire)) {
S
starlord 已提交
112 113 114
        return Status::OK();
    }

115
    // LOG_ENGINE_TRACE_ << "DB service start";
116
    initialized_.store(true, std::memory_order_release);
S
starlord 已提交
117

G
groot 已提交
118 119 120 121 122 123 124 125 126 127
    // server may be closed unexpected, these un-merge files need to be merged when server restart
    // and soft-delete files need to be deleted when server restart
    std::set<std::string> merge_collection_ids;
    std::vector<meta::CollectionSchema> collection_schema_array;
    meta_ptr_->AllCollections(collection_schema_array);
    for (auto& schema : collection_schema_array) {
        merge_collection_ids.insert(schema.collection_id_);
    }
    StartMergeTask(merge_collection_ids, true);

128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152
    // 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) {
G
groot 已提交
153 154
            // background wal thread
            bg_wal_thread_ = std::thread(&DBImpl::BackgroundWalThread, this);
155 156 157 158
        }
    } else {
        // for distribute version, some nodes are read only
        if (options_.mode_ != DBOptions::MODE::CLUSTER_READONLY) {
G
groot 已提交
159 160
            // background flush thread
            bg_flush_thread_ = std::thread(&DBImpl::BackgroundFlushThread, this);
161
        }
Z
update  
zhiru 已提交
162
    }
S
starlord 已提交
163

G
groot 已提交
164 165 166 167 168 169 170
    // for distribute version, some nodes are read only
    if (options_.mode_ != DBOptions::MODE::CLUSTER_READONLY) {
        // background build index thread
        bg_index_thread_ = std::thread(&DBImpl::BackgroundIndexThread, this);
    }

    // background metric thread
G
groot 已提交
171 172 173
    if (options_.metric_enable_) {
        bg_metric_thread_ = std::thread(&DBImpl::BackgroundMetricThread, this);
    }
G
groot 已提交
174

S
starlord 已提交
175 176 177
    return Status::OK();
}

S
starlord 已提交
178 179
Status
DBImpl::Stop() {
180
    if (!initialized_.load(std::memory_order_acquire)) {
S
starlord 已提交
181 182
        return Status::OK();
    }
183

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

186 187
    if (options_.mode_ != DBOptions::MODE::CLUSTER_READONLY) {
        if (options_.wal_enable_) {
G
groot 已提交
188 189
            // wait wal thread finish
            swn_wal_.Notify();
190 191
            bg_wal_thread_.join();
        } else {
G
groot 已提交
192
            // flush all without merge
193 194 195 196
            wal::MXLogRecord record;
            record.type = wal::MXLogType::Flush;
            ExecWalRecord(record);

G
groot 已提交
197 198 199
            // wait flush thread finish
            swn_flush_.Notify();
            bg_flush_thread_.join();
200
        }
S
starlord 已提交
201

202 203
        WaitMergeFileFinish();

G
groot 已提交
204 205 206
        swn_index_.Notify();
        bg_index_thread_.join();

207
        meta_ptr_->CleanUpShadowFiles();
S
starlord 已提交
208 209
    }

G
groot 已提交
210
    // wait metric thread exit
G
groot 已提交
211 212 213 214
    if (options_.metric_enable_) {
        swn_metric_.Notify();
        bg_metric_thread_.join();
    }
G
groot 已提交
215

216
    // LOG_ENGINE_TRACE_ << "DB service stop";
S
starlord 已提交
217
    return Status::OK();
X
Xu Peng 已提交
218 219
}

S
starlord 已提交
220 221
Status
DBImpl::DropAll() {
S
starlord 已提交
222 223 224
    return meta_ptr_->DropAll();
}

S
starlord 已提交
225
Status
226
DBImpl::CreateCollection(meta::CollectionSchema& collection_schema) {
227
    if (!initialized_.load(std::memory_order_acquire)) {
G
groot 已提交
228
        return SHUTDOWN_ERROR;
S
starlord 已提交
229 230
    }

231
    meta::CollectionSchema temp_schema = collection_schema;
B
bigbraver 已提交
232
    temp_schema.index_file_size_ *= MB;  // store as MB
233
    if (options_.wal_enable_) {
234
        temp_schema.flush_lsn_ = wal_mgr_->CreateCollection(collection_schema.collection_id_);
235 236
    }

237
    return meta_ptr_->CreateCollection(temp_schema);
238 239
}

240 241 242 243 244 245 246
Status
DBImpl::CreateHybridCollection(meta::CollectionSchema& collection_schema, meta::hybrid::FieldsSchema& fields_schema) {
    if (!initialized_.load(std::memory_order_acquire)) {
        return SHUTDOWN_ERROR;
    }

    meta::CollectionSchema temp_schema = collection_schema;
Y
yukun 已提交
247
    temp_schema.index_file_size_ *= MB;
248
    if (options_.wal_enable_) {
Y
yukun 已提交
249
        temp_schema.flush_lsn_ = wal_mgr_->CreateHybridCollection(collection_schema.collection_id_);
250 251 252 253 254 255 256 257 258 259 260 261 262 263 264 265
    }

    return meta_ptr_->CreateHybridCollection(temp_schema, fields_schema);
}

Status
DBImpl::DescribeHybridCollection(meta::CollectionSchema& collection_schema,
                                 milvus::engine::meta::hybrid::FieldsSchema& fields_schema) {
    if (!initialized_.load(std::memory_order_acquire)) {
        return SHUTDOWN_ERROR;
    }

    auto stat = meta_ptr_->DescribeHybridCollection(collection_schema, fields_schema);
    return stat;
}

S
starlord 已提交
266
Status
267
DBImpl::DropCollection(const std::string& collection_id) {
268
    if (!initialized_.load(std::memory_order_acquire)) {
G
groot 已提交
269
        return SHUTDOWN_ERROR;
S
starlord 已提交
270 271
    }

272 273 274 275
    // dates partly delete files of the collection but currently we don't support
    LOG_ENGINE_DEBUG_ << "Prepare to delete collection " << collection_id;

    Status status;
276
    if (options_.wal_enable_) {
277
        wal_mgr_->DropCollection(collection_id);
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
    status = mem_mgr_->EraseMemVector(collection_id);      // not allow insert
    status = meta_ptr_->DropCollections({collection_id});  // soft delete collection
    index_failed_checker_.CleanFailedIndexFileOfCollection(collection_id);

    std::vector<meta::CollectionSchema> partition_array;
    status = meta_ptr_->ShowPartitions(collection_id, partition_array);
    std::vector<std::string> partition_id_array;
    for (auto& schema : partition_array) {
        if (options_.wal_enable_) {
            wal_mgr_->DropCollection(schema.collection_id_);
        }
        status = mem_mgr_->EraseMemVector(schema.collection_id_);
        index_failed_checker_.CleanFailedIndexFileOfCollection(schema.collection_id_);
        partition_id_array.push_back(schema.collection_id_);
    }

    status = meta_ptr_->DropCollections(partition_id_array);
    fiu_do_on("DBImpl.DropCollection.failed", status = Status(DB_ERROR, ""));
    if (!status.ok()) {
        return status;
    }

    return Status::OK();
G
groot 已提交
303 304
}

S
starlord 已提交
305
Status
306
DBImpl::DescribeCollection(meta::CollectionSchema& collection_schema) {
307
    if (!initialized_.load(std::memory_order_acquire)) {
G
groot 已提交
308
        return SHUTDOWN_ERROR;
S
starlord 已提交
309 310
    }

311
    auto stat = meta_ptr_->DescribeCollection(collection_schema);
B
bigbraver 已提交
312
    collection_schema.index_file_size_ /= MB;  // return as MB
S
starlord 已提交
313
    return stat;
314 315
}

S
starlord 已提交
316
Status
317
DBImpl::HasCollection(const std::string& collection_id, bool& has_or_not) {
318
    if (!initialized_.load(std::memory_order_acquire)) {
G
groot 已提交
319
        return SHUTDOWN_ERROR;
S
starlord 已提交
320 321
    }

G
groot 已提交
322
    return meta_ptr_->HasCollection(collection_id, has_or_not, false);
323 324
}

325
Status
G
groot 已提交
326
DBImpl::HasNativeCollection(const std::string& collection_id, bool& has_or_not) {
327 328 329 330
    if (!initialized_.load(std::memory_order_acquire)) {
        return SHUTDOWN_ERROR;
    }

G
groot 已提交
331
    return meta_ptr_->HasCollection(collection_id, has_or_not, true);
332 333
}

S
starlord 已提交
334
Status
335
DBImpl::AllCollections(std::vector<meta::CollectionSchema>& collection_schema_array) {
336
    if (!initialized_.load(std::memory_order_acquire)) {
G
groot 已提交
337
        return SHUTDOWN_ERROR;
S
starlord 已提交
338 339
    }

340 341
    std::vector<meta::CollectionSchema> all_collections;
    auto status = meta_ptr_->AllCollections(all_collections);
342

343 344 345 346 347
    // only return real collections, dont return partition collections
    collection_schema_array.clear();
    for (auto& schema : all_collections) {
        if (schema.owner_collection_.empty()) {
            collection_schema_array.push_back(schema);
348 349 350 351
        }
    }

    return status;
G
groot 已提交
352 353
}

354
Status
355
DBImpl::GetCollectionInfo(const std::string& collection_id, std::string& collection_info) {
356 357 358 359 360
    if (!initialized_.load(std::memory_order_acquire)) {
        return SHUTDOWN_ERROR;
    }

    // step1: get all partition ids
J
Jin Hai 已提交
361 362
    std::vector<meta::CollectionSchema> partition_array;
    auto status = meta_ptr_->ShowPartitions(collection_id, partition_array);
363

J
Jin Hai 已提交
364 365
    std::vector<int> file_types{meta::SegmentSchema::FILE_TYPE::RAW, meta::SegmentSchema::FILE_TYPE::TO_INDEX,
                                meta::SegmentSchema::FILE_TYPE::INDEX};
366

367 368 369
    milvus::json json_info;
    milvus::json json_partitions;
    size_t total_row_count = 0;
370

371
    auto get_info = [&](const std::string& col_id, const std::string& tag) {
G
groot 已提交
372 373
        meta::FilesHolder files_holder;
        status = meta_ptr_->FilesByType(col_id, file_types, files_holder);
374
        if (!status.ok()) {
J
Jin Hai 已提交
375
            std::string err_msg = "Failed to get collection info: " + status.ToString();
376
            LOG_ENGINE_ERROR_ << err_msg;
377 378 379
            return Status(DB_ERROR, err_msg);
        }

380 381 382 383 384
        milvus::json json_partition;
        json_partition[JSON_PARTITION_TAG] = tag;

        milvus::json json_segments;
        size_t row_count = 0;
G
groot 已提交
385
        milvus::engine::meta::SegmentsSchema& collection_files = files_holder.HoldFiles();
386
        for (auto& file : collection_files) {
387 388 389 390 391 392 393 394 395
            milvus::json json_segment;
            json_segment[JSON_SEGMENT_NAME] = file.segment_id_;
            json_segment[JSON_ROW_COUNT] = file.row_count_;
            json_segment[JSON_INDEX_NAME] = utils::GetIndexName(file.engine_type_);
            json_segment[JSON_DATA_SIZE] = (int64_t)file.file_size_;
            json_segments.push_back(json_segment);

            row_count += file.row_count_;
            total_row_count += file.row_count_;
396 397
        }

398 399 400 401 402 403 404
        json_partition[JSON_ROW_COUNT] = row_count;
        json_partition[JSON_SEGMENTS] = json_segments;

        json_partitions.push_back(json_partition);

        return Status::OK();
    };
405

406 407 408 409
    // step2: get default partition info
    status = get_info(collection_id, milvus::engine::DEFAULT_PARTITON_TAG);
    if (!status.ok()) {
        return status;
410 411
    }

412 413 414 415 416 417 418 419 420 421 422 423 424
    // step3: get partitions info
    for (auto& schema : partition_array) {
        status = get_info(schema.collection_id_, schema.partition_tag_);
        if (!status.ok()) {
            return status;
        }
    }

    json_info[JSON_ROW_COUNT] = total_row_count;
    json_info[JSON_PARTITIONS] = json_partitions;

    collection_info = json_info.dump();

425 426 427
    return Status::OK();
}

S
starlord 已提交
428
Status
G
groot 已提交
429 430
DBImpl::PreloadCollection(const std::shared_ptr<server::Context>& context, const std::string& collection_id,
                          bool force) {
431
    if (!initialized_.load(std::memory_order_acquire)) {
G
groot 已提交
432
        return SHUTDOWN_ERROR;
S
starlord 已提交
433 434
    }

J
Jin Hai 已提交
435
    // step 1: get all collection files from parent collection
G
groot 已提交
436
    meta::FilesHolder files_holder;
G
groot 已提交
437
#if 0
G
groot 已提交
438
    auto status = meta_ptr_->FilesToSearch(collection_id, files_holder);
Y
Yu Kun 已提交
439 440 441
    if (!status.ok()) {
        return status;
    }
Y
Yu Kun 已提交
442

443
    // step 2: get files from partition collections
J
Jin Hai 已提交
444 445
    std::vector<meta::CollectionSchema> partition_array;
    status = meta_ptr_->ShowPartitions(collection_id, partition_array);
G
groot 已提交
446
    for (auto& schema : partition_array) {
G
groot 已提交
447
        status = meta_ptr_->FilesToSearch(schema.collection_id_, files_holder);
G
groot 已提交
448
    }
G
groot 已提交
449 450 451 452 453 454 455 456 457 458 459 460 461 462 463 464 465 466 467
#else
    auto status = meta_ptr_->FilesToSearch(collection_id, files_holder);
    if (!status.ok()) {
        return status;
    }

    std::vector<meta::CollectionSchema> partition_array;
    status = meta_ptr_->ShowPartitions(collection_id, partition_array);

    std::set<std::string> partition_ids;
    for (auto& schema : partition_array) {
        partition_ids.insert(schema.collection_id_);
    }

    status = meta_ptr_->FilesToSearchEx(collection_id, partition_ids, files_holder);
    if (!status.ok()) {
        return status;
    }
#endif
G
groot 已提交
468

Y
Yu Kun 已提交
469 470
    int64_t size = 0;
    int64_t cache_total = cache::CpuCacheMgr::GetInstance()->CacheCapacity();
Y
Yu Kun 已提交
471 472
    int64_t cache_usage = cache::CpuCacheMgr::GetInstance()->CacheUsage();
    int64_t available_size = cache_total - cache_usage;
Y
Yu Kun 已提交
473

474
    // step 3: load file one by one
G
groot 已提交
475
    milvus::engine::meta::SegmentsSchema& files_array = files_holder.HoldFiles();
476 477
    LOG_ENGINE_DEBUG_ << "Begin pre-load collection:" + collection_id + ", totally " << files_array.size()
                      << " files need to be pre-loaded";
J
Jin Hai 已提交
478
    TimeRecorderAuto rc("Pre-load collection:" + collection_id);
G
groot 已提交
479
    for (auto& file : files_array) {
G
groot 已提交
480 481 482 483 484 485
        // client break the connection, no need to continue
        if (context && context->IsConnectionBroken()) {
            LOG_ENGINE_DEBUG_ << "Client connection broken, stop load collection";
            break;
        }

486
        EngineType engine_type;
J
Jin Hai 已提交
487 488 489
        if (file.file_type_ == meta::SegmentSchema::FILE_TYPE::RAW ||
            file.file_type_ == meta::SegmentSchema::FILE_TYPE::TO_INDEX ||
            file.file_type_ == meta::SegmentSchema::FILE_TYPE::BACKUP) {
490 491
            engine_type =
                utils::IsBinaryMetricType(file.metric_type_) ? EngineType::FAISS_BIN_IDMAP : EngineType::FAISS_IDMAP;
492 493 494
        } else {
            engine_type = (EngineType)file.engine_type_;
        }
495 496 497 498

        auto json = milvus::json::parse(file.index_params_);
        ExecutionEnginePtr engine =
            EngineFactory::Build(file.dimension_, file.location_, engine_type, (MetricType)file.metric_type_, json);
499
        fiu_do_on("DBImpl.PreloadCollection.null_engine", engine = nullptr);
G
groot 已提交
500
        if (engine == nullptr) {
501
            LOG_ENGINE_ERROR_ << "Invalid engine type";
G
groot 已提交
502 503
            return Status(DB_ERROR, "Invalid engine type");
        }
Y
Yu Kun 已提交
504

505
        fiu_do_on("DBImpl.PreloadCollection.exceed_cache", size = available_size + 1);
506 507

        try {
508
            fiu_do_on("DBImpl.PreloadCollection.engine_throw_exception", throw std::exception());
509 510
            std::string msg = "Pre-loaded file: " + file.file_id_ + " size: " + std::to_string(file.file_size_);
            TimeRecorderAuto rc_1(msg);
511 512 513 514
            status = engine->Load(true);
            if (!status.ok()) {
                return status;
            }
515 516

            size += engine->Size();
G
groot 已提交
517
            if (!force && size > available_size) {
518
                LOG_ENGINE_DEBUG_ << "Pre-load cancelled since cache is almost full";
519
                return Status(SERVER_CACHE_FULL, "Cache is full");
Y
Yu Kun 已提交
520
            }
521
        } catch (std::exception& ex) {
J
Jin Hai 已提交
522
            std::string msg = "Pre-load collection encounter exception: " + std::string(ex.what());
523
            LOG_ENGINE_ERROR_ << msg;
524
            return Status(DB_ERROR, msg);
Y
Yu Kun 已提交
525 526
        }
    }
G
groot 已提交
527

Y
Yu Kun 已提交
528
    return Status::OK();
Y
Yu Kun 已提交
529 530
}

531 532 533 534 535 536 537 538 539 540 541 542 543 544 545 546 547 548 549 550 551 552 553 554 555 556 557 558 559 560 561 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
Status
DBImpl::ReLoadSegmentsDeletedDocs(const std::string& collection_id, const std::vector<int64_t>& segment_ids) {
    if (!initialized_.load(std::memory_order_acquire)) {
        return SHUTDOWN_ERROR;
    }

    meta::FilesHolder files_holder;
    std::vector<size_t> file_ids;
    for (auto& id : segment_ids) {
        file_ids.emplace_back(id);
    }

    auto status = meta_ptr_->FilesByID(file_ids, files_holder);
    if (!status.ok()) {
        std::string err_msg = "Failed get file holders by ids: " + status.ToString();
        LOG_ENGINE_ERROR_ << err_msg;
        return Status(DB_ERROR, err_msg);
    }

    milvus::engine::meta::SegmentsSchema hold_files = files_holder.HoldFiles();

    for (auto& file : hold_files) {
        std::string segment_dir;
        utils::GetParentPath(file.location_, segment_dir);

        auto data_obj_ptr = cache::CpuCacheMgr::GetInstance()->GetIndex(file.location_);
        auto index = std::static_pointer_cast<knowhere::VecIndex>(data_obj_ptr);
        if (nullptr == index) {
            LOG_ENGINE_WARNING_ << "Index " << file.location_ << " not found";
            continue;
        }

        segment::SegmentReader segment_reader(segment_dir);

        segment::DeletedDocsPtr delete_docs = std::make_shared<segment::DeletedDocs>();
        segment_reader.LoadDeletedDocs(delete_docs);
        auto& docs_offsets = delete_docs->GetDeletedDocs();

        faiss::ConcurrentBitsetPtr blacklist = index->GetBlacklist();
        if (nullptr == blacklist) {
            LOG_ENGINE_WARNING_ << "Index " << file.location_ << " is empty";
            faiss::ConcurrentBitsetPtr concurrent_bitset_ptr =
                std::make_shared<faiss::ConcurrentBitset>(index->Count());
            index->SetBlacklist(concurrent_bitset_ptr);
            blacklist = concurrent_bitset_ptr;
        }

        for (auto& i : docs_offsets) {
            if (!blacklist->test(i)) {
                blacklist->set(i);
            }
        }
    }

    return Status::OK();
}

S
starlord 已提交
588
Status
589
DBImpl::UpdateCollectionFlag(const std::string& collection_id, int64_t flag) {
590
    if (!initialized_.load(std::memory_order_acquire)) {
G
groot 已提交
591
        return SHUTDOWN_ERROR;
S
starlord 已提交
592 593
    }

594
    return meta_ptr_->UpdateCollectionFlag(collection_id, flag);
S
starlord 已提交
595 596
}

S
starlord 已提交
597
Status
598
DBImpl::GetCollectionRowCount(const std::string& collection_id, uint64_t& row_count) {
599
    if (!initialized_.load(std::memory_order_acquire)) {
G
groot 已提交
600 601 602
        return SHUTDOWN_ERROR;
    }

603
    return GetCollectionRowCountRecursively(collection_id, row_count);
G
groot 已提交
604 605 606
}

Status
J
Jin Hai 已提交
607
DBImpl::CreatePartition(const std::string& collection_id, const std::string& partition_name,
G
groot 已提交
608
                        const std::string& partition_tag) {
609
    if (!initialized_.load(std::memory_order_acquire)) {
G
groot 已提交
610 611 612
        return SHUTDOWN_ERROR;
    }

613
    uint64_t lsn = 0;
G
groot 已提交
614 615 616 617 618
    if (options_.wal_enable_) {
        lsn = wal_mgr_->CreatePartition(collection_id, partition_tag);
    } else {
        meta_ptr_->GetCollectionFlushLSN(collection_id, lsn);
    }
J
Jin Hai 已提交
619
    return meta_ptr_->CreatePartition(collection_id, partition_name, partition_tag, lsn);
G
groot 已提交
620 621
}

G
groot 已提交
622 623 624 625 626 627 628 629 630 631 632 633 634 635 636 637 638 639 640
Status
DBImpl::HasPartition(const std::string& collection_id, const std::string& tag, bool& has_or_not) {
    if (!initialized_.load(std::memory_order_acquire)) {
        return SHUTDOWN_ERROR;
    }

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

    if (valid_tag == milvus::engine::DEFAULT_PARTITON_TAG) {
        has_or_not = true;
        return Status::OK();
    }

    return meta_ptr_->HasPartition(collection_id, valid_tag, has_or_not);
}

G
groot 已提交
641 642
Status
DBImpl::DropPartition(const std::string& partition_name) {
643
    if (!initialized_.load(std::memory_order_acquire)) {
G
groot 已提交
644
        return SHUTDOWN_ERROR;
S
starlord 已提交
645 646
    }

647
    mem_mgr_->EraseMemVector(partition_name);                // not allow insert
J
Jin Hai 已提交
648
    auto status = meta_ptr_->DropPartition(partition_name);  // soft delete collection
649
    if (!status.ok()) {
650
        LOG_ENGINE_ERROR_ << status.message();
651 652
        return status;
    }
G
groot 已提交
653

J
Jin Hai 已提交
654
    // scheduler will determine when to delete collection files
G
groot 已提交
655 656 657 658 659 660
    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 已提交
661 662
}

S
starlord 已提交
663
Status
J
Jin Hai 已提交
664
DBImpl::DropPartitionByTag(const std::string& collection_id, const std::string& partition_tag) {
665
    if (!initialized_.load(std::memory_order_acquire)) {
G
groot 已提交
666 667 668 669
        return SHUTDOWN_ERROR;
    }

    std::string partition_name;
J
Jin Hai 已提交
670
    auto status = meta_ptr_->GetPartitionName(collection_id, partition_tag, partition_name);
671
    if (!status.ok()) {
672
        LOG_ENGINE_ERROR_ << status.message();
673 674 675
        return status;
    }

G
groot 已提交
676 677 678 679
    if (options_.wal_enable_) {
        wal_mgr_->DropPartition(collection_id, partition_tag);
    }

G
groot 已提交
680 681 682 683
    return DropPartition(partition_name);
}

Status
J
Jin Hai 已提交
684
DBImpl::ShowPartitions(const std::string& collection_id, std::vector<meta::CollectionSchema>& partition_schema_array) {
685
    if (!initialized_.load(std::memory_order_acquire)) {
G
groot 已提交
686 687 688
        return SHUTDOWN_ERROR;
    }

J
Jin Hai 已提交
689
    return meta_ptr_->ShowPartitions(collection_id, partition_schema_array);
G
groot 已提交
690 691 692
}

Status
J
Jin Hai 已提交
693
DBImpl::InsertVectors(const std::string& collection_id, const std::string& partition_tag, VectorsData& vectors) {
694
    //    LOG_ENGINE_DEBUG_ << "Insert " << n << " vectors to cache";
695
    if (!initialized_.load(std::memory_order_acquire)) {
G
groot 已提交
696
        return SHUTDOWN_ERROR;
S
starlord 已提交
697
    }
Y
yu yunfeng 已提交
698

J
Jin Hai 已提交
699
    // insert vectors into target collection
700 701
    // (zhiru): generate ids
    if (vectors.id_array_.empty()) {
J
Jin Hai 已提交
702 703 704
        SafeIDGenerator& id_generator = SafeIDGenerator::GetInstance();
        Status status = id_generator.GetNextIDNumbers(vectors.vector_count_, vectors.id_array_);
        if (!status.ok()) {
705
            LOG_ENGINE_ERROR_ << LogOut("[%s][%ld] Get next id number fail: %s", "insert", 0, status.message().c_str());
J
Jin Hai 已提交
706 707
            return status;
        }
708 709
    }

710
    Status status;
711
    if (options_.wal_enable_) {
712 713
        std::string target_collection_name;
        status = GetPartitionByTag(collection_id, partition_tag, target_collection_name);
G
groot 已提交
714
        if (!status.ok()) {
715
            LOG_ENGINE_ERROR_ << LogOut("[%s][%ld] Get partition fail: %s", "insert", 0, status.message().c_str());
G
groot 已提交
716 717
            return status;
        }
718 719

        if (!vectors.float_data_.empty()) {
J
Jin Hai 已提交
720
            wal_mgr_->Insert(collection_id, partition_tag, vectors.id_array_, vectors.float_data_);
721
        } else if (!vectors.binary_data_.empty()) {
J
Jin Hai 已提交
722
            wal_mgr_->Insert(collection_id, partition_tag, vectors.id_array_, vectors.binary_data_);
723
        }
G
groot 已提交
724
        swn_wal_.Notify();
725 726 727
    } else {
        wal::MXLogRecord record;
        record.lsn = 0;  // need to get from meta ?
J
Jin Hai 已提交
728
        record.collection_id = collection_id;
729 730 731 732 733 734 735 736 737 738 739 740 741 742 743 744
        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 已提交
745 746
    }

747 748 749
    return status;
}

750
Status
Y
yukun 已提交
751 752 753 754 755
CopyToAttr(std::vector<uint8_t>& record, uint64_t row_num, const std::vector<std::string>& field_names,
           std::unordered_map<std::string, meta::hybrid::DataType>& attr_types,
           std::unordered_map<std::string, std::vector<uint8_t>>& attr_datas,
           std::unordered_map<std::string, uint64_t>& attr_nbytes,
           std::unordered_map<std::string, uint64_t>& attr_data_size) {
756
    uint64_t offset = 0;
Y
yukun 已提交
757 758
    for (auto name : field_names) {
        switch (attr_types.at(name)) {
759 760
            case meta::hybrid::DataType::INT8: {
                std::vector<uint8_t> data;
Y
yukun 已提交
761
                data.resize(row_num * sizeof(int8_t));
762

Y
yukun 已提交
763 764
                std::vector<int64_t> attr_value(row_num, 0);
                memcpy(attr_value.data(), record.data() + offset, row_num * sizeof(int64_t));
765

Y
yukun 已提交
766 767
                std::vector<int8_t> raw_value(row_num, 0);
                for (uint64_t i = 0; i < row_num; ++i) {
768 769 770
                    raw_value[i] = attr_value[i];
                }

Y
yukun 已提交
771 772
                memcpy(data.data(), raw_value.data(), row_num * sizeof(int8_t));
                attr_datas.insert(std::make_pair(name, data));
773

Y
yukun 已提交
774 775 776
                attr_nbytes.insert(std::make_pair(name, sizeof(int8_t)));
                attr_data_size.insert(std::make_pair(name, row_num * sizeof(int8_t)));
                offset += row_num * sizeof(int64_t);
777 778 779 780
                break;
            }
            case meta::hybrid::DataType::INT16: {
                std::vector<uint8_t> data;
Y
yukun 已提交
781
                data.resize(row_num * sizeof(int16_t));
782

Y
yukun 已提交
783 784
                std::vector<int64_t> attr_value(row_num, 0);
                memcpy(attr_value.data(), record.data() + offset, row_num * sizeof(int64_t));
785

Y
yukun 已提交
786 787
                std::vector<int16_t> raw_value(row_num, 0);
                for (uint64_t i = 0; i < row_num; ++i) {
788 789 790
                    raw_value[i] = attr_value[i];
                }

Y
yukun 已提交
791 792
                memcpy(data.data(), raw_value.data(), row_num * sizeof(int16_t));
                attr_datas.insert(std::make_pair(name, data));
793

Y
yukun 已提交
794 795 796
                attr_nbytes.insert(std::make_pair(name, sizeof(int16_t)));
                attr_data_size.insert(std::make_pair(name, row_num * sizeof(int16_t)));
                offset += row_num * sizeof(int64_t);
797 798 799 800
                break;
            }
            case meta::hybrid::DataType::INT32: {
                std::vector<uint8_t> data;
Y
yukun 已提交
801
                data.resize(row_num * sizeof(int32_t));
802

Y
yukun 已提交
803 804
                std::vector<int64_t> attr_value(row_num, 0);
                memcpy(attr_value.data(), record.data() + offset, row_num * sizeof(int64_t));
805

Y
yukun 已提交
806 807
                std::vector<int32_t> raw_value(row_num, 0);
                for (uint64_t i = 0; i < row_num; ++i) {
808 809 810
                    raw_value[i] = attr_value[i];
                }

Y
yukun 已提交
811 812
                memcpy(data.data(), raw_value.data(), row_num * sizeof(int32_t));
                attr_datas.insert(std::make_pair(name, data));
813

Y
yukun 已提交
814 815 816
                attr_nbytes.insert(std::make_pair(name, sizeof(int32_t)));
                attr_data_size.insert(std::make_pair(name, row_num * sizeof(int32_t)));
                offset += row_num * sizeof(int64_t);
817 818 819 820
                break;
            }
            case meta::hybrid::DataType::INT64: {
                std::vector<uint8_t> data;
Y
yukun 已提交
821 822 823 824 825 826
                data.resize(row_num * sizeof(int64_t));
                memcpy(data.data(), record.data() + offset, row_num * sizeof(int64_t));
                attr_datas.insert(std::make_pair(name, data));

                std::vector<int64_t> test_data(row_num);
                memcpy(test_data.data(), record.data(), row_num * sizeof(int64_t));
827

Y
yukun 已提交
828 829 830
                attr_nbytes.insert(std::make_pair(name, sizeof(int64_t)));
                attr_data_size.insert(std::make_pair(name, row_num * sizeof(int64_t)));
                offset += row_num * sizeof(int64_t);
831 832 833 834
                break;
            }
            case meta::hybrid::DataType::FLOAT: {
                std::vector<uint8_t> data;
Y
yukun 已提交
835
                data.resize(row_num * sizeof(float));
836

Y
yukun 已提交
837 838
                std::vector<double> attr_value(row_num, 0);
                memcpy(attr_value.data(), record.data() + offset, row_num * sizeof(double));
839

Y
yukun 已提交
840 841
                std::vector<float> raw_value(row_num, 0);
                for (uint64_t i = 0; i < row_num; ++i) {
842 843 844
                    raw_value[i] = attr_value[i];
                }

Y
yukun 已提交
845 846
                memcpy(data.data(), raw_value.data(), row_num * sizeof(float));
                attr_datas.insert(std::make_pair(name, data));
847

Y
yukun 已提交
848 849 850
                attr_nbytes.insert(std::make_pair(name, sizeof(float)));
                attr_data_size.insert(std::make_pair(name, row_num * sizeof(float)));
                offset += row_num * sizeof(double);
851 852 853 854
                break;
            }
            case meta::hybrid::DataType::DOUBLE: {
                std::vector<uint8_t> data;
Y
yukun 已提交
855 856 857
                data.resize(row_num * sizeof(double));
                memcpy(data.data(), record.data() + offset, row_num * sizeof(double));
                attr_datas.insert(std::make_pair(name, data));
858

Y
yukun 已提交
859 860 861
                attr_nbytes.insert(std::make_pair(name, sizeof(double)));
                attr_data_size.insert(std::make_pair(name, row_num * sizeof(double)));
                offset += row_num * sizeof(double);
862 863
                break;
            }
864 865
            default:
                break;
866 867
        }
    }
Y
yukun 已提交
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
    return Status::OK();
}

Status
DBImpl::InsertEntities(const std::string& collection_id, const std::string& partition_tag,
                       const std::vector<std::string>& field_names, Entity& entity,
                       std::unordered_map<std::string, meta::hybrid::DataType>& attr_types) {
    if (!initialized_.load(std::memory_order_acquire)) {
        return SHUTDOWN_ERROR;
    }

    // Generate id
    if (entity.id_array_.empty()) {
        SafeIDGenerator& id_generator = SafeIDGenerator::GetInstance();
        Status status = id_generator.GetNextIDNumbers(entity.entity_count_, entity.id_array_);
        if (!status.ok()) {
            return status;
        }
    }

    Status status;
    std::unordered_map<std::string, std::vector<uint8_t>> attr_data;
    std::unordered_map<std::string, uint64_t> attr_nbytes;
    std::unordered_map<std::string, uint64_t> attr_data_size;
    status = CopyToAttr(entity.attr_value_, entity.entity_count_, field_names, attr_types, attr_data, attr_nbytes,
                        attr_data_size);
    if (!status.ok()) {
        return status;
    }

    wal::MXLogRecord record;
    record.lsn = 0;
    record.collection_id = collection_id;
    record.partition_tag = partition_tag;
    record.ids = entity.id_array_.data();
    record.length = entity.entity_count_;

    auto vector_it = entity.vector_data_.begin();
    if (vector_it->second.binary_data_.empty()) {
        record.type = wal::MXLogType::Entity;
        record.data = vector_it->second.float_data_.data();
        record.data_size = vector_it->second.float_data_.size() * sizeof(float);
        record.attr_data = attr_data;
        record.attr_nbytes = attr_nbytes;
        record.attr_data_size = attr_data_size;
    } else {
        //        record.type = wal::MXLogType::InsertBinary;
        //        record.data = entities.vector_data_[0].binary_data_.data();
        //        record.length = entities.vector_data_[0].binary_data_.size() * sizeof(uint8_t);
    }
918 919

    status = ExecWalRecord(record);
Y
yukun 已提交
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 951 952 953 954 955 956 957 958 959 960 961 962 963 964 965

#if 0
    if (options_.wal_enable_) {
        std::string target_collection_name;
        status = GetPartitionByTag(collection_id, partition_tag, target_collection_name);
        if (!status.ok()) {
            LOG_ENGINE_ERROR_ << LogOut("[%s][%ld] Get partition fail: %s", "insert", 0, status.message().c_str());
            return status;
        }

        auto vector_it = entity.vector_data_.begin();
        if (!vector_it->second.binary_data_.empty()) {
            wal_mgr_->InsertEntities(collection_id, partition_tag, entity.id_array_, vector_it->second.binary_data_,
                                     attr_nbytes, attr_data);
        } else if (!vector_it->second.float_data_.empty()) {
            wal_mgr_->InsertEntities(collection_id, partition_tag, entity.id_array_, vector_it->second.float_data_,
                                     attr_nbytes, attr_data);
        }
        swn_wal_.Notify();
    } else {
        // insert entities: collection_name is field id
        wal::MXLogRecord record;
        record.lsn = 0;
        record.collection_id = collection_id;
        record.partition_tag = partition_tag;
        record.ids = entity.id_array_.data();
        record.length = entity.entity_count_;

        auto vector_it = entity.vector_data_.begin();
        if (vector_it->second.binary_data_.empty()) {
            record.type = wal::MXLogType::Entity;
            record.data = vector_it->second.float_data_.data();
            record.data_size = vector_it->second.float_data_.size() * sizeof(float);
            record.attr_data = attr_data;
            record.attr_nbytes = attr_nbytes;
            record.attr_data_size = attr_data_size;
        } else {
            //        record.type = wal::MXLogType::InsertBinary;
            //        record.data = entities.vector_data_[0].binary_data_.data();
            //        record.length = entities.vector_data_[0].binary_data_.size() * sizeof(uint8_t);
        }

        status = ExecWalRecord(record);
    }
#endif

966 967 968
    return status;
}

969
Status
J
Jin Hai 已提交
970
DBImpl::DeleteVector(const std::string& collection_id, IDNumber vector_id) {
971 972
    IDNumbers ids;
    ids.push_back(vector_id);
J
Jin Hai 已提交
973
    return DeleteVectors(collection_id, ids);
974 975 976
}

Status
J
Jin Hai 已提交
977
DBImpl::DeleteVectors(const std::string& collection_id, IDNumbers vector_ids) {
978 979 980 981 982 983
    if (!initialized_.load(std::memory_order_acquire)) {
        return SHUTDOWN_ERROR;
    }

    Status status;
    if (options_.wal_enable_) {
J
Jin Hai 已提交
984
        wal_mgr_->DeleteById(collection_id, vector_ids);
G
groot 已提交
985
        swn_wal_.Notify();
986 987 988 989
    } else {
        wal::MXLogRecord record;
        record.lsn = 0;  // need to get from meta ?
        record.type = wal::MXLogType::Delete;
J
Jin Hai 已提交
990
        record.collection_id = collection_id;
991 992 993 994 995 996 997 998 999 1000
        record.ids = vector_ids.data();
        record.length = vector_ids.size();

        status = ExecWalRecord(record);
    }

    return status;
}

Status
J
Jin Hai 已提交
1001
DBImpl::Flush(const std::string& collection_id) {
1002 1003 1004 1005 1006
    if (!initialized_.load(std::memory_order_acquire)) {
        return SHUTDOWN_ERROR;
    }

    Status status;
1007 1008
    bool has_collection;
    status = HasCollection(collection_id, has_collection);
1009 1010 1011
    if (!status.ok()) {
        return status;
    }
1012
    if (!has_collection) {
1013
        LOG_ENGINE_ERROR_ << "Collection to flush does not exist: " << collection_id;
J
Jin Hai 已提交
1014
        return Status(DB_NOT_FOUND, "Collection to flush does not exist");
1015 1016
    }

1017
    LOG_ENGINE_DEBUG_ << "Begin flush collection: " << collection_id;
1018 1019

    if (options_.wal_enable_) {
1020
        LOG_ENGINE_DEBUG_ << "WAL flush";
J
Jin Hai 已提交
1021
        auto lsn = wal_mgr_->Flush(collection_id);
1022
        if (lsn != 0) {
G
groot 已提交
1023 1024
            swn_wal_.Notify();
            flush_req_swn_.Wait();
1025 1026
        }
    } else {
1027
        LOG_ENGINE_DEBUG_ << "MemTable flush";
G
groot 已提交
1028
        InternalFlush(collection_id);
1029 1030
    }

1031
    LOG_ENGINE_DEBUG_ << "End flush collection: " << collection_id;
1032 1033 1034 1035 1036 1037 1038 1039 1040 1041

    return status;
}

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

1042
    LOG_ENGINE_DEBUG_ << "Begin flush all collections";
1043 1044 1045

    Status status;
    if (options_.wal_enable_) {
1046
        LOG_ENGINE_DEBUG_ << "WAL flush";
1047 1048
        auto lsn = wal_mgr_->Flush();
        if (lsn != 0) {
G
groot 已提交
1049 1050
            swn_wal_.Notify();
            flush_req_swn_.Wait();
1051 1052
        }
    } else {
1053
        LOG_ENGINE_DEBUG_ << "MemTable flush";
G
groot 已提交
1054
        InternalFlush();
1055 1056
    }

1057
    LOG_ENGINE_DEBUG_ << "End flush all collections";
1058 1059 1060 1061 1062

    return status;
}

Status
G
groot 已提交
1063
DBImpl::Compact(const std::shared_ptr<server::Context>& context, const std::string& collection_id, double threshold) {
1064 1065 1066 1067
    if (!initialized_.load(std::memory_order_acquire)) {
        return SHUTDOWN_ERROR;
    }

1068 1069 1070
    engine::meta::CollectionSchema collection_schema;
    collection_schema.collection_id_ = collection_id;
    auto status = DescribeCollection(collection_schema);
1071 1072
    if (!status.ok()) {
        if (status.code() == DB_NOT_FOUND) {
1073
            LOG_ENGINE_ERROR_ << "Collection to compact does not exist: " << collection_id;
J
Jin Hai 已提交
1074
            return Status(DB_NOT_FOUND, "Collection to compact does not exist");
1075 1076 1077 1078
        } else {
            return status;
        }
    } else {
1079
        if (!collection_schema.owner_collection_.empty()) {
1080
            LOG_ENGINE_ERROR_ << "Collection to compact does not exist: " << collection_id;
J
Jin Hai 已提交
1081
            return Status(DB_NOT_FOUND, "Collection to compact does not exist");
1082 1083 1084
        }
    }

1085
    LOG_ENGINE_DEBUG_ << "Before compacting, wait for build index thread to finish...";
1086

G
groot 已提交
1087 1088 1089
    std::vector<meta::CollectionSchema> collection_array;
    status = meta_ptr_->ShowPartitions(collection_id, collection_array);
    collection_array.push_back(collection_schema);
1090

Z
update  
Zhiru Zhu 已提交
1091
    const std::lock_guard<std::mutex> index_lock(build_index_mutex_);
Z
Zhiru Zhu 已提交
1092
    const std::lock_guard<std::mutex> merge_lock(flush_merge_compact_mutex_);
Z
Zhiru Zhu 已提交
1093

1094
    LOG_ENGINE_DEBUG_ << "Compacting collection: " << collection_id;
Z
Zhiru Zhu 已提交
1095

1096
    // Get files to compact from meta.
J
Jin Hai 已提交
1097 1098
    std::vector<int> file_types{meta::SegmentSchema::FILE_TYPE::RAW, meta::SegmentSchema::FILE_TYPE::TO_INDEX,
                                meta::SegmentSchema::FILE_TYPE::BACKUP};
G
groot 已提交
1099
    meta::FilesHolder files_holder;
G
groot 已提交
1100
    status = meta_ptr_->FilesByTypeEx(collection_array, file_types, files_holder);
1101 1102
    if (!status.ok()) {
        std::string err_msg = "Failed to get files to compact: " + status.message();
1103
        LOG_ENGINE_ERROR_ << err_msg;
1104 1105 1106
        return Status(DB_ERROR, err_msg);
    }

G
groot 已提交
1107
    LOG_ENGINE_DEBUG_ << "Found " << files_holder.HoldFiles().size() << " segment to compact";
Z
Zhiru Zhu 已提交
1108 1109

    Status compact_status;
G
groot 已提交
1110 1111
    // attention: here is a copy, not reference, since files_holder.UnmarkFile will change the array internal
    milvus::engine::meta::SegmentsSchema files_to_compact = files_holder.HoldFiles();
Z
Zhiru Zhu 已提交
1112
    for (auto iter = files_to_compact.begin(); iter != files_to_compact.end();) {
G
groot 已提交
1113 1114 1115 1116 1117 1118
        // client break the connection, no need to continue
        if (context && context->IsConnectionBroken()) {
            LOG_ENGINE_DEBUG_ << "Client connection broken, stop compact operation";
            break;
        }

J
Jin Hai 已提交
1119
        meta::SegmentSchema file = *iter;
G
groot 已提交
1120 1121
        iter = files_to_compact.erase(iter);

Z
Zhiru Zhu 已提交
1122 1123 1124
        // Check if the segment needs compacting
        std::string segment_dir;
        utils::GetParentPath(file.location_, segment_dir);
1125

Z
Zhiru Zhu 已提交
1126
        segment::SegmentReader segment_reader(segment_dir);
Z
Zhiru Zhu 已提交
1127 1128
        size_t deleted_docs_size;
        status = segment_reader.ReadDeletedDocsSize(deleted_docs_size);
Z
Zhiru Zhu 已提交
1129
        if (!status.ok()) {
G
groot 已提交
1130
            files_holder.UnmarkFile(file);
G
groot 已提交
1131
            continue;  // skip this file and try compact next one
Z
Zhiru Zhu 已提交
1132 1133
        }

J
Jin Hai 已提交
1134
        meta::SegmentsSchema files_to_update;
Z
Zhiru Zhu 已提交
1135
        if (deleted_docs_size != 0) {
G
groot 已提交
1136
            compact_status = CompactFile(file, threshold, files_to_update);
Z
Zhiru Zhu 已提交
1137 1138

            if (!compact_status.ok()) {
1139 1140
                LOG_ENGINE_ERROR_ << "Compact failed for segment " << file.segment_id_ << ": "
                                  << compact_status.message();
G
groot 已提交
1141
                files_holder.UnmarkFile(file);
G
groot 已提交
1142
                continue;  // skip this file and try compact next one
Z
Zhiru Zhu 已提交
1143 1144
            }
        } else {
G
groot 已提交
1145
            files_holder.UnmarkFile(file);
1146
            LOG_ENGINE_DEBUG_ << "Segment " << file.segment_id_ << " has no deleted data. No need to compact";
G
groot 已提交
1147
            continue;  // skip this file and try compact next one
1148
        }
Z
Zhiru Zhu 已提交
1149

1150
        LOG_ENGINE_DEBUG_ << "Updating meta after compaction...";
1151
        status = meta_ptr_->UpdateCollectionFiles(files_to_update);
G
groot 已提交
1152
        files_holder.UnmarkFile(file);
G
groot 已提交
1153 1154 1155 1156
        if (!status.ok()) {
            compact_status = status;
            break;  // meta error, could not go on
        }
Z
Zhiru Zhu 已提交
1157 1158
    }

G
groot 已提交
1159
    if (compact_status.ok()) {
1160
        LOG_ENGINE_DEBUG_ << "Finished compacting collection: " << collection_id;
G
groot 已提交
1161
    }
1162

G
groot 已提交
1163
    return compact_status;
1164 1165 1166
}

Status
G
groot 已提交
1167 1168
DBImpl::CompactFile(const meta::SegmentSchema& file, double threshold, meta::SegmentsSchema& files_to_update) {
    LOG_ENGINE_DEBUG_ << "Compacting segment " << file.segment_id_ << " for collection: " << file.collection_id_;
1169

G
groot 已提交
1170 1171 1172 1173 1174 1175 1176 1177 1178 1179
    std::string segment_dir_to_merge;
    utils::GetParentPath(file.location_, segment_dir_to_merge);

    // no need to compact if deleted vectors are too few(less than threashold)
    if (file.row_count_ > 0 && threshold > 0.0) {
        segment::SegmentReader segment_reader_to_merge(segment_dir_to_merge);
        segment::DeletedDocsPtr deleted_docs_ptr;
        auto status = segment_reader_to_merge.LoadDeletedDocs(deleted_docs_ptr);
        if (status.ok()) {
            auto delete_items = deleted_docs_ptr->GetDeletedDocs();
G
groot 已提交
1180
            double delete_rate = (double)delete_items.size() / (double)(delete_items.size() + file.row_count_);
G
groot 已提交
1181 1182 1183 1184 1185 1186 1187 1188
            if (delete_rate < threshold) {
                LOG_ENGINE_DEBUG_ << "Delete rate less than " << threshold << ", no need to compact for"
                                  << segment_dir_to_merge;
                return Status::OK();
            }
        }
    }

J
Jin Hai 已提交
1189 1190
    // Create new collection file
    meta::SegmentSchema compacted_file;
G
groot 已提交
1191
    compacted_file.collection_id_ = file.collection_id_;
J
Jin Hai 已提交
1192
    compacted_file.file_type_ = meta::SegmentSchema::NEW_MERGE;  // TODO: use NEW_MERGE for now
G
groot 已提交
1193
    auto status = meta_ptr_->CreateCollectionFile(compacted_file);
1194 1195

    if (!status.ok()) {
1196
        LOG_ENGINE_ERROR_ << "Failed to create collection file: " << status.message();
1197 1198 1199
        return status;
    }

J
Jin Hai 已提交
1200
    // Compact (merge) file to the newly created collection file
1201 1202 1203 1204
    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);

1205
    LOG_ENGINE_DEBUG_ << "Compacting begin...";
1206 1207 1208
    segment_writer_ptr->Merge(segment_dir_to_merge, compacted_file.file_id_);

    // Serialize
1209
    LOG_ENGINE_DEBUG_ << "Serializing compacted segment...";
1210 1211
    status = segment_writer_ptr->Serialize();
    if (!status.ok()) {
1212
        LOG_ENGINE_ERROR_ << "Failed to serialize compacted segment: " << status.message();
J
Jin Hai 已提交
1213
        compacted_file.file_type_ = meta::SegmentSchema::TO_DELETE;
1214
        auto mark_status = meta_ptr_->UpdateCollectionFile(compacted_file);
1215
        if (mark_status.ok()) {
1216
            LOG_ENGINE_DEBUG_ << "Mark file: " << compacted_file.file_id_ << " to to_delete";
1217
        }
G
groot 已提交
1218

1219 1220 1221
        return status;
    }

G
groot 已提交
1222
    // Update compacted file state, if origin file is backup or to_index, set compacted file to to_index
1223 1224 1225 1226 1227 1228
    compacted_file.file_size_ = segment_writer_ptr->Size();
    compacted_file.row_count_ = segment_writer_ptr->VectorCount();
    if ((file.file_type_ == (int32_t)meta::SegmentSchema::BACKUP ||
         file.file_type_ == (int32_t)meta::SegmentSchema::TO_INDEX) &&
        (compacted_file.row_count_ > meta::BUILD_INDEX_THRESHOLD)) {
        compacted_file.file_type_ = meta::SegmentSchema::TO_INDEX;
1229
    } else {
J
Jin Hai 已提交
1230
        compacted_file.file_type_ = meta::SegmentSchema::RAW;
1231 1232 1233
    }

    if (compacted_file.row_count_ == 0) {
1234
        LOG_ENGINE_DEBUG_ << "Compacted segment is empty. Mark it as TO_DELETE";
J
Jin Hai 已提交
1235
        compacted_file.file_type_ = meta::SegmentSchema::TO_DELETE;
1236 1237
    }

Z
Zhiru Zhu 已提交
1238
    files_to_update.emplace_back(compacted_file);
Z
Zhiru Zhu 已提交
1239

Z
Zhiru Zhu 已提交
1240 1241
    // Set all files in segment to TO_DELETE
    auto& segment_id = file.segment_id_;
G
groot 已提交
1242 1243
    meta::FilesHolder files_holder;
    status = meta_ptr_->GetCollectionFilesBySegmentId(segment_id, files_holder);
Z
Zhiru Zhu 已提交
1244 1245 1246
    if (!status.ok()) {
        return status;
    }
G
groot 已提交
1247 1248

    milvus::engine::meta::SegmentsSchema& segment_files = files_holder.HoldFiles();
Z
Zhiru Zhu 已提交
1249
    for (auto& f : segment_files) {
J
Jin Hai 已提交
1250
        f.file_type_ = meta::SegmentSchema::FILE_TYPE::TO_DELETE;
Z
Zhiru Zhu 已提交
1251 1252
        files_to_update.emplace_back(f);
    }
G
groot 已提交
1253
    files_holder.ReleaseFiles();
1254

1255 1256 1257
    LOG_ENGINE_DEBUG_ << "Compacted segment " << compacted_file.segment_id_ << " from "
                      << std::to_string(file.file_size_) << " bytes to " << std::to_string(compacted_file.file_size_)
                      << " bytes";
1258 1259 1260 1261 1262 1263 1264 1265 1266

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

    return status;
}

Status
G
groot 已提交
1267
DBImpl::GetVectorsByID(const engine::meta::CollectionSchema& collection, const IDNumbers& id_array,
1268
                       std::vector<engine::VectorsData>& vectors) {
1269 1270 1271 1272
    if (!initialized_.load(std::memory_order_acquire)) {
        return SHUTDOWN_ERROR;
    }

G
groot 已提交
1273
    meta::FilesHolder files_holder;
J
Jin Hai 已提交
1274 1275
    std::vector<int> file_types{meta::SegmentSchema::FILE_TYPE::RAW, meta::SegmentSchema::FILE_TYPE::TO_INDEX,
                                meta::SegmentSchema::FILE_TYPE::BACKUP};
1276

G
groot 已提交
1277 1278 1279 1280 1281
    std::vector<meta::CollectionSchema> collection_array;
    auto status = meta_ptr_->ShowPartitions(collection.collection_id_, collection_array);

    collection_array.push_back(collection);
    status = meta_ptr_->FilesByTypeEx(collection_array, file_types, files_holder);
1282
    if (!status.ok()) {
G
groot 已提交
1283
        std::string err_msg = "Failed to get files for GetVectorByID: " + status.message();
1284
        LOG_ENGINE_ERROR_ << err_msg;
1285 1286 1287
        return status;
    }

G
groot 已提交
1288
    if (files_holder.HoldFiles().empty()) {
1289
        LOG_ENGINE_DEBUG_ << "No files to get vector by id from";
J
Jin Hai 已提交
1290
        return Status(DB_NOT_FOUND, "Collection is empty");
1291 1292 1293
    }

    cache::CpuCacheMgr::GetInstance()->PrintInfo();
G
groot 已提交
1294
    status = GetVectorsByIdHelper(id_array, vectors, files_holder);
1295 1296
    cache::CpuCacheMgr::GetInstance()->PrintInfo();

G
groot 已提交
1297 1298 1299 1300 1301
    if (vectors.empty()) {
        std::string msg = "Vectors not found in collection " + collection.collection_id_;
        LOG_ENGINE_DEBUG_ << msg;
    }

1302 1303 1304 1305
    return status;
}

Status
J
Jin Hai 已提交
1306
DBImpl::GetVectorIDs(const std::string& collection_id, const std::string& segment_id, IDNumbers& vector_ids) {
1307 1308 1309 1310
    if (!initialized_.load(std::memory_order_acquire)) {
        return SHUTDOWN_ERROR;
    }

J
Jin Hai 已提交
1311
    // step 1: check collection existence
1312 1313 1314
    bool has_collection;
    auto status = HasCollection(collection_id, has_collection);
    if (!has_collection) {
1315
        LOG_ENGINE_ERROR_ << "Collection " << collection_id << " does not exist: ";
J
Jin Hai 已提交
1316
        return Status(DB_NOT_FOUND, "Collection does not exist");
1317 1318 1319 1320 1321 1322
    }
    if (!status.ok()) {
        return status;
    }

    //  step 2: find segment
G
groot 已提交
1323 1324
    meta::FilesHolder files_holder;
    status = meta_ptr_->GetCollectionFilesBySegmentId(segment_id, files_holder);
1325 1326 1327 1328
    if (!status.ok()) {
        return status;
    }

G
groot 已提交
1329
    milvus::engine::meta::SegmentsSchema& collection_files = files_holder.HoldFiles();
1330
    if (collection_files.empty()) {
1331 1332 1333
        return Status(DB_NOT_FOUND, "Segment does not exist");
    }

J
Jin Hai 已提交
1334
    // check the segment is belong to this collection
1335
    if (collection_files[0].collection_id_ != collection_id) {
J
Jin Hai 已提交
1336
        // the segment could be in a partition under this collection
1337 1338 1339 1340
        meta::CollectionSchema collection_schema;
        collection_schema.collection_id_ = collection_files[0].collection_id_;
        status = DescribeCollection(collection_schema);
        if (collection_schema.owner_collection_ != collection_id) {
J
Jin Hai 已提交
1341
            return Status(DB_NOT_FOUND, "Segment does not belong to this collection");
1342 1343 1344 1345 1346
        }
    }

    // step 3: load segment ids and delete offset
    std::string segment_dir;
1347
    engine::utils::GetParentPath(collection_files[0].location_, segment_dir);
1348 1349 1350 1351 1352 1353 1354 1355 1356 1357 1358 1359 1360 1361 1362 1363 1364 1365 1366 1367 1368 1369 1370 1371 1372
    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 已提交
1373

G
groot 已提交
1374
    return status;
X
Xu Peng 已提交
1375 1376
}

1377
Status
G
groot 已提交
1378 1379
DBImpl::GetVectorsByIdHelper(const IDNumbers& id_array, std::vector<engine::VectorsData>& vectors,
                             meta::FilesHolder& files_holder) {
G
groot 已提交
1380 1381
    // attention: this is a copy, not a reference, since the files_holder.UnMarkFile will change the array internal
    milvus::engine::meta::SegmentsSchema files = files_holder.HoldFiles();
1382
    LOG_ENGINE_DEBUG_ << "Getting vector by id in " << files.size() << " files, id count = " << id_array.size();
J
Jin Hai 已提交
1383

1384 1385 1386 1387 1388 1389 1390
    // sometimes not all of id_array can be found, we need to return empty vector for id not found
    // for example:
    // id_array = [1, -1, 2, -1, 3]
    // vectors should return [valid_vector, empty_vector, valid_vector, empty_vector, valid_vector]
    // the ID2RAW is to ensure returned vector sequence is consist with id_array
    using ID2VECTOR = std::map<int64_t, VectorsData>;
    ID2VECTOR map_id2vector;
1391

1392 1393 1394
    vectors.clear();

    IDNumbers temp_ids = id_array;
1395
    for (auto& file : files) {
G
groot 已提交
1396 1397 1398
        if (temp_ids.empty()) {
            break;  // all vectors found, no need to continue
        }
1399 1400 1401 1402 1403 1404 1405
        // 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);

1406 1407 1408 1409 1410 1411 1412 1413 1414 1415 1416
        for (IDNumbers::iterator it = temp_ids.begin(); it != temp_ids.end();) {
            int64_t vector_id = *it;
            // each id must has a VectorsData
            // if vector not found for an id, its VectorsData's vector_count = 0, else 1
            VectorsData& vector_ref = map_id2vector[vector_id];

            // 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<segment::doc_id_t> uids;
                auto status = segment_reader.LoadUids(uids);
1417 1418 1419
                if (!status.ok()) {
                    return status;
                }
1420 1421 1422 1423 1424 1425 1426 1427

                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);
1428
                    if (!status.ok()) {
J
Jin Hai 已提交
1429
                        LOG_ENGINE_ERROR_ << status.message();
1430 1431
                        return status;
                    }
1432 1433 1434 1435 1436 1437 1438 1439 1440 1441 1442 1443 1444 1445 1446 1447 1448 1449 1450 1451 1452 1453 1454 1455 1456 1457
                    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
                        bool is_binary = utils::IsBinaryMetricType(file.metric_type_);
                        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()) {
                            LOG_ENGINE_ERROR_ << status.message();
                            return status;
                        }

                        vector_ref.vector_count_ = 1;
                        if (is_binary) {
                            vector_ref.binary_data_.swap(raw_vector);
                        } else {
                            std::vector<float> float_vector;
                            float_vector.resize(file.dimension_);
                            memcpy(float_vector.data(), raw_vector.data(), single_vector_bytes);
                            vector_ref.float_data_.swap(float_vector);
                        }
                        temp_ids.erase(it);
                        continue;
1458 1459 1460
                    }
                }
            }
1461 1462 1463

            it++;
        }
G
groot 已提交
1464 1465 1466

        // unmark file, allow the file to be deleted
        files_holder.UnmarkFile(file);
1467 1468 1469 1470 1471 1472 1473 1474
    }

    for (auto id : id_array) {
        VectorsData& vector_ref = map_id2vector[id];

        VectorsData data;
        data.vector_count_ = vector_ref.vector_count_;
        if (data.vector_count_ > 0) {
G
groot 已提交
1475 1476
            data.float_data_ = vector_ref.float_data_;    // copy data since there could be duplicated id
            data.binary_data_ = vector_ref.binary_data_;  // copy data since there could be duplicated id
1477
        }
1478
        vectors.emplace_back(data);
1479 1480 1481 1482 1483
    }

    return Status::OK();
}

S
starlord 已提交
1484
Status
G
groot 已提交
1485 1486
DBImpl::CreateIndex(const std::shared_ptr<server::Context>& context, const std::string& collection_id,
                    const CollectionIndex& index) {
1487
    if (!initialized_.load(std::memory_order_acquire)) {
G
groot 已提交
1488 1489 1490
        return SHUTDOWN_ERROR;
    }

G
groot 已提交
1491
    // step 1: wait merge file thread finished to avoid duplicate data bug
1492
    auto status = Flush();
G
groot 已提交
1493
    WaitMergeFileFinish();  // let merge file thread finish
G
groot 已提交
1494 1495 1496 1497 1498 1499 1500 1501

    // merge all files for this collection, including its partitions
    std::set<std::string> merge_collection_ids = {collection_id};
    std::vector<meta::CollectionSchema> partition_array;
    status = meta_ptr_->ShowPartitions(collection_id, partition_array);
    for (auto& schema : partition_array) {
        merge_collection_ids.insert(schema.collection_id_);
    }
G
groot 已提交
1502 1503
    StartMergeTask(merge_collection_ids, true);  // start force-merge task
    WaitMergeFileFinish();                       // let force-merge file thread finish
G
groot 已提交
1504

S
starlord 已提交
1505 1506 1507
    {
        std::unique_lock<std::mutex> lock(build_index_mutex_);

G
groot 已提交
1508
        // step 2: check index difference
1509
        CollectionIndex old_index;
J
Jin Hai 已提交
1510
        status = DescribeIndex(collection_id, old_index);
S
starlord 已提交
1511
        if (!status.ok()) {
1512
            LOG_ENGINE_ERROR_ << "Failed to get collection index info for collection: " << collection_id;
S
starlord 已提交
1513 1514 1515
            return status;
        }

G
groot 已提交
1516
        // step 3: update index info
1517 1518
        CollectionIndex new_index = index;
        new_index.metric_type_ = old_index.metric_type_;  // dont change metric type, it was defined by CreateCollection
S
starlord 已提交
1519
        if (!utils::IsSameIndex(old_index, new_index)) {
1520
            status = UpdateCollectionIndexRecursively(collection_id, new_index);
S
starlord 已提交
1521 1522 1523 1524 1525 1526
            if (!status.ok()) {
                return status;
            }
        }
    }

S
starlord 已提交
1527
    // step 4: wait and build index
1528
    status = index_failed_checker_.CleanFailedIndexFileOfCollection(collection_id);
G
groot 已提交
1529
    status = WaitCollectionIndexRecursively(context, collection_id, index);
S
starlord 已提交
1530

G
groot 已提交
1531
    return status;
S
starlord 已提交
1532 1533
}

S
starlord 已提交
1534
Status
1535
DBImpl::DescribeIndex(const std::string& collection_id, CollectionIndex& index) {
1536
    if (!initialized_.load(std::memory_order_acquire)) {
G
groot 已提交
1537 1538 1539
        return SHUTDOWN_ERROR;
    }

1540
    return meta_ptr_->DescribeCollectionIndex(collection_id, index);
S
starlord 已提交
1541 1542
}

S
starlord 已提交
1543
Status
J
Jin Hai 已提交
1544
DBImpl::DropIndex(const std::string& collection_id) {
1545
    if (!initialized_.load(std::memory_order_acquire)) {
G
groot 已提交
1546 1547 1548
        return SHUTDOWN_ERROR;
    }

1549
    LOG_ENGINE_DEBUG_ << "Drop index for collection: " << collection_id;
G
groot 已提交
1550
    auto status = DropCollectionIndexRecursively(collection_id);
G
groot 已提交
1551 1552
    std::set<std::string> merge_collection_ids = {collection_id};
    StartMergeTask(merge_collection_ids, true);  // merge small files after drop index
G
groot 已提交
1553
    return status;
S
starlord 已提交
1554 1555
}

S
starlord 已提交
1556
Status
1557 1558 1559
DBImpl::QueryByIDs(const std::shared_ptr<server::Context>& context, const std::string& collection_id,
                   const std::vector<std::string>& partition_tags, uint64_t k, const milvus::json& extra_params,
                   const IDNumbers& id_array, ResultIds& result_ids, ResultDistances& result_distances) {
1560
    if (!initialized_.load(std::memory_order_acquire)) {
G
groot 已提交
1561
        return SHUTDOWN_ERROR;
S
starlord 已提交
1562 1563
    }

1564 1565 1566 1567 1568 1569 1570 1571 1572 1573 1574 1575 1576 1577 1578 1579 1580 1581 1582 1583 1584 1585 1586 1587 1588 1589 1590 1591 1592 1593
    if (id_array.empty()) {
        return Status(DB_ERROR, "Empty id array during query by id");
    }

    TimeRecorder rc("Query by id in collection:" + collection_id);

    // get collection schema
    engine::meta::CollectionSchema collection_schema;
    collection_schema.collection_id_ = collection_id;
    auto status = DescribeCollection(collection_schema);
    if (!status.ok()) {
        if (status.code() == DB_NOT_FOUND) {
            std::string msg = "Collection to search does not exist: " + collection_id;
            LOG_ENGINE_ERROR_ << msg;
            return Status(DB_NOT_FOUND, msg);
        } else {
            return status;
        }
    } else {
        if (!collection_schema.owner_collection_.empty()) {
            std::string msg = "Collection to search does not exist: " + collection_id;
            LOG_ENGINE_ERROR_ << msg;
            return Status(DB_NOT_FOUND, msg);
        }
    }

    rc.RecordSection("get collection schema");

    // get target vectors data
    std::vector<milvus::engine::VectorsData> vectors;
G
groot 已提交
1594
    status = GetVectorsByID(collection_schema, id_array, vectors);
1595 1596 1597 1598 1599 1600 1601 1602 1603 1604 1605 1606 1607 1608 1609 1610 1611 1612 1613 1614 1615 1616 1617 1618 1619 1620 1621 1622 1623 1624 1625 1626 1627 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 1655 1656 1657 1658 1659 1660 1661 1662 1663 1664 1665 1666 1667 1668 1669 1670 1671 1672 1673 1674 1675 1676 1677 1678 1679 1680
    if (!status.ok()) {
        std::string msg = "Failed to get vector data for collection: " + collection_id;
        LOG_ENGINE_ERROR_ << msg;
        return status;
    }

    // some vectors could not be found, no need to search them
    uint64_t valid_count = 0;
    bool is_binary = utils::IsBinaryMetricType(collection_schema.metric_type_);
    for (auto& vector : vectors) {
        if (vector.vector_count_ > 0) {
            valid_count++;
        }
    }

    // copy valid vectors data for search input
    uint64_t dimension = collection_schema.dimension_;
    VectorsData valid_vectors;
    valid_vectors.vector_count_ = valid_count;
    if (is_binary) {
        valid_vectors.binary_data_.resize(valid_count * dimension / 8);
    } else {
        valid_vectors.float_data_.resize(valid_count * dimension * sizeof(float));
    }

    int64_t valid_index = 0;
    for (size_t i = 0; i < vectors.size(); i++) {
        if (vectors[i].vector_count_ == 0) {
            continue;
        }
        if (is_binary) {
            memcpy(valid_vectors.binary_data_.data() + valid_index * dimension / 8, vectors[i].binary_data_.data(),
                   vectors[i].binary_data_.size());
        } else {
            memcpy(valid_vectors.float_data_.data() + valid_index * dimension, vectors[i].float_data_.data(),
                   vectors[i].float_data_.size() * sizeof(float));
        }
        valid_index++;
    }

    rc.RecordSection("construct query input");

    // search valid vectors
    ResultIds valid_result_ids;
    ResultDistances valid_result_distances;
    status = Query(context, collection_id, partition_tags, k, extra_params, valid_vectors, valid_result_ids,
                   valid_result_distances);
    if (!status.ok()) {
        std::string msg = "Failed to query by id in collection " + collection_id + ", error: " + status.message();
        LOG_ENGINE_ERROR_ << msg;
        return status;
    }

    if (valid_result_ids.size() != valid_count * k || valid_result_distances.size() != valid_count * k) {
        std::string msg = "Failed to query by id in collection " + collection_id + ", result doesn't match id count";
        return Status(DB_ERROR, msg);
    }

    rc.RecordSection("query vealid vectors");

    // construct result
    if (valid_count == id_array.size()) {
        result_ids.swap(valid_result_ids);
        result_distances.swap(valid_result_distances);
    } else {
        result_ids.resize(vectors.size() * k);
        result_distances.resize(vectors.size() * k);
        int64_t valid_index = 0;
        for (uint64_t i = 0; i < vectors.size(); i++) {
            if (vectors[i].vector_count_ > 0) {
                memcpy(result_ids.data() + i * k, valid_result_ids.data() + valid_index * k, k * sizeof(int64_t));
                memcpy(result_distances.data() + i * k, valid_result_distances.data() + valid_index * k,
                       k * sizeof(float));
                valid_index++;
            } else {
                memset(result_ids.data() + i * k, -1, k * sizeof(int64_t));
                for (uint64_t j = i * k; j < i * k + k; j++) {
                    result_distances[j] = std::numeric_limits<float>::max();
                }
            }
        }
    }

    rc.RecordSection("construct result");

    return status;
X
Xu Peng 已提交
1681 1682
}

1683 1684 1685 1686 1687 1688 1689 1690 1691 1692 1693 1694 1695
Status
DBImpl::HybridQuery(const std::shared_ptr<server::Context>& context, const std::string& collection_id,
                    const std::vector<std::string>& partition_tags,
                    context::HybridSearchContextPtr hybrid_search_context, query::GeneralQueryPtr general_query,
                    std::unordered_map<std::string, engine::meta::hybrid::DataType>& attr_type, uint64_t& nq,
                    ResultIds& result_ids, ResultDistances& result_distances) {
    auto query_ctx = context->Child("Query");

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

    Status status;
G
groot 已提交
1696
    meta::FilesHolder files_holder;
1697 1698 1699
    if (partition_tags.empty()) {
        // no partition tag specified, means search in whole table
        // get all table files from parent table
G
groot 已提交
1700
        status = meta_ptr_->FilesToSearch(collection_id, files_holder);
1701 1702 1703 1704 1705 1706 1707 1708 1709 1710
        if (!status.ok()) {
            return status;
        }

        std::vector<meta::CollectionSchema> partition_array;
        status = meta_ptr_->ShowPartitions(collection_id, partition_array);
        if (!status.ok()) {
            return status;
        }
        for (auto& schema : partition_array) {
G
groot 已提交
1711
            status = meta_ptr_->FilesToSearch(schema.collection_id_, files_holder);
1712
            if (!status.ok()) {
G
groot 已提交
1713
                return Status(DB_ERROR, "get files to search failed in HybridQuery");
1714 1715 1716
            }
        }

G
groot 已提交
1717 1718
        if (files_holder.HoldFiles().empty()) {
            return Status::OK();  // no files to search
1719 1720 1721 1722 1723 1724 1725
        }
    } else {
        // get files from specified partitions
        std::set<std::string> partition_name_array;
        GetPartitionsByTags(collection_id, partition_tags, partition_name_array);

        for (auto& partition_name : partition_name_array) {
G
groot 已提交
1726
            status = meta_ptr_->FilesToSearch(partition_name, files_holder);
1727
            if (!status.ok()) {
G
groot 已提交
1728
                return Status(DB_ERROR, "get files to search failed in HybridQuery");
1729 1730 1731
            }
        }

G
groot 已提交
1732
        if (files_holder.HoldFiles().empty()) {
1733 1734 1735 1736 1737
            return Status::OK();
        }
    }

    cache::CpuCacheMgr::GetInstance()->PrintInfo();  // print cache info before query
G
groot 已提交
1738
    status = HybridQueryAsync(query_ctx, collection_id, files_holder, hybrid_search_context, general_query, attr_type,
1739 1740 1741 1742 1743 1744 1745 1746 1747 1748 1749
                              nq, result_ids, result_distances);
    if (!status.ok()) {
        return status;
    }
    cache::CpuCacheMgr::GetInstance()->PrintInfo();  // print cache info after query

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

    return status;
}

S
starlord 已提交
1750
Status
J
Jin Hai 已提交
1751
DBImpl::Query(const std::shared_ptr<server::Context>& context, const std::string& collection_id,
1752 1753
              const std::vector<std::string>& partition_tags, uint64_t k, const milvus::json& extra_params,
              const VectorsData& vectors, ResultIds& result_ids, ResultDistances& result_distances) {
1754
    milvus::server::ContextChild tracer(context, "Query");
Z
Zhiru Zhu 已提交
1755

1756
    if (!initialized_.load(std::memory_order_acquire)) {
G
groot 已提交
1757
        return SHUTDOWN_ERROR;
S
starlord 已提交
1758 1759
    }

G
groot 已提交
1760
    Status status;
G
groot 已提交
1761
    meta::FilesHolder files_holder;
G
groot 已提交
1762
    if (partition_tags.empty()) {
G
groot 已提交
1763
#if 0
J
Jin Hai 已提交
1764 1765
        // no partition tag specified, means search in whole collection
        // get all collection files from parent collection
G
groot 已提交
1766
        status = meta_ptr_->FilesToSearch(collection_id, files_holder);
G
groot 已提交
1767 1768 1769 1770
        if (!status.ok()) {
            return status;
        }

J
Jin Hai 已提交
1771 1772
        std::vector<meta::CollectionSchema> partition_array;
        status = meta_ptr_->ShowPartitions(collection_id, partition_array);
G
groot 已提交
1773
        for (auto& schema : partition_array) {
G
groot 已提交
1774
            status = meta_ptr_->FilesToSearch(schema.collection_id_, files_holder);
1775
        }
G
groot 已提交
1776 1777 1778 1779 1780 1781 1782 1783 1784 1785 1786 1787 1788 1789 1790 1791 1792 1793 1794 1795 1796
#else
        // no partition tag specified, means search in whole collection
        // get files from root collection
        status = meta_ptr_->FilesToSearch(collection_id, files_holder);
        if (!status.ok()) {
            return status;
        }

        // get files from partitions
        std::set<std::string> partition_ids;
        std::vector<meta::CollectionSchema> partition_array;
        status = meta_ptr_->ShowPartitions(collection_id, partition_array);
        for (auto& id : partition_array) {
            partition_ids.insert(id.collection_id_);
        }

        status = meta_ptr_->FilesToSearchEx(collection_id, partition_ids, files_holder);
        if (!status.ok()) {
            return status;
        }
#endif
1797

G
groot 已提交
1798 1799
        if (files_holder.HoldFiles().empty()) {
            return Status::OK();  // no files to search
G
groot 已提交
1800 1801
        }
    } else {
G
groot 已提交
1802
#if 0
G
groot 已提交
1803 1804
        // get files from specified partitions
        std::set<std::string> partition_name_array;
J
Jin Hai 已提交
1805
        status = GetPartitionsByTags(collection_id, partition_tags, partition_name_array);
T
Tinkerrr 已提交
1806 1807 1808
        if (!status.ok()) {
            return status;  // didn't match any partition.
        }
G
groot 已提交
1809 1810

        for (auto& partition_name : partition_name_array) {
G
groot 已提交
1811
            status = meta_ptr_->FilesToSearch(partition_name, files_holder);
1812
        }
G
groot 已提交
1813 1814 1815 1816 1817 1818
#else
        std::set<std::string> partition_name_array;
        status = GetPartitionsByTags(collection_id, partition_tags, partition_name_array);
        if (!status.ok()) {
            return status;  // didn't match any partition.
        }
1819

G
groot 已提交
1820 1821 1822 1823 1824 1825 1826
        std::set<std::string> partition_ids;
        for (auto& partition_name : partition_name_array) {
            partition_ids.insert(partition_name);
        }

        status = meta_ptr_->FilesToSearchEx(collection_id, partition_ids, files_holder);
#endif
G
groot 已提交
1827 1828
        if (files_holder.HoldFiles().empty()) {
            return Status::OK();  // no files to search
1829 1830 1831
        }
    }

S
starlord 已提交
1832
    cache::CpuCacheMgr::GetInstance()->PrintInfo();  // print cache info before query
G
groot 已提交
1833
    status = QueryAsync(tracer.Context(), files_holder, k, extra_params, vectors, result_ids, result_distances);
S
starlord 已提交
1834
    cache::CpuCacheMgr::GetInstance()->PrintInfo();  // print cache info after query
Z
Zhiru Zhu 已提交
1835

S
starlord 已提交
1836
    return status;
G
groot 已提交
1837
}
X
Xu Peng 已提交
1838

S
starlord 已提交
1839
Status
1840 1841 1842
DBImpl::QueryByFileID(const std::shared_ptr<server::Context>& context, const std::vector<std::string>& file_ids,
                      uint64_t k, const milvus::json& extra_params, const VectorsData& vectors, ResultIds& result_ids,
                      ResultDistances& result_distances) {
1843
    milvus::server::ContextChild tracer(context, "Query by file id");
Z
Zhiru Zhu 已提交
1844

1845
    if (!initialized_.load(std::memory_order_acquire)) {
G
groot 已提交
1846
        return SHUTDOWN_ERROR;
S
starlord 已提交
1847 1848
    }

S
starlord 已提交
1849
    // get specified files
1850
    std::vector<size_t> ids;
Y
Yu Kun 已提交
1851
    for (auto& id : file_ids) {
1852
        std::string::size_type sz;
J
jinhai 已提交
1853
        ids.push_back(std::stoul(id, &sz));
1854 1855
    }

G
groot 已提交
1856 1857
    meta::FilesHolder files_holder;
    auto status = meta_ptr_->FilesByID(ids, files_holder);
1858 1859
    if (!status.ok()) {
        return status;
1860 1861
    }

G
groot 已提交
1862
    milvus::engine::meta::SegmentsSchema& search_files = files_holder.HoldFiles();
1863
    if (search_files.empty()) {
S
starlord 已提交
1864
        return Status(DB_ERROR, "Invalid file id");
G
groot 已提交
1865 1866
    }

S
starlord 已提交
1867
    cache::CpuCacheMgr::GetInstance()->PrintInfo();  // print cache info before query
G
groot 已提交
1868
    status = QueryAsync(tracer.Context(), files_holder, k, extra_params, vectors, result_ids, result_distances);
S
starlord 已提交
1869
    cache::CpuCacheMgr::GetInstance()->PrintInfo();  // print cache info after query
Z
Zhiru Zhu 已提交
1870

S
starlord 已提交
1871
    return status;
1872 1873
}

S
starlord 已提交
1874
Status
Y
Yu Kun 已提交
1875
DBImpl::Size(uint64_t& result) {
1876
    if (!initialized_.load(std::memory_order_acquire)) {
G
groot 已提交
1877
        return SHUTDOWN_ERROR;
S
starlord 已提交
1878 1879
    }

S
starlord 已提交
1880
    return meta_ptr_->Size(result);
S
starlord 已提交
1881 1882 1883
}

///////////////////////////////////////////////////////////////////////////////////////////////////////////////////
S
starlord 已提交
1884
// internal methods
S
starlord 已提交
1885
///////////////////////////////////////////////////////////////////////////////////////////////////////////////////
S
starlord 已提交
1886
Status
G
groot 已提交
1887
DBImpl::QueryAsync(const std::shared_ptr<server::Context>& context, meta::FilesHolder& files_holder, uint64_t k,
1888 1889
                   const milvus::json& extra_params, const VectorsData& vectors, ResultIds& result_ids,
                   ResultDistances& result_distances) {
1890
    milvus::server::ContextChild tracer(context, "Query Async");
G
groot 已提交
1891
    server::CollectQueryMetrics metrics(vectors.vector_count_);
Y
Yu Kun 已提交
1892

G
groot 已提交
1893
    milvus::engine::meta::SegmentsSchema& files = files_holder.HoldFiles();
G
groot 已提交
1894 1895 1896
    if (files.size() > milvus::scheduler::TASK_TABLE_MAX_COUNT) {
        std::string msg =
            "Search files count exceed scheduler limit: " + std::to_string(milvus::scheduler::TASK_TABLE_MAX_COUNT);
1897
        LOG_ENGINE_ERROR_ << msg;
G
groot 已提交
1898 1899 1900
        return Status(DB_ERROR, msg);
    }

S
starlord 已提交
1901
    TimeRecorder rc("");
G
groot 已提交
1902

1903
    // step 1: construct search job
1904
    LOG_ENGINE_DEBUG_ << LogOut("Engine query begin, index file count: %ld", files.size());
1905
    scheduler::SearchJobPtr job = std::make_shared<scheduler::SearchJob>(tracer.Context(), k, extra_params, vectors);
Y
Yu Kun 已提交
1906
    for (auto& file : files) {
J
Jin Hai 已提交
1907
        scheduler::SegmentSchemaPtr file_ptr = std::make_shared<meta::SegmentSchema>(file);
W
wxyu 已提交
1908
        job->AddIndexFile(file_ptr);
G
groot 已提交
1909 1910
    }

1911 1912 1913
    // Suspend builder
    SuspendIfFirst();

1914
    // step 2: put search job to scheduler and wait result
S
starlord 已提交
1915
    scheduler::JobMgrInst::GetInstance()->Put(job);
W
wxyu 已提交
1916
    job->WaitResult();
1917

1918 1919 1920
    // Resume builder
    ResumeIfLast();

G
groot 已提交
1921
    files_holder.ReleaseFiles();
W
wxyu 已提交
1922 1923
    if (!job->GetStatus().ok()) {
        return job->GetStatus();
1924
    }
G
groot 已提交
1925

1926
    // step 3: construct results
G
groot 已提交
1927 1928
    result_ids = job->GetResultIds();
    result_distances = job->GetResultDistances();
S
starlord 已提交
1929
    rc.ElapseFromBegin("Engine query totally cost");
G
groot 已提交
1930 1931 1932 1933

    return Status::OK();
}

1934 1935
Status
DBImpl::HybridQueryAsync(const std::shared_ptr<server::Context>& context, const std::string& table_id,
G
groot 已提交
1936
                         meta::FilesHolder& files_holder, context::HybridSearchContextPtr hybrid_search_context,
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
                         query::GeneralQueryPtr general_query,
                         std::unordered_map<std::string, engine::meta::hybrid::DataType>& attr_type, uint64_t& nq,
                         ResultIds& result_ids, ResultDistances& result_distances) {
    auto query_async_ctx = context->Child("Query Async");

#if 0
    // Construct tasks
    for (auto file : files) {
        std::unordered_map<std::string, engine::DataType> types;
        auto it = attr_type.begin();
        for (; it != attr_type.end(); it++) {
            types.insert(std::make_pair(it->first, (engine::DataType)it->second));
        }

        auto file_ptr = std::make_shared<meta::TableFileSchema>(file);
        search::TaskPtr
            task = std::make_shared<search::Task>(context, file_ptr, general_query, types, hybrid_search_context);
        search::TaskInst::GetInstance().load_queue().push(task);
        search::TaskInst::GetInstance().load_cv().notify_one();
        hybrid_search_context->tasks_.emplace_back(task);
    }

#endif

    //#if 0
    TimeRecorder rc("");

    // step 1: construct search job
    VectorsData vectors;
G
groot 已提交
1966 1967
    milvus::engine::meta::SegmentsSchema& files = files_holder.HoldFiles();
    LOG_ENGINE_DEBUG_ << LogOut("Engine query begin, index file count: %ld", files_holder.HoldFiles().size());
1968 1969 1970 1971 1972 1973 1974 1975 1976 1977 1978
    scheduler::SearchJobPtr job =
        std::make_shared<scheduler::SearchJob>(query_async_ctx, general_query, attr_type, vectors);
    for (auto& file : files) {
        scheduler::SegmentSchemaPtr file_ptr = std::make_shared<meta::SegmentSchema>(file);
        job->AddIndexFile(file_ptr);
    }

    // step 2: put search job to scheduler and wait result
    scheduler::JobMgrInst::GetInstance()->Put(job);
    job->WaitResult();

G
groot 已提交
1979
    files_holder.ReleaseFiles();
1980 1981 1982 1983 1984 1985 1986 1987 1988 1989 1990 1991 1992 1993 1994 1995
    if (!job->GetStatus().ok()) {
        return job->GetStatus();
    }

    // step 3: construct results
    nq = job->vector_count();
    result_ids = job->GetResultIds();
    result_distances = job->GetResultDistances();
    rc.ElapseFromBegin("Engine query totally cost");

    query_async_ctx->GetTraceContext()->GetSpan()->Finish();
    //#endif

    return Status::OK();
}

S
starlord 已提交
1996
void
G
groot 已提交
1997
DBImpl::BackgroundIndexThread() {
G
groot 已提交
1998
    SetThreadName("index_thread");
Y
yu yunfeng 已提交
1999
    server::SystemInfo::GetInstance().Init();
X
Xu Peng 已提交
2000
    while (true) {
2001
        if (!initialized_.load(std::memory_order_acquire)) {
2002 2003
            WaitMergeFileFinish();
            WaitBuildIndexFinish();
S
starlord 已提交
2004

2005
            LOG_ENGINE_DEBUG_ << "DB background thread exit";
G
groot 已提交
2006 2007
            break;
        }
X
Xu Peng 已提交
2008

G
groot 已提交
2009
        swn_index_.Wait_For(std::chrono::seconds(BACKGROUND_INDEX_INTERVAL));
X
Xu Peng 已提交
2010

G
groot 已提交
2011
        WaitMergeFileFinish();
G
groot 已提交
2012 2013
        StartBuildIndexTask();
    }
X
Xu Peng 已提交
2014 2015
}

S
starlord 已提交
2016 2017
void
DBImpl::WaitMergeFileFinish() {
2018
    //    LOG_ENGINE_DEBUG_ << "Begin WaitMergeFileFinish";
2019 2020
    std::lock_guard<std::mutex> lck(merge_result_mutex_);
    for (auto& iter : merge_thread_results_) {
2021 2022
        iter.wait();
    }
2023
    //    LOG_ENGINE_DEBUG_ << "End WaitMergeFileFinish";
2024 2025
}

S
starlord 已提交
2026 2027
void
DBImpl::WaitBuildIndexFinish() {
2028
    //    LOG_ENGINE_DEBUG_ << "Begin WaitBuildIndexFinish";
2029
    std::lock_guard<std::mutex> lck(index_result_mutex_);
Y
Yu Kun 已提交
2030
    for (auto& iter : index_thread_results_) {
2031 2032
        iter.wait();
    }
2033
    //    LOG_ENGINE_DEBUG_ << "End WaitBuildIndexFinish";
2034 2035
}

S
starlord 已提交
2036 2037
void
DBImpl::StartMetricTask() {
G
groot 已提交
2038
    server::Metrics::GetInstance().KeepingAliveCounterIncrement(BACKGROUND_METRIC_INTERVAL);
G
groot 已提交
2039 2040
    int64_t cache_usage = cache::CpuCacheMgr::GetInstance()->CacheUsage();
    int64_t cache_total = cache::CpuCacheMgr::GetInstance()->CacheCapacity();
S
shengjh 已提交
2041 2042
    fiu_do_on("DBImpl.StartMetricTask.InvalidTotalCache", cache_total = 0);

J
JinHai-CN 已提交
2043 2044 2045 2046 2047 2048 2049
    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 已提交
2050
    server::Metrics::GetInstance().GpuCacheUsageGaugeSet();
G
groot 已提交
2051 2052 2053 2054 2055 2056 2057 2058
    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 已提交
2059

K
kun yu 已提交
2060
    server::Metrics::GetInstance().CPUCoreUsagePercentSet();
K
kun yu 已提交
2061 2062
    server::Metrics::GetInstance().GPUTemperature();
    server::Metrics::GetInstance().CPUTemperature();
2063
    server::Metrics::GetInstance().PushToGateway();
G
groot 已提交
2064 2065
}

S
starlord 已提交
2066
void
G
groot 已提交
2067
DBImpl::StartMergeTask(const std::set<std::string>& merge_collection_ids, bool force_merge_all) {
2068
    // LOG_ENGINE_DEBUG_ << "Begin StartMergeTask";
2069
    // merge task has been finished?
2070
    {
2071 2072
        std::lock_guard<std::mutex> lck(merge_result_mutex_);
        if (!merge_thread_results_.empty()) {
2073
            std::chrono::milliseconds span(10);
2074 2075
            if (merge_thread_results_.back().wait_for(span) == std::future_status::ready) {
                merge_thread_results_.pop_back();
2076
            }
G
groot 已提交
2077 2078
        }
    }
X
Xu Peng 已提交
2079

2080
    // add new merge task
2081
    {
2082 2083
        std::lock_guard<std::mutex> lck(merge_result_mutex_);
        if (merge_thread_results_.empty()) {
2084
            // start merge file thread
2085
            merge_thread_results_.push_back(
G
groot 已提交
2086
                merge_thread_pool_.enqueue(&DBImpl::BackgroundMerge, this, merge_collection_ids, force_merge_all));
2087
        }
G
groot 已提交
2088
    }
2089

2090
    // LOG_ENGINE_DEBUG_ << "End StartMergeTask";
X
Xu Peng 已提交
2091 2092
}

2093
Status
G
groot 已提交
2094
DBImpl::MergeHybridFiles(const std::string& collection_id, meta::FilesHolder& files_holder) {
2095 2096 2097 2098 2099 2100 2101 2102 2103 2104 2105 2106 2107 2108 2109 2110 2111 2112 2113 2114 2115 2116 2117 2118 2119 2120 2121
    // const std::lock_guard<std::mutex> lock(flush_merge_compact_mutex_);

    LOG_ENGINE_DEBUG_ << "Merge files for collection: " << collection_id;

    // step 1: create table file
    meta::SegmentSchema table_file;
    table_file.collection_id_ = collection_id;
    table_file.file_type_ = meta::SegmentSchema::NEW_MERGE;
    Status status = meta_ptr_->CreateHybridCollectionFile(table_file);

    if (!status.ok()) {
        LOG_ENGINE_ERROR_ << "Failed to create collection: " << status.ToString();
        return status;
    }

    // step 2: merge files
    /*
    ExecutionEnginePtr index =
        EngineFactory::Build(table_file.dimension_, table_file.location_, (EngineType)table_file.engine_type_,
                             (MetricType)table_file.metric_type_, table_file.nlist_);
*/
    meta::SegmentsSchema updated;

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

G
groot 已提交
2122 2123
    // attention: here is a copy, not reference, since files_holder.UnmarkFile will change the array internal
    milvus::engine::meta::SegmentsSchema files = files_holder.HoldFiles();
2124 2125 2126 2127 2128
    for (auto& file : files) {
        server::CollectMergeFilesMetrics metrics;
        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_);
G
groot 已提交
2129 2130 2131

        files_holder.UnmarkFile(file);

2132 2133 2134
        auto file_schema = file;
        file_schema.file_type_ = meta::SegmentSchema::TO_DELETE;
        updated.push_back(file_schema);
C
Cai Yudong 已提交
2135
        int64_t size = segment_writer_ptr->Size();
2136 2137 2138 2139 2140 2141 2142 2143 2144 2145 2146 2147 2148 2149 2150 2151 2152 2153 2154 2155 2156 2157 2158 2159 2160 2161 2162 2163 2164 2165 2166 2167
        if (size >= file_schema.index_file_size_) {
            break;
        }
    }

    // step 3: serialize to disk
    try {
        status = segment_writer_ptr->Serialize();
        fiu_do_on("DBImpl.MergeFiles.Serialize_ThrowException", throw std::exception());
        fiu_do_on("DBImpl.MergeFiles.Serialize_ErrorStatus", status = Status(DB_ERROR, ""));
    } catch (std::exception& ex) {
        std::string msg = "Serialize merged index encounter exception: " + std::string(ex.what());
        LOG_ENGINE_ERROR_ << msg;
        status = Status(DB_ERROR, msg);
    }

    if (!status.ok()) {
        LOG_ENGINE_ERROR_ << "Failed to persist merged segment: " << new_segment_dir << ". Error: " << status.message();

        // if failed to serialize merge file to disk
        // typical error: out of disk space, out of memory or permission denied
        table_file.file_type_ = meta::SegmentSchema::TO_DELETE;
        status = meta_ptr_->UpdateCollectionFile(table_file);
        LOG_ENGINE_DEBUG_ << "Failed to update file to index, mark file: " << table_file.file_id_ << " to to_delete";

        return status;
    }

    // step 4: 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
    if (!utils::IsRawIndexType(table_file.engine_type_)) {
C
Cai Yudong 已提交
2168
        table_file.file_type_ = (segment_writer_ptr->Size() >= (size_t)(table_file.index_file_size_))
2169 2170 2171 2172 2173 2174 2175 2176 2177 2178 2179 2180 2181 2182 2183 2184 2185 2186 2187
                                    ? meta::SegmentSchema::TO_INDEX
                                    : meta::SegmentSchema::RAW;
    } else {
        table_file.file_type_ = meta::SegmentSchema::RAW;
    }
    table_file.file_size_ = segment_writer_ptr->Size();
    table_file.row_count_ = segment_writer_ptr->VectorCount();
    updated.push_back(table_file);
    status = meta_ptr_->UpdateCollectionFiles(updated);
    LOG_ENGINE_DEBUG_ << "New merged segment " << table_file.segment_id_ << " of size " << segment_writer_ptr->Size()
                      << " bytes";

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

    return status;
}

S
starlord 已提交
2188
void
G
groot 已提交
2189
DBImpl::BackgroundMerge(std::set<std::string> collection_ids, bool force_merge_all) {
2190
    // LOG_ENGINE_TRACE_ << " Background merge thread start";
S
starlord 已提交
2191

G
groot 已提交
2192
    Status status;
2193
    for (auto& collection_id : collection_ids) {
G
groot 已提交
2194 2195
        const std::lock_guard<std::mutex> lock(flush_merge_compact_mutex_);

G
groot 已提交
2196 2197 2198 2199 2200
        auto old_strategy = merge_mgr_ptr_->Strategy();
        if (force_merge_all) {
            merge_mgr_ptr_->UseStrategy(MergeStrategyType::ADAPTIVE);
        }

G
groot 已提交
2201
        auto status = merge_mgr_ptr_->MergeFiles(collection_id);
G
groot 已提交
2202
        merge_mgr_ptr_->UseStrategy(old_strategy);
G
groot 已提交
2203
        if (!status.ok()) {
G
groot 已提交
2204 2205
            LOG_ENGINE_ERROR_ << "Failed to get merge files for collection: " << collection_id
                              << " reason:" << status.message();
G
groot 已提交
2206
        }
S
starlord 已提交
2207

2208
        if (!initialized_.load(std::memory_order_acquire)) {
G
groot 已提交
2209
            LOG_ENGINE_DEBUG_ << "Server will shutdown, skip merge action for collection: " << collection_id;
S
starlord 已提交
2210 2211
            break;
        }
G
groot 已提交
2212
    }
X
Xu Peng 已提交
2213

G
groot 已提交
2214
    //    meta_ptr_->Archive();
Z
update  
zhiru 已提交
2215

2216
    {
G
groot 已提交
2217
        uint64_t timeout = (options_.file_cleanup_timeout_ >= 0) ? options_.file_cleanup_timeout_ : 10;
G
groot 已提交
2218
        uint64_t ttl = timeout * meta::SECOND;  // default: file will be hard-deleted few seconds after soft-deleted
2219
        meta_ptr_->CleanUpFilesWithTTL(ttl);
Z
update  
zhiru 已提交
2220
    }
S
starlord 已提交
2221

2222
    // LOG_ENGINE_TRACE_ << " Background merge thread exit";
G
groot 已提交
2223
}
X
Xu Peng 已提交
2224

S
starlord 已提交
2225
void
G
groot 已提交
2226
DBImpl::StartBuildIndexTask() {
S
starlord 已提交
2227
    // build index has been finished?
2228 2229 2230 2231 2232 2233 2234
    {
        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 已提交
2235 2236 2237
        }
    }

S
starlord 已提交
2238
    // add new build index task
2239 2240 2241
    {
        std::lock_guard<std::mutex> lck(index_result_mutex_);
        if (index_thread_results_.empty()) {
S
starlord 已提交
2242
            index_thread_results_.push_back(index_thread_pool_.enqueue(&DBImpl::BackgroundBuildIndex, this));
2243
        }
G
groot 已提交
2244
    }
X
Xu Peng 已提交
2245 2246
}

S
starlord 已提交
2247 2248
void
DBImpl::BackgroundBuildIndex() {
P
peng.xu 已提交
2249
    std::unique_lock<std::mutex> lock(build_index_mutex_);
G
groot 已提交
2250 2251 2252
    meta::FilesHolder files_holder;
    meta_ptr_->FilesToIndex(files_holder);

G
groot 已提交
2253
    milvus::engine::meta::SegmentsSchema to_index_files = files_holder.HoldFiles();
2254
    Status status = index_failed_checker_.IgnoreFailedIndexFiles(to_index_files);
2255

2256
    if (!to_index_files.empty()) {
G
groot 已提交
2257
        LOG_ENGINE_DEBUG_ << "Background build index thread begin " << to_index_files.size() << " files";
2258

2259
        // step 2: put build index task to scheduler
J
Jin Hai 已提交
2260
        std::vector<std::pair<scheduler::BuildIndexJobPtr, scheduler::SegmentSchemaPtr>> job2file_map;
2261
        for (auto& file : to_index_files) {
G
groot 已提交
2262
            scheduler::BuildIndexJobPtr job = std::make_shared<scheduler::BuildIndexJob>(meta_ptr_, options_);
J
Jin Hai 已提交
2263
            scheduler::SegmentSchemaPtr file_ptr = std::make_shared<meta::SegmentSchema>(file);
2264
            job->AddToIndexFiles(file_ptr);
G
groot 已提交
2265
            scheduler::JobMgrInst::GetInstance()->Put(job);
G
groot 已提交
2266
            job2file_map.push_back(std::make_pair(job, file_ptr));
2267
        }
G
groot 已提交
2268

G
groot 已提交
2269
        // step 3: wait build index finished and mark failed files
2270
        int64_t completed = 0;
G
groot 已提交
2271 2272
        for (auto iter = job2file_map.begin(); iter != job2file_map.end(); ++iter) {
            scheduler::BuildIndexJobPtr job = iter->first;
J
Jin Hai 已提交
2273
            meta::SegmentSchema& file_schema = *(iter->second.get());
G
groot 已提交
2274
            job->WaitBuildIndexFinish();
2275
            LOG_ENGINE_INFO_ << "Build Index Progress: " << ++completed << " of " << job2file_map.size();
G
groot 已提交
2276 2277
            if (!job->GetStatus().ok()) {
                Status status = job->GetStatus();
2278
                LOG_ENGINE_ERROR_ << "Building index job " << job->id() << " failed: " << status.ToString();
G
groot 已提交
2279

2280
                index_failed_checker_.MarkFailedIndexFile(file_schema, status.message());
G
groot 已提交
2281
            } else {
2282
                LOG_ENGINE_DEBUG_ << "Building index job " << job->id() << " succeed.";
G
groot 已提交
2283 2284

                index_failed_checker_.MarkSucceedIndexFile(file_schema);
G
groot 已提交
2285
            }
G
groot 已提交
2286 2287
            status = files_holder.UnmarkFile(file_schema);
            LOG_ENGINE_DEBUG_ << "Finish build index file " << file_schema.file_id_;
2288
        }
G
groot 已提交
2289

2290
        LOG_ENGINE_DEBUG_ << "Background build index thread finished";
G
groot 已提交
2291
        index_req_swn_.Notify();  // notify CreateIndex check circle
Y
Yu Kun 已提交
2292
    }
X
Xu Peng 已提交
2293 2294
}

G
groot 已提交
2295
Status
J
Jin Hai 已提交
2296
DBImpl::GetFilesToBuildIndex(const std::string& collection_id, const std::vector<int>& file_types,
G
groot 已提交
2297 2298 2299
                             meta::FilesHolder& files_holder) {
    files_holder.ReleaseFiles();
    auto status = meta_ptr_->FilesByType(collection_id, file_types, files_holder);
G
groot 已提交
2300

G
groot 已提交
2301 2302 2303 2304 2305 2306 2307 2308 2309 2310 2311
    // attention: here is a copy, not reference, since files_holder.UnmarkFile will change the array internal
    milvus::engine::meta::SegmentsSchema files = files_holder.HoldFiles();
    for (const milvus::engine::meta::SegmentSchema& file : files) {
        if (file.file_type_ == static_cast<int>(meta::SegmentSchema::RAW) &&
            file.row_count_ < meta::BUILD_INDEX_THRESHOLD) {
            // skip build index for files that row count less than certain threshold
            files_holder.UnmarkFile(file);
        } else if (index_failed_checker_.IsFailedIndexFile(file)) {
            // skip build index for files that failed before
            files_holder.UnmarkFile(file);
        }
G
groot 已提交
2312 2313 2314 2315 2316
    }

    return Status::OK();
}

2317
Status
J
Jin Hai 已提交
2318 2319
DBImpl::GetPartitionByTag(const std::string& collection_id, const std::string& partition_tag,
                          std::string& partition_name) {
2320 2321 2322
    Status status;

    if (partition_tag.empty()) {
J
Jin Hai 已提交
2323
        partition_name = collection_id;
2324 2325 2326 2327 2328 2329 2330 2331

    } 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) {
J
Jin Hai 已提交
2332
            partition_name = collection_id;
2333 2334 2335
            return status;
        }

J
Jin Hai 已提交
2336
        status = meta_ptr_->GetPartitionName(collection_id, partition_tag, partition_name);
2337
        if (!status.ok()) {
2338
            LOG_ENGINE_ERROR_ << status.message();
2339 2340 2341 2342 2343 2344
        }
    }

    return status;
}

G
groot 已提交
2345
Status
J
Jin Hai 已提交
2346
DBImpl::GetPartitionsByTags(const std::string& collection_id, const std::vector<std::string>& partition_tags,
G
groot 已提交
2347
                            std::set<std::string>& partition_name_array) {
J
Jin Hai 已提交
2348 2349
    std::vector<meta::CollectionSchema> partition_array;
    auto status = meta_ptr_->ShowPartitions(collection_id, partition_array);
G
groot 已提交
2350 2351

    for (auto& tag : partition_tags) {
2352 2353 2354 2355
        // 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);
2356 2357

        if (valid_tag == milvus::engine::DEFAULT_PARTITON_TAG) {
J
Jin Hai 已提交
2358
            partition_name_array.insert(collection_id);
2359 2360 2361
            return status;
        }

G
groot 已提交
2362
        for (auto& schema : partition_array) {
2363
            if (server::StringHelpFunctions::IsRegexMatch(schema.partition_tag_, valid_tag)) {
J
Jin Hai 已提交
2364
                partition_name_array.insert(schema.collection_id_);
G
groot 已提交
2365 2366 2367 2368
            }
        }
    }

T
Tinkerrr 已提交
2369
    if (partition_name_array.empty()) {
G
groot 已提交
2370
        return Status(DB_PARTITION_NOT_FOUND, "The specified partiton does not exist");
T
Tinkerrr 已提交
2371 2372
    }

G
groot 已提交
2373 2374 2375 2376
    return Status::OK();
}

Status
2377
DBImpl::UpdateCollectionIndexRecursively(const std::string& collection_id, const CollectionIndex& index) {
J
Jin Hai 已提交
2378
    DropIndex(collection_id);
G
groot 已提交
2379

2380 2381
    auto status = meta_ptr_->UpdateCollectionIndex(collection_id, index);
    fiu_do_on("DBImpl.UpdateCollectionIndexRecursively.fail_update_collection_index",
S
shengjh 已提交
2382
              status = Status(DB_META_TRANSACTION_FAILED, ""));
G
groot 已提交
2383
    if (!status.ok()) {
2384
        LOG_ENGINE_ERROR_ << "Failed to update collection index info for collection: " << collection_id;
G
groot 已提交
2385 2386 2387
        return status;
    }

J
Jin Hai 已提交
2388 2389
    std::vector<meta::CollectionSchema> partition_array;
    status = meta_ptr_->ShowPartitions(collection_id, partition_array);
A
AzAz 已提交
2390 2391 2392
    if (!status.ok()) {
        return status;
    }
G
groot 已提交
2393
    for (auto& schema : partition_array) {
2394
        status = UpdateCollectionIndexRecursively(schema.collection_id_, index);
G
groot 已提交
2395 2396 2397 2398 2399 2400 2401 2402 2403
        if (!status.ok()) {
            return status;
        }
    }

    return Status::OK();
}

Status
G
groot 已提交
2404 2405
DBImpl::WaitCollectionIndexRecursively(const std::shared_ptr<server::Context>& context,
                                       const std::string& collection_id, const CollectionIndex& index) {
G
groot 已提交
2406 2407 2408
    // 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;
2409
    if (utils::IsRawIndexType(index.engine_type_)) {
G
groot 已提交
2410
        file_types = {
J
Jin Hai 已提交
2411 2412
            static_cast<int32_t>(meta::SegmentSchema::NEW),
            static_cast<int32_t>(meta::SegmentSchema::NEW_MERGE),
G
groot 已提交
2413 2414 2415
        };
    } else {
        file_types = {
J
Jin Hai 已提交
2416 2417 2418
            static_cast<int32_t>(meta::SegmentSchema::RAW),       static_cast<int32_t>(meta::SegmentSchema::NEW),
            static_cast<int32_t>(meta::SegmentSchema::NEW_MERGE), static_cast<int32_t>(meta::SegmentSchema::NEW_INDEX),
            static_cast<int32_t>(meta::SegmentSchema::TO_INDEX),
G
groot 已提交
2419 2420 2421 2422
        };
    }

    // get files to build index
G
groot 已提交
2423 2424 2425 2426
    {
        meta::FilesHolder files_holder;
        auto status = GetFilesToBuildIndex(collection_id, file_types, files_holder);
        int times = 1;
G
groot 已提交
2427
        uint64_t repeat = 0;
G
groot 已提交
2428
        while (!files_holder.HoldFiles().empty()) {
G
groot 已提交
2429 2430 2431 2432 2433 2434
            if (repeat % WAIT_BUILD_INDEX_INTERVAL == 0) {
                LOG_ENGINE_DEBUG_ << files_holder.HoldFiles().size() << " non-index files detected! Will build index "
                                  << times;
                if (!utils::IsRawIndexType(index.engine_type_)) {
                    status = meta_ptr_->UpdateCollectionFilesToIndex(collection_id);
                }
G
groot 已提交
2435
            }
G
groot 已提交
2436

2437
            auto ret = index_req_swn_.Wait_For(std::chrono::seconds(1));
G
groot 已提交
2438
            // client break the connection, no need to block, check every 1 second
G
groot 已提交
2439
            if (context && context->IsConnectionBroken()) {
G
groot 已提交
2440 2441 2442 2443
                LOG_ENGINE_DEBUG_ << "Client connection broken, build index in background";
                break;  // just break, not return, continue to update partitions files to to_index
            }

2444
            // check to_index files every 5 seconds or background index thread finished
G
groot 已提交
2445
            repeat++;
2446
            if ((ret == std::cv_status::no_timeout) || (repeat % WAIT_BUILD_INDEX_INTERVAL == 0)) {
G
groot 已提交
2447 2448 2449
                GetFilesToBuildIndex(collection_id, file_types, files_holder);
                ++times;
            }
G
groot 已提交
2450 2451 2452 2453
        }
    }

    // build index for partition
J
Jin Hai 已提交
2454
    std::vector<meta::CollectionSchema> partition_array;
G
groot 已提交
2455
    auto status = meta_ptr_->ShowPartitions(collection_id, partition_array);
G
groot 已提交
2456
    for (auto& schema : partition_array) {
G
groot 已提交
2457
        status = WaitCollectionIndexRecursively(context, schema.collection_id_, index);
2458
        fiu_do_on("DBImpl.WaitCollectionIndexRecursively.fail_build_collection_Index_for_partition",
S
shengjh 已提交
2459
                  status = Status(DB_ERROR, ""));
G
groot 已提交
2460 2461 2462 2463 2464
        if (!status.ok()) {
            return status;
        }
    }

G
groot 已提交
2465
    // failed to build index for some files, return error
2466
    std::string err_msg;
2467 2468
    index_failed_checker_.GetErrMsgForCollection(collection_id, err_msg);
    fiu_do_on("DBImpl.WaitCollectionIndexRecursively.not_empty_err_msg", err_msg.append("fiu"));
2469 2470
    if (!err_msg.empty()) {
        return Status(DB_ERROR, err_msg);
G
groot 已提交
2471 2472
    }

G
groot 已提交
2473 2474
    LOG_ENGINE_DEBUG_ << "WaitCollectionIndexRecursively finished";

G
groot 已提交
2475 2476 2477 2478
    return Status::OK();
}

Status
2479
DBImpl::DropCollectionIndexRecursively(const std::string& collection_id) {
2480
    LOG_ENGINE_DEBUG_ << "Drop index for collection: " << collection_id;
2481 2482
    index_failed_checker_.CleanFailedIndexFileOfCollection(collection_id);
    auto status = meta_ptr_->DropCollectionIndex(collection_id);
G
groot 已提交
2483 2484 2485 2486 2487
    if (!status.ok()) {
        return status;
    }

    // drop partition index
J
Jin Hai 已提交
2488 2489
    std::vector<meta::CollectionSchema> partition_array;
    status = meta_ptr_->ShowPartitions(collection_id, partition_array);
G
groot 已提交
2490
    for (auto& schema : partition_array) {
2491 2492
        status = DropCollectionIndexRecursively(schema.collection_id_);
        fiu_do_on("DBImpl.DropCollectionIndexRecursively.fail_drop_collection_Index_for_partition",
S
shengjh 已提交
2493
                  status = Status(DB_ERROR, ""));
G
groot 已提交
2494 2495 2496 2497 2498 2499 2500 2501 2502
        if (!status.ok()) {
            return status;
        }
    }

    return Status::OK();
}

Status
2503
DBImpl::GetCollectionRowCountRecursively(const std::string& collection_id, uint64_t& row_count) {
G
groot 已提交
2504
    row_count = 0;
J
Jin Hai 已提交
2505
    auto status = meta_ptr_->Count(collection_id, row_count);
G
groot 已提交
2506 2507 2508 2509 2510
    if (!status.ok()) {
        return status;
    }

    // get partition row count
J
Jin Hai 已提交
2511 2512
    std::vector<meta::CollectionSchema> partition_array;
    status = meta_ptr_->ShowPartitions(collection_id, partition_array);
G
groot 已提交
2513
    for (auto& schema : partition_array) {
G
groot 已提交
2514
        uint64_t partition_row_count = 0;
2515 2516
        status = GetCollectionRowCountRecursively(schema.collection_id_, partition_row_count);
        fiu_do_on("DBImpl.GetCollectionRowCountRecursively.fail_get_collection_rowcount_for_partition",
S
shengjh 已提交
2517
                  status = Status(DB_ERROR, ""));
G
groot 已提交
2518 2519 2520 2521 2522 2523 2524 2525 2526 2527
        if (!status.ok()) {
            return status;
        }

        row_count += partition_row_count;
    }

    return Status::OK();
}

2528 2529 2530 2531
Status
DBImpl::ExecWalRecord(const wal::MXLogRecord& record) {
    fiu_return_on("DBImpl.ExexWalRecord.return", Status(););

G
groot 已提交
2532 2533
    auto collections_flushed = [&](const std::string collection_id,
                                   const std::set<std::string>& target_collection_names) -> uint64_t {
2534 2535
        uint64_t max_lsn = 0;
        if (options_.wal_enable_) {
G
groot 已提交
2536 2537
            uint64_t lsn = 0;
            for (auto& collection : target_collection_names) {
2538
                meta_ptr_->GetCollectionFlushLSN(collection, lsn);
2539 2540 2541 2542
                if (lsn > max_lsn) {
                    max_lsn = lsn;
                }
            }
G
groot 已提交
2543
            wal_mgr_->CollectionFlushed(collection_id, lsn);
2544 2545
        }

G
groot 已提交
2546
        std::set<std::string> merge_collection_ids;
G
groot 已提交
2547
        for (auto& collection : target_collection_names) {
G
groot 已提交
2548
            merge_collection_ids.insert(collection);
2549
        }
G
groot 已提交
2550
        StartMergeTask(merge_collection_ids);
2551 2552 2553
        return max_lsn;
    };

G
groot 已提交
2554 2555 2556 2557 2558 2559 2560 2561
    auto partition_flushed = [&](const std::string& collection_id, const std::string& partition,
                                 const std::string& target_collection_name) {
        if (options_.wal_enable_) {
            uint64_t lsn = 0;
            meta_ptr_->GetCollectionFlushLSN(target_collection_name, lsn);
            wal_mgr_->PartitionFlushed(collection_id, partition, lsn);
        }

G
groot 已提交
2562 2563
        std::set<std::string> merge_collection_ids = {target_collection_name};
        StartMergeTask(merge_collection_ids);
G
groot 已提交
2564 2565
    };

2566 2567 2568
    Status status;

    switch (record.type) {
2569 2570 2571 2572
        case wal::MXLogType::Entity: {
            std::string target_collection_name;
            status = GetPartitionByTag(record.collection_id, record.partition_tag, target_collection_name);
            if (!status.ok()) {
2573
                LOG_WAL_ERROR_ << LogOut("[%s][%ld] ", "insert", 0) << "Get partition fail: " << status.message();
2574 2575 2576
                return status;
            }

2577
            std::set<std::string> flushed_collections;
2578 2579 2580
            status = mem_mgr_->InsertEntities(target_collection_name, record.length, record.ids,
                                              (record.data_size / record.length / sizeof(float)),
                                              (const float*)record.data, record.attr_nbytes, record.attr_data_size,
2581
                                              record.attr_data, record.lsn, flushed_collections);
G
groot 已提交
2582 2583 2584
            if (!flushed_collections.empty()) {
                partition_flushed(record.collection_id, record.partition_tag, target_collection_name);
            }
2585 2586 2587 2588

            milvus::server::CollectInsertMetrics metrics(record.length, status);
            break;
        }
2589
        case wal::MXLogType::InsertBinary: {
2590 2591
            std::string target_collection_name;
            status = GetPartitionByTag(record.collection_id, record.partition_tag, target_collection_name);
2592
            if (!status.ok()) {
2593
                LOG_WAL_ERROR_ << LogOut("[%s][%ld] ", "insert", 0) << "Get partition fail: " << status.message();
2594 2595 2596
                return status;
            }

2597 2598
            std::set<std::string> flushed_collections;
            status = mem_mgr_->InsertVectors(target_collection_name, record.length, record.ids,
2599
                                             (record.data_size / record.length / sizeof(uint8_t)),
2600
                                             (const u_int8_t*)record.data, record.lsn, flushed_collections);
2601
            // even though !status.ok, run
G
groot 已提交
2602 2603 2604
            if (!flushed_collections.empty()) {
                partition_flushed(record.collection_id, record.partition_tag, target_collection_name);
            }
2605 2606 2607 2608 2609 2610 2611

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

        case wal::MXLogType::InsertVector: {
2612 2613
            std::string target_collection_name;
            status = GetPartitionByTag(record.collection_id, record.partition_tag, target_collection_name);
2614
            if (!status.ok()) {
2615
                LOG_WAL_ERROR_ << LogOut("[%s][%ld] ", "insert", 0) << "Get partition fail: " << status.message();
2616 2617 2618
                return status;
            }

2619 2620
            std::set<std::string> flushed_collections;
            status = mem_mgr_->InsertVectors(target_collection_name, record.length, record.ids,
2621
                                             (record.data_size / record.length / sizeof(float)),
2622
                                             (const float*)record.data, record.lsn, flushed_collections);
2623
            // even though !status.ok, run
G
groot 已提交
2624 2625 2626
            if (!flushed_collections.empty()) {
                partition_flushed(record.collection_id, record.partition_tag, target_collection_name);
            }
2627 2628 2629 2630 2631 2632 2633

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

        case wal::MXLogType::Delete: {
J
Jin Hai 已提交
2634 2635
            std::vector<meta::CollectionSchema> partition_array;
            status = meta_ptr_->ShowPartitions(record.collection_id, partition_array);
2636 2637 2638 2639
            if (!status.ok()) {
                return status;
            }

2640
            std::vector<std::string> collection_ids{record.collection_id};
2641
            for (auto& partition : partition_array) {
2642 2643
                auto& partition_collection_id = partition.collection_id_;
                collection_ids.emplace_back(partition_collection_id);
2644 2645 2646
            }

            if (record.length == 1) {
2647
                for (auto& collection_id : collection_ids) {
J
Jin Hai 已提交
2648
                    status = mem_mgr_->DeleteVector(collection_id, *record.ids, record.lsn);
2649 2650 2651 2652 2653
                    if (!status.ok()) {
                        return status;
                    }
                }
            } else {
2654
                for (auto& collection_id : collection_ids) {
J
Jin Hai 已提交
2655
                    status = mem_mgr_->DeleteVectors(collection_id, record.length, record.ids, record.lsn);
2656 2657 2658 2659 2660 2661 2662 2663 2664
                    if (!status.ok()) {
                        return status;
                    }
                }
            }
            break;
        }

        case wal::MXLogType::Flush: {
J
Jin Hai 已提交
2665 2666 2667 2668
            if (!record.collection_id.empty()) {
                // flush one collection
                std::vector<meta::CollectionSchema> partition_array;
                status = meta_ptr_->ShowPartitions(record.collection_id, partition_array);
2669 2670 2671 2672
                if (!status.ok()) {
                    return status;
                }

2673
                std::vector<std::string> collection_ids{record.collection_id};
2674
                for (auto& partition : partition_array) {
2675 2676
                    auto& partition_collection_id = partition.collection_id_;
                    collection_ids.emplace_back(partition_collection_id);
2677 2678
                }

2679 2680
                std::set<std::string> flushed_collections;
                for (auto& collection_id : collection_ids) {
2681
                    const std::lock_guard<std::mutex> lock(flush_merge_compact_mutex_);
J
Jin Hai 已提交
2682
                    status = mem_mgr_->Flush(collection_id);
2683 2684 2685
                    if (!status.ok()) {
                        break;
                    }
2686
                    flushed_collections.insert(collection_id);
2687 2688
                }

G
groot 已提交
2689
                collections_flushed(record.collection_id, flushed_collections);
2690 2691

            } else {
2692 2693
                // flush all collections
                std::set<std::string> collection_ids;
2694 2695
                {
                    const std::lock_guard<std::mutex> lock(flush_merge_compact_mutex_);
2696
                    status = mem_mgr_->Flush(collection_ids);
2697 2698
                }

G
groot 已提交
2699
                uint64_t lsn = collections_flushed("", collection_ids);
2700 2701 2702 2703 2704 2705
                if (options_.wal_enable_) {
                    wal_mgr_->RemoveOldFiles(lsn);
                }
            }
            break;
        }
C
Cai Yudong 已提交
2706 2707 2708

        default:
            break;
2709 2710 2711 2712 2713 2714
    }

    return status;
}

void
G
groot 已提交
2715 2716 2717 2718 2719 2720 2721 2722 2723
DBImpl::InternalFlush(const std::string& collection_id) {
    wal::MXLogRecord record;
    record.type = wal::MXLogType::Flush;
    record.collection_id = collection_id;
    ExecWalRecord(record);
}

void
DBImpl::BackgroundWalThread() {
2724
    SetThreadName("wal_thread");
2725 2726
    server::SystemInfo::GetInstance().Init();

2727
    std::chrono::system_clock::time_point next_auto_flush_time;
2728
    auto get_next_auto_flush_time = [&]() {
2729
        return std::chrono::system_clock::now() + std::chrono::seconds(options_.auto_flush_interval_);
2730
    };
2731 2732 2733
    if (options_.auto_flush_interval_ > 0) {
        next_auto_flush_time = get_next_auto_flush_time();
    }
2734 2735

    while (true) {
2736 2737
        if (options_.auto_flush_interval_ > 0) {
            if (std::chrono::system_clock::now() >= next_auto_flush_time) {
G
groot 已提交
2738
                InternalFlush();
2739 2740
                next_auto_flush_time = get_next_auto_flush_time();
            }
2741 2742
        }

G
groot 已提交
2743
        wal::MXLogRecord record;
2744 2745
        auto error_code = wal_mgr_->GetNextRecord(record);
        if (error_code != WAL_SUCCESS) {
2746
            LOG_ENGINE_ERROR_ << "WAL background GetNextRecord error";
2747 2748 2749 2750 2751 2752
            break;
        }

        if (record.type != wal::MXLogType::None) {
            ExecWalRecord(record);
            if (record.type == wal::MXLogType::Flush) {
G
groot 已提交
2753 2754
                // notify flush request to return
                flush_req_swn_.Notify();
2755 2756

                // if user flush all manually, update auto flush also
J
Jin Hai 已提交
2757
                if (record.collection_id.empty() && options_.auto_flush_interval_ > 0) {
2758 2759 2760 2761 2762 2763
                    next_auto_flush_time = get_next_auto_flush_time();
                }
            }

        } else {
            if (!initialized_.load(std::memory_order_acquire)) {
G
groot 已提交
2764 2765
                InternalFlush();
                flush_req_swn_.Notify();
2766 2767
                WaitMergeFileFinish();
                WaitBuildIndexFinish();
2768
                LOG_ENGINE_DEBUG_ << "WAL background thread exit";
2769 2770 2771
                break;
            }

2772
            if (options_.auto_flush_interval_ > 0) {
G
groot 已提交
2773
                swn_wal_.Wait_Until(next_auto_flush_time);
2774
            } else {
G
groot 已提交
2775
                swn_wal_.Wait();
2776
            }
2777 2778 2779 2780
        }
    }
}

G
groot 已提交
2781 2782
void
DBImpl::BackgroundFlushThread() {
2783
    SetThreadName("flush_thread");
G
groot 已提交
2784 2785 2786
    server::SystemInfo::GetInstance().Init();
    while (true) {
        if (!initialized_.load(std::memory_order_acquire)) {
2787
            LOG_ENGINE_DEBUG_ << "DB background flush thread exit";
G
groot 已提交
2788 2789 2790 2791 2792 2793 2794 2795 2796 2797 2798 2799 2800 2801
            break;
        }

        InternalFlush();
        if (options_.auto_flush_interval_ > 0) {
            swn_flush_.Wait_For(std::chrono::seconds(options_.auto_flush_interval_));
        } else {
            swn_flush_.Wait();
        }
    }
}

void
DBImpl::BackgroundMetricThread() {
G
groot 已提交
2802
    SetThreadName("metric_thread");
G
groot 已提交
2803 2804 2805
    server::SystemInfo::GetInstance().Init();
    while (true) {
        if (!initialized_.load(std::memory_order_acquire)) {
2806
            LOG_ENGINE_DEBUG_ << "DB background metric thread exit";
G
groot 已提交
2807 2808 2809 2810 2811
            break;
        }

        swn_metric_.Wait_For(std::chrono::seconds(BACKGROUND_METRIC_INTERVAL));
        StartMetricTask();
G
groot 已提交
2812
        meta::FilesHolder::PrintInfo();
G
groot 已提交
2813 2814 2815
    }
}

2816 2817 2818 2819 2820
void
DBImpl::OnCacheInsertDataChanged(bool value) {
    options_.insert_cache_immediately_ = value;
}

2821 2822 2823 2824 2825
void
DBImpl::OnUseBlasThresholdChanged(int64_t threshold) {
    faiss::distance_compute_blas_threshold = threshold;
}

2826 2827 2828 2829 2830 2831 2832 2833 2834 2835 2836 2837 2838 2839 2840 2841 2842 2843
void
DBImpl::SuspendIfFirst() {
    std::lock_guard<std::mutex> lock(suspend_build_mutex_);
    if (++live_search_num_ == 1) {
        LOG_ENGINE_TRACE_ << "live_search_num_: " << live_search_num_;
        knowhere::BuilderSuspend();
    }
}

void
DBImpl::ResumeIfLast() {
    std::lock_guard<std::mutex> lock(suspend_build_mutex_);
    if (--live_search_num_ == 0) {
        LOG_ENGINE_TRACE_ << "live_search_num_: " << live_search_num_;
        knowhere::BuildResume();
    }
}

S
starlord 已提交
2844 2845
}  // namespace engine
}  // namespace milvus