DBImpl.cpp 100.2 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

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

G
groot 已提交
154 155 156 157 158 159 160 161 162
    // 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
    bg_metric_thread_ = std::thread(&DBImpl::BackgroundMetricThread, this);

S
starlord 已提交
163 164 165
    return Status::OK();
}

S
starlord 已提交
166 167
Status
DBImpl::Stop() {
168
    if (!initialized_.load(std::memory_order_acquire)) {
S
starlord 已提交
169 170
        return Status::OK();
    }
171

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

174 175
    if (options_.mode_ != DBOptions::MODE::CLUSTER_READONLY) {
        if (options_.wal_enable_) {
G
groot 已提交
176 177
            // wait wal thread finish
            swn_wal_.Notify();
178 179
            bg_wal_thread_.join();
        } else {
G
groot 已提交
180
            // flush all without merge
181 182 183 184
            wal::MXLogRecord record;
            record.type = wal::MXLogType::Flush;
            ExecWalRecord(record);

G
groot 已提交
185 186 187
            // wait flush thread finish
            swn_flush_.Notify();
            bg_flush_thread_.join();
188
        }
S
starlord 已提交
189

190 191
        WaitMergeFileFinish();

G
groot 已提交
192 193 194
        swn_index_.Notify();
        bg_index_thread_.join();

195
        meta_ptr_->CleanUpShadowFiles();
S
starlord 已提交
196 197
    }

G
groot 已提交
198 199 200 201
    // wait metric thread exit
    swn_metric_.Notify();
    bg_metric_thread_.join();

202
    // LOG_ENGINE_TRACE_ << "DB service stop";
S
starlord 已提交
203
    return Status::OK();
X
Xu Peng 已提交
204 205
}

S
starlord 已提交
206 207
Status
DBImpl::DropAll() {
S
starlord 已提交
208 209 210
    return meta_ptr_->DropAll();
}

S
starlord 已提交
211
Status
212
DBImpl::CreateCollection(meta::CollectionSchema& collection_schema) {
213
    if (!initialized_.load(std::memory_order_acquire)) {
G
groot 已提交
214
        return SHUTDOWN_ERROR;
S
starlord 已提交
215 216
    }

217
    meta::CollectionSchema temp_schema = collection_schema;
B
bigbraver 已提交
218
    temp_schema.index_file_size_ *= MB;  // store as MB
219
    if (options_.wal_enable_) {
220
        temp_schema.flush_lsn_ = wal_mgr_->CreateCollection(collection_schema.collection_id_);
221 222
    }

223
    return meta_ptr_->CreateCollection(temp_schema);
224 225
}

226 227 228 229 230 231 232
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 已提交
233
    temp_schema.index_file_size_ *= MB;
234
    if (options_.wal_enable_) {
Y
yukun 已提交
235
        temp_schema.flush_lsn_ = wal_mgr_->CreateHybridCollection(collection_schema.collection_id_);
236 237 238 239 240 241 242 243 244 245 246 247 248 249 250 251
    }

    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 已提交
252
Status
253
DBImpl::DropCollection(const std::string& collection_id) {
254
    if (!initialized_.load(std::memory_order_acquire)) {
G
groot 已提交
255
        return SHUTDOWN_ERROR;
S
starlord 已提交
256 257
    }

258
    if (options_.wal_enable_) {
259
        wal_mgr_->DropCollection(collection_id);
260 261
    }

262
    return DropCollectionRecursively(collection_id);
G
groot 已提交
263 264
}

S
starlord 已提交
265
Status
266
DBImpl::DescribeCollection(meta::CollectionSchema& collection_schema) {
267
    if (!initialized_.load(std::memory_order_acquire)) {
G
groot 已提交
268
        return SHUTDOWN_ERROR;
S
starlord 已提交
269 270
    }

271
    auto stat = meta_ptr_->DescribeCollection(collection_schema);
B
bigbraver 已提交
272
    collection_schema.index_file_size_ /= MB;  // return as MB
S
starlord 已提交
273
    return stat;
274 275
}

S
starlord 已提交
276
Status
277
DBImpl::HasCollection(const std::string& collection_id, bool& has_or_not) {
278
    if (!initialized_.load(std::memory_order_acquire)) {
G
groot 已提交
279
        return SHUTDOWN_ERROR;
S
starlord 已提交
280 281
    }

G
groot 已提交
282
    return meta_ptr_->HasCollection(collection_id, has_or_not, false);
283 284
}

285
Status
G
groot 已提交
286
DBImpl::HasNativeCollection(const std::string& collection_id, bool& has_or_not) {
287 288 289 290
    if (!initialized_.load(std::memory_order_acquire)) {
        return SHUTDOWN_ERROR;
    }

G
groot 已提交
291
    return meta_ptr_->HasCollection(collection_id, has_or_not, true);
292 293
}

S
starlord 已提交
294
Status
295
DBImpl::AllCollections(std::vector<meta::CollectionSchema>& collection_schema_array) {
296
    if (!initialized_.load(std::memory_order_acquire)) {
G
groot 已提交
297
        return SHUTDOWN_ERROR;
S
starlord 已提交
298 299
    }

300 301
    std::vector<meta::CollectionSchema> all_collections;
    auto status = meta_ptr_->AllCollections(all_collections);
302

303 304 305 306 307
    // 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);
308 309 310 311
        }
    }

    return status;
G
groot 已提交
312 313
}

314
Status
315
DBImpl::GetCollectionInfo(const std::string& collection_id, std::string& collection_info) {
316 317 318 319 320
    if (!initialized_.load(std::memory_order_acquire)) {
        return SHUTDOWN_ERROR;
    }

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

J
Jin Hai 已提交
324 325
    std::vector<int> file_types{meta::SegmentSchema::FILE_TYPE::RAW, meta::SegmentSchema::FILE_TYPE::TO_INDEX,
                                meta::SegmentSchema::FILE_TYPE::INDEX};
326

327 328 329
    milvus::json json_info;
    milvus::json json_partitions;
    size_t total_row_count = 0;
330

331
    auto get_info = [&](const std::string& col_id, const std::string& tag) {
G
groot 已提交
332 333
        meta::FilesHolder files_holder;
        status = meta_ptr_->FilesByType(col_id, file_types, files_holder);
334
        if (!status.ok()) {
J
Jin Hai 已提交
335
            std::string err_msg = "Failed to get collection info: " + status.ToString();
336
            LOG_ENGINE_ERROR_ << err_msg;
337 338 339
            return Status(DB_ERROR, err_msg);
        }

340 341 342 343 344
        milvus::json json_partition;
        json_partition[JSON_PARTITION_TAG] = tag;

        milvus::json json_segments;
        size_t row_count = 0;
G
groot 已提交
345
        milvus::engine::meta::SegmentsSchema& collection_files = files_holder.HoldFiles();
346
        for (auto& file : collection_files) {
347 348 349 350 351 352 353 354 355
            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_;
356 357
        }

358 359 360 361 362 363 364
        json_partition[JSON_ROW_COUNT] = row_count;
        json_partition[JSON_SEGMENTS] = json_segments;

        json_partitions.push_back(json_partition);

        return Status::OK();
    };
365

366 367 368 369
    // step2: get default partition info
    status = get_info(collection_id, milvus::engine::DEFAULT_PARTITON_TAG);
    if (!status.ok()) {
        return status;
370 371
    }

372 373 374 375 376 377 378 379 380 381 382 383 384
    // 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();

385 386 387
    return Status::OK();
}

S
starlord 已提交
388
Status
389
DBImpl::PreloadCollection(const std::string& collection_id) {
390
    if (!initialized_.load(std::memory_order_acquire)) {
G
groot 已提交
391
        return SHUTDOWN_ERROR;
S
starlord 已提交
392 393
    }

J
Jin Hai 已提交
394
    // step 1: get all collection files from parent collection
G
groot 已提交
395 396
    meta::FilesHolder files_holder;
    auto status = meta_ptr_->FilesToSearch(collection_id, files_holder);
Y
Yu Kun 已提交
397 398 399
    if (!status.ok()) {
        return status;
    }
Y
Yu Kun 已提交
400

401
    // step 2: get files from partition collections
J
Jin Hai 已提交
402 403
    std::vector<meta::CollectionSchema> partition_array;
    status = meta_ptr_->ShowPartitions(collection_id, partition_array);
G
groot 已提交
404
    for (auto& schema : partition_array) {
G
groot 已提交
405
        status = meta_ptr_->FilesToSearch(schema.collection_id_, files_holder);
G
groot 已提交
406 407
    }

Y
Yu Kun 已提交
408 409
    int64_t size = 0;
    int64_t cache_total = cache::CpuCacheMgr::GetInstance()->CacheCapacity();
Y
Yu Kun 已提交
410 411
    int64_t cache_usage = cache::CpuCacheMgr::GetInstance()->CacheUsage();
    int64_t available_size = cache_total - cache_usage;
Y
Yu Kun 已提交
412

413
    // step 3: load file one by one
G
groot 已提交
414
    milvus::engine::meta::SegmentsSchema& files_array = files_holder.HoldFiles();
415 416
    LOG_ENGINE_DEBUG_ << "Begin pre-load collection:" + collection_id + ", totally " << files_array.size()
                      << " files need to be pre-loaded";
J
Jin Hai 已提交
417
    TimeRecorderAuto rc("Pre-load collection:" + collection_id);
G
groot 已提交
418
    for (auto& file : files_array) {
419
        EngineType engine_type;
J
Jin Hai 已提交
420 421 422
        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) {
423 424
            engine_type =
                utils::IsBinaryMetricType(file.metric_type_) ? EngineType::FAISS_BIN_IDMAP : EngineType::FAISS_IDMAP;
425 426 427
        } else {
            engine_type = (EngineType)file.engine_type_;
        }
428 429 430 431

        auto json = milvus::json::parse(file.index_params_);
        ExecutionEnginePtr engine =
            EngineFactory::Build(file.dimension_, file.location_, engine_type, (MetricType)file.metric_type_, json);
432
        fiu_do_on("DBImpl.PreloadCollection.null_engine", engine = nullptr);
G
groot 已提交
433
        if (engine == nullptr) {
434
            LOG_ENGINE_ERROR_ << "Invalid engine type";
G
groot 已提交
435 436
            return Status(DB_ERROR, "Invalid engine type");
        }
Y
Yu Kun 已提交
437

438
        fiu_do_on("DBImpl.PreloadCollection.exceed_cache", size = available_size + 1);
439 440

        try {
441
            fiu_do_on("DBImpl.PreloadCollection.engine_throw_exception", throw std::exception());
442 443
            std::string msg = "Pre-loaded file: " + file.file_id_ + " size: " + std::to_string(file.file_size_);
            TimeRecorderAuto rc_1(msg);
444 445 446 447
            status = engine->Load(true);
            if (!status.ok()) {
                return status;
            }
448 449 450

            size += engine->Size();
            if (size > available_size) {
451
                LOG_ENGINE_DEBUG_ << "Pre-load cancelled since cache is almost full";
452
                return Status(SERVER_CACHE_FULL, "Cache is full");
Y
Yu Kun 已提交
453
            }
454
        } catch (std::exception& ex) {
J
Jin Hai 已提交
455
            std::string msg = "Pre-load collection encounter exception: " + std::string(ex.what());
456
            LOG_ENGINE_ERROR_ << msg;
457
            return Status(DB_ERROR, msg);
Y
Yu Kun 已提交
458 459
        }
    }
G
groot 已提交
460

Y
Yu Kun 已提交
461
    return Status::OK();
Y
Yu Kun 已提交
462 463
}

S
starlord 已提交
464
Status
465
DBImpl::UpdateCollectionFlag(const std::string& collection_id, int64_t flag) {
466
    if (!initialized_.load(std::memory_order_acquire)) {
G
groot 已提交
467
        return SHUTDOWN_ERROR;
S
starlord 已提交
468 469
    }

470
    return meta_ptr_->UpdateCollectionFlag(collection_id, flag);
S
starlord 已提交
471 472
}

S
starlord 已提交
473
Status
474
DBImpl::GetCollectionRowCount(const std::string& collection_id, uint64_t& row_count) {
475
    if (!initialized_.load(std::memory_order_acquire)) {
G
groot 已提交
476 477 478
        return SHUTDOWN_ERROR;
    }

479
    return GetCollectionRowCountRecursively(collection_id, row_count);
G
groot 已提交
480 481 482
}

Status
J
Jin Hai 已提交
483
DBImpl::CreatePartition(const std::string& collection_id, const std::string& partition_name,
G
groot 已提交
484
                        const std::string& partition_tag) {
485
    if (!initialized_.load(std::memory_order_acquire)) {
G
groot 已提交
486 487 488
        return SHUTDOWN_ERROR;
    }

489
    uint64_t lsn = 0;
G
groot 已提交
490 491 492 493 494
    if (options_.wal_enable_) {
        lsn = wal_mgr_->CreatePartition(collection_id, partition_tag);
    } else {
        meta_ptr_->GetCollectionFlushLSN(collection_id, lsn);
    }
J
Jin Hai 已提交
495
    return meta_ptr_->CreatePartition(collection_id, partition_name, partition_tag, lsn);
G
groot 已提交
496 497
}

G
groot 已提交
498 499 500 501 502 503 504 505 506 507 508 509 510 511 512 513 514 515 516
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 已提交
517 518
Status
DBImpl::DropPartition(const std::string& partition_name) {
519
    if (!initialized_.load(std::memory_order_acquire)) {
G
groot 已提交
520
        return SHUTDOWN_ERROR;
S
starlord 已提交
521 522
    }

523
    mem_mgr_->EraseMemVector(partition_name);                // not allow insert
J
Jin Hai 已提交
524
    auto status = meta_ptr_->DropPartition(partition_name);  // soft delete collection
525
    if (!status.ok()) {
526
        LOG_ENGINE_ERROR_ << status.message();
527 528
        return status;
    }
G
groot 已提交
529

J
Jin Hai 已提交
530
    // scheduler will determine when to delete collection files
G
groot 已提交
531 532 533 534 535 536
    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 已提交
537 538
}

S
starlord 已提交
539
Status
J
Jin Hai 已提交
540
DBImpl::DropPartitionByTag(const std::string& collection_id, const std::string& partition_tag) {
541
    if (!initialized_.load(std::memory_order_acquire)) {
G
groot 已提交
542 543 544 545
        return SHUTDOWN_ERROR;
    }

    std::string partition_name;
J
Jin Hai 已提交
546
    auto status = meta_ptr_->GetPartitionName(collection_id, partition_tag, partition_name);
547
    if (!status.ok()) {
548
        LOG_ENGINE_ERROR_ << status.message();
549 550 551
        return status;
    }

G
groot 已提交
552 553 554 555
    if (options_.wal_enable_) {
        wal_mgr_->DropPartition(collection_id, partition_tag);
    }

G
groot 已提交
556 557 558 559
    return DropPartition(partition_name);
}

Status
J
Jin Hai 已提交
560
DBImpl::ShowPartitions(const std::string& collection_id, std::vector<meta::CollectionSchema>& partition_schema_array) {
561
    if (!initialized_.load(std::memory_order_acquire)) {
G
groot 已提交
562 563 564
        return SHUTDOWN_ERROR;
    }

J
Jin Hai 已提交
565
    return meta_ptr_->ShowPartitions(collection_id, partition_schema_array);
G
groot 已提交
566 567 568
}

Status
J
Jin Hai 已提交
569
DBImpl::InsertVectors(const std::string& collection_id, const std::string& partition_tag, VectorsData& vectors) {
570
    //    LOG_ENGINE_DEBUG_ << "Insert " << n << " vectors to cache";
571
    if (!initialized_.load(std::memory_order_acquire)) {
G
groot 已提交
572
        return SHUTDOWN_ERROR;
S
starlord 已提交
573
    }
Y
yu yunfeng 已提交
574

J
Jin Hai 已提交
575
    // insert vectors into target collection
576 577
    // (zhiru): generate ids
    if (vectors.id_array_.empty()) {
J
Jin Hai 已提交
578 579 580
        SafeIDGenerator& id_generator = SafeIDGenerator::GetInstance();
        Status status = id_generator.GetNextIDNumbers(vectors.vector_count_, vectors.id_array_);
        if (!status.ok()) {
581
            LOG_ENGINE_ERROR_ << LogOut("[%s][%ld] Get next id number fail: %s", "insert", 0, status.message().c_str());
J
Jin Hai 已提交
582 583
            return status;
        }
584 585
    }

586
    Status status;
587
    if (options_.wal_enable_) {
588 589
        std::string target_collection_name;
        status = GetPartitionByTag(collection_id, partition_tag, target_collection_name);
G
groot 已提交
590
        if (!status.ok()) {
591
            LOG_ENGINE_ERROR_ << LogOut("[%s][%ld] Get partition fail: %s", "insert", 0, status.message().c_str());
G
groot 已提交
592 593
            return status;
        }
594 595

        if (!vectors.float_data_.empty()) {
J
Jin Hai 已提交
596
            wal_mgr_->Insert(collection_id, partition_tag, vectors.id_array_, vectors.float_data_);
597
        } else if (!vectors.binary_data_.empty()) {
J
Jin Hai 已提交
598
            wal_mgr_->Insert(collection_id, partition_tag, vectors.id_array_, vectors.binary_data_);
599
        }
G
groot 已提交
600
        swn_wal_.Notify();
601 602 603
    } else {
        wal::MXLogRecord record;
        record.lsn = 0;  // need to get from meta ?
J
Jin Hai 已提交
604
        record.collection_id = collection_id;
605 606 607 608 609 610 611 612 613 614 615 616 617 618 619 620
        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 已提交
621 622
    }

623 624 625
    return status;
}

626
Status
Y
yukun 已提交
627 628 629 630 631
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) {
632
    uint64_t offset = 0;
Y
yukun 已提交
633 634
    for (auto name : field_names) {
        switch (attr_types.at(name)) {
635 636
            case meta::hybrid::DataType::INT8: {
                std::vector<uint8_t> data;
Y
yukun 已提交
637
                data.resize(row_num * sizeof(int8_t));
638

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

Y
yukun 已提交
642 643
                std::vector<int8_t> raw_value(row_num, 0);
                for (uint64_t i = 0; i < row_num; ++i) {
644 645 646
                    raw_value[i] = attr_value[i];
                }

Y
yukun 已提交
647 648
                memcpy(data.data(), raw_value.data(), row_num * sizeof(int8_t));
                attr_datas.insert(std::make_pair(name, data));
649

Y
yukun 已提交
650 651 652
                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);
653 654 655 656
                break;
            }
            case meta::hybrid::DataType::INT16: {
                std::vector<uint8_t> data;
Y
yukun 已提交
657
                data.resize(row_num * sizeof(int16_t));
658

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

Y
yukun 已提交
662 663
                std::vector<int16_t> raw_value(row_num, 0);
                for (uint64_t i = 0; i < row_num; ++i) {
664 665 666
                    raw_value[i] = attr_value[i];
                }

Y
yukun 已提交
667 668
                memcpy(data.data(), raw_value.data(), row_num * sizeof(int16_t));
                attr_datas.insert(std::make_pair(name, data));
669

Y
yukun 已提交
670 671 672
                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);
673 674 675 676
                break;
            }
            case meta::hybrid::DataType::INT32: {
                std::vector<uint8_t> data;
Y
yukun 已提交
677
                data.resize(row_num * sizeof(int32_t));
678

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

Y
yukun 已提交
682 683
                std::vector<int32_t> raw_value(row_num, 0);
                for (uint64_t i = 0; i < row_num; ++i) {
684 685 686
                    raw_value[i] = attr_value[i];
                }

Y
yukun 已提交
687 688
                memcpy(data.data(), raw_value.data(), row_num * sizeof(int32_t));
                attr_datas.insert(std::make_pair(name, data));
689

Y
yukun 已提交
690 691 692
                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);
693 694 695 696
                break;
            }
            case meta::hybrid::DataType::INT64: {
                std::vector<uint8_t> data;
Y
yukun 已提交
697 698 699 700 701 702
                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));
703

Y
yukun 已提交
704 705 706
                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);
707 708 709 710
                break;
            }
            case meta::hybrid::DataType::FLOAT: {
                std::vector<uint8_t> data;
Y
yukun 已提交
711
                data.resize(row_num * sizeof(float));
712

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

Y
yukun 已提交
716 717
                std::vector<float> raw_value(row_num, 0);
                for (uint64_t i = 0; i < row_num; ++i) {
718 719 720
                    raw_value[i] = attr_value[i];
                }

Y
yukun 已提交
721 722
                memcpy(data.data(), raw_value.data(), row_num * sizeof(float));
                attr_datas.insert(std::make_pair(name, data));
723

Y
yukun 已提交
724 725 726
                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);
727 728 729 730
                break;
            }
            case meta::hybrid::DataType::DOUBLE: {
                std::vector<uint8_t> data;
Y
yukun 已提交
731 732 733
                data.resize(row_num * sizeof(double));
                memcpy(data.data(), record.data() + offset, row_num * sizeof(double));
                attr_datas.insert(std::make_pair(name, data));
734

Y
yukun 已提交
735 736 737
                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);
738 739
                break;
            }
740 741
            default:
                break;
742 743
        }
    }
Y
yukun 已提交
744 745 746 747 748 749 750 751 752 753 754 755 756 757 758 759 760 761 762 763 764 765 766 767 768 769 770 771 772 773 774 775 776 777 778 779 780 781 782 783 784 785 786 787 788 789 790 791 792 793
    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);
    }
794 795

    status = ExecWalRecord(record);
Y
yukun 已提交
796 797 798 799 800 801 802 803 804 805 806 807 808 809 810 811 812 813 814 815 816 817 818 819 820 821 822 823 824 825 826 827 828 829 830 831 832 833 834 835 836 837 838 839 840 841

#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

842 843 844
    return status;
}

845
Status
J
Jin Hai 已提交
846
DBImpl::DeleteVector(const std::string& collection_id, IDNumber vector_id) {
847 848
    IDNumbers ids;
    ids.push_back(vector_id);
J
Jin Hai 已提交
849
    return DeleteVectors(collection_id, ids);
850 851 852
}

Status
J
Jin Hai 已提交
853
DBImpl::DeleteVectors(const std::string& collection_id, IDNumbers vector_ids) {
854 855 856 857 858 859
    if (!initialized_.load(std::memory_order_acquire)) {
        return SHUTDOWN_ERROR;
    }

    Status status;
    if (options_.wal_enable_) {
J
Jin Hai 已提交
860
        wal_mgr_->DeleteById(collection_id, vector_ids);
G
groot 已提交
861
        swn_wal_.Notify();
862 863 864 865
    } else {
        wal::MXLogRecord record;
        record.lsn = 0;  // need to get from meta ?
        record.type = wal::MXLogType::Delete;
J
Jin Hai 已提交
866
        record.collection_id = collection_id;
867 868 869 870 871 872 873 874 875 876
        record.ids = vector_ids.data();
        record.length = vector_ids.size();

        status = ExecWalRecord(record);
    }

    return status;
}

Status
J
Jin Hai 已提交
877
DBImpl::Flush(const std::string& collection_id) {
878 879 880 881 882
    if (!initialized_.load(std::memory_order_acquire)) {
        return SHUTDOWN_ERROR;
    }

    Status status;
883 884
    bool has_collection;
    status = HasCollection(collection_id, has_collection);
885 886 887
    if (!status.ok()) {
        return status;
    }
888
    if (!has_collection) {
889
        LOG_ENGINE_ERROR_ << "Collection to flush does not exist: " << collection_id;
J
Jin Hai 已提交
890
        return Status(DB_NOT_FOUND, "Collection to flush does not exist");
891 892
    }

893
    LOG_ENGINE_DEBUG_ << "Begin flush collection: " << collection_id;
894 895

    if (options_.wal_enable_) {
896
        LOG_ENGINE_DEBUG_ << "WAL flush";
J
Jin Hai 已提交
897
        auto lsn = wal_mgr_->Flush(collection_id);
898
        if (lsn != 0) {
G
groot 已提交
899 900
            swn_wal_.Notify();
            flush_req_swn_.Wait();
901
        }
G
groot 已提交
902
        StartMergeTask();
903
    } else {
904
        LOG_ENGINE_DEBUG_ << "MemTable flush";
G
groot 已提交
905
        InternalFlush(collection_id);
906 907
    }

908
    LOG_ENGINE_DEBUG_ << "End flush collection: " << collection_id;
909 910 911 912 913 914 915 916 917 918

    return status;
}

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

919
    LOG_ENGINE_DEBUG_ << "Begin flush all collections";
920 921 922

    Status status;
    if (options_.wal_enable_) {
923
        LOG_ENGINE_DEBUG_ << "WAL flush";
924 925
        auto lsn = wal_mgr_->Flush();
        if (lsn != 0) {
G
groot 已提交
926 927
            swn_wal_.Notify();
            flush_req_swn_.Wait();
928
        }
G
groot 已提交
929
        StartMergeTask();
930
    } else {
931
        LOG_ENGINE_DEBUG_ << "MemTable flush";
G
groot 已提交
932
        InternalFlush();
933 934
    }

935
    LOG_ENGINE_DEBUG_ << "End flush all collections";
936 937 938 939 940

    return status;
}

Status
G
groot 已提交
941
DBImpl::Compact(const std::string& collection_id, double threshold) {
942 943 944 945
    if (!initialized_.load(std::memory_order_acquire)) {
        return SHUTDOWN_ERROR;
    }

946 947 948
    engine::meta::CollectionSchema collection_schema;
    collection_schema.collection_id_ = collection_id;
    auto status = DescribeCollection(collection_schema);
949 950
    if (!status.ok()) {
        if (status.code() == DB_NOT_FOUND) {
951
            LOG_ENGINE_ERROR_ << "Collection to compact does not exist: " << collection_id;
J
Jin Hai 已提交
952
            return Status(DB_NOT_FOUND, "Collection to compact does not exist");
953 954 955 956
        } else {
            return status;
        }
    } else {
957
        if (!collection_schema.owner_collection_.empty()) {
958
            LOG_ENGINE_ERROR_ << "Collection to compact does not exist: " << collection_id;
J
Jin Hai 已提交
959
            return Status(DB_NOT_FOUND, "Collection to compact does not exist");
960 961 962
        }
    }

963
    LOG_ENGINE_DEBUG_ << "Before compacting, wait for build index thread to finish...";
964

Z
update  
Zhiru Zhu 已提交
965
    // WaitBuildIndexFinish();
966

Z
update  
Zhiru Zhu 已提交
967
    const std::lock_guard<std::mutex> index_lock(build_index_mutex_);
Z
Zhiru Zhu 已提交
968
    const std::lock_guard<std::mutex> merge_lock(flush_merge_compact_mutex_);
Z
Zhiru Zhu 已提交
969

970
    LOG_ENGINE_DEBUG_ << "Compacting collection: " << collection_id;
Z
Zhiru Zhu 已提交
971

972
    // Get files to compact from meta.
J
Jin Hai 已提交
973 974
    std::vector<int> file_types{meta::SegmentSchema::FILE_TYPE::RAW, meta::SegmentSchema::FILE_TYPE::TO_INDEX,
                                meta::SegmentSchema::FILE_TYPE::BACKUP};
G
groot 已提交
975 976
    meta::FilesHolder files_holder;
    status = meta_ptr_->FilesByType(collection_id, file_types, files_holder);
977 978
    if (!status.ok()) {
        std::string err_msg = "Failed to get files to compact: " + status.message();
979
        LOG_ENGINE_ERROR_ << err_msg;
980 981 982
        return Status(DB_ERROR, err_msg);
    }

G
groot 已提交
983
    LOG_ENGINE_DEBUG_ << "Found " << files_holder.HoldFiles().size() << " segment to compact";
Z
Zhiru Zhu 已提交
984 985

    Status compact_status;
G
groot 已提交
986 987
    // 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 已提交
988
    for (auto iter = files_to_compact.begin(); iter != files_to_compact.end();) {
J
Jin Hai 已提交
989
        meta::SegmentSchema file = *iter;
G
groot 已提交
990 991
        iter = files_to_compact.erase(iter);

Z
Zhiru Zhu 已提交
992 993 994
        // Check if the segment needs compacting
        std::string segment_dir;
        utils::GetParentPath(file.location_, segment_dir);
995

Z
Zhiru Zhu 已提交
996
        segment::SegmentReader segment_reader(segment_dir);
Z
Zhiru Zhu 已提交
997 998
        size_t deleted_docs_size;
        status = segment_reader.ReadDeletedDocsSize(deleted_docs_size);
Z
Zhiru Zhu 已提交
999
        if (!status.ok()) {
G
groot 已提交
1000
            files_holder.UnmarkFile(file);
G
groot 已提交
1001
            continue;  // skip this file and try compact next one
Z
Zhiru Zhu 已提交
1002 1003
        }

J
Jin Hai 已提交
1004
        meta::SegmentsSchema files_to_update;
Z
Zhiru Zhu 已提交
1005
        if (deleted_docs_size != 0) {
G
groot 已提交
1006
            compact_status = CompactFile(collection_id, threshold, file, files_to_update);
Z
Zhiru Zhu 已提交
1007 1008

            if (!compact_status.ok()) {
1009 1010
                LOG_ENGINE_ERROR_ << "Compact failed for segment " << file.segment_id_ << ": "
                                  << compact_status.message();
G
groot 已提交
1011
                files_holder.UnmarkFile(file);
G
groot 已提交
1012
                continue;  // skip this file and try compact next one
Z
Zhiru Zhu 已提交
1013 1014
            }
        } else {
G
groot 已提交
1015
            files_holder.UnmarkFile(file);
1016
            LOG_ENGINE_DEBUG_ << "Segment " << file.segment_id_ << " has no deleted data. No need to compact";
G
groot 已提交
1017
            continue;  // skip this file and try compact next one
1018
        }
Z
Zhiru Zhu 已提交
1019

1020
        LOG_ENGINE_DEBUG_ << "Updating meta after compaction...";
1021
        status = meta_ptr_->UpdateCollectionFiles(files_to_update);
G
groot 已提交
1022
        files_holder.UnmarkFile(file);
G
groot 已提交
1023 1024 1025 1026
        if (!status.ok()) {
            compact_status = status;
            break;  // meta error, could not go on
        }
Z
Zhiru Zhu 已提交
1027 1028
    }

G
groot 已提交
1029
    if (compact_status.ok()) {
1030
        LOG_ENGINE_DEBUG_ << "Finished compacting collection: " << collection_id;
G
groot 已提交
1031
    }
1032

G
groot 已提交
1033
    return compact_status;
1034 1035 1036
}

Status
G
groot 已提交
1037
DBImpl::CompactFile(const std::string& collection_id, double threshold, const meta::SegmentSchema& file,
J
Jin Hai 已提交
1038
                    meta::SegmentsSchema& files_to_update) {
1039
    LOG_ENGINE_DEBUG_ << "Compacting segment " << file.segment_id_ << " for collection: " << collection_id;
1040

G
groot 已提交
1041 1042 1043 1044 1045 1046 1047 1048 1049 1050 1051 1052 1053 1054 1055 1056 1057 1058 1059
    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();
            double delete_rate = (double)delete_items.size() / (double)file.row_count_;
            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 已提交
1060 1061 1062
    // Create new collection file
    meta::SegmentSchema compacted_file;
    compacted_file.collection_id_ = collection_id;
1063
    // compacted_file.date_ = date;
J
Jin Hai 已提交
1064
    compacted_file.file_type_ = meta::SegmentSchema::NEW_MERGE;  // TODO: use NEW_MERGE for now
G
groot 已提交
1065
    auto status = meta_ptr_->CreateCollectionFile(compacted_file);
1066 1067

    if (!status.ok()) {
1068
        LOG_ENGINE_ERROR_ << "Failed to create collection file: " << status.message();
1069 1070 1071
        return status;
    }

J
Jin Hai 已提交
1072
    // Compact (merge) file to the newly created collection file
1073 1074 1075 1076 1077

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

1078
    LOG_ENGINE_DEBUG_ << "Compacting begin...";
1079 1080 1081
    segment_writer_ptr->Merge(segment_dir_to_merge, compacted_file.file_id_);

    // Serialize
1082
    LOG_ENGINE_DEBUG_ << "Serializing compacted segment...";
1083 1084
    status = segment_writer_ptr->Serialize();
    if (!status.ok()) {
1085
        LOG_ENGINE_ERROR_ << "Failed to serialize compacted segment: " << status.message();
J
Jin Hai 已提交
1086
        compacted_file.file_type_ = meta::SegmentSchema::TO_DELETE;
1087
        auto mark_status = meta_ptr_->UpdateCollectionFile(compacted_file);
1088
        if (mark_status.ok()) {
1089
            LOG_ENGINE_DEBUG_ << "Mark file: " << compacted_file.file_id_ << " to to_delete";
1090
        }
G
groot 已提交
1091

1092 1093 1094
        return status;
    }

1095 1096 1097 1098 1099 1100 1101
    // Update compacted file state, if origin file is backup or to_index, set compected file to to_index
    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;
1102
    } else {
J
Jin Hai 已提交
1103
        compacted_file.file_type_ = meta::SegmentSchema::RAW;
1104 1105 1106
    }

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

Z
Zhiru Zhu 已提交
1111
    files_to_update.emplace_back(compacted_file);
Z
Zhiru Zhu 已提交
1112

Z
Zhiru Zhu 已提交
1113 1114
    // Set all files in segment to TO_DELETE
    auto& segment_id = file.segment_id_;
G
groot 已提交
1115 1116
    meta::FilesHolder files_holder;
    status = meta_ptr_->GetCollectionFilesBySegmentId(segment_id, files_holder);
Z
Zhiru Zhu 已提交
1117 1118 1119
    if (!status.ok()) {
        return status;
    }
G
groot 已提交
1120 1121

    milvus::engine::meta::SegmentsSchema& segment_files = files_holder.HoldFiles();
Z
Zhiru Zhu 已提交
1122
    for (auto& f : segment_files) {
J
Jin Hai 已提交
1123
        f.file_type_ = meta::SegmentSchema::FILE_TYPE::TO_DELETE;
Z
Zhiru Zhu 已提交
1124 1125
        files_to_update.emplace_back(f);
    }
G
groot 已提交
1126
    files_holder.ReleaseFiles();
1127

1128 1129 1130
    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";
1131 1132 1133 1134 1135 1136 1137 1138 1139

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

    return status;
}

Status
1140 1141
DBImpl::GetVectorsByID(const std::string& collection_id, const IDNumbers& id_array,
                       std::vector<engine::VectorsData>& vectors) {
1142 1143 1144 1145
    if (!initialized_.load(std::memory_order_acquire)) {
        return SHUTDOWN_ERROR;
    }

1146 1147 1148
    bool has_collection;
    auto status = HasCollection(collection_id, has_collection);
    if (!has_collection) {
1149
        LOG_ENGINE_ERROR_ << "Collection " << collection_id << " does not exist: ";
J
Jin Hai 已提交
1150
        return Status(DB_NOT_FOUND, "Collection does not exist");
1151 1152 1153 1154 1155
    }
    if (!status.ok()) {
        return status;
    }

G
groot 已提交
1156
    meta::FilesHolder files_holder;
J
Jin Hai 已提交
1157 1158
    std::vector<int> file_types{meta::SegmentSchema::FILE_TYPE::RAW, meta::SegmentSchema::FILE_TYPE::TO_INDEX,
                                meta::SegmentSchema::FILE_TYPE::BACKUP};
1159

G
groot 已提交
1160
    status = meta_ptr_->FilesByType(collection_id, file_types, files_holder);
1161
    if (!status.ok()) {
1162
        std::string err_msg = "Failed to get files for GetVectorsByID: " + status.message();
1163
        LOG_ENGINE_ERROR_ << err_msg;
1164 1165 1166
        return status;
    }

J
Jin Hai 已提交
1167 1168
    std::vector<meta::CollectionSchema> partition_array;
    status = meta_ptr_->ShowPartitions(collection_id, partition_array);
1169
    for (auto& schema : partition_array) {
G
groot 已提交
1170
        status = meta_ptr_->FilesByType(schema.collection_id_, file_types, files_holder);
1171 1172
        if (!status.ok()) {
            std::string err_msg = "Failed to get files for GetVectorByID: " + status.message();
1173
            LOG_ENGINE_ERROR_ << err_msg;
1174 1175 1176 1177
            return status;
        }
    }

G
groot 已提交
1178
    if (files_holder.HoldFiles().empty()) {
1179
        LOG_ENGINE_DEBUG_ << "No files to get vector by id from";
J
Jin Hai 已提交
1180
        return Status(DB_NOT_FOUND, "Collection is empty");
1181 1182 1183
    }

    cache::CpuCacheMgr::GetInstance()->PrintInfo();
G
groot 已提交
1184
    status = GetVectorsByIdHelper(collection_id, id_array, vectors, files_holder);
1185 1186 1187 1188 1189 1190
    cache::CpuCacheMgr::GetInstance()->PrintInfo();

    return status;
}

Status
J
Jin Hai 已提交
1191
DBImpl::GetVectorIDs(const std::string& collection_id, const std::string& segment_id, IDNumbers& vector_ids) {
1192 1193 1194 1195
    if (!initialized_.load(std::memory_order_acquire)) {
        return SHUTDOWN_ERROR;
    }

J
Jin Hai 已提交
1196
    // step 1: check collection existence
1197 1198 1199
    bool has_collection;
    auto status = HasCollection(collection_id, has_collection);
    if (!has_collection) {
1200
        LOG_ENGINE_ERROR_ << "Collection " << collection_id << " does not exist: ";
J
Jin Hai 已提交
1201
        return Status(DB_NOT_FOUND, "Collection does not exist");
1202 1203 1204 1205 1206 1207
    }
    if (!status.ok()) {
        return status;
    }

    //  step 2: find segment
G
groot 已提交
1208 1209
    meta::FilesHolder files_holder;
    status = meta_ptr_->GetCollectionFilesBySegmentId(segment_id, files_holder);
1210 1211 1212 1213
    if (!status.ok()) {
        return status;
    }

G
groot 已提交
1214
    milvus::engine::meta::SegmentsSchema& collection_files = files_holder.HoldFiles();
1215
    if (collection_files.empty()) {
1216 1217 1218
        return Status(DB_NOT_FOUND, "Segment does not exist");
    }

J
Jin Hai 已提交
1219
    // check the segment is belong to this collection
1220
    if (collection_files[0].collection_id_ != collection_id) {
J
Jin Hai 已提交
1221
        // the segment could be in a partition under this collection
1222 1223 1224 1225
        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 已提交
1226
            return Status(DB_NOT_FOUND, "Segment does not belong to this collection");
1227 1228 1229 1230 1231
        }
    }

    // step 3: load segment ids and delete offset
    std::string segment_dir;
1232
    engine::utils::GetParentPath(collection_files[0].location_, segment_dir);
1233 1234 1235 1236 1237 1238 1239 1240 1241 1242 1243 1244 1245 1246 1247 1248 1249 1250 1251 1252 1253 1254 1255 1256 1257
    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 已提交
1258

G
groot 已提交
1259
    return status;
X
Xu Peng 已提交
1260 1261
}

1262
Status
1263
DBImpl::GetVectorsByIdHelper(const std::string& collection_id, const IDNumbers& id_array,
G
groot 已提交
1264 1265 1266
                             std::vector<engine::VectorsData>& vectors, meta::FilesHolder& files_holder) {
    // 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();
1267
    LOG_ENGINE_DEBUG_ << "Getting vector by id in " << files.size() << " files, id count = " << id_array.size();
J
Jin Hai 已提交
1268

1269 1270 1271 1272 1273 1274 1275
    // 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;
1276

1277 1278 1279
    vectors.clear();

    IDNumbers temp_ids = id_array;
1280 1281 1282 1283 1284 1285 1286 1287
    for (auto& file : files) {
        // Load bloom filter
        std::string segment_dir;
        engine::utils::GetParentPath(file.location_, segment_dir);
        segment::SegmentReader segment_reader(segment_dir);
        segment::IdBloomFilterPtr id_bloom_filter_ptr;
        segment_reader.LoadBloomFilter(id_bloom_filter_ptr);

1288 1289 1290 1291 1292 1293 1294 1295 1296 1297 1298
        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);
1299 1300 1301
                if (!status.ok()) {
                    return status;
                }
1302 1303 1304 1305 1306 1307 1308 1309

                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);
1310
                    if (!status.ok()) {
J
Jin Hai 已提交
1311
                        LOG_ENGINE_ERROR_ << status.message();
1312 1313
                        return status;
                    }
1314 1315 1316 1317 1318 1319 1320 1321 1322 1323 1324 1325 1326 1327 1328 1329 1330 1331 1332 1333 1334 1335 1336 1337 1338 1339
                    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;
1340 1341 1342
                    }
                }
            }
1343 1344 1345

            it++;
        }
G
groot 已提交
1346 1347 1348

        // unmark file, allow the file to be deleted
        files_holder.UnmarkFile(file);
1349 1350 1351 1352 1353 1354 1355 1356
    }

    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 已提交
1357 1358
            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
1359
        }
1360
        vectors.emplace_back(data);
1361 1362
    }

1363 1364
    if (vectors.empty()) {
        std::string msg = "Vectors not found in collection " + collection_id;
J
Jin Hai 已提交
1365 1366 1367
        LOG_ENGINE_DEBUG_ << msg;
    }

1368 1369 1370
    return Status::OK();
}

S
starlord 已提交
1371
Status
G
groot 已提交
1372 1373
DBImpl::CreateIndex(const std::shared_ptr<server::Context>& context, const std::string& collection_id,
                    const CollectionIndex& index) {
1374
    if (!initialized_.load(std::memory_order_acquire)) {
G
groot 已提交
1375 1376 1377
        return SHUTDOWN_ERROR;
    }

G
groot 已提交
1378
    // serialize memory data
1379 1380
    //    std::set<std::string> sync_collection_ids;
    //    auto status = SyncMemData(sync_collection_ids);
1381
    auto status = Flush();
G
groot 已提交
1382

S
starlord 已提交
1383 1384 1385
    {
        std::unique_lock<std::mutex> lock(build_index_mutex_);

S
starlord 已提交
1386
        // step 1: check index difference
1387
        CollectionIndex old_index;
J
Jin Hai 已提交
1388
        status = DescribeIndex(collection_id, old_index);
S
starlord 已提交
1389
        if (!status.ok()) {
1390
            LOG_ENGINE_ERROR_ << "Failed to get collection index info for collection: " << collection_id;
S
starlord 已提交
1391 1392 1393
            return status;
        }

S
starlord 已提交
1394
        // step 2: update index info
1395 1396
        CollectionIndex new_index = index;
        new_index.metric_type_ = old_index.metric_type_;  // dont change metric type, it was defined by CreateCollection
S
starlord 已提交
1397
        if (!utils::IsSameIndex(old_index, new_index)) {
1398
            status = UpdateCollectionIndexRecursively(collection_id, new_index);
S
starlord 已提交
1399 1400 1401 1402 1403 1404
            if (!status.ok()) {
                return status;
            }
        }
    }

G
groot 已提交
1405 1406 1407 1408
    // step 3: wait merge file thread finished to avoid duplicate data bug
    WaitMergeFileFinish();  // let merge file thread finish
    StartMergeTask(true);   // start force-merge task
    WaitMergeFileFinish();  // let force-merge file thread finish
1409

S
starlord 已提交
1410
    // step 4: wait and build index
1411
    status = index_failed_checker_.CleanFailedIndexFileOfCollection(collection_id);
G
groot 已提交
1412
    status = WaitCollectionIndexRecursively(context, collection_id, index);
S
starlord 已提交
1413

G
groot 已提交
1414
    return status;
S
starlord 已提交
1415 1416
}

S
starlord 已提交
1417
Status
1418
DBImpl::DescribeIndex(const std::string& collection_id, CollectionIndex& index) {
1419
    if (!initialized_.load(std::memory_order_acquire)) {
G
groot 已提交
1420 1421 1422
        return SHUTDOWN_ERROR;
    }

1423
    return meta_ptr_->DescribeCollectionIndex(collection_id, index);
S
starlord 已提交
1424 1425
}

S
starlord 已提交
1426
Status
J
Jin Hai 已提交
1427
DBImpl::DropIndex(const std::string& collection_id) {
1428
    if (!initialized_.load(std::memory_order_acquire)) {
G
groot 已提交
1429 1430 1431
        return SHUTDOWN_ERROR;
    }

1432
    LOG_ENGINE_DEBUG_ << "Drop index for collection: " << collection_id;
G
groot 已提交
1433 1434 1435
    auto status = DropCollectionIndexRecursively(collection_id);
    StartMergeTask();  // merge small files after drop index
    return status;
S
starlord 已提交
1436 1437
}

S
starlord 已提交
1438
Status
1439 1440 1441
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) {
1442
    if (!initialized_.load(std::memory_order_acquire)) {
G
groot 已提交
1443
        return SHUTDOWN_ERROR;
S
starlord 已提交
1444 1445
    }

1446 1447 1448 1449 1450 1451 1452 1453 1454 1455 1456 1457 1458 1459 1460 1461 1462 1463 1464 1465 1466 1467 1468 1469 1470 1471 1472 1473 1474 1475 1476 1477 1478 1479 1480 1481 1482 1483 1484 1485 1486 1487 1488 1489 1490 1491 1492 1493 1494 1495 1496 1497 1498 1499 1500 1501 1502 1503 1504 1505 1506 1507 1508 1509 1510 1511 1512 1513 1514 1515 1516 1517 1518 1519 1520 1521 1522 1523 1524 1525 1526 1527 1528 1529 1530 1531 1532 1533 1534 1535 1536 1537 1538 1539 1540 1541 1542 1543 1544 1545 1546 1547 1548 1549 1550 1551 1552 1553 1554 1555 1556 1557 1558 1559 1560 1561 1562
    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;
    status = GetVectorsByID(collection_id, id_array, vectors);
    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 已提交
1563 1564
}

1565 1566 1567 1568 1569 1570 1571 1572 1573 1574 1575 1576 1577
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 已提交
1578
    meta::FilesHolder files_holder;
1579 1580 1581
    if (partition_tags.empty()) {
        // no partition tag specified, means search in whole table
        // get all table files from parent table
G
groot 已提交
1582
        status = meta_ptr_->FilesToSearch(collection_id, files_holder);
1583 1584 1585 1586 1587 1588 1589 1590 1591 1592
        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 已提交
1593
            status = meta_ptr_->FilesToSearch(schema.collection_id_, files_holder);
1594
            if (!status.ok()) {
G
groot 已提交
1595
                return Status(DB_ERROR, "get files to search failed in HybridQuery");
1596 1597 1598
            }
        }

G
groot 已提交
1599 1600
        if (files_holder.HoldFiles().empty()) {
            return Status::OK();  // no files to search
1601 1602 1603 1604 1605 1606 1607
        }
    } 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 已提交
1608
            status = meta_ptr_->FilesToSearch(partition_name, files_holder);
1609
            if (!status.ok()) {
G
groot 已提交
1610
                return Status(DB_ERROR, "get files to search failed in HybridQuery");
1611 1612 1613
            }
        }

G
groot 已提交
1614
        if (files_holder.HoldFiles().empty()) {
1615 1616 1617 1618 1619
            return Status::OK();
        }
    }

    cache::CpuCacheMgr::GetInstance()->PrintInfo();  // print cache info before query
G
groot 已提交
1620
    status = HybridQueryAsync(query_ctx, collection_id, files_holder, hybrid_search_context, general_query, attr_type,
1621 1622 1623 1624 1625 1626 1627 1628 1629 1630 1631
                              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 已提交
1632
Status
J
Jin Hai 已提交
1633
DBImpl::Query(const std::shared_ptr<server::Context>& context, const std::string& collection_id,
1634 1635
              const std::vector<std::string>& partition_tags, uint64_t k, const milvus::json& extra_params,
              const VectorsData& vectors, ResultIds& result_ids, ResultDistances& result_distances) {
1636
    milvus::server::ContextChild tracer(context, "Query");
Z
Zhiru Zhu 已提交
1637

1638
    if (!initialized_.load(std::memory_order_acquire)) {
G
groot 已提交
1639
        return SHUTDOWN_ERROR;
S
starlord 已提交
1640 1641
    }

G
groot 已提交
1642
    Status status;
G
groot 已提交
1643
    meta::FilesHolder files_holder;
G
groot 已提交
1644
    if (partition_tags.empty()) {
J
Jin Hai 已提交
1645 1646
        // no partition tag specified, means search in whole collection
        // get all collection files from parent collection
G
groot 已提交
1647
        status = meta_ptr_->FilesToSearch(collection_id, files_holder);
G
groot 已提交
1648 1649 1650 1651
        if (!status.ok()) {
            return status;
        }

J
Jin Hai 已提交
1652 1653
        std::vector<meta::CollectionSchema> partition_array;
        status = meta_ptr_->ShowPartitions(collection_id, partition_array);
G
groot 已提交
1654
        for (auto& schema : partition_array) {
G
groot 已提交
1655
            status = meta_ptr_->FilesToSearch(schema.collection_id_, files_holder);
1656 1657
        }

G
groot 已提交
1658 1659
        if (files_holder.HoldFiles().empty()) {
            return Status::OK();  // no files to search
G
groot 已提交
1660 1661 1662 1663
        }
    } else {
        // get files from specified partitions
        std::set<std::string> partition_name_array;
J
Jin Hai 已提交
1664
        status = GetPartitionsByTags(collection_id, partition_tags, partition_name_array);
T
Tinkerrr 已提交
1665 1666 1667
        if (!status.ok()) {
            return status;  // didn't match any partition.
        }
G
groot 已提交
1668 1669

        for (auto& partition_name : partition_name_array) {
G
groot 已提交
1670
            status = meta_ptr_->FilesToSearch(partition_name, files_holder);
1671 1672
        }

G
groot 已提交
1673 1674
        if (files_holder.HoldFiles().empty()) {
            return Status::OK();  // no files to search
1675 1676 1677
        }
    }

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

S
starlord 已提交
1682
    return status;
G
groot 已提交
1683
}
X
Xu Peng 已提交
1684

S
starlord 已提交
1685
Status
1686 1687 1688
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) {
1689
    milvus::server::ContextChild tracer(context, "Query by file id");
Z
Zhiru Zhu 已提交
1690

1691
    if (!initialized_.load(std::memory_order_acquire)) {
G
groot 已提交
1692
        return SHUTDOWN_ERROR;
S
starlord 已提交
1693 1694
    }

S
starlord 已提交
1695
    // get specified files
1696
    std::vector<size_t> ids;
Y
Yu Kun 已提交
1697
    for (auto& id : file_ids) {
1698
        std::string::size_type sz;
J
jinhai 已提交
1699
        ids.push_back(std::stoul(id, &sz));
1700 1701
    }

G
groot 已提交
1702 1703
    meta::FilesHolder files_holder;
    auto status = meta_ptr_->FilesByID(ids, files_holder);
1704 1705
    if (!status.ok()) {
        return status;
1706 1707
    }

G
groot 已提交
1708
    milvus::engine::meta::SegmentsSchema& search_files = files_holder.HoldFiles();
1709
    if (search_files.empty()) {
S
starlord 已提交
1710
        return Status(DB_ERROR, "Invalid file id");
G
groot 已提交
1711 1712
    }

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

S
starlord 已提交
1717
    return status;
1718 1719
}

S
starlord 已提交
1720
Status
Y
Yu Kun 已提交
1721
DBImpl::Size(uint64_t& result) {
1722
    if (!initialized_.load(std::memory_order_acquire)) {
G
groot 已提交
1723
        return SHUTDOWN_ERROR;
S
starlord 已提交
1724 1725
    }

S
starlord 已提交
1726
    return meta_ptr_->Size(result);
S
starlord 已提交
1727 1728 1729
}

///////////////////////////////////////////////////////////////////////////////////////////////////////////////////
S
starlord 已提交
1730
// internal methods
S
starlord 已提交
1731
///////////////////////////////////////////////////////////////////////////////////////////////////////////////////
S
starlord 已提交
1732
Status
G
groot 已提交
1733
DBImpl::QueryAsync(const std::shared_ptr<server::Context>& context, meta::FilesHolder& files_holder, uint64_t k,
1734 1735
                   const milvus::json& extra_params, const VectorsData& vectors, ResultIds& result_ids,
                   ResultDistances& result_distances) {
1736
    milvus::server::ContextChild tracer(context, "Query Async");
G
groot 已提交
1737
    server::CollectQueryMetrics metrics(vectors.vector_count_);
Y
Yu Kun 已提交
1738

G
groot 已提交
1739
    milvus::engine::meta::SegmentsSchema& files = files_holder.HoldFiles();
G
groot 已提交
1740 1741 1742
    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);
1743
        LOG_ENGINE_ERROR_ << msg;
G
groot 已提交
1744 1745 1746
        return Status(DB_ERROR, msg);
    }

S
starlord 已提交
1747
    TimeRecorder rc("");
G
groot 已提交
1748

1749
    // step 1: construct search job
1750
    LOG_ENGINE_DEBUG_ << LogOut("Engine query begin, index file count: %ld", files.size());
1751
    scheduler::SearchJobPtr job = std::make_shared<scheduler::SearchJob>(tracer.Context(), k, extra_params, vectors);
Y
Yu Kun 已提交
1752
    for (auto& file : files) {
J
Jin Hai 已提交
1753
        scheduler::SegmentSchemaPtr file_ptr = std::make_shared<meta::SegmentSchema>(file);
W
wxyu 已提交
1754
        job->AddIndexFile(file_ptr);
G
groot 已提交
1755 1756
    }

1757 1758 1759
    // Suspend builder
    SuspendIfFirst();

1760
    // step 2: put search job to scheduler and wait result
S
starlord 已提交
1761
    scheduler::JobMgrInst::GetInstance()->Put(job);
W
wxyu 已提交
1762
    job->WaitResult();
1763

1764 1765 1766
    // Resume builder
    ResumeIfLast();

G
groot 已提交
1767
    files_holder.ReleaseFiles();
W
wxyu 已提交
1768 1769
    if (!job->GetStatus().ok()) {
        return job->GetStatus();
1770
    }
G
groot 已提交
1771

1772
    // step 3: construct results
G
groot 已提交
1773 1774
    result_ids = job->GetResultIds();
    result_distances = job->GetResultDistances();
S
starlord 已提交
1775
    rc.ElapseFromBegin("Engine query totally cost");
G
groot 已提交
1776 1777 1778 1779

    return Status::OK();
}

1780 1781
Status
DBImpl::HybridQueryAsync(const std::shared_ptr<server::Context>& context, const std::string& table_id,
G
groot 已提交
1782
                         meta::FilesHolder& files_holder, context::HybridSearchContextPtr hybrid_search_context,
1783 1784 1785 1786 1787 1788 1789 1790 1791 1792 1793 1794 1795 1796 1797 1798 1799 1800 1801 1802 1803 1804 1805 1806 1807 1808 1809 1810 1811
                         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 已提交
1812 1813
    milvus::engine::meta::SegmentsSchema& files = files_holder.HoldFiles();
    LOG_ENGINE_DEBUG_ << LogOut("Engine query begin, index file count: %ld", files_holder.HoldFiles().size());
1814 1815 1816 1817 1818 1819 1820 1821 1822 1823 1824
    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 已提交
1825
    files_holder.ReleaseFiles();
1826 1827 1828 1829 1830 1831 1832 1833 1834 1835 1836 1837 1838 1839 1840 1841
    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 已提交
1842
void
G
groot 已提交
1843
DBImpl::BackgroundIndexThread() {
Y
yu yunfeng 已提交
1844
    server::SystemInfo::GetInstance().Init();
X
Xu Peng 已提交
1845
    while (true) {
1846
        if (!initialized_.load(std::memory_order_acquire)) {
1847 1848
            WaitMergeFileFinish();
            WaitBuildIndexFinish();
S
starlord 已提交
1849

1850
            LOG_ENGINE_DEBUG_ << "DB background thread exit";
G
groot 已提交
1851 1852
            break;
        }
X
Xu Peng 已提交
1853

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

G
groot 已提交
1856
        WaitMergeFileFinish();
G
groot 已提交
1857 1858
        StartBuildIndexTask();
    }
X
Xu Peng 已提交
1859 1860
}

S
starlord 已提交
1861 1862
void
DBImpl::WaitMergeFileFinish() {
1863
    //    LOG_ENGINE_DEBUG_ << "Begin WaitMergeFileFinish";
1864 1865
    std::lock_guard<std::mutex> lck(merge_result_mutex_);
    for (auto& iter : merge_thread_results_) {
1866 1867
        iter.wait();
    }
1868
    //    LOG_ENGINE_DEBUG_ << "End WaitMergeFileFinish";
1869 1870
}

S
starlord 已提交
1871 1872
void
DBImpl::WaitBuildIndexFinish() {
1873
    //    LOG_ENGINE_DEBUG_ << "Begin WaitBuildIndexFinish";
1874
    std::lock_guard<std::mutex> lck(index_result_mutex_);
Y
Yu Kun 已提交
1875
    for (auto& iter : index_thread_results_) {
1876 1877
        iter.wait();
    }
1878
    //    LOG_ENGINE_DEBUG_ << "End WaitBuildIndexFinish";
1879 1880
}

S
starlord 已提交
1881 1882
void
DBImpl::StartMetricTask() {
G
groot 已提交
1883
    server::Metrics::GetInstance().KeepingAliveCounterIncrement(BACKGROUND_METRIC_INTERVAL);
G
groot 已提交
1884 1885
    int64_t cache_usage = cache::CpuCacheMgr::GetInstance()->CacheUsage();
    int64_t cache_total = cache::CpuCacheMgr::GetInstance()->CacheCapacity();
S
shengjh 已提交
1886 1887
    fiu_do_on("DBImpl.StartMetricTask.InvalidTotalCache", cache_total = 0);

J
JinHai-CN 已提交
1888 1889 1890 1891 1892 1893 1894
    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 已提交
1895
    server::Metrics::GetInstance().GpuCacheUsageGaugeSet();
G
groot 已提交
1896 1897 1898 1899 1900 1901 1902 1903
    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 已提交
1904

K
kun yu 已提交
1905
    server::Metrics::GetInstance().CPUCoreUsagePercentSet();
K
kun yu 已提交
1906 1907
    server::Metrics::GetInstance().GPUTemperature();
    server::Metrics::GetInstance().CPUTemperature();
1908
    server::Metrics::GetInstance().PushToGateway();
G
groot 已提交
1909 1910
}

S
starlord 已提交
1911
void
G
groot 已提交
1912
DBImpl::StartMergeTask(bool force_merge_all) {
1913
    // LOG_ENGINE_DEBUG_ << "Begin StartMergeTask";
1914
    // merge task has been finished?
1915
    {
1916 1917
        std::lock_guard<std::mutex> lck(merge_result_mutex_);
        if (!merge_thread_results_.empty()) {
1918
            std::chrono::milliseconds span(10);
1919 1920
            if (merge_thread_results_.back().wait_for(span) == std::future_status::ready) {
                merge_thread_results_.pop_back();
1921
            }
G
groot 已提交
1922 1923
        }
    }
X
Xu Peng 已提交
1924

1925
    // add new merge task
1926
    {
1927 1928
        std::lock_guard<std::mutex> lck(merge_result_mutex_);
        if (merge_thread_results_.empty()) {
1929 1930
            // collect merge files for all collections(if merge_collection_ids_ is empty) for two reasons:
            // 1. other collections may still has un-merged files
1931
            // 2. server may be closed unexpected, these un-merge files need to be merged when server restart
1932 1933 1934 1935 1936
            if (merge_collection_ids_.empty()) {
                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_);
1937 1938 1939 1940
                }
            }

            // start merge file thread
1941
            merge_thread_results_.push_back(
G
groot 已提交
1942
                merge_thread_pool_.enqueue(&DBImpl::BackgroundMerge, this, merge_collection_ids_, force_merge_all));
1943
            merge_collection_ids_.clear();
1944
        }
G
groot 已提交
1945
    }
1946

1947
    // LOG_ENGINE_DEBUG_ << "End StartMergeTask";
X
Xu Peng 已提交
1948 1949
}

1950
Status
G
groot 已提交
1951
DBImpl::MergeHybridFiles(const std::string& collection_id, meta::FilesHolder& files_holder) {
1952 1953 1954 1955 1956 1957 1958 1959 1960 1961 1962 1963 1964 1965 1966 1967 1968 1969 1970 1971 1972 1973 1974 1975 1976 1977 1978
    // 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 已提交
1979 1980
    // attention: here is a copy, not reference, since files_holder.UnmarkFile will change the array internal
    milvus::engine::meta::SegmentsSchema files = files_holder.HoldFiles();
1981 1982 1983 1984 1985
    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 已提交
1986 1987 1988

        files_holder.UnmarkFile(file);

1989 1990 1991
        auto file_schema = file;
        file_schema.file_type_ = meta::SegmentSchema::TO_DELETE;
        updated.push_back(file_schema);
C
Cai Yudong 已提交
1992
        int64_t size = segment_writer_ptr->Size();
1993 1994 1995 1996 1997 1998 1999 2000 2001 2002 2003 2004 2005 2006 2007 2008 2009 2010 2011 2012 2013 2014 2015 2016 2017 2018 2019 2020 2021 2022 2023 2024
        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 已提交
2025
        table_file.file_type_ = (segment_writer_ptr->Size() >= (size_t)(table_file.index_file_size_))
2026 2027 2028 2029 2030 2031 2032 2033 2034 2035 2036 2037 2038 2039 2040 2041 2042 2043 2044
                                    ? 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 已提交
2045
void
G
groot 已提交
2046
DBImpl::BackgroundMerge(std::set<std::string> collection_ids, bool force_merge_all) {
2047
    // LOG_ENGINE_TRACE_ << " Background merge thread start";
S
starlord 已提交
2048

G
groot 已提交
2049
    Status status;
2050
    for (auto& collection_id : collection_ids) {
G
groot 已提交
2051 2052
        const std::lock_guard<std::mutex> lock(flush_merge_compact_mutex_);

G
groot 已提交
2053 2054 2055 2056 2057
        auto old_strategy = merge_mgr_ptr_->Strategy();
        if (force_merge_all) {
            merge_mgr_ptr_->UseStrategy(MergeStrategyType::ADAPTIVE);
        }

G
groot 已提交
2058
        auto status = merge_mgr_ptr_->MergeFiles(collection_id);
G
groot 已提交
2059
        merge_mgr_ptr_->UseStrategy(old_strategy);
G
groot 已提交
2060
        if (!status.ok()) {
G
groot 已提交
2061 2062
            LOG_ENGINE_ERROR_ << "Failed to get merge files for collection: " << collection_id
                              << " reason:" << status.message();
G
groot 已提交
2063
        }
S
starlord 已提交
2064

2065
        if (!initialized_.load(std::memory_order_acquire)) {
G
groot 已提交
2066
            LOG_ENGINE_DEBUG_ << "Server will shutdown, skip merge action for collection: " << collection_id;
S
starlord 已提交
2067 2068
            break;
        }
G
groot 已提交
2069
    }
X
Xu Peng 已提交
2070

G
groot 已提交
2071
    meta_ptr_->Archive();
Z
update  
zhiru 已提交
2072

2073
    {
G
groot 已提交
2074
        uint64_t timeout = (options_.file_cleanup_timeout_ >= 0) ? options_.file_cleanup_timeout_ : 10;
G
groot 已提交
2075
        uint64_t ttl = timeout * meta::SECOND;  // default: file will be hard-deleted few seconds after soft-deleted
2076
        meta_ptr_->CleanUpFilesWithTTL(ttl);
Z
update  
zhiru 已提交
2077
    }
S
starlord 已提交
2078

2079
    // LOG_ENGINE_TRACE_ << " Background merge thread exit";
G
groot 已提交
2080
}
X
Xu Peng 已提交
2081

S
starlord 已提交
2082
void
G
groot 已提交
2083
DBImpl::StartBuildIndexTask() {
S
starlord 已提交
2084
    // build index has been finished?
2085 2086 2087 2088 2089 2090 2091
    {
        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 已提交
2092 2093 2094
        }
    }

S
starlord 已提交
2095
    // add new build index task
2096 2097 2098
    {
        std::lock_guard<std::mutex> lck(index_result_mutex_);
        if (index_thread_results_.empty()) {
S
starlord 已提交
2099
            index_thread_results_.push_back(index_thread_pool_.enqueue(&DBImpl::BackgroundBuildIndex, this));
2100
        }
G
groot 已提交
2101
    }
X
Xu Peng 已提交
2102 2103
}

S
starlord 已提交
2104 2105
void
DBImpl::BackgroundBuildIndex() {
P
peng.xu 已提交
2106
    std::unique_lock<std::mutex> lock(build_index_mutex_);
G
groot 已提交
2107 2108 2109 2110
    meta::FilesHolder files_holder;
    meta_ptr_->FilesToIndex(files_holder);

    milvus::engine::meta::SegmentsSchema& to_index_files = files_holder.HoldFiles();
2111
    Status status = index_failed_checker_.IgnoreFailedIndexFiles(to_index_files);
2112

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

2116
        // step 2: put build index task to scheduler
J
Jin Hai 已提交
2117
        std::vector<std::pair<scheduler::BuildIndexJobPtr, scheduler::SegmentSchemaPtr>> job2file_map;
2118
        for (auto& file : to_index_files) {
G
groot 已提交
2119
            scheduler::BuildIndexJobPtr job = std::make_shared<scheduler::BuildIndexJob>(meta_ptr_, options_);
J
Jin Hai 已提交
2120
            scheduler::SegmentSchemaPtr file_ptr = std::make_shared<meta::SegmentSchema>(file);
2121
            job->AddToIndexFiles(file_ptr);
G
groot 已提交
2122
            scheduler::JobMgrInst::GetInstance()->Put(job);
G
groot 已提交
2123
            job2file_map.push_back(std::make_pair(job, file_ptr));
2124
        }
G
groot 已提交
2125

G
groot 已提交
2126
        // step 3: wait build index finished and mark failed files
G
groot 已提交
2127 2128
        for (auto iter = job2file_map.begin(); iter != job2file_map.end(); ++iter) {
            scheduler::BuildIndexJobPtr job = iter->first;
J
Jin Hai 已提交
2129
            meta::SegmentSchema& file_schema = *(iter->second.get());
G
groot 已提交
2130 2131 2132
            job->WaitBuildIndexFinish();
            if (!job->GetStatus().ok()) {
                Status status = job->GetStatus();
2133
                LOG_ENGINE_ERROR_ << "Building index job " << job->id() << " failed: " << status.ToString();
G
groot 已提交
2134

2135
                index_failed_checker_.MarkFailedIndexFile(file_schema, status.message());
G
groot 已提交
2136
            } else {
2137
                LOG_ENGINE_DEBUG_ << "Building index job " << job->id() << " succeed.";
G
groot 已提交
2138 2139

                index_failed_checker_.MarkSucceedIndexFile(file_schema);
G
groot 已提交
2140
            }
G
groot 已提交
2141 2142
            status = files_holder.UnmarkFile(file_schema);
            LOG_ENGINE_DEBUG_ << "Finish build index file " << file_schema.file_id_;
2143
        }
G
groot 已提交
2144

2145
        LOG_ENGINE_DEBUG_ << "Background build index thread finished";
G
groot 已提交
2146
        index_req_swn_.Notify();  // notify CreateIndex check circle
Y
Yu Kun 已提交
2147
    }
X
Xu Peng 已提交
2148 2149
}

G
groot 已提交
2150
Status
J
Jin Hai 已提交
2151
DBImpl::GetFilesToBuildIndex(const std::string& collection_id, const std::vector<int>& file_types,
G
groot 已提交
2152 2153 2154
                             meta::FilesHolder& files_holder) {
    files_holder.ReleaseFiles();
    auto status = meta_ptr_->FilesByType(collection_id, file_types, files_holder);
G
groot 已提交
2155

G
groot 已提交
2156 2157 2158 2159 2160 2161 2162 2163 2164 2165 2166
    // 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 已提交
2167 2168 2169 2170 2171
    }

    return Status::OK();
}

2172
Status
J
Jin Hai 已提交
2173 2174
DBImpl::GetPartitionByTag(const std::string& collection_id, const std::string& partition_tag,
                          std::string& partition_name) {
2175 2176 2177
    Status status;

    if (partition_tag.empty()) {
J
Jin Hai 已提交
2178
        partition_name = collection_id;
2179 2180 2181 2182 2183 2184 2185 2186

    } 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 已提交
2187
            partition_name = collection_id;
2188 2189 2190
            return status;
        }

J
Jin Hai 已提交
2191
        status = meta_ptr_->GetPartitionName(collection_id, partition_tag, partition_name);
2192
        if (!status.ok()) {
2193
            LOG_ENGINE_ERROR_ << status.message();
2194 2195 2196 2197 2198 2199
        }
    }

    return status;
}

G
groot 已提交
2200
Status
J
Jin Hai 已提交
2201
DBImpl::GetPartitionsByTags(const std::string& collection_id, const std::vector<std::string>& partition_tags,
G
groot 已提交
2202
                            std::set<std::string>& partition_name_array) {
J
Jin Hai 已提交
2203 2204
    std::vector<meta::CollectionSchema> partition_array;
    auto status = meta_ptr_->ShowPartitions(collection_id, partition_array);
G
groot 已提交
2205 2206

    for (auto& tag : partition_tags) {
2207 2208 2209 2210
        // 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);
2211 2212

        if (valid_tag == milvus::engine::DEFAULT_PARTITON_TAG) {
J
Jin Hai 已提交
2213
            partition_name_array.insert(collection_id);
2214 2215 2216
            return status;
        }

G
groot 已提交
2217
        for (auto& schema : partition_array) {
2218
            if (server::StringHelpFunctions::IsRegexMatch(schema.partition_tag_, valid_tag)) {
J
Jin Hai 已提交
2219
                partition_name_array.insert(schema.collection_id_);
G
groot 已提交
2220 2221 2222 2223
            }
        }
    }

T
Tinkerrr 已提交
2224
    if (partition_name_array.empty()) {
G
groot 已提交
2225
        return Status(DB_PARTITION_NOT_FOUND, "The specified partiton does not exist");
T
Tinkerrr 已提交
2226 2227
    }

G
groot 已提交
2228 2229 2230 2231
    return Status::OK();
}

Status
2232
DBImpl::DropCollectionRecursively(const std::string& collection_id) {
J
Jin Hai 已提交
2233
    // dates partly delete files of the collection but currently we don't support
2234
    LOG_ENGINE_DEBUG_ << "Prepare to delete collection " << collection_id;
G
groot 已提交
2235 2236

    Status status;
2237
    if (options_.wal_enable_) {
2238
        wal_mgr_->DropCollection(collection_id);
G
groot 已提交
2239 2240
    }

2241 2242 2243
    status = mem_mgr_->EraseMemVector(collection_id);   // not allow insert
    status = meta_ptr_->DropCollection(collection_id);  // soft delete collection
    index_failed_checker_.CleanFailedIndexFileOfCollection(collection_id);
2244

J
Jin Hai 已提交
2245
    // scheduler will determine when to delete collection files
2246
    auto nres = scheduler::ResMgrInst::GetInstance()->GetNumOfComputeResource();
J
Jin Hai 已提交
2247
    scheduler::DeleteJobPtr job = std::make_shared<scheduler::DeleteJob>(collection_id, meta_ptr_, nres);
2248 2249 2250
    scheduler::JobMgrInst::GetInstance()->Put(job);
    job->WaitAndDelete();

J
Jin Hai 已提交
2251 2252
    std::vector<meta::CollectionSchema> partition_array;
    status = meta_ptr_->ShowPartitions(collection_id, partition_array);
G
groot 已提交
2253
    for (auto& schema : partition_array) {
2254 2255
        status = DropCollectionRecursively(schema.collection_id_);
        fiu_do_on("DBImpl.DropCollectionRecursively.failed", status = Status(DB_ERROR, ""));
G
groot 已提交
2256 2257 2258 2259 2260 2261 2262 2263 2264
        if (!status.ok()) {
            return status;
        }
    }

    return Status::OK();
}

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

2268 2269
    auto status = meta_ptr_->UpdateCollectionIndex(collection_id, index);
    fiu_do_on("DBImpl.UpdateCollectionIndexRecursively.fail_update_collection_index",
S
shengjh 已提交
2270
              status = Status(DB_META_TRANSACTION_FAILED, ""));
G
groot 已提交
2271
    if (!status.ok()) {
2272
        LOG_ENGINE_ERROR_ << "Failed to update collection index info for collection: " << collection_id;
G
groot 已提交
2273 2274 2275
        return status;
    }

J
Jin Hai 已提交
2276 2277
    std::vector<meta::CollectionSchema> partition_array;
    status = meta_ptr_->ShowPartitions(collection_id, partition_array);
A
AzAz 已提交
2278 2279 2280
    if (!status.ok()) {
        return status;
    }
G
groot 已提交
2281
    for (auto& schema : partition_array) {
2282
        status = UpdateCollectionIndexRecursively(schema.collection_id_, index);
G
groot 已提交
2283 2284 2285 2286 2287 2288 2289 2290 2291
        if (!status.ok()) {
            return status;
        }
    }

    return Status::OK();
}

Status
G
groot 已提交
2292 2293
DBImpl::WaitCollectionIndexRecursively(const std::shared_ptr<server::Context>& context,
                                       const std::string& collection_id, const CollectionIndex& index) {
G
groot 已提交
2294 2295 2296
    // 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;
2297
    if (utils::IsRawIndexType(index.engine_type_)) {
G
groot 已提交
2298
        file_types = {
J
Jin Hai 已提交
2299 2300
            static_cast<int32_t>(meta::SegmentSchema::NEW),
            static_cast<int32_t>(meta::SegmentSchema::NEW_MERGE),
G
groot 已提交
2301 2302 2303
        };
    } else {
        file_types = {
J
Jin Hai 已提交
2304 2305 2306
            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 已提交
2307 2308 2309 2310
        };
    }

    // get files to build index
G
groot 已提交
2311 2312 2313 2314
    {
        meta::FilesHolder files_holder;
        auto status = GetFilesToBuildIndex(collection_id, file_types, files_holder);
        int times = 1;
G
groot 已提交
2315
        uint64_t repeat = 0;
G
groot 已提交
2316
        while (!files_holder.HoldFiles().empty()) {
G
groot 已提交
2317 2318 2319 2320 2321 2322
            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 已提交
2323
            }
G
groot 已提交
2324

G
groot 已提交
2325 2326 2327 2328 2329 2330 2331 2332 2333 2334 2335 2336 2337 2338
            index_req_swn_.Wait_For(std::chrono::seconds(1));

            // client break the connection, no need to block, check every 1 second
            if (context->IsConnectionBroken()) {
                LOG_ENGINE_DEBUG_ << "Client connection broken, build index in background";
                break;  // just break, not return, continue to update partitions files to to_index
            }

            // check to_index files every 5 seconds
            repeat++;
            if (repeat % WAIT_BUILD_INDEX_INTERVAL == 0) {
                GetFilesToBuildIndex(collection_id, file_types, files_holder);
                ++times;
            }
G
groot 已提交
2339 2340 2341 2342
        }
    }

    // build index for partition
J
Jin Hai 已提交
2343
    std::vector<meta::CollectionSchema> partition_array;
G
groot 已提交
2344
    auto status = meta_ptr_->ShowPartitions(collection_id, partition_array);
G
groot 已提交
2345
    for (auto& schema : partition_array) {
G
groot 已提交
2346
        status = WaitCollectionIndexRecursively(context, schema.collection_id_, index);
2347
        fiu_do_on("DBImpl.WaitCollectionIndexRecursively.fail_build_collection_Index_for_partition",
S
shengjh 已提交
2348
                  status = Status(DB_ERROR, ""));
G
groot 已提交
2349 2350 2351 2352 2353
        if (!status.ok()) {
            return status;
        }
    }

G
groot 已提交
2354
    // failed to build index for some files, return error
2355
    std::string err_msg;
2356 2357
    index_failed_checker_.GetErrMsgForCollection(collection_id, err_msg);
    fiu_do_on("DBImpl.WaitCollectionIndexRecursively.not_empty_err_msg", err_msg.append("fiu"));
2358 2359
    if (!err_msg.empty()) {
        return Status(DB_ERROR, err_msg);
G
groot 已提交
2360 2361
    }

G
groot 已提交
2362 2363
    LOG_ENGINE_DEBUG_ << "WaitCollectionIndexRecursively finished";

G
groot 已提交
2364 2365 2366 2367
    return Status::OK();
}

Status
2368
DBImpl::DropCollectionIndexRecursively(const std::string& collection_id) {
2369
    LOG_ENGINE_DEBUG_ << "Drop index for collection: " << collection_id;
2370 2371
    index_failed_checker_.CleanFailedIndexFileOfCollection(collection_id);
    auto status = meta_ptr_->DropCollectionIndex(collection_id);
G
groot 已提交
2372 2373 2374 2375 2376
    if (!status.ok()) {
        return status;
    }

    // drop partition index
J
Jin Hai 已提交
2377 2378
    std::vector<meta::CollectionSchema> partition_array;
    status = meta_ptr_->ShowPartitions(collection_id, partition_array);
G
groot 已提交
2379
    for (auto& schema : partition_array) {
2380 2381
        status = DropCollectionIndexRecursively(schema.collection_id_);
        fiu_do_on("DBImpl.DropCollectionIndexRecursively.fail_drop_collection_Index_for_partition",
S
shengjh 已提交
2382
                  status = Status(DB_ERROR, ""));
G
groot 已提交
2383 2384 2385 2386 2387 2388 2389 2390 2391
        if (!status.ok()) {
            return status;
        }
    }

    return Status::OK();
}

Status
2392
DBImpl::GetCollectionRowCountRecursively(const std::string& collection_id, uint64_t& row_count) {
G
groot 已提交
2393
    row_count = 0;
J
Jin Hai 已提交
2394
    auto status = meta_ptr_->Count(collection_id, row_count);
G
groot 已提交
2395 2396 2397 2398 2399
    if (!status.ok()) {
        return status;
    }

    // get partition row count
J
Jin Hai 已提交
2400 2401
    std::vector<meta::CollectionSchema> partition_array;
    status = meta_ptr_->ShowPartitions(collection_id, partition_array);
G
groot 已提交
2402
    for (auto& schema : partition_array) {
G
groot 已提交
2403
        uint64_t partition_row_count = 0;
2404 2405
        status = GetCollectionRowCountRecursively(schema.collection_id_, partition_row_count);
        fiu_do_on("DBImpl.GetCollectionRowCountRecursively.fail_get_collection_rowcount_for_partition",
S
shengjh 已提交
2406
                  status = Status(DB_ERROR, ""));
G
groot 已提交
2407 2408 2409 2410 2411 2412 2413 2414 2415 2416
        if (!status.ok()) {
            return status;
        }

        row_count += partition_row_count;
    }

    return Status::OK();
}

2417 2418 2419 2420
Status
DBImpl::ExecWalRecord(const wal::MXLogRecord& record) {
    fiu_return_on("DBImpl.ExexWalRecord.return", Status(););

G
groot 已提交
2421 2422
    auto collections_flushed = [&](const std::string collection_id,
                                   const std::set<std::string>& target_collection_names) -> uint64_t {
2423 2424
        uint64_t max_lsn = 0;
        if (options_.wal_enable_) {
G
groot 已提交
2425 2426
            uint64_t lsn = 0;
            for (auto& collection : target_collection_names) {
2427
                meta_ptr_->GetCollectionFlushLSN(collection, lsn);
2428 2429 2430 2431
                if (lsn > max_lsn) {
                    max_lsn = lsn;
                }
            }
G
groot 已提交
2432
            wal_mgr_->CollectionFlushed(collection_id, lsn);
2433 2434 2435
        }

        std::lock_guard<std::mutex> lck(merge_result_mutex_);
G
groot 已提交
2436
        for (auto& collection : target_collection_names) {
2437
            merge_collection_ids_.insert(collection);
2438 2439 2440 2441
        }
        return max_lsn;
    };

G
groot 已提交
2442 2443 2444 2445 2446 2447 2448 2449 2450 2451 2452 2453
    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);
        }

        std::lock_guard<std::mutex> lck(merge_result_mutex_);
        merge_collection_ids_.insert(target_collection_name);
    };

2454 2455 2456
    Status status;

    switch (record.type) {
2457 2458 2459 2460
        case wal::MXLogType::Entity: {
            std::string target_collection_name;
            status = GetPartitionByTag(record.collection_id, record.partition_tag, target_collection_name);
            if (!status.ok()) {
2461
                LOG_WAL_ERROR_ << LogOut("[%s][%ld] ", "insert", 0) << "Get partition fail: " << status.message();
2462 2463 2464
                return status;
            }

2465
            std::set<std::string> flushed_collections;
2466 2467 2468
            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,
2469
                                              record.attr_data, record.lsn, flushed_collections);
G
groot 已提交
2470 2471 2472
            if (!flushed_collections.empty()) {
                partition_flushed(record.collection_id, record.partition_tag, target_collection_name);
            }
2473 2474 2475 2476

            milvus::server::CollectInsertMetrics metrics(record.length, status);
            break;
        }
2477
        case wal::MXLogType::InsertBinary: {
2478 2479
            std::string target_collection_name;
            status = GetPartitionByTag(record.collection_id, record.partition_tag, target_collection_name);
2480
            if (!status.ok()) {
2481
                LOG_WAL_ERROR_ << LogOut("[%s][%ld] ", "insert", 0) << "Get partition fail: " << status.message();
2482 2483 2484
                return status;
            }

2485 2486
            std::set<std::string> flushed_collections;
            status = mem_mgr_->InsertVectors(target_collection_name, record.length, record.ids,
2487
                                             (record.data_size / record.length / sizeof(uint8_t)),
2488
                                             (const u_int8_t*)record.data, record.lsn, flushed_collections);
2489
            // even though !status.ok, run
G
groot 已提交
2490 2491 2492
            if (!flushed_collections.empty()) {
                partition_flushed(record.collection_id, record.partition_tag, target_collection_name);
            }
2493 2494 2495 2496 2497 2498 2499

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

        case wal::MXLogType::InsertVector: {
2500 2501
            std::string target_collection_name;
            status = GetPartitionByTag(record.collection_id, record.partition_tag, target_collection_name);
2502
            if (!status.ok()) {
2503
                LOG_WAL_ERROR_ << LogOut("[%s][%ld] ", "insert", 0) << "Get partition fail: " << status.message();
2504 2505 2506
                return status;
            }

2507 2508
            std::set<std::string> flushed_collections;
            status = mem_mgr_->InsertVectors(target_collection_name, record.length, record.ids,
2509
                                             (record.data_size / record.length / sizeof(float)),
2510
                                             (const float*)record.data, record.lsn, flushed_collections);
2511
            // even though !status.ok, run
G
groot 已提交
2512 2513 2514
            if (!flushed_collections.empty()) {
                partition_flushed(record.collection_id, record.partition_tag, target_collection_name);
            }
2515 2516 2517 2518 2519 2520 2521

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

        case wal::MXLogType::Delete: {
J
Jin Hai 已提交
2522 2523
            std::vector<meta::CollectionSchema> partition_array;
            status = meta_ptr_->ShowPartitions(record.collection_id, partition_array);
2524 2525 2526 2527
            if (!status.ok()) {
                return status;
            }

2528
            std::vector<std::string> collection_ids{record.collection_id};
2529
            for (auto& partition : partition_array) {
2530 2531
                auto& partition_collection_id = partition.collection_id_;
                collection_ids.emplace_back(partition_collection_id);
2532 2533 2534
            }

            if (record.length == 1) {
2535
                for (auto& collection_id : collection_ids) {
J
Jin Hai 已提交
2536
                    status = mem_mgr_->DeleteVector(collection_id, *record.ids, record.lsn);
2537 2538 2539 2540 2541
                    if (!status.ok()) {
                        return status;
                    }
                }
            } else {
2542
                for (auto& collection_id : collection_ids) {
J
Jin Hai 已提交
2543
                    status = mem_mgr_->DeleteVectors(collection_id, record.length, record.ids, record.lsn);
2544 2545 2546 2547 2548 2549 2550 2551 2552
                    if (!status.ok()) {
                        return status;
                    }
                }
            }
            break;
        }

        case wal::MXLogType::Flush: {
J
Jin Hai 已提交
2553 2554 2555 2556
            if (!record.collection_id.empty()) {
                // flush one collection
                std::vector<meta::CollectionSchema> partition_array;
                status = meta_ptr_->ShowPartitions(record.collection_id, partition_array);
2557 2558 2559 2560
                if (!status.ok()) {
                    return status;
                }

2561
                std::vector<std::string> collection_ids{record.collection_id};
2562
                for (auto& partition : partition_array) {
2563 2564
                    auto& partition_collection_id = partition.collection_id_;
                    collection_ids.emplace_back(partition_collection_id);
2565 2566
                }

2567 2568
                std::set<std::string> flushed_collections;
                for (auto& collection_id : collection_ids) {
2569
                    const std::lock_guard<std::mutex> lock(flush_merge_compact_mutex_);
J
Jin Hai 已提交
2570
                    status = mem_mgr_->Flush(collection_id);
2571 2572 2573
                    if (!status.ok()) {
                        break;
                    }
2574
                    flushed_collections.insert(collection_id);
2575 2576
                }

G
groot 已提交
2577
                collections_flushed(record.collection_id, flushed_collections);
2578 2579

            } else {
2580 2581
                // flush all collections
                std::set<std::string> collection_ids;
2582 2583
                {
                    const std::lock_guard<std::mutex> lock(flush_merge_compact_mutex_);
2584
                    status = mem_mgr_->Flush(collection_ids);
2585 2586
                }

G
groot 已提交
2587
                uint64_t lsn = collections_flushed("", collection_ids);
2588 2589 2590 2591 2592 2593
                if (options_.wal_enable_) {
                    wal_mgr_->RemoveOldFiles(lsn);
                }
            }
            break;
        }
C
Cai Yudong 已提交
2594 2595 2596

        default:
            break;
2597 2598 2599 2600 2601 2602
    }

    return status;
}

void
G
groot 已提交
2603 2604 2605 2606 2607 2608 2609 2610 2611 2612 2613
DBImpl::InternalFlush(const std::string& collection_id) {
    wal::MXLogRecord record;
    record.type = wal::MXLogType::Flush;
    record.collection_id = collection_id;
    ExecWalRecord(record);

    StartMergeTask();
}

void
DBImpl::BackgroundWalThread() {
2614
    SetThreadName("wal_thread");
2615 2616
    server::SystemInfo::GetInstance().Init();

2617
    std::chrono::system_clock::time_point next_auto_flush_time;
2618
    auto get_next_auto_flush_time = [&]() {
2619
        return std::chrono::system_clock::now() + std::chrono::seconds(options_.auto_flush_interval_);
2620
    };
2621 2622 2623
    if (options_.auto_flush_interval_ > 0) {
        next_auto_flush_time = get_next_auto_flush_time();
    }
2624 2625

    while (true) {
2626 2627
        if (options_.auto_flush_interval_ > 0) {
            if (std::chrono::system_clock::now() >= next_auto_flush_time) {
G
groot 已提交
2628
                InternalFlush();
2629 2630
                next_auto_flush_time = get_next_auto_flush_time();
            }
2631 2632
        }

G
groot 已提交
2633
        wal::MXLogRecord record;
2634 2635
        auto error_code = wal_mgr_->GetNextRecord(record);
        if (error_code != WAL_SUCCESS) {
2636
            LOG_ENGINE_ERROR_ << "WAL background GetNextRecord error";
2637 2638 2639 2640 2641 2642
            break;
        }

        if (record.type != wal::MXLogType::None) {
            ExecWalRecord(record);
            if (record.type == wal::MXLogType::Flush) {
G
groot 已提交
2643 2644
                // notify flush request to return
                flush_req_swn_.Notify();
2645 2646

                // if user flush all manually, update auto flush also
J
Jin Hai 已提交
2647
                if (record.collection_id.empty() && options_.auto_flush_interval_ > 0) {
2648 2649 2650 2651 2652 2653
                    next_auto_flush_time = get_next_auto_flush_time();
                }
            }

        } else {
            if (!initialized_.load(std::memory_order_acquire)) {
G
groot 已提交
2654 2655
                InternalFlush();
                flush_req_swn_.Notify();
2656 2657
                WaitMergeFileFinish();
                WaitBuildIndexFinish();
2658
                LOG_ENGINE_DEBUG_ << "WAL background thread exit";
2659 2660 2661
                break;
            }

2662
            if (options_.auto_flush_interval_ > 0) {
G
groot 已提交
2663
                swn_wal_.Wait_Until(next_auto_flush_time);
2664
            } else {
G
groot 已提交
2665
                swn_wal_.Wait();
2666
            }
2667 2668 2669 2670
        }
    }
}

G
groot 已提交
2671 2672
void
DBImpl::BackgroundFlushThread() {
2673
    SetThreadName("flush_thread");
G
groot 已提交
2674 2675 2676
    server::SystemInfo::GetInstance().Init();
    while (true) {
        if (!initialized_.load(std::memory_order_acquire)) {
2677
            LOG_ENGINE_DEBUG_ << "DB background flush thread exit";
G
groot 已提交
2678 2679 2680 2681 2682 2683 2684 2685 2686 2687 2688 2689 2690 2691 2692 2693 2694
            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() {
    server::SystemInfo::GetInstance().Init();
    while (true) {
        if (!initialized_.load(std::memory_order_acquire)) {
2695
            LOG_ENGINE_DEBUG_ << "DB background metric thread exit";
G
groot 已提交
2696 2697 2698 2699 2700
            break;
        }

        swn_metric_.Wait_For(std::chrono::seconds(BACKGROUND_METRIC_INTERVAL));
        StartMetricTask();
G
groot 已提交
2701
        meta::FilesHolder::PrintInfo();
G
groot 已提交
2702 2703 2704
    }
}

2705 2706 2707 2708 2709
void
DBImpl::OnCacheInsertDataChanged(bool value) {
    options_.insert_cache_immediately_ = value;
}

2710 2711 2712 2713 2714
void
DBImpl::OnUseBlasThresholdChanged(int64_t threshold) {
    faiss::distance_compute_blas_threshold = threshold;
}

2715 2716 2717 2718 2719 2720 2721 2722 2723 2724 2725 2726 2727 2728 2729 2730 2731 2732
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 已提交
2733 2734
}  // namespace engine
}  // namespace milvus