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

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

55 56
#include "search/TaskInst.h"

J
jinhai 已提交
57
namespace milvus {
X
Xu Peng 已提交
58
namespace engine {
X
Xu Peng 已提交
59

G
groot 已提交
60
namespace {
G
groot 已提交
61 62
constexpr uint64_t BACKGROUND_METRIC_INTERVAL = 1;
constexpr uint64_t BACKGROUND_INDEX_INTERVAL = 1;
G
groot 已提交
63
constexpr uint64_t WAIT_BUILD_INDEX_INTERVAL = 5;
G
groot 已提交
64

65 66 67 68 69 70 71 72
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 已提交
73
static const Status SHUTDOWN_ERROR = Status(DB_ERROR, "Milvus server is shutdown!");
G
groot 已提交
74

S
starlord 已提交
75
}  // namespace
G
groot 已提交
76

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

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

91 92
    SetIdentity("DBImpl");
    AddCacheInsertDataListener();
93
    AddUseBlasThresholdListener();
94

S
starlord 已提交
95 96 97 98 99 100 101
    Start();
}

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

S
starlord 已提交
102
///////////////////////////////////////////////////////////////////////////////////////////////////////////////////
S
starlord 已提交
103
// external api
S
starlord 已提交
104
///////////////////////////////////////////////////////////////////////////////////////////////////////////////////
S
starlord 已提交
105 106
Status
DBImpl::Start() {
107
    if (initialized_.load(std::memory_order_acquire)) {
S
starlord 已提交
108 109 110
        return Status::OK();
    }

111
    // LOG_ENGINE_TRACE_ << "DB service start";
112
    initialized_.store(true, std::memory_order_release);
S
starlord 已提交
113

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

G
groot 已提交
150 151 152 153 154 155 156 157 158
    // 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 已提交
159 160 161
    return Status::OK();
}

S
starlord 已提交
162 163
Status
DBImpl::Stop() {
164
    if (!initialized_.load(std::memory_order_acquire)) {
S
starlord 已提交
165 166
        return Status::OK();
    }
167

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

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

G
groot 已提交
181 182 183
            // wait flush thread finish
            swn_flush_.Notify();
            bg_flush_thread_.join();
184
        }
S
starlord 已提交
185

186 187
        WaitMergeFileFinish();

G
groot 已提交
188 189 190
        swn_index_.Notify();
        bg_index_thread_.join();

191
        meta_ptr_->CleanUpShadowFiles();
S
starlord 已提交
192 193
    }

G
groot 已提交
194 195 196 197
    // wait metric thread exit
    swn_metric_.Notify();
    bg_metric_thread_.join();

198
    // LOG_ENGINE_TRACE_ << "DB service stop";
S
starlord 已提交
199
    return Status::OK();
X
Xu Peng 已提交
200 201
}

S
starlord 已提交
202 203
Status
DBImpl::DropAll() {
S
starlord 已提交
204 205 206
    return meta_ptr_->DropAll();
}

S
starlord 已提交
207
Status
208
DBImpl::CreateCollection(meta::CollectionSchema& collection_schema) {
209
    if (!initialized_.load(std::memory_order_acquire)) {
G
groot 已提交
210
        return SHUTDOWN_ERROR;
S
starlord 已提交
211 212
    }

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

219
    return meta_ptr_->CreateCollection(temp_schema);
220 221
}

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

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

254
    if (options_.wal_enable_) {
255
        wal_mgr_->DropCollection(collection_id);
256 257
    }

258
    return DropCollectionRecursively(collection_id);
G
groot 已提交
259 260
}

S
starlord 已提交
261
Status
262
DBImpl::DescribeCollection(meta::CollectionSchema& collection_schema) {
263
    if (!initialized_.load(std::memory_order_acquire)) {
G
groot 已提交
264
        return SHUTDOWN_ERROR;
S
starlord 已提交
265 266
    }

267
    auto stat = meta_ptr_->DescribeCollection(collection_schema);
B
bigbraver 已提交
268
    collection_schema.index_file_size_ /= MB;  // return as MB
S
starlord 已提交
269
    return stat;
270 271
}

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

278
    return meta_ptr_->HasCollection(collection_id, has_or_not);
279 280
}

281
Status
282
DBImpl::HasNativeCollection(const std::string& collection_id, bool& has_or_not_) {
283 284 285 286
    if (!initialized_.load(std::memory_order_acquire)) {
        return SHUTDOWN_ERROR;
    }

287 288 289
    engine::meta::CollectionSchema collection_schema;
    collection_schema.collection_id_ = collection_id;
    auto status = DescribeCollection(collection_schema);
290 291 292 293
    if (!status.ok()) {
        has_or_not_ = false;
        return status;
    } else {
294
        if (!collection_schema.owner_collection_.empty()) {
295 296 297 298 299 300 301 302 303
            has_or_not_ = false;
            return Status(DB_NOT_FOUND, "");
        }

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

S
starlord 已提交
304
Status
305
DBImpl::AllCollections(std::vector<meta::CollectionSchema>& collection_schema_array) {
306
    if (!initialized_.load(std::memory_order_acquire)) {
G
groot 已提交
307
        return SHUTDOWN_ERROR;
S
starlord 已提交
308 309
    }

310 311
    std::vector<meta::CollectionSchema> all_collections;
    auto status = meta_ptr_->AllCollections(all_collections);
312

313 314 315 316 317
    // 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);
318 319 320 321
        }
    }

    return status;
G
groot 已提交
322 323
}

324
Status
325
DBImpl::GetCollectionInfo(const std::string& collection_id, std::string& collection_info) {
326 327 328 329 330
    if (!initialized_.load(std::memory_order_acquire)) {
        return SHUTDOWN_ERROR;
    }

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

J
Jin Hai 已提交
334 335
    std::vector<int> file_types{meta::SegmentSchema::FILE_TYPE::RAW, meta::SegmentSchema::FILE_TYPE::TO_INDEX,
                                meta::SegmentSchema::FILE_TYPE::INDEX};
336

337 338 339
    milvus::json json_info;
    milvus::json json_partitions;
    size_t total_row_count = 0;
340

341
    auto get_info = [&](const std::string& col_id, const std::string& tag) {
G
groot 已提交
342 343
        meta::FilesHolder files_holder;
        status = meta_ptr_->FilesByType(col_id, file_types, files_holder);
344
        if (!status.ok()) {
J
Jin Hai 已提交
345
            std::string err_msg = "Failed to get collection info: " + status.ToString();
346
            LOG_ENGINE_ERROR_ << err_msg;
347 348 349
            return Status(DB_ERROR, err_msg);
        }

350 351 352 353 354
        milvus::json json_partition;
        json_partition[JSON_PARTITION_TAG] = tag;

        milvus::json json_segments;
        size_t row_count = 0;
G
groot 已提交
355
        milvus::engine::meta::SegmentsSchema& collection_files = files_holder.HoldFiles();
356
        for (auto& file : collection_files) {
357 358 359 360 361 362 363 364 365
            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_;
366 367
        }

368 369 370 371 372 373 374
        json_partition[JSON_ROW_COUNT] = row_count;
        json_partition[JSON_SEGMENTS] = json_segments;

        json_partitions.push_back(json_partition);

        return Status::OK();
    };
375

376 377 378 379
    // step2: get default partition info
    status = get_info(collection_id, milvus::engine::DEFAULT_PARTITON_TAG);
    if (!status.ok()) {
        return status;
380 381
    }

382 383 384 385 386 387 388 389 390 391 392 393 394
    // 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();

395 396 397
    return Status::OK();
}

S
starlord 已提交
398
Status
399
DBImpl::PreloadCollection(const std::string& collection_id) {
400
    if (!initialized_.load(std::memory_order_acquire)) {
G
groot 已提交
401
        return SHUTDOWN_ERROR;
S
starlord 已提交
402 403
    }

J
Jin Hai 已提交
404
    // step 1: get all collection files from parent collection
G
groot 已提交
405 406
    meta::FilesHolder files_holder;
    auto status = meta_ptr_->FilesToSearch(collection_id, files_holder);
Y
Yu Kun 已提交
407 408 409
    if (!status.ok()) {
        return status;
    }
Y
Yu Kun 已提交
410

411
    // step 2: get files from partition collections
J
Jin Hai 已提交
412 413
    std::vector<meta::CollectionSchema> partition_array;
    status = meta_ptr_->ShowPartitions(collection_id, partition_array);
G
groot 已提交
414
    for (auto& schema : partition_array) {
G
groot 已提交
415
        status = meta_ptr_->FilesToSearch(schema.collection_id_, files_holder);
G
groot 已提交
416 417
    }

Y
Yu Kun 已提交
418 419
    int64_t size = 0;
    int64_t cache_total = cache::CpuCacheMgr::GetInstance()->CacheCapacity();
Y
Yu Kun 已提交
420 421
    int64_t cache_usage = cache::CpuCacheMgr::GetInstance()->CacheUsage();
    int64_t available_size = cache_total - cache_usage;
Y
Yu Kun 已提交
422

423
    // step 3: load file one by one
G
groot 已提交
424
    milvus::engine::meta::SegmentsSchema& files_array = files_holder.HoldFiles();
425 426
    LOG_ENGINE_DEBUG_ << "Begin pre-load collection:" + collection_id + ", totally " << files_array.size()
                      << " files need to be pre-loaded";
J
Jin Hai 已提交
427
    TimeRecorderAuto rc("Pre-load collection:" + collection_id);
G
groot 已提交
428
    for (auto& file : files_array) {
429
        EngineType engine_type;
J
Jin Hai 已提交
430 431 432
        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) {
433 434
            engine_type =
                utils::IsBinaryMetricType(file.metric_type_) ? EngineType::FAISS_BIN_IDMAP : EngineType::FAISS_IDMAP;
435 436 437
        } else {
            engine_type = (EngineType)file.engine_type_;
        }
438 439 440 441

        auto json = milvus::json::parse(file.index_params_);
        ExecutionEnginePtr engine =
            EngineFactory::Build(file.dimension_, file.location_, engine_type, (MetricType)file.metric_type_, json);
442
        fiu_do_on("DBImpl.PreloadCollection.null_engine", engine = nullptr);
G
groot 已提交
443
        if (engine == nullptr) {
444
            LOG_ENGINE_ERROR_ << "Invalid engine type";
G
groot 已提交
445 446
            return Status(DB_ERROR, "Invalid engine type");
        }
Y
Yu Kun 已提交
447

448
        fiu_do_on("DBImpl.PreloadCollection.exceed_cache", size = available_size + 1);
449 450

        try {
451
            fiu_do_on("DBImpl.PreloadCollection.engine_throw_exception", throw std::exception());
452 453
            std::string msg = "Pre-loaded file: " + file.file_id_ + " size: " + std::to_string(file.file_size_);
            TimeRecorderAuto rc_1(msg);
454 455 456 457
            status = engine->Load(true);
            if (!status.ok()) {
                return status;
            }
458 459 460

            size += engine->Size();
            if (size > available_size) {
461
                LOG_ENGINE_DEBUG_ << "Pre-load cancelled since cache is almost full";
462
                return Status(SERVER_CACHE_FULL, "Cache is full");
Y
Yu Kun 已提交
463
            }
464
        } catch (std::exception& ex) {
J
Jin Hai 已提交
465
            std::string msg = "Pre-load collection encounter exception: " + std::string(ex.what());
466
            LOG_ENGINE_ERROR_ << msg;
467
            return Status(DB_ERROR, msg);
Y
Yu Kun 已提交
468 469
        }
    }
G
groot 已提交
470

Y
Yu Kun 已提交
471
    return Status::OK();
Y
Yu Kun 已提交
472 473
}

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

480
    return meta_ptr_->UpdateCollectionFlag(collection_id, flag);
S
starlord 已提交
481 482
}

S
starlord 已提交
483
Status
484
DBImpl::GetCollectionRowCount(const std::string& collection_id, uint64_t& row_count) {
485
    if (!initialized_.load(std::memory_order_acquire)) {
G
groot 已提交
486 487 488
        return SHUTDOWN_ERROR;
    }

489
    return GetCollectionRowCountRecursively(collection_id, row_count);
G
groot 已提交
490 491 492
}

Status
J
Jin Hai 已提交
493
DBImpl::CreatePartition(const std::string& collection_id, const std::string& partition_name,
G
groot 已提交
494
                        const std::string& partition_tag) {
495
    if (!initialized_.load(std::memory_order_acquire)) {
G
groot 已提交
496 497 498
        return SHUTDOWN_ERROR;
    }

499
    uint64_t lsn = 0;
500
    meta_ptr_->GetCollectionFlushLSN(collection_id, lsn);
J
Jin Hai 已提交
501
    return meta_ptr_->CreatePartition(collection_id, partition_name, partition_tag, lsn);
G
groot 已提交
502 503 504 505
}

Status
DBImpl::DropPartition(const std::string& partition_name) {
506
    if (!initialized_.load(std::memory_order_acquire)) {
G
groot 已提交
507
        return SHUTDOWN_ERROR;
S
starlord 已提交
508 509
    }

510
    mem_mgr_->EraseMemVector(partition_name);                // not allow insert
J
Jin Hai 已提交
511
    auto status = meta_ptr_->DropPartition(partition_name);  // soft delete collection
512
    if (!status.ok()) {
513
        LOG_ENGINE_ERROR_ << status.message();
514 515
        return status;
    }
G
groot 已提交
516

J
Jin Hai 已提交
517
    // scheduler will determine when to delete collection files
G
groot 已提交
518 519 520 521 522 523
    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 已提交
524 525
}

S
starlord 已提交
526
Status
J
Jin Hai 已提交
527
DBImpl::DropPartitionByTag(const std::string& collection_id, const std::string& partition_tag) {
528
    if (!initialized_.load(std::memory_order_acquire)) {
G
groot 已提交
529 530 531 532
        return SHUTDOWN_ERROR;
    }

    std::string partition_name;
J
Jin Hai 已提交
533
    auto status = meta_ptr_->GetPartitionName(collection_id, partition_tag, partition_name);
534
    if (!status.ok()) {
535
        LOG_ENGINE_ERROR_ << status.message();
536 537 538
        return status;
    }

G
groot 已提交
539 540 541 542
    return DropPartition(partition_name);
}

Status
J
Jin Hai 已提交
543
DBImpl::ShowPartitions(const std::string& collection_id, std::vector<meta::CollectionSchema>& partition_schema_array) {
544
    if (!initialized_.load(std::memory_order_acquire)) {
G
groot 已提交
545 546 547
        return SHUTDOWN_ERROR;
    }

J
Jin Hai 已提交
548
    return meta_ptr_->ShowPartitions(collection_id, partition_schema_array);
G
groot 已提交
549 550 551
}

Status
J
Jin Hai 已提交
552
DBImpl::InsertVectors(const std::string& collection_id, const std::string& partition_tag, VectorsData& vectors) {
553
    //    LOG_ENGINE_DEBUG_ << "Insert " << n << " vectors to cache";
554
    if (!initialized_.load(std::memory_order_acquire)) {
G
groot 已提交
555
        return SHUTDOWN_ERROR;
S
starlord 已提交
556
    }
Y
yu yunfeng 已提交
557

J
Jin Hai 已提交
558
    // insert vectors into target collection
559 560
    // (zhiru): generate ids
    if (vectors.id_array_.empty()) {
J
Jin Hai 已提交
561 562 563
        SafeIDGenerator& id_generator = SafeIDGenerator::GetInstance();
        Status status = id_generator.GetNextIDNumbers(vectors.vector_count_, vectors.id_array_);
        if (!status.ok()) {
564
            LOG_ENGINE_ERROR_ << LogOut("[%s][%ld] Get next id number fail: %s", "insert", 0, status.message().c_str());
J
Jin Hai 已提交
565 566
            return status;
        }
567 568
    }

569
    Status status;
570
    if (options_.wal_enable_) {
571 572
        std::string target_collection_name;
        status = GetPartitionByTag(collection_id, partition_tag, target_collection_name);
G
groot 已提交
573
        if (!status.ok()) {
574
            LOG_ENGINE_ERROR_ << LogOut("[%s][%ld] Get partition fail: %s", "insert", 0, status.message().c_str());
G
groot 已提交
575 576
            return status;
        }
577 578

        if (!vectors.float_data_.empty()) {
J
Jin Hai 已提交
579
            wal_mgr_->Insert(collection_id, partition_tag, vectors.id_array_, vectors.float_data_);
580
        } else if (!vectors.binary_data_.empty()) {
J
Jin Hai 已提交
581
            wal_mgr_->Insert(collection_id, partition_tag, vectors.id_array_, vectors.binary_data_);
582
        }
G
groot 已提交
583
        swn_wal_.Notify();
584 585 586
    } else {
        wal::MXLogRecord record;
        record.lsn = 0;  // need to get from meta ?
J
Jin Hai 已提交
587
        record.collection_id = collection_id;
588 589 590 591 592 593 594 595 596 597 598 599 600 601 602 603
        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 已提交
604 605
    }

606 607 608
    return status;
}

609
Status
Y
yukun 已提交
610 611 612 613 614
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) {
615
    uint64_t offset = 0;
Y
yukun 已提交
616 617
    for (auto name : field_names) {
        switch (attr_types.at(name)) {
618 619
            case meta::hybrid::DataType::INT8: {
                std::vector<uint8_t> data;
Y
yukun 已提交
620
                data.resize(row_num * sizeof(int8_t));
621

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

Y
yukun 已提交
625 626
                std::vector<int8_t> raw_value(row_num, 0);
                for (uint64_t i = 0; i < row_num; ++i) {
627 628 629
                    raw_value[i] = attr_value[i];
                }

Y
yukun 已提交
630 631
                memcpy(data.data(), raw_value.data(), row_num * sizeof(int8_t));
                attr_datas.insert(std::make_pair(name, data));
632

Y
yukun 已提交
633 634 635
                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);
636 637 638 639
                break;
            }
            case meta::hybrid::DataType::INT16: {
                std::vector<uint8_t> data;
Y
yukun 已提交
640
                data.resize(row_num * sizeof(int16_t));
641

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

Y
yukun 已提交
645 646
                std::vector<int16_t> raw_value(row_num, 0);
                for (uint64_t i = 0; i < row_num; ++i) {
647 648 649
                    raw_value[i] = attr_value[i];
                }

Y
yukun 已提交
650 651
                memcpy(data.data(), raw_value.data(), row_num * sizeof(int16_t));
                attr_datas.insert(std::make_pair(name, data));
652

Y
yukun 已提交
653 654 655
                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);
656 657 658 659
                break;
            }
            case meta::hybrid::DataType::INT32: {
                std::vector<uint8_t> data;
Y
yukun 已提交
660
                data.resize(row_num * sizeof(int32_t));
661

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

Y
yukun 已提交
665 666
                std::vector<int32_t> raw_value(row_num, 0);
                for (uint64_t i = 0; i < row_num; ++i) {
667 668 669
                    raw_value[i] = attr_value[i];
                }

Y
yukun 已提交
670 671
                memcpy(data.data(), raw_value.data(), row_num * sizeof(int32_t));
                attr_datas.insert(std::make_pair(name, data));
672

Y
yukun 已提交
673 674 675
                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);
676 677 678 679
                break;
            }
            case meta::hybrid::DataType::INT64: {
                std::vector<uint8_t> data;
Y
yukun 已提交
680 681 682 683 684 685
                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));
686

Y
yukun 已提交
687 688 689
                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);
690 691 692 693
                break;
            }
            case meta::hybrid::DataType::FLOAT: {
                std::vector<uint8_t> data;
Y
yukun 已提交
694
                data.resize(row_num * sizeof(float));
695

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

Y
yukun 已提交
699 700
                std::vector<float> raw_value(row_num, 0);
                for (uint64_t i = 0; i < row_num; ++i) {
701 702 703
                    raw_value[i] = attr_value[i];
                }

Y
yukun 已提交
704 705
                memcpy(data.data(), raw_value.data(), row_num * sizeof(float));
                attr_datas.insert(std::make_pair(name, data));
706

Y
yukun 已提交
707 708 709
                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);
710 711 712 713
                break;
            }
            case meta::hybrid::DataType::DOUBLE: {
                std::vector<uint8_t> data;
Y
yukun 已提交
714 715 716
                data.resize(row_num * sizeof(double));
                memcpy(data.data(), record.data() + offset, row_num * sizeof(double));
                attr_datas.insert(std::make_pair(name, data));
717

Y
yukun 已提交
718 719 720
                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);
721 722
                break;
            }
723 724
            default:
                break;
725 726
        }
    }
Y
yukun 已提交
727 728 729 730 731 732 733 734 735 736 737 738 739 740 741 742 743 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
    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);
    }
777 778

    status = ExecWalRecord(record);
Y
yukun 已提交
779 780 781 782 783 784 785 786 787 788 789 790 791 792 793 794 795 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

#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

825 826 827
    return status;
}

828
Status
J
Jin Hai 已提交
829
DBImpl::DeleteVector(const std::string& collection_id, IDNumber vector_id) {
830 831
    IDNumbers ids;
    ids.push_back(vector_id);
J
Jin Hai 已提交
832
    return DeleteVectors(collection_id, ids);
833 834 835
}

Status
J
Jin Hai 已提交
836
DBImpl::DeleteVectors(const std::string& collection_id, IDNumbers vector_ids) {
837 838 839 840 841 842
    if (!initialized_.load(std::memory_order_acquire)) {
        return SHUTDOWN_ERROR;
    }

    Status status;
    if (options_.wal_enable_) {
J
Jin Hai 已提交
843
        wal_mgr_->DeleteById(collection_id, vector_ids);
G
groot 已提交
844
        swn_wal_.Notify();
845 846 847 848
    } else {
        wal::MXLogRecord record;
        record.lsn = 0;  // need to get from meta ?
        record.type = wal::MXLogType::Delete;
J
Jin Hai 已提交
849
        record.collection_id = collection_id;
850 851 852 853 854 855 856 857 858 859
        record.ids = vector_ids.data();
        record.length = vector_ids.size();

        status = ExecWalRecord(record);
    }

    return status;
}

Status
J
Jin Hai 已提交
860
DBImpl::Flush(const std::string& collection_id) {
861 862 863 864 865
    if (!initialized_.load(std::memory_order_acquire)) {
        return SHUTDOWN_ERROR;
    }

    Status status;
866 867
    bool has_collection;
    status = HasCollection(collection_id, has_collection);
868 869 870
    if (!status.ok()) {
        return status;
    }
871
    if (!has_collection) {
872
        LOG_ENGINE_ERROR_ << "Collection to flush does not exist: " << collection_id;
J
Jin Hai 已提交
873
        return Status(DB_NOT_FOUND, "Collection to flush does not exist");
874 875
    }

876
    LOG_ENGINE_DEBUG_ << "Begin flush collection: " << collection_id;
877 878

    if (options_.wal_enable_) {
879
        LOG_ENGINE_DEBUG_ << "WAL flush";
J
Jin Hai 已提交
880
        auto lsn = wal_mgr_->Flush(collection_id);
881
        if (lsn != 0) {
G
groot 已提交
882 883
            swn_wal_.Notify();
            flush_req_swn_.Wait();
884 885 886
        }

    } else {
887
        LOG_ENGINE_DEBUG_ << "MemTable flush";
G
groot 已提交
888
        InternalFlush(collection_id);
889 890
    }

891
    LOG_ENGINE_DEBUG_ << "End flush collection: " << collection_id;
892 893 894 895 896 897 898 899 900 901

    return status;
}

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

902
    LOG_ENGINE_DEBUG_ << "Begin flush all collections";
903 904 905

    Status status;
    if (options_.wal_enable_) {
906
        LOG_ENGINE_DEBUG_ << "WAL flush";
907 908
        auto lsn = wal_mgr_->Flush();
        if (lsn != 0) {
G
groot 已提交
909 910
            swn_wal_.Notify();
            flush_req_swn_.Wait();
911 912
        }
    } else {
913
        LOG_ENGINE_DEBUG_ << "MemTable flush";
G
groot 已提交
914
        InternalFlush();
915 916
    }

917
    LOG_ENGINE_DEBUG_ << "End flush all collections";
918 919 920 921 922

    return status;
}

Status
G
groot 已提交
923
DBImpl::Compact(const std::string& collection_id, double threshold) {
924 925 926 927
    if (!initialized_.load(std::memory_order_acquire)) {
        return SHUTDOWN_ERROR;
    }

928 929 930
    engine::meta::CollectionSchema collection_schema;
    collection_schema.collection_id_ = collection_id;
    auto status = DescribeCollection(collection_schema);
931 932
    if (!status.ok()) {
        if (status.code() == DB_NOT_FOUND) {
933
            LOG_ENGINE_ERROR_ << "Collection to compact does not exist: " << collection_id;
J
Jin Hai 已提交
934
            return Status(DB_NOT_FOUND, "Collection to compact does not exist");
935 936 937 938
        } else {
            return status;
        }
    } else {
939
        if (!collection_schema.owner_collection_.empty()) {
940
            LOG_ENGINE_ERROR_ << "Collection to compact does not exist: " << collection_id;
J
Jin Hai 已提交
941
            return Status(DB_NOT_FOUND, "Collection to compact does not exist");
942 943 944
        }
    }

945
    LOG_ENGINE_DEBUG_ << "Before compacting, wait for build index thread to finish...";
946

Z
update  
Zhiru Zhu 已提交
947
    // WaitBuildIndexFinish();
948

Z
update  
Zhiru Zhu 已提交
949
    const std::lock_guard<std::mutex> index_lock(build_index_mutex_);
Z
Zhiru Zhu 已提交
950
    const std::lock_guard<std::mutex> merge_lock(flush_merge_compact_mutex_);
Z
Zhiru Zhu 已提交
951

952
    LOG_ENGINE_DEBUG_ << "Compacting collection: " << collection_id;
Z
Zhiru Zhu 已提交
953

954
    // Get files to compact from meta.
J
Jin Hai 已提交
955 956
    std::vector<int> file_types{meta::SegmentSchema::FILE_TYPE::RAW, meta::SegmentSchema::FILE_TYPE::TO_INDEX,
                                meta::SegmentSchema::FILE_TYPE::BACKUP};
G
groot 已提交
957 958
    meta::FilesHolder files_holder;
    status = meta_ptr_->FilesByType(collection_id, file_types, files_holder);
959 960
    if (!status.ok()) {
        std::string err_msg = "Failed to get files to compact: " + status.message();
961
        LOG_ENGINE_ERROR_ << err_msg;
962 963 964
        return Status(DB_ERROR, err_msg);
    }

G
groot 已提交
965
    LOG_ENGINE_DEBUG_ << "Found " << files_holder.HoldFiles().size() << " segment to compact";
Z
Zhiru Zhu 已提交
966 967

    Status compact_status;
G
groot 已提交
968 969
    // 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 已提交
970
    for (auto iter = files_to_compact.begin(); iter != files_to_compact.end();) {
J
Jin Hai 已提交
971
        meta::SegmentSchema file = *iter;
G
groot 已提交
972 973
        iter = files_to_compact.erase(iter);

Z
Zhiru Zhu 已提交
974 975 976
        // Check if the segment needs compacting
        std::string segment_dir;
        utils::GetParentPath(file.location_, segment_dir);
977

Z
Zhiru Zhu 已提交
978
        segment::SegmentReader segment_reader(segment_dir);
Z
Zhiru Zhu 已提交
979 980
        size_t deleted_docs_size;
        status = segment_reader.ReadDeletedDocsSize(deleted_docs_size);
Z
Zhiru Zhu 已提交
981
        if (!status.ok()) {
G
groot 已提交
982
            files_holder.UnmarkFile(file);
G
groot 已提交
983
            continue;  // skip this file and try compact next one
Z
Zhiru Zhu 已提交
984 985
        }

J
Jin Hai 已提交
986
        meta::SegmentsSchema files_to_update;
Z
Zhiru Zhu 已提交
987
        if (deleted_docs_size != 0) {
G
groot 已提交
988
            compact_status = CompactFile(collection_id, threshold, file, files_to_update);
Z
Zhiru Zhu 已提交
989 990

            if (!compact_status.ok()) {
991 992
                LOG_ENGINE_ERROR_ << "Compact failed for segment " << file.segment_id_ << ": "
                                  << compact_status.message();
G
groot 已提交
993
                files_holder.UnmarkFile(file);
G
groot 已提交
994
                continue;  // skip this file and try compact next one
Z
Zhiru Zhu 已提交
995 996
            }
        } else {
G
groot 已提交
997
            files_holder.UnmarkFile(file);
998
            LOG_ENGINE_DEBUG_ << "Segment " << file.segment_id_ << " has no deleted data. No need to compact";
G
groot 已提交
999
            continue;  // skip this file and try compact next one
1000
        }
Z
Zhiru Zhu 已提交
1001

1002
        LOG_ENGINE_DEBUG_ << "Updating meta after compaction...";
1003
        status = meta_ptr_->UpdateCollectionFiles(files_to_update);
G
groot 已提交
1004
        files_holder.UnmarkFile(file);
G
groot 已提交
1005 1006 1007 1008
        if (!status.ok()) {
            compact_status = status;
            break;  // meta error, could not go on
        }
Z
Zhiru Zhu 已提交
1009 1010
    }

G
groot 已提交
1011
    if (compact_status.ok()) {
1012
        LOG_ENGINE_DEBUG_ << "Finished compacting collection: " << collection_id;
G
groot 已提交
1013
    }
1014

G
groot 已提交
1015
    return compact_status;
1016 1017 1018
}

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

G
groot 已提交
1023 1024 1025 1026 1027 1028 1029 1030 1031 1032 1033 1034 1035 1036 1037 1038 1039 1040 1041
    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 已提交
1042 1043 1044
    // Create new collection file
    meta::SegmentSchema compacted_file;
    compacted_file.collection_id_ = collection_id;
1045
    // compacted_file.date_ = date;
J
Jin Hai 已提交
1046
    compacted_file.file_type_ = meta::SegmentSchema::NEW_MERGE;  // TODO: use NEW_MERGE for now
G
groot 已提交
1047
    auto status = meta_ptr_->CreateCollectionFile(compacted_file);
1048 1049

    if (!status.ok()) {
1050
        LOG_ENGINE_ERROR_ << "Failed to create collection file: " << status.message();
1051 1052 1053
        return status;
    }

J
Jin Hai 已提交
1054
    // Compact (merge) file to the newly created collection file
1055 1056 1057 1058 1059

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

1060
    LOG_ENGINE_DEBUG_ << "Compacting begin...";
1061 1062 1063
    segment_writer_ptr->Merge(segment_dir_to_merge, compacted_file.file_id_);

    // Serialize
1064
    LOG_ENGINE_DEBUG_ << "Serializing compacted segment...";
1065 1066
    status = segment_writer_ptr->Serialize();
    if (!status.ok()) {
1067
        LOG_ENGINE_ERROR_ << "Failed to serialize compacted segment: " << status.message();
J
Jin Hai 已提交
1068
        compacted_file.file_type_ = meta::SegmentSchema::TO_DELETE;
1069
        auto mark_status = meta_ptr_->UpdateCollectionFile(compacted_file);
1070
        if (mark_status.ok()) {
1071
            LOG_ENGINE_DEBUG_ << "Mark file: " << compacted_file.file_id_ << " to to_delete";
1072
        }
G
groot 已提交
1073

1074 1075 1076
        return status;
    }

1077 1078 1079 1080 1081 1082 1083
    // 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;
1084
    } else {
J
Jin Hai 已提交
1085
        compacted_file.file_type_ = meta::SegmentSchema::RAW;
1086 1087 1088
    }

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

Z
Zhiru Zhu 已提交
1093
    files_to_update.emplace_back(compacted_file);
Z
Zhiru Zhu 已提交
1094

Z
Zhiru Zhu 已提交
1095 1096
    // Set all files in segment to TO_DELETE
    auto& segment_id = file.segment_id_;
G
groot 已提交
1097 1098
    meta::FilesHolder files_holder;
    status = meta_ptr_->GetCollectionFilesBySegmentId(segment_id, files_holder);
Z
Zhiru Zhu 已提交
1099 1100 1101
    if (!status.ok()) {
        return status;
    }
G
groot 已提交
1102 1103

    milvus::engine::meta::SegmentsSchema& segment_files = files_holder.HoldFiles();
Z
Zhiru Zhu 已提交
1104
    for (auto& f : segment_files) {
J
Jin Hai 已提交
1105
        f.file_type_ = meta::SegmentSchema::FILE_TYPE::TO_DELETE;
Z
Zhiru Zhu 已提交
1106 1107
        files_to_update.emplace_back(f);
    }
G
groot 已提交
1108
    files_holder.ReleaseFiles();
1109

1110 1111 1112
    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";
1113 1114 1115 1116 1117 1118 1119 1120 1121

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

    return status;
}

Status
1122 1123
DBImpl::GetVectorsByID(const std::string& collection_id, const IDNumbers& id_array,
                       std::vector<engine::VectorsData>& vectors) {
1124 1125 1126 1127
    if (!initialized_.load(std::memory_order_acquire)) {
        return SHUTDOWN_ERROR;
    }

1128 1129 1130
    bool has_collection;
    auto status = HasCollection(collection_id, has_collection);
    if (!has_collection) {
1131
        LOG_ENGINE_ERROR_ << "Collection " << collection_id << " does not exist: ";
J
Jin Hai 已提交
1132
        return Status(DB_NOT_FOUND, "Collection does not exist");
1133 1134 1135 1136 1137
    }
    if (!status.ok()) {
        return status;
    }

G
groot 已提交
1138
    meta::FilesHolder files_holder;
J
Jin Hai 已提交
1139 1140
    std::vector<int> file_types{meta::SegmentSchema::FILE_TYPE::RAW, meta::SegmentSchema::FILE_TYPE::TO_INDEX,
                                meta::SegmentSchema::FILE_TYPE::BACKUP};
1141

G
groot 已提交
1142
    status = meta_ptr_->FilesByType(collection_id, file_types, files_holder);
1143
    if (!status.ok()) {
1144
        std::string err_msg = "Failed to get files for GetVectorsByID: " + status.message();
1145
        LOG_ENGINE_ERROR_ << err_msg;
1146 1147 1148
        return status;
    }

J
Jin Hai 已提交
1149 1150
    std::vector<meta::CollectionSchema> partition_array;
    status = meta_ptr_->ShowPartitions(collection_id, partition_array);
1151
    for (auto& schema : partition_array) {
G
groot 已提交
1152
        status = meta_ptr_->FilesByType(schema.collection_id_, file_types, files_holder);
1153 1154
        if (!status.ok()) {
            std::string err_msg = "Failed to get files for GetVectorByID: " + status.message();
1155
            LOG_ENGINE_ERROR_ << err_msg;
1156 1157 1158 1159
            return status;
        }
    }

G
groot 已提交
1160
    if (files_holder.HoldFiles().empty()) {
1161
        LOG_ENGINE_DEBUG_ << "No files to get vector by id from";
J
Jin Hai 已提交
1162
        return Status(DB_NOT_FOUND, "Collection is empty");
1163 1164 1165
    }

    cache::CpuCacheMgr::GetInstance()->PrintInfo();
G
groot 已提交
1166
    status = GetVectorsByIdHelper(collection_id, id_array, vectors, files_holder);
1167 1168 1169 1170 1171 1172
    cache::CpuCacheMgr::GetInstance()->PrintInfo();

    return status;
}

Status
J
Jin Hai 已提交
1173
DBImpl::GetVectorIDs(const std::string& collection_id, const std::string& segment_id, IDNumbers& vector_ids) {
1174 1175 1176 1177
    if (!initialized_.load(std::memory_order_acquire)) {
        return SHUTDOWN_ERROR;
    }

J
Jin Hai 已提交
1178
    // step 1: check collection existence
1179 1180 1181
    bool has_collection;
    auto status = HasCollection(collection_id, has_collection);
    if (!has_collection) {
1182
        LOG_ENGINE_ERROR_ << "Collection " << collection_id << " does not exist: ";
J
Jin Hai 已提交
1183
        return Status(DB_NOT_FOUND, "Collection does not exist");
1184 1185 1186 1187 1188 1189
    }
    if (!status.ok()) {
        return status;
    }

    //  step 2: find segment
G
groot 已提交
1190 1191
    meta::FilesHolder files_holder;
    status = meta_ptr_->GetCollectionFilesBySegmentId(segment_id, files_holder);
1192 1193 1194 1195
    if (!status.ok()) {
        return status;
    }

G
groot 已提交
1196
    milvus::engine::meta::SegmentsSchema& collection_files = files_holder.HoldFiles();
1197
    if (collection_files.empty()) {
1198 1199 1200
        return Status(DB_NOT_FOUND, "Segment does not exist");
    }

J
Jin Hai 已提交
1201
    // check the segment is belong to this collection
1202
    if (collection_files[0].collection_id_ != collection_id) {
J
Jin Hai 已提交
1203
        // the segment could be in a partition under this collection
1204 1205 1206 1207
        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 已提交
1208
            return Status(DB_NOT_FOUND, "Segment does not belong to this collection");
1209 1210 1211 1212 1213
        }
    }

    // step 3: load segment ids and delete offset
    std::string segment_dir;
1214
    engine::utils::GetParentPath(collection_files[0].location_, segment_dir);
1215 1216 1217 1218 1219 1220 1221 1222 1223 1224 1225 1226 1227 1228 1229 1230 1231 1232 1233 1234 1235 1236 1237 1238 1239
    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 已提交
1240

G
groot 已提交
1241
    return status;
X
Xu Peng 已提交
1242 1243
}

1244
Status
1245
DBImpl::GetVectorsByIdHelper(const std::string& collection_id, const IDNumbers& id_array,
G
groot 已提交
1246 1247 1248
                             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();
1249
    LOG_ENGINE_DEBUG_ << "Getting vector by id in " << files.size() << " files, id count = " << id_array.size();
J
Jin Hai 已提交
1250

1251 1252 1253 1254 1255 1256 1257
    // 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;
1258

1259 1260 1261
    vectors.clear();

    IDNumbers temp_ids = id_array;
1262 1263 1264 1265 1266 1267 1268 1269
    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);

1270 1271 1272 1273 1274 1275 1276 1277 1278 1279 1280
        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);
1281 1282 1283
                if (!status.ok()) {
                    return status;
                }
1284 1285 1286 1287 1288 1289 1290 1291

                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);
1292
                    if (!status.ok()) {
J
Jin Hai 已提交
1293
                        LOG_ENGINE_ERROR_ << status.message();
1294 1295
                        return status;
                    }
1296 1297 1298 1299 1300 1301 1302 1303 1304 1305 1306 1307 1308 1309 1310 1311 1312 1313 1314 1315 1316 1317 1318 1319 1320 1321
                    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;
1322 1323 1324
                    }
                }
            }
1325 1326 1327

            it++;
        }
G
groot 已提交
1328 1329 1330

        // unmark file, allow the file to be deleted
        files_holder.UnmarkFile(file);
1331 1332 1333 1334 1335 1336 1337 1338
    }

    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 已提交
1339 1340
            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
1341
        }
1342
        vectors.emplace_back(data);
1343 1344
    }

1345 1346
    if (vectors.empty()) {
        std::string msg = "Vectors not found in collection " + collection_id;
J
Jin Hai 已提交
1347 1348 1349
        LOG_ENGINE_DEBUG_ << msg;
    }

1350 1351 1352
    return Status::OK();
}

S
starlord 已提交
1353
Status
G
groot 已提交
1354 1355
DBImpl::CreateIndex(const std::shared_ptr<server::Context>& context, const std::string& collection_id,
                    const CollectionIndex& index) {
1356
    if (!initialized_.load(std::memory_order_acquire)) {
G
groot 已提交
1357 1358 1359
        return SHUTDOWN_ERROR;
    }

G
groot 已提交
1360
    // serialize memory data
1361 1362
    //    std::set<std::string> sync_collection_ids;
    //    auto status = SyncMemData(sync_collection_ids);
1363
    auto status = Flush();
G
groot 已提交
1364

S
starlord 已提交
1365 1366 1367
    {
        std::unique_lock<std::mutex> lock(build_index_mutex_);

S
starlord 已提交
1368
        // step 1: check index difference
1369
        CollectionIndex old_index;
J
Jin Hai 已提交
1370
        status = DescribeIndex(collection_id, old_index);
S
starlord 已提交
1371
        if (!status.ok()) {
1372
            LOG_ENGINE_ERROR_ << "Failed to get collection index info for collection: " << collection_id;
S
starlord 已提交
1373 1374 1375
            return status;
        }

S
starlord 已提交
1376
        // step 2: update index info
1377 1378
        CollectionIndex new_index = index;
        new_index.metric_type_ = old_index.metric_type_;  // dont change metric type, it was defined by CreateCollection
S
starlord 已提交
1379
        if (!utils::IsSameIndex(old_index, new_index)) {
1380
            status = UpdateCollectionIndexRecursively(collection_id, new_index);
S
starlord 已提交
1381 1382 1383 1384 1385 1386
            if (!status.ok()) {
                return status;
            }
        }
    }

S
starlord 已提交
1387 1388
    // step 3: let merge file thread finish
    // to avoid duplicate data bug
1389 1390
    WaitMergeFileFinish();

S
starlord 已提交
1391
    // step 4: wait and build index
1392
    status = index_failed_checker_.CleanFailedIndexFileOfCollection(collection_id);
G
groot 已提交
1393
    status = WaitCollectionIndexRecursively(context, collection_id, index);
S
starlord 已提交
1394

G
groot 已提交
1395
    return status;
S
starlord 已提交
1396 1397
}

S
starlord 已提交
1398
Status
1399
DBImpl::DescribeIndex(const std::string& collection_id, CollectionIndex& index) {
1400
    if (!initialized_.load(std::memory_order_acquire)) {
G
groot 已提交
1401 1402 1403
        return SHUTDOWN_ERROR;
    }

1404
    return meta_ptr_->DescribeCollectionIndex(collection_id, index);
S
starlord 已提交
1405 1406
}

S
starlord 已提交
1407
Status
J
Jin Hai 已提交
1408
DBImpl::DropIndex(const std::string& collection_id) {
1409
    if (!initialized_.load(std::memory_order_acquire)) {
G
groot 已提交
1410 1411 1412
        return SHUTDOWN_ERROR;
    }

1413
    LOG_ENGINE_DEBUG_ << "Drop index for collection: " << collection_id;
1414
    return DropCollectionIndexRecursively(collection_id);
S
starlord 已提交
1415 1416
}

S
starlord 已提交
1417
Status
1418 1419 1420
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) {
1421
    if (!initialized_.load(std::memory_order_acquire)) {
G
groot 已提交
1422
        return SHUTDOWN_ERROR;
S
starlord 已提交
1423 1424
    }

1425 1426 1427 1428 1429 1430 1431 1432 1433 1434 1435 1436 1437 1438 1439 1440 1441 1442 1443 1444 1445 1446 1447 1448 1449 1450 1451 1452 1453 1454 1455 1456 1457 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
    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 已提交
1542 1543
}

1544 1545 1546 1547 1548 1549 1550 1551 1552 1553 1554 1555 1556
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 已提交
1557
    meta::FilesHolder files_holder;
1558 1559 1560
    if (partition_tags.empty()) {
        // no partition tag specified, means search in whole table
        // get all table files from parent table
G
groot 已提交
1561
        status = meta_ptr_->FilesToSearch(collection_id, files_holder);
1562 1563 1564 1565 1566 1567 1568 1569 1570 1571
        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 已提交
1572
            status = meta_ptr_->FilesToSearch(schema.collection_id_, files_holder);
1573
            if (!status.ok()) {
G
groot 已提交
1574
                return Status(DB_ERROR, "get files to search failed in HybridQuery");
1575 1576 1577
            }
        }

G
groot 已提交
1578 1579
        if (files_holder.HoldFiles().empty()) {
            return Status::OK();  // no files to search
1580 1581 1582 1583 1584 1585 1586
        }
    } 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 已提交
1587
            status = meta_ptr_->FilesToSearch(partition_name, files_holder);
1588
            if (!status.ok()) {
G
groot 已提交
1589
                return Status(DB_ERROR, "get files to search failed in HybridQuery");
1590 1591 1592
            }
        }

G
groot 已提交
1593
        if (files_holder.HoldFiles().empty()) {
1594 1595 1596 1597 1598
            return Status::OK();
        }
    }

    cache::CpuCacheMgr::GetInstance()->PrintInfo();  // print cache info before query
G
groot 已提交
1599
    status = HybridQueryAsync(query_ctx, collection_id, files_holder, hybrid_search_context, general_query, attr_type,
1600 1601 1602 1603 1604 1605 1606 1607 1608 1609 1610
                              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 已提交
1611
Status
J
Jin Hai 已提交
1612
DBImpl::Query(const std::shared_ptr<server::Context>& context, const std::string& collection_id,
1613 1614
              const std::vector<std::string>& partition_tags, uint64_t k, const milvus::json& extra_params,
              const VectorsData& vectors, ResultIds& result_ids, ResultDistances& result_distances) {
1615
    milvus::server::ContextChild tracer(context, "Query");
Z
Zhiru Zhu 已提交
1616

1617
    if (!initialized_.load(std::memory_order_acquire)) {
G
groot 已提交
1618
        return SHUTDOWN_ERROR;
S
starlord 已提交
1619 1620
    }

G
groot 已提交
1621
    Status status;
G
groot 已提交
1622
    meta::FilesHolder files_holder;
G
groot 已提交
1623
    if (partition_tags.empty()) {
J
Jin Hai 已提交
1624 1625
        // no partition tag specified, means search in whole collection
        // get all collection files from parent collection
G
groot 已提交
1626
        status = meta_ptr_->FilesToSearch(collection_id, files_holder);
G
groot 已提交
1627 1628 1629 1630
        if (!status.ok()) {
            return status;
        }

J
Jin Hai 已提交
1631 1632
        std::vector<meta::CollectionSchema> partition_array;
        status = meta_ptr_->ShowPartitions(collection_id, partition_array);
G
groot 已提交
1633
        for (auto& schema : partition_array) {
G
groot 已提交
1634
            status = meta_ptr_->FilesToSearch(schema.collection_id_, files_holder);
1635 1636
        }

G
groot 已提交
1637 1638
        if (files_holder.HoldFiles().empty()) {
            return Status::OK();  // no files to search
G
groot 已提交
1639 1640 1641 1642
        }
    } else {
        // get files from specified partitions
        std::set<std::string> partition_name_array;
J
Jin Hai 已提交
1643
        status = GetPartitionsByTags(collection_id, partition_tags, partition_name_array);
T
Tinkerrr 已提交
1644 1645 1646
        if (!status.ok()) {
            return status;  // didn't match any partition.
        }
G
groot 已提交
1647 1648

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

G
groot 已提交
1652 1653
        if (files_holder.HoldFiles().empty()) {
            return Status::OK();  // no files to search
1654 1655 1656
        }
    }

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

S
starlord 已提交
1661
    return status;
G
groot 已提交
1662
}
X
Xu Peng 已提交
1663

S
starlord 已提交
1664
Status
1665 1666 1667
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) {
1668
    milvus::server::ContextChild tracer(context, "Query by file id");
Z
Zhiru Zhu 已提交
1669

1670
    if (!initialized_.load(std::memory_order_acquire)) {
G
groot 已提交
1671
        return SHUTDOWN_ERROR;
S
starlord 已提交
1672 1673
    }

S
starlord 已提交
1674
    // get specified files
1675
    std::vector<size_t> ids;
Y
Yu Kun 已提交
1676
    for (auto& id : file_ids) {
1677
        std::string::size_type sz;
J
jinhai 已提交
1678
        ids.push_back(std::stoul(id, &sz));
1679 1680
    }

G
groot 已提交
1681 1682
    meta::FilesHolder files_holder;
    auto status = meta_ptr_->FilesByID(ids, files_holder);
1683 1684
    if (!status.ok()) {
        return status;
1685 1686
    }

G
groot 已提交
1687
    milvus::engine::meta::SegmentsSchema& search_files = files_holder.HoldFiles();
1688
    if (search_files.empty()) {
S
starlord 已提交
1689
        return Status(DB_ERROR, "Invalid file id");
G
groot 已提交
1690 1691
    }

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

S
starlord 已提交
1696
    return status;
1697 1698
}

S
starlord 已提交
1699
Status
Y
Yu Kun 已提交
1700
DBImpl::Size(uint64_t& result) {
1701
    if (!initialized_.load(std::memory_order_acquire)) {
G
groot 已提交
1702
        return SHUTDOWN_ERROR;
S
starlord 已提交
1703 1704
    }

S
starlord 已提交
1705
    return meta_ptr_->Size(result);
S
starlord 已提交
1706 1707 1708
}

///////////////////////////////////////////////////////////////////////////////////////////////////////////////////
S
starlord 已提交
1709
// internal methods
S
starlord 已提交
1710
///////////////////////////////////////////////////////////////////////////////////////////////////////////////////
S
starlord 已提交
1711
Status
G
groot 已提交
1712
DBImpl::QueryAsync(const std::shared_ptr<server::Context>& context, meta::FilesHolder& files_holder, uint64_t k,
1713 1714
                   const milvus::json& extra_params, const VectorsData& vectors, ResultIds& result_ids,
                   ResultDistances& result_distances) {
1715
    milvus::server::ContextChild tracer(context, "Query Async");
G
groot 已提交
1716
    server::CollectQueryMetrics metrics(vectors.vector_count_);
Y
Yu Kun 已提交
1717

G
groot 已提交
1718
    milvus::engine::meta::SegmentsSchema& files = files_holder.HoldFiles();
G
groot 已提交
1719 1720 1721
    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);
1722
        LOG_ENGINE_ERROR_ << msg;
G
groot 已提交
1723 1724 1725
        return Status(DB_ERROR, msg);
    }

S
starlord 已提交
1726
    TimeRecorder rc("");
G
groot 已提交
1727

1728
    // step 1: construct search job
1729
    LOG_ENGINE_DEBUG_ << LogOut("Engine query begin, index file count: %ld", files.size());
1730
    scheduler::SearchJobPtr job = std::make_shared<scheduler::SearchJob>(tracer.Context(), k, extra_params, vectors);
Y
Yu Kun 已提交
1731
    for (auto& file : files) {
J
Jin Hai 已提交
1732
        scheduler::SegmentSchemaPtr file_ptr = std::make_shared<meta::SegmentSchema>(file);
W
wxyu 已提交
1733
        job->AddIndexFile(file_ptr);
G
groot 已提交
1734 1735
    }

1736
    // step 2: put search job to scheduler and wait result
S
starlord 已提交
1737
    scheduler::JobMgrInst::GetInstance()->Put(job);
W
wxyu 已提交
1738
    job->WaitResult();
1739

G
groot 已提交
1740
    files_holder.ReleaseFiles();
W
wxyu 已提交
1741 1742
    if (!job->GetStatus().ok()) {
        return job->GetStatus();
1743
    }
G
groot 已提交
1744

1745
    // step 3: construct results
G
groot 已提交
1746 1747
    result_ids = job->GetResultIds();
    result_distances = job->GetResultDistances();
S
starlord 已提交
1748
    rc.ElapseFromBegin("Engine query totally cost");
G
groot 已提交
1749 1750 1751 1752

    return Status::OK();
}

1753 1754
Status
DBImpl::HybridQueryAsync(const std::shared_ptr<server::Context>& context, const std::string& table_id,
G
groot 已提交
1755
                         meta::FilesHolder& files_holder, context::HybridSearchContextPtr hybrid_search_context,
1756 1757 1758 1759 1760 1761 1762 1763 1764 1765 1766 1767 1768 1769 1770 1771 1772 1773 1774 1775 1776 1777 1778 1779 1780 1781 1782 1783 1784
                         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 已提交
1785 1786
    milvus::engine::meta::SegmentsSchema& files = files_holder.HoldFiles();
    LOG_ENGINE_DEBUG_ << LogOut("Engine query begin, index file count: %ld", files_holder.HoldFiles().size());
1787 1788 1789 1790 1791 1792 1793 1794 1795 1796 1797
    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 已提交
1798
    files_holder.ReleaseFiles();
1799 1800 1801 1802 1803 1804 1805 1806 1807 1808 1809 1810 1811 1812 1813 1814
    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 已提交
1815
void
G
groot 已提交
1816
DBImpl::BackgroundIndexThread() {
Y
yu yunfeng 已提交
1817
    server::SystemInfo::GetInstance().Init();
X
Xu Peng 已提交
1818
    while (true) {
1819
        if (!initialized_.load(std::memory_order_acquire)) {
1820 1821
            WaitMergeFileFinish();
            WaitBuildIndexFinish();
S
starlord 已提交
1822

1823
            LOG_ENGINE_DEBUG_ << "DB background thread exit";
G
groot 已提交
1824 1825
            break;
        }
X
Xu Peng 已提交
1826

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

G
groot 已提交
1829
        WaitMergeFileFinish();
G
groot 已提交
1830 1831
        StartBuildIndexTask();
    }
X
Xu Peng 已提交
1832 1833
}

S
starlord 已提交
1834 1835
void
DBImpl::WaitMergeFileFinish() {
1836
    //    LOG_ENGINE_DEBUG_ << "Begin WaitMergeFileFinish";
1837 1838
    std::lock_guard<std::mutex> lck(merge_result_mutex_);
    for (auto& iter : merge_thread_results_) {
1839 1840
        iter.wait();
    }
1841
    //    LOG_ENGINE_DEBUG_ << "End WaitMergeFileFinish";
1842 1843
}

S
starlord 已提交
1844 1845
void
DBImpl::WaitBuildIndexFinish() {
1846
    //    LOG_ENGINE_DEBUG_ << "Begin WaitBuildIndexFinish";
1847
    std::lock_guard<std::mutex> lck(index_result_mutex_);
Y
Yu Kun 已提交
1848
    for (auto& iter : index_thread_results_) {
1849 1850
        iter.wait();
    }
1851
    //    LOG_ENGINE_DEBUG_ << "End WaitBuildIndexFinish";
1852 1853
}

S
starlord 已提交
1854 1855
void
DBImpl::StartMetricTask() {
G
groot 已提交
1856
    server::Metrics::GetInstance().KeepingAliveCounterIncrement(BACKGROUND_METRIC_INTERVAL);
G
groot 已提交
1857 1858
    int64_t cache_usage = cache::CpuCacheMgr::GetInstance()->CacheUsage();
    int64_t cache_total = cache::CpuCacheMgr::GetInstance()->CacheCapacity();
S
shengjh 已提交
1859 1860
    fiu_do_on("DBImpl.StartMetricTask.InvalidTotalCache", cache_total = 0);

J
JinHai-CN 已提交
1861 1862 1863 1864 1865 1866 1867
    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 已提交
1868
    server::Metrics::GetInstance().GpuCacheUsageGaugeSet();
G
groot 已提交
1869 1870 1871 1872 1873 1874 1875 1876
    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 已提交
1877

K
kun yu 已提交
1878
    server::Metrics::GetInstance().CPUCoreUsagePercentSet();
K
kun yu 已提交
1879 1880
    server::Metrics::GetInstance().GPUTemperature();
    server::Metrics::GetInstance().CPUTemperature();
1881
    server::Metrics::GetInstance().PushToGateway();
G
groot 已提交
1882 1883
}

S
starlord 已提交
1884
void
1885
DBImpl::StartMergeTask() {
1886
    // LOG_ENGINE_DEBUG_ << "Begin StartMergeTask";
1887
    // merge task has been finished?
1888
    {
1889 1890
        std::lock_guard<std::mutex> lck(merge_result_mutex_);
        if (!merge_thread_results_.empty()) {
1891
            std::chrono::milliseconds span(10);
1892 1893
            if (merge_thread_results_.back().wait_for(span) == std::future_status::ready) {
                merge_thread_results_.pop_back();
1894
            }
G
groot 已提交
1895 1896
        }
    }
X
Xu Peng 已提交
1897

1898
    // add new merge task
1899
    {
1900 1901
        std::lock_guard<std::mutex> lck(merge_result_mutex_);
        if (merge_thread_results_.empty()) {
1902 1903
            // collect merge files for all collections(if merge_collection_ids_ is empty) for two reasons:
            // 1. other collections may still has un-merged files
1904
            // 2. server may be closed unexpected, these un-merge files need to be merged when server restart
1905 1906 1907 1908 1909
            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_);
1910 1911 1912 1913
                }
            }

            // start merge file thread
1914
            merge_thread_results_.push_back(
1915 1916
                merge_thread_pool_.enqueue(&DBImpl::BackgroundMerge, this, merge_collection_ids_));
            merge_collection_ids_.clear();
1917
        }
G
groot 已提交
1918
    }
1919

1920
    // LOG_ENGINE_DEBUG_ << "End StartMergeTask";
X
Xu Peng 已提交
1921 1922
}

S
starlord 已提交
1923
Status
G
groot 已提交
1924
DBImpl::MergeFiles(const std::string& collection_id, meta::FilesHolder& files_holder) {
Z
Zhiru Zhu 已提交
1925
    // const std::lock_guard<std::mutex> lock(flush_merge_compact_mutex_);
1926

1927
    LOG_ENGINE_DEBUG_ << "Merge files for collection: " << collection_id;
S
starlord 已提交
1928

J
Jin Hai 已提交
1929
    // step 1: create collection file
1930 1931 1932 1933
    meta::SegmentSchema collection_file;
    collection_file.collection_id_ = collection_id;
    collection_file.file_type_ = meta::SegmentSchema::NEW_MERGE;
    Status status = meta_ptr_->CreateCollectionFile(collection_file);
1934
    if (!status.ok()) {
1935
        LOG_ENGINE_ERROR_ << "Failed to create collection: " << status.ToString();
1936 1937 1938
        return status;
    }

S
starlord 已提交
1939
    // step 2: merge files
1940
    /*
G
groot 已提交
1941
    ExecutionEnginePtr index =
1942 1943
        EngineFactory::Build(collection_file.dimension_, collection_file.location_,
    (EngineType)collection_file.engine_type_, (MetricType)collection_file.metric_type_, collection_file.nlist_);
1944
*/
J
Jin Hai 已提交
1945
    meta::SegmentsSchema updated;
1946 1947

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

G
groot 已提交
1951 1952
    // attention: here is a copy, not reference, since files_holder.UnmarkFile will change the array internal
    milvus::engine::meta::SegmentsSchema files = files_holder.HoldFiles();
Y
Yu Kun 已提交
1953
    for (auto& file : files) {
Y
Yu Kun 已提交
1954
        server::CollectMergeFilesMetrics metrics;
1955 1956
        std::string segment_dir_to_merge;
        utils::GetParentPath(file.location_, segment_dir_to_merge);
1957
        segment_writer_ptr->Merge(segment_dir_to_merge, collection_file.file_id_);
G
groot 已提交
1958 1959 1960

        files_holder.UnmarkFile(file);

1961
        auto file_schema = file;
J
Jin Hai 已提交
1962
        file_schema.file_type_ = meta::SegmentSchema::TO_DELETE;
1963
        updated.push_back(file_schema);
1964 1965
        auto size = segment_writer_ptr->Size();
        if (size >= file_schema.index_file_size_) {
S
starlord 已提交
1966
            break;
S
starlord 已提交
1967
        }
1968 1969
    }

S
starlord 已提交
1970
    // step 3: serialize to disk
S
starlord 已提交
1971
    try {
1972
        status = segment_writer_ptr->Serialize();
S
shengjh 已提交
1973 1974
        fiu_do_on("DBImpl.MergeFiles.Serialize_ThrowException", throw std::exception());
        fiu_do_on("DBImpl.MergeFiles.Serialize_ErrorStatus", status = Status(DB_ERROR, ""));
Y
Yu Kun 已提交
1975
    } catch (std::exception& ex) {
S
starlord 已提交
1976
        std::string msg = "Serialize merged index encounter exception: " + std::string(ex.what());
1977
        LOG_ENGINE_ERROR_ << msg;
G
groot 已提交
1978 1979
        status = Status(DB_ERROR, msg);
    }
Y
yu yunfeng 已提交
1980

G
groot 已提交
1981
    if (!status.ok()) {
1982
        LOG_ENGINE_ERROR_ << "Failed to persist merged segment: " << new_segment_dir << ". Error: " << status.message();
1983

G
groot 已提交
1984
        // if failed to serialize merge file to disk
1985
        // typical error: out of disk space, out of memory or permission denied
1986 1987
        collection_file.file_type_ = meta::SegmentSchema::TO_DELETE;
        status = meta_ptr_->UpdateCollectionFile(collection_file);
1988 1989
        LOG_ENGINE_DEBUG_ << "Failed to update file to index, mark file: " << collection_file.file_id_
                          << " to to_delete";
X
Xu Peng 已提交
1990

G
groot 已提交
1991
        return status;
S
starlord 已提交
1992 1993
    }

J
Jin Hai 已提交
1994
    // step 4: update collection files state
1995
    // if index type isn't IDMAP, set file type to TO_INDEX if file size exceed index_file_size
S
starlord 已提交
1996
    // else set file type to RAW, no need to build index
1997 1998 1999 2000
    if (!utils::IsRawIndexType(collection_file.engine_type_)) {
        collection_file.file_type_ = (segment_writer_ptr->Size() >= collection_file.index_file_size_)
                                         ? meta::SegmentSchema::TO_INDEX
                                         : meta::SegmentSchema::RAW;
2001
    } else {
2002
        collection_file.file_type_ = meta::SegmentSchema::RAW;
2003
    }
2004 2005 2006 2007
    collection_file.file_size_ = segment_writer_ptr->Size();
    collection_file.row_count_ = segment_writer_ptr->VectorCount();
    updated.push_back(collection_file);
    status = meta_ptr_->UpdateCollectionFiles(updated);
2008 2009
    LOG_ENGINE_DEBUG_ << "New merged segment " << collection_file.segment_id_ << " of size "
                      << segment_writer_ptr->Size() << " bytes";
2010

S
starlord 已提交
2011
    if (options_.insert_cache_immediately_) {
2012
        segment_writer_ptr->Cache();
S
starlord 已提交
2013
    }
X
Xu Peng 已提交
2014

2015 2016 2017
    return status;
}

2018
Status
G
groot 已提交
2019
DBImpl::MergeHybridFiles(const std::string& collection_id, meta::FilesHolder& files_holder) {
2020 2021 2022 2023 2024 2025 2026 2027 2028 2029 2030 2031 2032 2033 2034 2035 2036 2037 2038 2039 2040 2041 2042 2043 2044 2045 2046
    // 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 已提交
2047 2048
    // attention: here is a copy, not reference, since files_holder.UnmarkFile will change the array internal
    milvus::engine::meta::SegmentsSchema files = files_holder.HoldFiles();
2049 2050 2051 2052 2053
    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 已提交
2054 2055 2056

        files_holder.UnmarkFile(file);

2057 2058 2059 2060 2061 2062 2063 2064 2065 2066 2067 2068 2069 2070 2071 2072 2073 2074 2075 2076 2077 2078 2079 2080 2081 2082 2083 2084 2085 2086 2087 2088 2089 2090 2091 2092 2093 2094 2095 2096 2097 2098 2099 2100 2101 2102 2103 2104 2105 2106 2107 2108 2109 2110 2111 2112
        auto file_schema = file;
        file_schema.file_type_ = meta::SegmentSchema::TO_DELETE;
        updated.push_back(file_schema);
        auto size = segment_writer_ptr->Size();
        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_)) {
        table_file.file_type_ = (segment_writer_ptr->Size() >= table_file.index_file_size_)
                                    ? 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 已提交
2113
Status
J
Jin Hai 已提交
2114
DBImpl::BackgroundMergeFiles(const std::string& collection_id) {
Z
Zhiru Zhu 已提交
2115
    const std::lock_guard<std::mutex> lock(flush_merge_compact_mutex_);
2116

G
groot 已提交
2117 2118
    meta::FilesHolder files_holder;
    auto status = meta_ptr_->FilesToMerge(collection_id, files_holder);
X
Xu Peng 已提交
2119
    if (!status.ok()) {
2120
        LOG_ENGINE_ERROR_ << "Failed to get merge files for collection: " << collection_id;
X
Xu Peng 已提交
2121 2122
        return status;
    }
2123

G
groot 已提交
2124
    if (files_holder.HoldFiles().size() < options_.merge_trigger_number_) {
2125
        LOG_ENGINE_TRACE_ << "Files number not greater equal than merge trigger number, skip merge action";
2126 2127
        return Status::OK();
    }
2128

G
groot 已提交
2129
    MergeFiles(collection_id, files_holder);
G
groot 已提交
2130

2131
    if (!initialized_.load(std::memory_order_acquire)) {
2132
        LOG_ENGINE_DEBUG_ << "Server will shutdown, skip merge action for collection: " << collection_id;
2133
    }
X
Xu Peng 已提交
2134

G
groot 已提交
2135 2136
    return Status::OK();
}
2137

S
starlord 已提交
2138
void
2139
DBImpl::BackgroundMerge(std::set<std::string> collection_ids) {
2140
    // LOG_ENGINE_TRACE_ << " Background merge thread start";
S
starlord 已提交
2141

G
groot 已提交
2142
    Status status;
2143
    for (auto& collection_id : collection_ids) {
J
Jin Hai 已提交
2144
        status = BackgroundMergeFiles(collection_id);
G
groot 已提交
2145
        if (!status.ok()) {
2146
            LOG_ENGINE_ERROR_ << "Merge files for collection " << collection_id << " failed: " << status.ToString();
G
groot 已提交
2147
        }
S
starlord 已提交
2148

2149
        if (!initialized_.load(std::memory_order_acquire)) {
2150
            LOG_ENGINE_DEBUG_ << "Server will shutdown, skip merge action";
S
starlord 已提交
2151 2152
            break;
        }
G
groot 已提交
2153
    }
X
Xu Peng 已提交
2154

G
groot 已提交
2155
    meta_ptr_->Archive();
Z
update  
zhiru 已提交
2156

2157
    {
G
groot 已提交
2158 2159
        uint64_t timeout = (options_.file_cleanup_timeout_ > 0) ? options_.file_cleanup_timeout_ : 10;
        uint64_t ttl = timeout * meta::SECOND;  // default: file will be hard-deleted few seconds after soft-deleted
2160
        meta_ptr_->CleanUpFilesWithTTL(ttl);
Z
update  
zhiru 已提交
2161
    }
S
starlord 已提交
2162

2163
    // LOG_ENGINE_TRACE_ << " Background merge thread exit";
G
groot 已提交
2164
}
X
Xu Peng 已提交
2165

S
starlord 已提交
2166
void
G
groot 已提交
2167
DBImpl::StartBuildIndexTask() {
S
starlord 已提交
2168
    // build index has been finished?
2169 2170 2171 2172 2173 2174 2175
    {
        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 已提交
2176 2177 2178
        }
    }

S
starlord 已提交
2179
    // add new build index task
2180 2181 2182
    {
        std::lock_guard<std::mutex> lck(index_result_mutex_);
        if (index_thread_results_.empty()) {
S
starlord 已提交
2183
            index_thread_results_.push_back(index_thread_pool_.enqueue(&DBImpl::BackgroundBuildIndex, this));
2184
        }
G
groot 已提交
2185
    }
X
Xu Peng 已提交
2186 2187
}

S
starlord 已提交
2188 2189
void
DBImpl::BackgroundBuildIndex() {
P
peng.xu 已提交
2190
    std::unique_lock<std::mutex> lock(build_index_mutex_);
G
groot 已提交
2191 2192 2193 2194
    meta::FilesHolder files_holder;
    meta_ptr_->FilesToIndex(files_holder);

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

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

2200
        // step 2: put build index task to scheduler
J
Jin Hai 已提交
2201
        std::vector<std::pair<scheduler::BuildIndexJobPtr, scheduler::SegmentSchemaPtr>> job2file_map;
2202
        for (auto& file : to_index_files) {
G
groot 已提交
2203
            scheduler::BuildIndexJobPtr job = std::make_shared<scheduler::BuildIndexJob>(meta_ptr_, options_);
J
Jin Hai 已提交
2204
            scheduler::SegmentSchemaPtr file_ptr = std::make_shared<meta::SegmentSchema>(file);
2205
            job->AddToIndexFiles(file_ptr);
G
groot 已提交
2206
            scheduler::JobMgrInst::GetInstance()->Put(job);
G
groot 已提交
2207
            job2file_map.push_back(std::make_pair(job, file_ptr));
2208
        }
G
groot 已提交
2209

G
groot 已提交
2210
        // step 3: wait build index finished and mark failed files
G
groot 已提交
2211 2212
        for (auto iter = job2file_map.begin(); iter != job2file_map.end(); ++iter) {
            scheduler::BuildIndexJobPtr job = iter->first;
J
Jin Hai 已提交
2213
            meta::SegmentSchema& file_schema = *(iter->second.get());
G
groot 已提交
2214 2215 2216
            job->WaitBuildIndexFinish();
            if (!job->GetStatus().ok()) {
                Status status = job->GetStatus();
2217
                LOG_ENGINE_ERROR_ << "Building index job " << job->id() << " failed: " << status.ToString();
G
groot 已提交
2218

2219
                index_failed_checker_.MarkFailedIndexFile(file_schema, status.message());
G
groot 已提交
2220
            } else {
2221
                LOG_ENGINE_DEBUG_ << "Building index job " << job->id() << " succeed.";
G
groot 已提交
2222 2223

                index_failed_checker_.MarkSucceedIndexFile(file_schema);
G
groot 已提交
2224
            }
G
groot 已提交
2225 2226
            status = files_holder.UnmarkFile(file_schema);
            LOG_ENGINE_DEBUG_ << "Finish build index file " << file_schema.file_id_;
2227
        }
G
groot 已提交
2228

2229
        LOG_ENGINE_DEBUG_ << "Background build index thread finished";
G
groot 已提交
2230
        index_req_swn_.Notify();  // notify CreateIndex check circle
Y
Yu Kun 已提交
2231
    }
X
Xu Peng 已提交
2232 2233
}

G
groot 已提交
2234
Status
J
Jin Hai 已提交
2235
DBImpl::GetFilesToBuildIndex(const std::string& collection_id, const std::vector<int>& file_types,
G
groot 已提交
2236 2237 2238
                             meta::FilesHolder& files_holder) {
    files_holder.ReleaseFiles();
    auto status = meta_ptr_->FilesByType(collection_id, file_types, files_holder);
G
groot 已提交
2239

G
groot 已提交
2240 2241 2242 2243 2244 2245 2246 2247 2248 2249 2250
    // 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 已提交
2251 2252 2253 2254 2255
    }

    return Status::OK();
}

2256
Status
J
Jin Hai 已提交
2257 2258
DBImpl::GetPartitionByTag(const std::string& collection_id, const std::string& partition_tag,
                          std::string& partition_name) {
2259 2260 2261
    Status status;

    if (partition_tag.empty()) {
J
Jin Hai 已提交
2262
        partition_name = collection_id;
2263 2264 2265 2266 2267 2268 2269 2270

    } 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 已提交
2271
            partition_name = collection_id;
2272 2273 2274
            return status;
        }

J
Jin Hai 已提交
2275
        status = meta_ptr_->GetPartitionName(collection_id, partition_tag, partition_name);
2276
        if (!status.ok()) {
2277
            LOG_ENGINE_ERROR_ << status.message();
2278 2279 2280 2281 2282 2283
        }
    }

    return status;
}

G
groot 已提交
2284
Status
J
Jin Hai 已提交
2285
DBImpl::GetPartitionsByTags(const std::string& collection_id, const std::vector<std::string>& partition_tags,
G
groot 已提交
2286
                            std::set<std::string>& partition_name_array) {
J
Jin Hai 已提交
2287 2288
    std::vector<meta::CollectionSchema> partition_array;
    auto status = meta_ptr_->ShowPartitions(collection_id, partition_array);
G
groot 已提交
2289 2290

    for (auto& tag : partition_tags) {
2291 2292 2293 2294
        // 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);
2295 2296

        if (valid_tag == milvus::engine::DEFAULT_PARTITON_TAG) {
J
Jin Hai 已提交
2297
            partition_name_array.insert(collection_id);
2298 2299 2300
            return status;
        }

G
groot 已提交
2301
        for (auto& schema : partition_array) {
2302
            if (server::StringHelpFunctions::IsRegexMatch(schema.partition_tag_, valid_tag)) {
J
Jin Hai 已提交
2303
                partition_name_array.insert(schema.collection_id_);
G
groot 已提交
2304 2305 2306 2307
            }
        }
    }

T
Tinkerrr 已提交
2308
    if (partition_name_array.empty()) {
G
groot 已提交
2309
        return Status(DB_PARTITION_NOT_FOUND, "The specified partiton does not exist");
T
Tinkerrr 已提交
2310 2311
    }

G
groot 已提交
2312 2313 2314 2315
    return Status::OK();
}

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

    Status status;
2321
    if (options_.wal_enable_) {
2322
        wal_mgr_->DropCollection(collection_id);
G
groot 已提交
2323 2324
    }

2325 2326 2327
    status = mem_mgr_->EraseMemVector(collection_id);   // not allow insert
    status = meta_ptr_->DropCollection(collection_id);  // soft delete collection
    index_failed_checker_.CleanFailedIndexFileOfCollection(collection_id);
2328

J
Jin Hai 已提交
2329
    // scheduler will determine when to delete collection files
2330
    auto nres = scheduler::ResMgrInst::GetInstance()->GetNumOfComputeResource();
J
Jin Hai 已提交
2331
    scheduler::DeleteJobPtr job = std::make_shared<scheduler::DeleteJob>(collection_id, meta_ptr_, nres);
2332 2333 2334
    scheduler::JobMgrInst::GetInstance()->Put(job);
    job->WaitAndDelete();

J
Jin Hai 已提交
2335 2336
    std::vector<meta::CollectionSchema> partition_array;
    status = meta_ptr_->ShowPartitions(collection_id, partition_array);
G
groot 已提交
2337
    for (auto& schema : partition_array) {
2338 2339
        status = DropCollectionRecursively(schema.collection_id_);
        fiu_do_on("DBImpl.DropCollectionRecursively.failed", status = Status(DB_ERROR, ""));
G
groot 已提交
2340 2341 2342 2343 2344 2345 2346 2347 2348
        if (!status.ok()) {
            return status;
        }
    }

    return Status::OK();
}

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

2352 2353
    auto status = meta_ptr_->UpdateCollectionIndex(collection_id, index);
    fiu_do_on("DBImpl.UpdateCollectionIndexRecursively.fail_update_collection_index",
S
shengjh 已提交
2354
              status = Status(DB_META_TRANSACTION_FAILED, ""));
G
groot 已提交
2355
    if (!status.ok()) {
2356
        LOG_ENGINE_ERROR_ << "Failed to update collection index info for collection: " << collection_id;
G
groot 已提交
2357 2358 2359
        return status;
    }

J
Jin Hai 已提交
2360 2361
    std::vector<meta::CollectionSchema> partition_array;
    status = meta_ptr_->ShowPartitions(collection_id, partition_array);
A
AzAz 已提交
2362 2363 2364
    if (!status.ok()) {
        return status;
    }
G
groot 已提交
2365
    for (auto& schema : partition_array) {
2366
        status = UpdateCollectionIndexRecursively(schema.collection_id_, index);
G
groot 已提交
2367 2368 2369 2370 2371 2372 2373 2374 2375
        if (!status.ok()) {
            return status;
        }
    }

    return Status::OK();
}

Status
G
groot 已提交
2376 2377
DBImpl::WaitCollectionIndexRecursively(const std::shared_ptr<server::Context>& context,
                                       const std::string& collection_id, const CollectionIndex& index) {
G
groot 已提交
2378 2379 2380
    // 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;
2381
    if (utils::IsRawIndexType(index.engine_type_)) {
G
groot 已提交
2382
        file_types = {
J
Jin Hai 已提交
2383 2384
            static_cast<int32_t>(meta::SegmentSchema::NEW),
            static_cast<int32_t>(meta::SegmentSchema::NEW_MERGE),
G
groot 已提交
2385 2386 2387
        };
    } else {
        file_types = {
J
Jin Hai 已提交
2388 2389 2390
            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 已提交
2391 2392 2393 2394
        };
    }

    // get files to build index
G
groot 已提交
2395 2396 2397 2398
    {
        meta::FilesHolder files_holder;
        auto status = GetFilesToBuildIndex(collection_id, file_types, files_holder);
        int times = 1;
G
groot 已提交
2399
        uint64_t repeat = 0;
G
groot 已提交
2400
        while (!files_holder.HoldFiles().empty()) {
G
groot 已提交
2401 2402 2403 2404 2405 2406
            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 已提交
2407
            }
G
groot 已提交
2408

G
groot 已提交
2409 2410 2411 2412 2413 2414 2415 2416 2417 2418 2419 2420 2421 2422
            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 已提交
2423 2424 2425 2426
        }
    }

    // build index for partition
J
Jin Hai 已提交
2427
    std::vector<meta::CollectionSchema> partition_array;
G
groot 已提交
2428
    auto status = meta_ptr_->ShowPartitions(collection_id, partition_array);
G
groot 已提交
2429
    for (auto& schema : partition_array) {
G
groot 已提交
2430
        status = WaitCollectionIndexRecursively(context, schema.collection_id_, index);
2431
        fiu_do_on("DBImpl.WaitCollectionIndexRecursively.fail_build_collection_Index_for_partition",
S
shengjh 已提交
2432
                  status = Status(DB_ERROR, ""));
G
groot 已提交
2433 2434 2435 2436 2437
        if (!status.ok()) {
            return status;
        }
    }

G
groot 已提交
2438
    // failed to build index for some files, return error
2439
    std::string err_msg;
2440 2441
    index_failed_checker_.GetErrMsgForCollection(collection_id, err_msg);
    fiu_do_on("DBImpl.WaitCollectionIndexRecursively.not_empty_err_msg", err_msg.append("fiu"));
2442 2443
    if (!err_msg.empty()) {
        return Status(DB_ERROR, err_msg);
G
groot 已提交
2444 2445
    }

G
groot 已提交
2446 2447
    LOG_ENGINE_DEBUG_ << "WaitCollectionIndexRecursively finished";

G
groot 已提交
2448 2449 2450 2451
    return Status::OK();
}

Status
2452
DBImpl::DropCollectionIndexRecursively(const std::string& collection_id) {
2453
    LOG_ENGINE_DEBUG_ << "Drop index for collection: " << collection_id;
2454 2455
    index_failed_checker_.CleanFailedIndexFileOfCollection(collection_id);
    auto status = meta_ptr_->DropCollectionIndex(collection_id);
G
groot 已提交
2456 2457 2458 2459 2460
    if (!status.ok()) {
        return status;
    }

    // drop partition index
J
Jin Hai 已提交
2461 2462
    std::vector<meta::CollectionSchema> partition_array;
    status = meta_ptr_->ShowPartitions(collection_id, partition_array);
G
groot 已提交
2463
    for (auto& schema : partition_array) {
2464 2465
        status = DropCollectionIndexRecursively(schema.collection_id_);
        fiu_do_on("DBImpl.DropCollectionIndexRecursively.fail_drop_collection_Index_for_partition",
S
shengjh 已提交
2466
                  status = Status(DB_ERROR, ""));
G
groot 已提交
2467 2468 2469 2470 2471 2472 2473 2474 2475
        if (!status.ok()) {
            return status;
        }
    }

    return Status::OK();
}

Status
2476
DBImpl::GetCollectionRowCountRecursively(const std::string& collection_id, uint64_t& row_count) {
G
groot 已提交
2477
    row_count = 0;
J
Jin Hai 已提交
2478
    auto status = meta_ptr_->Count(collection_id, row_count);
G
groot 已提交
2479 2480 2481 2482 2483
    if (!status.ok()) {
        return status;
    }

    // get partition row count
J
Jin Hai 已提交
2484 2485
    std::vector<meta::CollectionSchema> partition_array;
    status = meta_ptr_->ShowPartitions(collection_id, partition_array);
G
groot 已提交
2486
    for (auto& schema : partition_array) {
G
groot 已提交
2487
        uint64_t partition_row_count = 0;
2488 2489
        status = GetCollectionRowCountRecursively(schema.collection_id_, partition_row_count);
        fiu_do_on("DBImpl.GetCollectionRowCountRecursively.fail_get_collection_rowcount_for_partition",
S
shengjh 已提交
2490
                  status = Status(DB_ERROR, ""));
G
groot 已提交
2491 2492 2493 2494 2495 2496 2497 2498 2499 2500
        if (!status.ok()) {
            return status;
        }

        row_count += partition_row_count;
    }

    return Status::OK();
}

2501 2502 2503 2504
Status
DBImpl::ExecWalRecord(const wal::MXLogRecord& record) {
    fiu_return_on("DBImpl.ExexWalRecord.return", Status(););

2505 2506
    auto collections_flushed = [&](const std::set<std::string>& collection_ids) -> uint64_t {
        if (collection_ids.empty()) {
2507 2508 2509 2510 2511
            return 0;
        }

        uint64_t max_lsn = 0;
        if (options_.wal_enable_) {
2512
            for (auto& collection : collection_ids) {
2513
                uint64_t lsn = 0;
2514 2515
                meta_ptr_->GetCollectionFlushLSN(collection, lsn);
                wal_mgr_->CollectionFlushed(collection, lsn);
2516 2517 2518 2519 2520 2521 2522
                if (lsn > max_lsn) {
                    max_lsn = lsn;
                }
            }
        }

        std::lock_guard<std::mutex> lck(merge_result_mutex_);
2523 2524
        for (auto& collection : collection_ids) {
            merge_collection_ids_.insert(collection);
2525 2526 2527 2528 2529 2530 2531
        }
        return max_lsn;
    };

    Status status;

    switch (record.type) {
2532 2533 2534 2535
        case wal::MXLogType::Entity: {
            std::string target_collection_name;
            status = GetPartitionByTag(record.collection_id, record.partition_tag, target_collection_name);
            if (!status.ok()) {
2536
                LOG_WAL_ERROR_ << LogOut("[%s][%ld] ", "insert", 0) << "Get partition fail: " << status.message();
2537 2538 2539
                return status;
            }

2540
            std::set<std::string> flushed_collections;
2541 2542 2543
            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,
2544 2545
                                              record.attr_data, record.lsn, flushed_collections);
            collections_flushed(flushed_collections);
2546 2547 2548 2549

            milvus::server::CollectInsertMetrics metrics(record.length, status);
            break;
        }
2550
        case wal::MXLogType::InsertBinary: {
2551 2552
            std::string target_collection_name;
            status = GetPartitionByTag(record.collection_id, record.partition_tag, target_collection_name);
2553
            if (!status.ok()) {
2554
                LOG_WAL_ERROR_ << LogOut("[%s][%ld] ", "insert", 0) << "Get partition fail: " << status.message();
2555 2556 2557
                return status;
            }

2558 2559
            std::set<std::string> flushed_collections;
            status = mem_mgr_->InsertVectors(target_collection_name, record.length, record.ids,
2560
                                             (record.data_size / record.length / sizeof(uint8_t)),
2561
                                             (const u_int8_t*)record.data, record.lsn, flushed_collections);
2562
            // even though !status.ok, run
2563
            collections_flushed(flushed_collections);
2564 2565 2566 2567 2568 2569 2570

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

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

2578 2579
            std::set<std::string> flushed_collections;
            status = mem_mgr_->InsertVectors(target_collection_name, record.length, record.ids,
2580
                                             (record.data_size / record.length / sizeof(float)),
2581
                                             (const float*)record.data, record.lsn, flushed_collections);
2582
            // even though !status.ok, run
2583
            collections_flushed(flushed_collections);
2584 2585 2586 2587 2588 2589 2590

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

        case wal::MXLogType::Delete: {
J
Jin Hai 已提交
2591 2592
            std::vector<meta::CollectionSchema> partition_array;
            status = meta_ptr_->ShowPartitions(record.collection_id, partition_array);
2593 2594 2595 2596
            if (!status.ok()) {
                return status;
            }

2597
            std::vector<std::string> collection_ids{record.collection_id};
2598
            for (auto& partition : partition_array) {
2599 2600
                auto& partition_collection_id = partition.collection_id_;
                collection_ids.emplace_back(partition_collection_id);
2601 2602 2603
            }

            if (record.length == 1) {
2604
                for (auto& collection_id : collection_ids) {
J
Jin Hai 已提交
2605
                    status = mem_mgr_->DeleteVector(collection_id, *record.ids, record.lsn);
2606 2607 2608 2609 2610
                    if (!status.ok()) {
                        return status;
                    }
                }
            } else {
2611
                for (auto& collection_id : collection_ids) {
J
Jin Hai 已提交
2612
                    status = mem_mgr_->DeleteVectors(collection_id, record.length, record.ids, record.lsn);
2613 2614 2615 2616 2617 2618 2619 2620 2621
                    if (!status.ok()) {
                        return status;
                    }
                }
            }
            break;
        }

        case wal::MXLogType::Flush: {
J
Jin Hai 已提交
2622 2623 2624 2625
            if (!record.collection_id.empty()) {
                // flush one collection
                std::vector<meta::CollectionSchema> partition_array;
                status = meta_ptr_->ShowPartitions(record.collection_id, partition_array);
2626 2627 2628 2629
                if (!status.ok()) {
                    return status;
                }

2630
                std::vector<std::string> collection_ids{record.collection_id};
2631
                for (auto& partition : partition_array) {
2632 2633
                    auto& partition_collection_id = partition.collection_id_;
                    collection_ids.emplace_back(partition_collection_id);
2634 2635
                }

2636 2637
                std::set<std::string> flushed_collections;
                for (auto& collection_id : collection_ids) {
2638
                    const std::lock_guard<std::mutex> lock(flush_merge_compact_mutex_);
J
Jin Hai 已提交
2639
                    status = mem_mgr_->Flush(collection_id);
2640 2641 2642
                    if (!status.ok()) {
                        break;
                    }
2643
                    flushed_collections.insert(collection_id);
2644 2645
                }

2646
                collections_flushed(flushed_collections);
2647 2648

            } else {
2649 2650
                // flush all collections
                std::set<std::string> collection_ids;
2651 2652
                {
                    const std::lock_guard<std::mutex> lock(flush_merge_compact_mutex_);
2653
                    status = mem_mgr_->Flush(collection_ids);
2654 2655
                }

2656
                uint64_t lsn = collections_flushed(collection_ids);
2657 2658 2659 2660 2661 2662 2663 2664 2665 2666 2667 2668
                if (options_.wal_enable_) {
                    wal_mgr_->RemoveOldFiles(lsn);
                }
            }
            break;
        }
    }

    return status;
}

void
G
groot 已提交
2669 2670 2671 2672 2673 2674 2675 2676 2677 2678 2679
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() {
2680
    SetThreadName("wal_thread");
2681 2682
    server::SystemInfo::GetInstance().Init();

2683
    std::chrono::system_clock::time_point next_auto_flush_time;
2684
    auto get_next_auto_flush_time = [&]() {
2685
        return std::chrono::system_clock::now() + std::chrono::seconds(options_.auto_flush_interval_);
2686
    };
2687 2688 2689
    if (options_.auto_flush_interval_ > 0) {
        next_auto_flush_time = get_next_auto_flush_time();
    }
2690 2691

    while (true) {
2692 2693
        if (options_.auto_flush_interval_ > 0) {
            if (std::chrono::system_clock::now() >= next_auto_flush_time) {
G
groot 已提交
2694
                InternalFlush();
2695 2696
                next_auto_flush_time = get_next_auto_flush_time();
            }
2697 2698
        }

G
groot 已提交
2699
        wal::MXLogRecord record;
2700 2701
        auto error_code = wal_mgr_->GetNextRecord(record);
        if (error_code != WAL_SUCCESS) {
2702
            LOG_ENGINE_ERROR_ << "WAL background GetNextRecord error";
2703 2704 2705 2706 2707 2708
            break;
        }

        if (record.type != wal::MXLogType::None) {
            ExecWalRecord(record);
            if (record.type == wal::MXLogType::Flush) {
G
groot 已提交
2709 2710
                // notify flush request to return
                flush_req_swn_.Notify();
2711 2712

                // if user flush all manually, update auto flush also
J
Jin Hai 已提交
2713
                if (record.collection_id.empty() && options_.auto_flush_interval_ > 0) {
2714 2715 2716 2717 2718 2719
                    next_auto_flush_time = get_next_auto_flush_time();
                }
            }

        } else {
            if (!initialized_.load(std::memory_order_acquire)) {
G
groot 已提交
2720 2721
                InternalFlush();
                flush_req_swn_.Notify();
2722 2723
                WaitMergeFileFinish();
                WaitBuildIndexFinish();
2724
                LOG_ENGINE_DEBUG_ << "WAL background thread exit";
2725 2726 2727
                break;
            }

2728
            if (options_.auto_flush_interval_ > 0) {
G
groot 已提交
2729
                swn_wal_.Wait_Until(next_auto_flush_time);
2730
            } else {
G
groot 已提交
2731
                swn_wal_.Wait();
2732
            }
2733 2734 2735 2736
        }
    }
}

G
groot 已提交
2737 2738
void
DBImpl::BackgroundFlushThread() {
2739
    SetThreadName("flush_thread");
G
groot 已提交
2740 2741 2742
    server::SystemInfo::GetInstance().Init();
    while (true) {
        if (!initialized_.load(std::memory_order_acquire)) {
2743
            LOG_ENGINE_DEBUG_ << "DB background flush thread exit";
G
groot 已提交
2744 2745 2746 2747 2748 2749 2750 2751 2752 2753 2754 2755 2756 2757 2758 2759 2760
            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)) {
2761
            LOG_ENGINE_DEBUG_ << "DB background metric thread exit";
G
groot 已提交
2762 2763 2764 2765 2766
            break;
        }

        swn_metric_.Wait_For(std::chrono::seconds(BACKGROUND_METRIC_INTERVAL));
        StartMetricTask();
G
groot 已提交
2767
        meta::FilesHolder::PrintInfo();
G
groot 已提交
2768 2769 2770
    }
}

2771 2772 2773 2774 2775
void
DBImpl::OnCacheInsertDataChanged(bool value) {
    options_.insert_cache_immediately_ = value;
}

2776 2777 2778 2779 2780
void
DBImpl::OnUseBlasThresholdChanged(int64_t threshold) {
    faiss::distance_compute_blas_threshold = threshold;
}

S
starlord 已提交
2781 2782
}  // namespace engine
}  // namespace milvus