DBImpl.cpp 98.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 139
    // 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 已提交
140 141
            // background wal thread
            bg_wal_thread_ = std::thread(&DBImpl::BackgroundWalThread, this);
142 143 144 145
        }
    } else {
        // for distribute version, some nodes are read only
        if (options_.mode_ != DBOptions::MODE::CLUSTER_READONLY) {
G
groot 已提交
146 147
            // background flush thread
            bg_flush_thread_ = std::thread(&DBImpl::BackgroundFlushThread, this);
148
        }
Z
update  
zhiru 已提交
149
    }
S
starlord 已提交
150

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

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

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

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

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

187 188
        WaitMergeFileFinish();

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

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

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

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

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

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

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

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

223 224 225 226 227 228 229 230 231 232 233 234 235 236 237 238 239 240 241 242 243 244 245 246 247
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;
    if (options_.wal_enable_) {
        // TODO(yukun): wal_mgr_->CreateHybridCollection()
    }

    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 454 455 456 457
            std::string msg = "Pre-loaded file: " + file.file_id_ + " size: " + std::to_string(file.file_size_);
            TimeRecorderAuto rc_1(msg);
            engine->Load(true);

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

Y
Yu Kun 已提交
468
    return Status::OK();
Y
Yu Kun 已提交
469 470
}

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

477
    return meta_ptr_->UpdateCollectionFlag(collection_id, flag);
S
starlord 已提交
478 479
}

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

486
    return GetCollectionRowCountRecursively(collection_id, row_count);
G
groot 已提交
487 488 489
}

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

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

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

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

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

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

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

G
groot 已提交
536 537 538 539
    return DropPartition(partition_name);
}

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

J
Jin Hai 已提交
545
    return meta_ptr_->ShowPartitions(collection_id, partition_schema_array);
G
groot 已提交
546 547 548
}

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

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

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

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

603 604 605
    return status;
}

606
Status
607 608
DBImpl::InsertEntities(const std::string& collection_id, const std::string& partition_tag,
                       const std::vector<std::string>& field_names, Entity& entity,
609 610 611 612 613 614 615 616 617 618 619 620 621 622 623 624 625 626 627 628 629 630 631 632 633 634 635 636 637 638 639 640 641 642
                       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;
    // 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);
    } 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);
    }

643 644 645
    uint64_t offset = 0;
    for (auto field_name : field_names) {
        switch (attr_types.at(field_name)) {
646 647 648 649
            case meta::hybrid::DataType::INT8: {
                std::vector<uint8_t> data;
                data.resize(entity.entity_count_ * sizeof(int8_t));

650 651 652 653 654 655 656 657 658 659 660 661 662 663
                std::vector<int64_t> attr_value(entity.entity_count_, 0);
                memcpy(attr_value.data(), entity.attr_value_.data() + offset, entity.entity_count_ * sizeof(int64_t));
                offset += entity.entity_count_ * sizeof(int64_t);

                std::vector<int8_t> raw_value(entity.entity_count_, 0);
                for (uint64_t i = 0; i < entity.entity_count_; ++i) {
                    raw_value[i] = attr_value[i];
                }

                memcpy(data.data(), raw_value.data(), entity.entity_count_ * sizeof(int8_t));
                record.attr_data.insert(std::make_pair(field_name, data));

                record.attr_nbytes.insert(std::make_pair(field_name, sizeof(int8_t)));
                record.attr_data_size.insert(std::make_pair(field_name, entity.entity_count_ * sizeof(int8_t)));
664 665 666 667 668 669
                break;
            }
            case meta::hybrid::DataType::INT16: {
                std::vector<uint8_t> data;
                data.resize(entity.entity_count_ * sizeof(int16_t));

670 671 672 673 674 675 676 677 678 679 680 681 682 683
                std::vector<int64_t> attr_value(entity.entity_count_, 0);
                memcpy(attr_value.data(), entity.attr_value_.data() + offset, entity.entity_count_ * sizeof(int64_t));
                offset += entity.entity_count_ * sizeof(int64_t);

                std::vector<int16_t> raw_value(entity.entity_count_, 0);
                for (uint64_t i = 0; i < entity.entity_count_; ++i) {
                    raw_value[i] = attr_value[i];
                }

                memcpy(data.data(), raw_value.data(), entity.entity_count_ * sizeof(int16_t));
                record.attr_data.insert(std::make_pair(field_name, data));

                record.attr_nbytes.insert(std::make_pair(field_name, sizeof(int16_t)));
                record.attr_data_size.insert(std::make_pair(field_name, entity.entity_count_ * sizeof(int16_t)));
684 685 686 687 688 689
                break;
            }
            case meta::hybrid::DataType::INT32: {
                std::vector<uint8_t> data;
                data.resize(entity.entity_count_ * sizeof(int32_t));

690 691 692 693 694 695 696 697 698 699 700 701 702 703
                std::vector<int64_t> attr_value(entity.entity_count_, 0);
                memcpy(attr_value.data(), entity.attr_value_.data() + offset, entity.entity_count_ * sizeof(int64_t));
                offset += entity.entity_count_ * sizeof(int64_t);

                std::vector<int32_t> raw_value(entity.entity_count_, 0);
                for (uint64_t i = 0; i < entity.entity_count_; ++i) {
                    raw_value[i] = attr_value[i];
                }

                memcpy(data.data(), raw_value.data(), entity.entity_count_ * sizeof(int32_t));
                record.attr_data.insert(std::make_pair(field_name, data));

                record.attr_nbytes.insert(std::make_pair(field_name, sizeof(int32_t)));
                record.attr_data_size.insert(std::make_pair(field_name, entity.entity_count_ * sizeof(int32_t)));
704 705 706 707 708
                break;
            }
            case meta::hybrid::DataType::INT64: {
                std::vector<uint8_t> data;
                data.resize(entity.entity_count_ * sizeof(int64_t));
709 710
                memcpy(data.data(), entity.attr_value_.data() + offset, entity.entity_count_ * sizeof(int64_t));
                record.attr_data.insert(std::make_pair(field_name, data));
711

712 713 714
                record.attr_nbytes.insert(std::make_pair(field_name, sizeof(int64_t)));
                record.attr_data_size.insert(std::make_pair(field_name, entity.entity_count_ * sizeof(int64_t)));
                offset += entity.entity_count_ * sizeof(int64_t);
715 716 717 718 719 720
                break;
            }
            case meta::hybrid::DataType::FLOAT: {
                std::vector<uint8_t> data;
                data.resize(entity.entity_count_ * sizeof(float));

721 722 723
                std::vector<double> attr_value(entity.entity_count_, 0);
                memcpy(attr_value.data(), entity.attr_value_.data() + offset, entity.entity_count_ * sizeof(double));
                offset += entity.entity_count_ * sizeof(double);
724

725 726 727 728 729 730 731 732 733 734
                std::vector<float> raw_value(entity.entity_count_, 0);
                for (uint64_t i = 0; i < entity.entity_count_; ++i) {
                    raw_value[i] = attr_value[i];
                }

                memcpy(data.data(), raw_value.data(), entity.entity_count_ * sizeof(float));
                record.attr_data.insert(std::make_pair(field_name, data));

                record.attr_nbytes.insert(std::make_pair(field_name, sizeof(float)));
                record.attr_data_size.insert(std::make_pair(field_name, entity.entity_count_ * sizeof(float)));
735 736 737 738 739
                break;
            }
            case meta::hybrid::DataType::DOUBLE: {
                std::vector<uint8_t> data;
                data.resize(entity.entity_count_ * sizeof(double));
740 741
                memcpy(data.data(), entity.attr_value_.data() + offset, entity.entity_count_ * sizeof(double));
                record.attr_data.insert(std::make_pair(field_name, data));
742

743 744 745
                record.attr_nbytes.insert(std::make_pair(field_name, sizeof(double)));
                record.attr_data_size.insert(std::make_pair(field_name, entity.entity_count_ * sizeof(double)));
                offset += entity.entity_count_ * sizeof(double);
746 747
                break;
            }
748 749
            default:
                break;
750 751 752 753 754 755 756
        }
    }

    status = ExecWalRecord(record);
    return status;
}

757
Status
J
Jin Hai 已提交
758
DBImpl::DeleteVector(const std::string& collection_id, IDNumber vector_id) {
759 760
    IDNumbers ids;
    ids.push_back(vector_id);
J
Jin Hai 已提交
761
    return DeleteVectors(collection_id, ids);
762 763 764
}

Status
J
Jin Hai 已提交
765
DBImpl::DeleteVectors(const std::string& collection_id, IDNumbers vector_ids) {
766 767 768 769 770 771
    if (!initialized_.load(std::memory_order_acquire)) {
        return SHUTDOWN_ERROR;
    }

    Status status;
    if (options_.wal_enable_) {
J
Jin Hai 已提交
772
        wal_mgr_->DeleteById(collection_id, vector_ids);
G
groot 已提交
773
        swn_wal_.Notify();
774 775 776 777
    } else {
        wal::MXLogRecord record;
        record.lsn = 0;  // need to get from meta ?
        record.type = wal::MXLogType::Delete;
J
Jin Hai 已提交
778
        record.collection_id = collection_id;
779 780 781 782 783 784 785 786 787 788
        record.ids = vector_ids.data();
        record.length = vector_ids.size();

        status = ExecWalRecord(record);
    }

    return status;
}

Status
J
Jin Hai 已提交
789
DBImpl::Flush(const std::string& collection_id) {
790 791 792 793 794
    if (!initialized_.load(std::memory_order_acquire)) {
        return SHUTDOWN_ERROR;
    }

    Status status;
795 796
    bool has_collection;
    status = HasCollection(collection_id, has_collection);
797 798 799
    if (!status.ok()) {
        return status;
    }
800
    if (!has_collection) {
801
        LOG_ENGINE_ERROR_ << "Collection to flush does not exist: " << collection_id;
J
Jin Hai 已提交
802
        return Status(DB_NOT_FOUND, "Collection to flush does not exist");
803 804
    }

805
    LOG_ENGINE_DEBUG_ << "Begin flush collection: " << collection_id;
806 807

    if (options_.wal_enable_) {
808
        LOG_ENGINE_DEBUG_ << "WAL flush";
J
Jin Hai 已提交
809
        auto lsn = wal_mgr_->Flush(collection_id);
810
        if (lsn != 0) {
G
groot 已提交
811 812
            swn_wal_.Notify();
            flush_req_swn_.Wait();
813 814 815
        }

    } else {
816
        LOG_ENGINE_DEBUG_ << "MemTable flush";
G
groot 已提交
817
        InternalFlush(collection_id);
818 819
    }

820
    LOG_ENGINE_DEBUG_ << "End flush collection: " << collection_id;
821 822 823 824 825 826 827 828 829 830

    return status;
}

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

831
    LOG_ENGINE_DEBUG_ << "Begin flush all collections";
832 833 834

    Status status;
    if (options_.wal_enable_) {
835
        LOG_ENGINE_DEBUG_ << "WAL flush";
836 837
        auto lsn = wal_mgr_->Flush();
        if (lsn != 0) {
G
groot 已提交
838 839
            swn_wal_.Notify();
            flush_req_swn_.Wait();
840 841
        }
    } else {
842
        LOG_ENGINE_DEBUG_ << "MemTable flush";
G
groot 已提交
843
        InternalFlush();
844 845
    }

846
    LOG_ENGINE_DEBUG_ << "End flush all collections";
847 848 849 850 851

    return status;
}

Status
J
Jin Hai 已提交
852
DBImpl::Compact(const std::string& collection_id) {
853 854 855 856
    if (!initialized_.load(std::memory_order_acquire)) {
        return SHUTDOWN_ERROR;
    }

857 858 859
    engine::meta::CollectionSchema collection_schema;
    collection_schema.collection_id_ = collection_id;
    auto status = DescribeCollection(collection_schema);
860 861
    if (!status.ok()) {
        if (status.code() == DB_NOT_FOUND) {
862
            LOG_ENGINE_ERROR_ << "Collection to compact does not exist: " << collection_id;
J
Jin Hai 已提交
863
            return Status(DB_NOT_FOUND, "Collection to compact does not exist");
864 865 866 867
        } else {
            return status;
        }
    } else {
868
        if (!collection_schema.owner_collection_.empty()) {
869
            LOG_ENGINE_ERROR_ << "Collection to compact does not exist: " << collection_id;
J
Jin Hai 已提交
870
            return Status(DB_NOT_FOUND, "Collection to compact does not exist");
871 872 873
        }
    }

874
    LOG_ENGINE_DEBUG_ << "Before compacting, wait for build index thread to finish...";
875

Z
update  
Zhiru Zhu 已提交
876
    // WaitBuildIndexFinish();
877

Z
update  
Zhiru Zhu 已提交
878
    const std::lock_guard<std::mutex> index_lock(build_index_mutex_);
Z
Zhiru Zhu 已提交
879
    const std::lock_guard<std::mutex> merge_lock(flush_merge_compact_mutex_);
Z
Zhiru Zhu 已提交
880

881
    LOG_ENGINE_DEBUG_ << "Compacting collection: " << collection_id;
Z
Zhiru Zhu 已提交
882

883
    // Get files to compact from meta.
J
Jin Hai 已提交
884 885
    std::vector<int> file_types{meta::SegmentSchema::FILE_TYPE::RAW, meta::SegmentSchema::FILE_TYPE::TO_INDEX,
                                meta::SegmentSchema::FILE_TYPE::BACKUP};
G
groot 已提交
886 887
    meta::FilesHolder files_holder;
    status = meta_ptr_->FilesByType(collection_id, file_types, files_holder);
888 889
    if (!status.ok()) {
        std::string err_msg = "Failed to get files to compact: " + status.message();
890
        LOG_ENGINE_ERROR_ << err_msg;
891 892 893
        return Status(DB_ERROR, err_msg);
    }

G
groot 已提交
894
    LOG_ENGINE_DEBUG_ << "Found " << files_holder.HoldFiles().size() << " segment to compact";
Z
Zhiru Zhu 已提交
895 896

    Status compact_status;
G
groot 已提交
897 898
    // 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 已提交
899
    for (auto iter = files_to_compact.begin(); iter != files_to_compact.end();) {
J
Jin Hai 已提交
900
        meta::SegmentSchema file = *iter;
G
groot 已提交
901 902
        iter = files_to_compact.erase(iter);

Z
Zhiru Zhu 已提交
903 904 905
        // Check if the segment needs compacting
        std::string segment_dir;
        utils::GetParentPath(file.location_, segment_dir);
906

Z
Zhiru Zhu 已提交
907
        segment::SegmentReader segment_reader(segment_dir);
Z
Zhiru Zhu 已提交
908 909
        size_t deleted_docs_size;
        status = segment_reader.ReadDeletedDocsSize(deleted_docs_size);
Z
Zhiru Zhu 已提交
910
        if (!status.ok()) {
G
groot 已提交
911
            files_holder.UnmarkFile(file);
G
groot 已提交
912
            continue;  // skip this file and try compact next one
Z
Zhiru Zhu 已提交
913 914
        }

J
Jin Hai 已提交
915
        meta::SegmentsSchema files_to_update;
Z
Zhiru Zhu 已提交
916
        if (deleted_docs_size != 0) {
J
Jin Hai 已提交
917
            compact_status = CompactFile(collection_id, file, files_to_update);
Z
Zhiru Zhu 已提交
918 919

            if (!compact_status.ok()) {
920 921
                LOG_ENGINE_ERROR_ << "Compact failed for segment " << file.segment_id_ << ": "
                                  << compact_status.message();
G
groot 已提交
922
                files_holder.UnmarkFile(file);
G
groot 已提交
923
                continue;  // skip this file and try compact next one
Z
Zhiru Zhu 已提交
924 925
            }
        } else {
G
groot 已提交
926
            files_holder.UnmarkFile(file);
927
            LOG_ENGINE_DEBUG_ << "Segment " << file.segment_id_ << " has no deleted data. No need to compact";
G
groot 已提交
928
            continue;  // skip this file and try compact next one
929
        }
Z
Zhiru Zhu 已提交
930

931
        LOG_ENGINE_DEBUG_ << "Updating meta after compaction...";
932
        status = meta_ptr_->UpdateCollectionFiles(files_to_update);
G
groot 已提交
933
        files_holder.UnmarkFile(file);
G
groot 已提交
934 935 936 937
        if (!status.ok()) {
            compact_status = status;
            break;  // meta error, could not go on
        }
Z
Zhiru Zhu 已提交
938 939
    }

G
groot 已提交
940
    if (compact_status.ok()) {
941
        LOG_ENGINE_DEBUG_ << "Finished compacting collection: " << collection_id;
G
groot 已提交
942
    }
943

G
groot 已提交
944
    return compact_status;
945 946 947
}

Status
J
Jin Hai 已提交
948 949
DBImpl::CompactFile(const std::string& collection_id, const meta::SegmentSchema& file,
                    meta::SegmentsSchema& files_to_update) {
950
    LOG_ENGINE_DEBUG_ << "Compacting segment " << file.segment_id_ << " for collection: " << collection_id;
951

J
Jin Hai 已提交
952 953 954
    // Create new collection file
    meta::SegmentSchema compacted_file;
    compacted_file.collection_id_ = collection_id;
955
    // compacted_file.date_ = date;
J
Jin Hai 已提交
956
    compacted_file.file_type_ = meta::SegmentSchema::NEW_MERGE;  // TODO: use NEW_MERGE for now
957
    Status status = meta_ptr_->CreateCollectionFile(compacted_file);
958 959

    if (!status.ok()) {
960
        LOG_ENGINE_ERROR_ << "Failed to create collection file: " << status.message();
961 962 963
        return status;
    }

J
Jin Hai 已提交
964
    // Compact (merge) file to the newly created collection file
965 966 967 968 969 970 971 972

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

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

973
    LOG_ENGINE_DEBUG_ << "Compacting begin...";
974 975 976
    segment_writer_ptr->Merge(segment_dir_to_merge, compacted_file.file_id_);

    // Serialize
977
    LOG_ENGINE_DEBUG_ << "Serializing compacted segment...";
978 979
    status = segment_writer_ptr->Serialize();
    if (!status.ok()) {
980
        LOG_ENGINE_ERROR_ << "Failed to serialize compacted segment: " << status.message();
J
Jin Hai 已提交
981
        compacted_file.file_type_ = meta::SegmentSchema::TO_DELETE;
982
        auto mark_status = meta_ptr_->UpdateCollectionFile(compacted_file);
983
        if (mark_status.ok()) {
984
            LOG_ENGINE_DEBUG_ << "Mark file: " << compacted_file.file_id_ << " to to_delete";
985 986 987 988
        }
        return status;
    }

989 990 991 992 993 994 995
    // 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;
996
    } else {
J
Jin Hai 已提交
997
        compacted_file.file_type_ = meta::SegmentSchema::RAW;
998 999 1000
    }

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

Z
Zhiru Zhu 已提交
1005
    files_to_update.emplace_back(compacted_file);
Z
Zhiru Zhu 已提交
1006

Z
Zhiru Zhu 已提交
1007 1008
    // Set all files in segment to TO_DELETE
    auto& segment_id = file.segment_id_;
G
groot 已提交
1009 1010
    meta::FilesHolder files_holder;
    status = meta_ptr_->GetCollectionFilesBySegmentId(segment_id, files_holder);
Z
Zhiru Zhu 已提交
1011 1012 1013
    if (!status.ok()) {
        return status;
    }
G
groot 已提交
1014 1015

    milvus::engine::meta::SegmentsSchema& segment_files = files_holder.HoldFiles();
Z
Zhiru Zhu 已提交
1016
    for (auto& f : segment_files) {
J
Jin Hai 已提交
1017
        f.file_type_ = meta::SegmentSchema::FILE_TYPE::TO_DELETE;
Z
Zhiru Zhu 已提交
1018 1019
        files_to_update.emplace_back(f);
    }
G
groot 已提交
1020
    files_holder.ReleaseFiles();
1021

1022 1023 1024
    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";
1025 1026 1027 1028 1029 1030 1031 1032 1033

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

    return status;
}

Status
1034 1035
DBImpl::GetVectorsByID(const std::string& collection_id, const IDNumbers& id_array,
                       std::vector<engine::VectorsData>& vectors) {
1036 1037 1038 1039
    if (!initialized_.load(std::memory_order_acquire)) {
        return SHUTDOWN_ERROR;
    }

1040 1041 1042
    bool has_collection;
    auto status = HasCollection(collection_id, has_collection);
    if (!has_collection) {
1043
        LOG_ENGINE_ERROR_ << "Collection " << collection_id << " does not exist: ";
J
Jin Hai 已提交
1044
        return Status(DB_NOT_FOUND, "Collection does not exist");
1045 1046 1047 1048 1049
    }
    if (!status.ok()) {
        return status;
    }

G
groot 已提交
1050
    meta::FilesHolder files_holder;
J
Jin Hai 已提交
1051 1052
    std::vector<int> file_types{meta::SegmentSchema::FILE_TYPE::RAW, meta::SegmentSchema::FILE_TYPE::TO_INDEX,
                                meta::SegmentSchema::FILE_TYPE::BACKUP};
1053

G
groot 已提交
1054
    status = meta_ptr_->FilesByType(collection_id, file_types, files_holder);
1055
    if (!status.ok()) {
1056
        std::string err_msg = "Failed to get files for GetVectorsByID: " + status.message();
1057
        LOG_ENGINE_ERROR_ << err_msg;
1058 1059 1060
        return status;
    }

J
Jin Hai 已提交
1061 1062
    std::vector<meta::CollectionSchema> partition_array;
    status = meta_ptr_->ShowPartitions(collection_id, partition_array);
1063
    for (auto& schema : partition_array) {
G
groot 已提交
1064
        status = meta_ptr_->FilesByType(schema.collection_id_, file_types, files_holder);
1065 1066
        if (!status.ok()) {
            std::string err_msg = "Failed to get files for GetVectorByID: " + status.message();
1067
            LOG_ENGINE_ERROR_ << err_msg;
1068 1069 1070 1071
            return status;
        }
    }

G
groot 已提交
1072
    if (files_holder.HoldFiles().empty()) {
1073
        LOG_ENGINE_DEBUG_ << "No files to get vector by id from";
J
Jin Hai 已提交
1074
        return Status(DB_NOT_FOUND, "Collection is empty");
1075 1076 1077
    }

    cache::CpuCacheMgr::GetInstance()->PrintInfo();
G
groot 已提交
1078
    status = GetVectorsByIdHelper(collection_id, id_array, vectors, files_holder);
1079 1080 1081 1082 1083 1084
    cache::CpuCacheMgr::GetInstance()->PrintInfo();

    return status;
}

Status
J
Jin Hai 已提交
1085
DBImpl::GetVectorIDs(const std::string& collection_id, const std::string& segment_id, IDNumbers& vector_ids) {
1086 1087 1088 1089
    if (!initialized_.load(std::memory_order_acquire)) {
        return SHUTDOWN_ERROR;
    }

J
Jin Hai 已提交
1090
    // step 1: check collection existence
1091 1092 1093
    bool has_collection;
    auto status = HasCollection(collection_id, has_collection);
    if (!has_collection) {
1094
        LOG_ENGINE_ERROR_ << "Collection " << collection_id << " does not exist: ";
J
Jin Hai 已提交
1095
        return Status(DB_NOT_FOUND, "Collection does not exist");
1096 1097 1098 1099 1100 1101
    }
    if (!status.ok()) {
        return status;
    }

    //  step 2: find segment
G
groot 已提交
1102 1103
    meta::FilesHolder files_holder;
    status = meta_ptr_->GetCollectionFilesBySegmentId(segment_id, files_holder);
1104 1105 1106 1107
    if (!status.ok()) {
        return status;
    }

G
groot 已提交
1108
    milvus::engine::meta::SegmentsSchema& collection_files = files_holder.HoldFiles();
1109
    if (collection_files.empty()) {
1110 1111 1112
        return Status(DB_NOT_FOUND, "Segment does not exist");
    }

J
Jin Hai 已提交
1113
    // check the segment is belong to this collection
1114
    if (collection_files[0].collection_id_ != collection_id) {
J
Jin Hai 已提交
1115
        // the segment could be in a partition under this collection
1116 1117 1118 1119
        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 已提交
1120
            return Status(DB_NOT_FOUND, "Segment does not belong to this collection");
1121 1122 1123 1124 1125
        }
    }

    // step 3: load segment ids and delete offset
    std::string segment_dir;
1126
    engine::utils::GetParentPath(collection_files[0].location_, segment_dir);
1127 1128 1129 1130 1131 1132 1133 1134 1135 1136 1137 1138 1139 1140 1141 1142 1143 1144 1145 1146 1147 1148 1149 1150 1151
    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 已提交
1152

G
groot 已提交
1153
    return status;
X
Xu Peng 已提交
1154 1155
}

1156
Status
1157
DBImpl::GetVectorsByIdHelper(const std::string& collection_id, const IDNumbers& id_array,
G
groot 已提交
1158 1159 1160
                             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();
1161
    LOG_ENGINE_DEBUG_ << "Getting vector by id in " << files.size() << " files, id count = " << id_array.size();
J
Jin Hai 已提交
1162

1163 1164 1165 1166 1167 1168 1169
    // 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;
1170

1171 1172 1173
    vectors.clear();

    IDNumbers temp_ids = id_array;
1174 1175 1176 1177 1178 1179 1180 1181
    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);

1182 1183 1184 1185 1186 1187 1188 1189 1190 1191 1192
        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);
1193 1194 1195
                if (!status.ok()) {
                    return status;
                }
1196 1197 1198 1199 1200 1201 1202 1203

                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);
1204
                    if (!status.ok()) {
J
Jin Hai 已提交
1205
                        LOG_ENGINE_ERROR_ << status.message();
1206 1207
                        return status;
                    }
1208 1209 1210 1211 1212 1213 1214 1215 1216 1217 1218 1219 1220 1221 1222 1223 1224 1225 1226 1227 1228 1229 1230 1231 1232 1233
                    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;
1234 1235 1236
                    }
                }
            }
1237 1238 1239

            it++;
        }
G
groot 已提交
1240 1241 1242

        // unmark file, allow the file to be deleted
        files_holder.UnmarkFile(file);
1243 1244 1245 1246 1247 1248 1249 1250
    }

    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 已提交
1251 1252
            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
1253
        }
1254
        vectors.emplace_back(data);
1255 1256
    }

1257 1258
    if (vectors.empty()) {
        std::string msg = "Vectors not found in collection " + collection_id;
J
Jin Hai 已提交
1259 1260 1261
        LOG_ENGINE_DEBUG_ << msg;
    }

1262 1263 1264
    return Status::OK();
}

S
starlord 已提交
1265
Status
1266
DBImpl::CreateIndex(const std::string& collection_id, const CollectionIndex& index) {
1267
    if (!initialized_.load(std::memory_order_acquire)) {
G
groot 已提交
1268 1269 1270
        return SHUTDOWN_ERROR;
    }

G
groot 已提交
1271
    // serialize memory data
1272 1273
    //    std::set<std::string> sync_collection_ids;
    //    auto status = SyncMemData(sync_collection_ids);
1274
    auto status = Flush();
G
groot 已提交
1275

S
starlord 已提交
1276 1277 1278
    {
        std::unique_lock<std::mutex> lock(build_index_mutex_);

S
starlord 已提交
1279
        // step 1: check index difference
1280
        CollectionIndex old_index;
J
Jin Hai 已提交
1281
        status = DescribeIndex(collection_id, old_index);
S
starlord 已提交
1282
        if (!status.ok()) {
1283
            LOG_ENGINE_ERROR_ << "Failed to get collection index info for collection: " << collection_id;
S
starlord 已提交
1284 1285 1286
            return status;
        }

S
starlord 已提交
1287
        // step 2: update index info
1288 1289
        CollectionIndex new_index = index;
        new_index.metric_type_ = old_index.metric_type_;  // dont change metric type, it was defined by CreateCollection
S
starlord 已提交
1290
        if (!utils::IsSameIndex(old_index, new_index)) {
1291
            status = UpdateCollectionIndexRecursively(collection_id, new_index);
S
starlord 已提交
1292 1293 1294 1295 1296 1297
            if (!status.ok()) {
                return status;
            }
        }
    }

S
starlord 已提交
1298 1299
    // step 3: let merge file thread finish
    // to avoid duplicate data bug
1300 1301
    WaitMergeFileFinish();

S
starlord 已提交
1302
    // step 4: wait and build index
1303 1304
    status = index_failed_checker_.CleanFailedIndexFileOfCollection(collection_id);
    status = WaitCollectionIndexRecursively(collection_id, index);
S
starlord 已提交
1305

G
groot 已提交
1306
    return status;
S
starlord 已提交
1307 1308
}

S
starlord 已提交
1309
Status
1310
DBImpl::DescribeIndex(const std::string& collection_id, CollectionIndex& index) {
1311
    if (!initialized_.load(std::memory_order_acquire)) {
G
groot 已提交
1312 1313 1314
        return SHUTDOWN_ERROR;
    }

1315
    return meta_ptr_->DescribeCollectionIndex(collection_id, index);
S
starlord 已提交
1316 1317
}

S
starlord 已提交
1318
Status
J
Jin Hai 已提交
1319
DBImpl::DropIndex(const std::string& collection_id) {
1320
    if (!initialized_.load(std::memory_order_acquire)) {
G
groot 已提交
1321 1322 1323
        return SHUTDOWN_ERROR;
    }

1324
    LOG_ENGINE_DEBUG_ << "Drop index for collection: " << collection_id;
1325
    return DropCollectionIndexRecursively(collection_id);
S
starlord 已提交
1326 1327
}

S
starlord 已提交
1328
Status
1329 1330 1331
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) {
1332
    if (!initialized_.load(std::memory_order_acquire)) {
G
groot 已提交
1333
        return SHUTDOWN_ERROR;
S
starlord 已提交
1334 1335
    }

1336 1337 1338 1339 1340 1341 1342 1343 1344 1345 1346 1347 1348 1349 1350 1351 1352 1353 1354 1355 1356 1357 1358 1359 1360 1361 1362 1363 1364 1365 1366 1367 1368 1369 1370 1371 1372 1373 1374 1375 1376 1377 1378 1379 1380 1381 1382 1383 1384 1385 1386 1387 1388 1389 1390 1391 1392 1393 1394 1395 1396 1397 1398 1399 1400 1401 1402 1403 1404 1405 1406 1407 1408 1409 1410 1411 1412 1413 1414 1415 1416 1417 1418 1419 1420 1421 1422 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
    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 已提交
1453 1454
}

1455 1456 1457 1458 1459 1460 1461 1462 1463 1464 1465 1466 1467
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 已提交
1468
    meta::FilesHolder files_holder;
1469 1470 1471
    if (partition_tags.empty()) {
        // no partition tag specified, means search in whole table
        // get all table files from parent table
G
groot 已提交
1472
        status = meta_ptr_->FilesToSearch(collection_id, files_holder);
1473 1474 1475 1476 1477 1478 1479 1480 1481 1482
        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 已提交
1483
            status = meta_ptr_->FilesToSearch(schema.collection_id_, files_holder);
1484
            if (!status.ok()) {
G
groot 已提交
1485
                return Status(DB_ERROR, "get files to search failed in HybridQuery");
1486 1487 1488
            }
        }

G
groot 已提交
1489 1490
        if (files_holder.HoldFiles().empty()) {
            return Status::OK();  // no files to search
1491 1492 1493 1494 1495 1496 1497
        }
    } 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 已提交
1498
            status = meta_ptr_->FilesToSearch(partition_name, files_holder);
1499
            if (!status.ok()) {
G
groot 已提交
1500
                return Status(DB_ERROR, "get files to search failed in HybridQuery");
1501 1502 1503
            }
        }

G
groot 已提交
1504
        if (files_holder.HoldFiles().empty()) {
1505 1506 1507 1508 1509
            return Status::OK();
        }
    }

    cache::CpuCacheMgr::GetInstance()->PrintInfo();  // print cache info before query
G
groot 已提交
1510
    status = HybridQueryAsync(query_ctx, collection_id, files_holder, hybrid_search_context, general_query, attr_type,
1511 1512 1513 1514 1515 1516 1517 1518 1519 1520 1521
                              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 已提交
1522
Status
J
Jin Hai 已提交
1523
DBImpl::Query(const std::shared_ptr<server::Context>& context, const std::string& collection_id,
1524 1525
              const std::vector<std::string>& partition_tags, uint64_t k, const milvus::json& extra_params,
              const VectorsData& vectors, ResultIds& result_ids, ResultDistances& result_distances) {
1526
    milvus::server::ContextChild tracer(context, "Query");
Z
Zhiru Zhu 已提交
1527

1528
    if (!initialized_.load(std::memory_order_acquire)) {
G
groot 已提交
1529
        return SHUTDOWN_ERROR;
S
starlord 已提交
1530 1531
    }

G
groot 已提交
1532
    Status status;
G
groot 已提交
1533
    meta::FilesHolder files_holder;
G
groot 已提交
1534
    if (partition_tags.empty()) {
J
Jin Hai 已提交
1535 1536
        // no partition tag specified, means search in whole collection
        // get all collection files from parent collection
G
groot 已提交
1537
        status = meta_ptr_->FilesToSearch(collection_id, files_holder);
G
groot 已提交
1538 1539 1540 1541
        if (!status.ok()) {
            return status;
        }

J
Jin Hai 已提交
1542 1543
        std::vector<meta::CollectionSchema> partition_array;
        status = meta_ptr_->ShowPartitions(collection_id, partition_array);
G
groot 已提交
1544
        for (auto& schema : partition_array) {
G
groot 已提交
1545
            status = meta_ptr_->FilesToSearch(schema.collection_id_, files_holder);
1546 1547
        }

G
groot 已提交
1548 1549
        if (files_holder.HoldFiles().empty()) {
            return Status::OK();  // no files to search
G
groot 已提交
1550 1551 1552 1553
        }
    } else {
        // get files from specified partitions
        std::set<std::string> partition_name_array;
J
Jin Hai 已提交
1554
        status = GetPartitionsByTags(collection_id, partition_tags, partition_name_array);
T
Tinkerrr 已提交
1555 1556 1557
        if (!status.ok()) {
            return status;  // didn't match any partition.
        }
G
groot 已提交
1558 1559

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

G
groot 已提交
1563 1564
        if (files_holder.HoldFiles().empty()) {
            return Status::OK();  // no files to search
1565 1566 1567
        }
    }

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

S
starlord 已提交
1572
    return status;
G
groot 已提交
1573
}
X
Xu Peng 已提交
1574

S
starlord 已提交
1575
Status
1576 1577 1578
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) {
1579
    milvus::server::ContextChild tracer(context, "Query by file id");
Z
Zhiru Zhu 已提交
1580

1581
    if (!initialized_.load(std::memory_order_acquire)) {
G
groot 已提交
1582
        return SHUTDOWN_ERROR;
S
starlord 已提交
1583 1584
    }

S
starlord 已提交
1585
    // get specified files
1586
    std::vector<size_t> ids;
Y
Yu Kun 已提交
1587
    for (auto& id : file_ids) {
1588
        std::string::size_type sz;
J
jinhai 已提交
1589
        ids.push_back(std::stoul(id, &sz));
1590 1591
    }

G
groot 已提交
1592 1593
    meta::FilesHolder files_holder;
    auto status = meta_ptr_->FilesByID(ids, files_holder);
1594 1595
    if (!status.ok()) {
        return status;
1596 1597
    }

G
groot 已提交
1598
    milvus::engine::meta::SegmentsSchema& search_files = files_holder.HoldFiles();
1599
    if (search_files.empty()) {
S
starlord 已提交
1600
        return Status(DB_ERROR, "Invalid file id");
G
groot 已提交
1601 1602
    }

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

S
starlord 已提交
1607
    return status;
1608 1609
}

S
starlord 已提交
1610
Status
Y
Yu Kun 已提交
1611
DBImpl::Size(uint64_t& result) {
1612
    if (!initialized_.load(std::memory_order_acquire)) {
G
groot 已提交
1613
        return SHUTDOWN_ERROR;
S
starlord 已提交
1614 1615
    }

S
starlord 已提交
1616
    return meta_ptr_->Size(result);
S
starlord 已提交
1617 1618 1619
}

///////////////////////////////////////////////////////////////////////////////////////////////////////////////////
S
starlord 已提交
1620
// internal methods
S
starlord 已提交
1621
///////////////////////////////////////////////////////////////////////////////////////////////////////////////////
S
starlord 已提交
1622
Status
G
groot 已提交
1623
DBImpl::QueryAsync(const std::shared_ptr<server::Context>& context, meta::FilesHolder& files_holder, uint64_t k,
1624 1625
                   const milvus::json& extra_params, const VectorsData& vectors, ResultIds& result_ids,
                   ResultDistances& result_distances) {
1626
    milvus::server::ContextChild tracer(context, "Query Async");
G
groot 已提交
1627
    server::CollectQueryMetrics metrics(vectors.vector_count_);
Y
Yu Kun 已提交
1628

G
groot 已提交
1629
    milvus::engine::meta::SegmentsSchema& files = files_holder.HoldFiles();
G
groot 已提交
1630 1631 1632
    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);
1633
        LOG_ENGINE_ERROR_ << msg;
G
groot 已提交
1634 1635 1636
        return Status(DB_ERROR, msg);
    }

S
starlord 已提交
1637
    TimeRecorder rc("");
G
groot 已提交
1638

1639
    // step 1: construct search job
1640
    LOG_ENGINE_DEBUG_ << LogOut("Engine query begin, index file count: %ld", files.size());
1641
    scheduler::SearchJobPtr job = std::make_shared<scheduler::SearchJob>(tracer.Context(), k, extra_params, vectors);
Y
Yu Kun 已提交
1642
    for (auto& file : files) {
J
Jin Hai 已提交
1643
        scheduler::SegmentSchemaPtr file_ptr = std::make_shared<meta::SegmentSchema>(file);
W
wxyu 已提交
1644
        job->AddIndexFile(file_ptr);
G
groot 已提交
1645 1646
    }

1647
    // step 2: put search job to scheduler and wait result
S
starlord 已提交
1648
    scheduler::JobMgrInst::GetInstance()->Put(job);
W
wxyu 已提交
1649
    job->WaitResult();
1650

G
groot 已提交
1651
    files_holder.ReleaseFiles();
W
wxyu 已提交
1652 1653
    if (!job->GetStatus().ok()) {
        return job->GetStatus();
1654
    }
G
groot 已提交
1655

1656
    // step 3: construct results
G
groot 已提交
1657 1658
    result_ids = job->GetResultIds();
    result_distances = job->GetResultDistances();
S
starlord 已提交
1659
    rc.ElapseFromBegin("Engine query totally cost");
G
groot 已提交
1660 1661 1662 1663

    return Status::OK();
}

1664 1665
Status
DBImpl::HybridQueryAsync(const std::shared_ptr<server::Context>& context, const std::string& table_id,
G
groot 已提交
1666
                         meta::FilesHolder& files_holder, context::HybridSearchContextPtr hybrid_search_context,
1667 1668 1669 1670 1671 1672 1673 1674 1675 1676 1677 1678 1679 1680 1681 1682 1683 1684 1685 1686 1687 1688 1689 1690 1691 1692 1693 1694 1695
                         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 已提交
1696 1697
    milvus::engine::meta::SegmentsSchema& files = files_holder.HoldFiles();
    LOG_ENGINE_DEBUG_ << LogOut("Engine query begin, index file count: %ld", files_holder.HoldFiles().size());
1698 1699 1700 1701 1702 1703 1704 1705 1706 1707 1708
    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 已提交
1709
    files_holder.ReleaseFiles();
1710 1711 1712 1713 1714 1715 1716 1717 1718 1719 1720 1721 1722 1723 1724 1725
    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 已提交
1726
void
G
groot 已提交
1727
DBImpl::BackgroundIndexThread() {
Y
yu yunfeng 已提交
1728
    server::SystemInfo::GetInstance().Init();
X
Xu Peng 已提交
1729
    while (true) {
1730
        if (!initialized_.load(std::memory_order_acquire)) {
1731 1732
            WaitMergeFileFinish();
            WaitBuildIndexFinish();
S
starlord 已提交
1733

1734
            LOG_ENGINE_DEBUG_ << "DB background thread exit";
G
groot 已提交
1735 1736
            break;
        }
X
Xu Peng 已提交
1737

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

G
groot 已提交
1740
        WaitMergeFileFinish();
G
groot 已提交
1741 1742
        StartBuildIndexTask();
    }
X
Xu Peng 已提交
1743 1744
}

S
starlord 已提交
1745 1746
void
DBImpl::WaitMergeFileFinish() {
1747
    //    LOG_ENGINE_DEBUG_ << "Begin WaitMergeFileFinish";
1748 1749
    std::lock_guard<std::mutex> lck(merge_result_mutex_);
    for (auto& iter : merge_thread_results_) {
1750 1751
        iter.wait();
    }
1752
    //    LOG_ENGINE_DEBUG_ << "End WaitMergeFileFinish";
1753 1754
}

S
starlord 已提交
1755 1756
void
DBImpl::WaitBuildIndexFinish() {
1757
    //    LOG_ENGINE_DEBUG_ << "Begin WaitBuildIndexFinish";
1758
    std::lock_guard<std::mutex> lck(index_result_mutex_);
Y
Yu Kun 已提交
1759
    for (auto& iter : index_thread_results_) {
1760 1761
        iter.wait();
    }
1762
    //    LOG_ENGINE_DEBUG_ << "End WaitBuildIndexFinish";
1763 1764
}

S
starlord 已提交
1765 1766
void
DBImpl::StartMetricTask() {
G
groot 已提交
1767
    server::Metrics::GetInstance().KeepingAliveCounterIncrement(BACKGROUND_METRIC_INTERVAL);
G
groot 已提交
1768 1769
    int64_t cache_usage = cache::CpuCacheMgr::GetInstance()->CacheUsage();
    int64_t cache_total = cache::CpuCacheMgr::GetInstance()->CacheCapacity();
S
shengjh 已提交
1770 1771
    fiu_do_on("DBImpl.StartMetricTask.InvalidTotalCache", cache_total = 0);

J
JinHai-CN 已提交
1772 1773 1774 1775 1776 1777 1778
    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 已提交
1779
    server::Metrics::GetInstance().GpuCacheUsageGaugeSet();
G
groot 已提交
1780 1781 1782 1783 1784 1785 1786 1787
    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 已提交
1788

K
kun yu 已提交
1789
    server::Metrics::GetInstance().CPUCoreUsagePercentSet();
K
kun yu 已提交
1790 1791
    server::Metrics::GetInstance().GPUTemperature();
    server::Metrics::GetInstance().CPUTemperature();
1792
    server::Metrics::GetInstance().PushToGateway();
G
groot 已提交
1793 1794
}

S
starlord 已提交
1795
void
1796
DBImpl::StartMergeTask() {
1797
    // LOG_ENGINE_DEBUG_ << "Begin StartMergeTask";
1798
    // merge task has been finished?
1799
    {
1800 1801
        std::lock_guard<std::mutex> lck(merge_result_mutex_);
        if (!merge_thread_results_.empty()) {
1802
            std::chrono::milliseconds span(10);
1803 1804
            if (merge_thread_results_.back().wait_for(span) == std::future_status::ready) {
                merge_thread_results_.pop_back();
1805
            }
G
groot 已提交
1806 1807
        }
    }
X
Xu Peng 已提交
1808

1809
    // add new merge task
1810
    {
1811 1812
        std::lock_guard<std::mutex> lck(merge_result_mutex_);
        if (merge_thread_results_.empty()) {
1813 1814
            // collect merge files for all collections(if merge_collection_ids_ is empty) for two reasons:
            // 1. other collections may still has un-merged files
1815
            // 2. server may be closed unexpected, these un-merge files need to be merged when server restart
1816 1817 1818 1819 1820
            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_);
1821 1822 1823 1824
                }
            }

            // start merge file thread
1825
            merge_thread_results_.push_back(
1826 1827
                merge_thread_pool_.enqueue(&DBImpl::BackgroundMerge, this, merge_collection_ids_));
            merge_collection_ids_.clear();
1828
        }
G
groot 已提交
1829
    }
1830

1831
    // LOG_ENGINE_DEBUG_ << "End StartMergeTask";
X
Xu Peng 已提交
1832 1833
}

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

1838
    LOG_ENGINE_DEBUG_ << "Merge files for collection: " << collection_id;
S
starlord 已提交
1839

J
Jin Hai 已提交
1840
    // step 1: create collection file
1841 1842 1843 1844
    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);
X
Xu Peng 已提交
1845

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

S
starlord 已提交
1851
    // step 2: merge files
1852
    /*
G
groot 已提交
1853
    ExecutionEnginePtr index =
1854 1855
        EngineFactory::Build(collection_file.dimension_, collection_file.location_,
    (EngineType)collection_file.engine_type_, (MetricType)collection_file.metric_type_, collection_file.nlist_);
1856
*/
J
Jin Hai 已提交
1857
    meta::SegmentsSchema updated;
1858 1859

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

G
groot 已提交
1863 1864
    // 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 已提交
1865
    for (auto& file : files) {
Y
Yu Kun 已提交
1866
        server::CollectMergeFilesMetrics metrics;
1867 1868
        std::string segment_dir_to_merge;
        utils::GetParentPath(file.location_, segment_dir_to_merge);
1869
        segment_writer_ptr->Merge(segment_dir_to_merge, collection_file.file_id_);
G
groot 已提交
1870 1871 1872

        files_holder.UnmarkFile(file);

1873
        auto file_schema = file;
J
Jin Hai 已提交
1874
        file_schema.file_type_ = meta::SegmentSchema::TO_DELETE;
1875
        updated.push_back(file_schema);
1876 1877
        auto size = segment_writer_ptr->Size();
        if (size >= file_schema.index_file_size_) {
S
starlord 已提交
1878
            break;
S
starlord 已提交
1879
        }
1880 1881
    }

S
starlord 已提交
1882
    // step 3: serialize to disk
S
starlord 已提交
1883
    try {
1884
        status = segment_writer_ptr->Serialize();
S
shengjh 已提交
1885 1886
        fiu_do_on("DBImpl.MergeFiles.Serialize_ThrowException", throw std::exception());
        fiu_do_on("DBImpl.MergeFiles.Serialize_ErrorStatus", status = Status(DB_ERROR, ""));
Y
Yu Kun 已提交
1887
    } catch (std::exception& ex) {
S
starlord 已提交
1888
        std::string msg = "Serialize merged index encounter exception: " + std::string(ex.what());
1889
        LOG_ENGINE_ERROR_ << msg;
G
groot 已提交
1890 1891
        status = Status(DB_ERROR, msg);
    }
Y
yu yunfeng 已提交
1892

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

G
groot 已提交
1896
        // if failed to serialize merge file to disk
1897
        // typical error: out of disk space, out of memory or permission denied
1898 1899
        collection_file.file_type_ = meta::SegmentSchema::TO_DELETE;
        status = meta_ptr_->UpdateCollectionFile(collection_file);
1900 1901
        LOG_ENGINE_DEBUG_ << "Failed to update file to index, mark file: " << collection_file.file_id_
                          << " to to_delete";
X
Xu Peng 已提交
1902

G
groot 已提交
1903
        return status;
S
starlord 已提交
1904 1905
    }

J
Jin Hai 已提交
1906
    // step 4: update collection files state
1907
    // if index type isn't IDMAP, set file type to TO_INDEX if file size exceed index_file_size
S
starlord 已提交
1908
    // else set file type to RAW, no need to build index
1909 1910 1911 1912
    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;
1913
    } else {
1914
        collection_file.file_type_ = meta::SegmentSchema::RAW;
1915
    }
1916 1917 1918 1919
    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);
1920 1921
    LOG_ENGINE_DEBUG_ << "New merged segment " << collection_file.segment_id_ << " of size "
                      << segment_writer_ptr->Size() << " bytes";
1922

S
starlord 已提交
1923
    if (options_.insert_cache_immediately_) {
1924
        segment_writer_ptr->Cache();
S
starlord 已提交
1925
    }
X
Xu Peng 已提交
1926

1927 1928 1929
    return status;
}

1930
Status
G
groot 已提交
1931
DBImpl::MergeHybridFiles(const std::string& collection_id, meta::FilesHolder& files_holder) {
1932 1933 1934 1935 1936 1937 1938 1939 1940 1941 1942 1943 1944 1945 1946 1947 1948 1949 1950 1951 1952 1953 1954 1955 1956 1957 1958
    // 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 已提交
1959 1960
    // attention: here is a copy, not reference, since files_holder.UnmarkFile will change the array internal
    milvus::engine::meta::SegmentsSchema files = files_holder.HoldFiles();
1961 1962 1963 1964 1965
    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 已提交
1966 1967 1968

        files_holder.UnmarkFile(file);

1969 1970 1971 1972 1973 1974 1975 1976 1977 1978 1979 1980 1981 1982 1983 1984 1985 1986 1987 1988 1989 1990 1991 1992 1993 1994 1995 1996 1997 1998 1999 2000 2001 2002 2003 2004 2005 2006 2007 2008 2009 2010 2011 2012 2013 2014 2015 2016 2017 2018 2019 2020 2021 2022 2023 2024
        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 已提交
2025
Status
J
Jin Hai 已提交
2026
DBImpl::BackgroundMergeFiles(const std::string& collection_id) {
Z
Zhiru Zhu 已提交
2027
    const std::lock_guard<std::mutex> lock(flush_merge_compact_mutex_);
2028

G
groot 已提交
2029 2030
    meta::FilesHolder files_holder;
    auto status = meta_ptr_->FilesToMerge(collection_id, files_holder);
X
Xu Peng 已提交
2031
    if (!status.ok()) {
2032
        LOG_ENGINE_ERROR_ << "Failed to get merge files for collection: " << collection_id;
X
Xu Peng 已提交
2033 2034
        return status;
    }
2035

G
groot 已提交
2036
    if (files_holder.HoldFiles().size() < options_.merge_trigger_number_) {
2037
        LOG_ENGINE_TRACE_ << "Files number not greater equal than merge trigger number, skip merge action";
2038 2039
        return Status::OK();
    }
2040

G
groot 已提交
2041
    MergeFiles(collection_id, files_holder);
G
groot 已提交
2042

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

G
groot 已提交
2047 2048
    return Status::OK();
}
2049

S
starlord 已提交
2050
void
2051
DBImpl::BackgroundMerge(std::set<std::string> collection_ids) {
2052
    // LOG_ENGINE_TRACE_ << " Background merge thread start";
S
starlord 已提交
2053

G
groot 已提交
2054
    Status status;
2055
    for (auto& collection_id : collection_ids) {
J
Jin Hai 已提交
2056
        status = BackgroundMergeFiles(collection_id);
G
groot 已提交
2057
        if (!status.ok()) {
2058
            LOG_ENGINE_ERROR_ << "Merge files for collection " << collection_id << " failed: " << status.ToString();
G
groot 已提交
2059
        }
S
starlord 已提交
2060

2061
        if (!initialized_.load(std::memory_order_acquire)) {
2062
            LOG_ENGINE_DEBUG_ << "Server will shutdown, skip merge action";
S
starlord 已提交
2063 2064
            break;
        }
G
groot 已提交
2065
    }
X
Xu Peng 已提交
2066

G
groot 已提交
2067
    meta_ptr_->Archive();
Z
update  
zhiru 已提交
2068

2069
    {
G
groot 已提交
2070
        uint64_t ttl = 10 * meta::SECOND;  // default: file will be hard-deleted few seconds after soft-deleted
2071
        if (options_.mode_ == DBOptions::MODE::CLUSTER_WRITABLE) {
2072
            ttl = meta::HOUR;
2073
        }
G
groot 已提交
2074

2075
        meta_ptr_->CleanUpFilesWithTTL(ttl);
Z
update  
zhiru 已提交
2076
    }
S
starlord 已提交
2077

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

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

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

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

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

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

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

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

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

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

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

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

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

    return Status::OK();
}

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

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

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

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

    return status;
}

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

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

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

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

T
Tinkerrr 已提交
2223 2224 2225 2226
    if (partition_name_array.empty()) {
        return Status(PARTITION_NOT_FOUND, "Cannot find the specified partitions");
    }

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

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

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

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

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

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

    return Status::OK();
}

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

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

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

    return Status::OK();
}

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

    // get files to build index
G
groot 已提交
2309 2310 2311 2312 2313 2314 2315 2316 2317 2318 2319
    {
        meta::FilesHolder files_holder;
        auto status = GetFilesToBuildIndex(collection_id, file_types, files_holder);
        int times = 1;

        while (!files_holder.HoldFiles().empty()) {
            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 已提交
2320

G
groot 已提交
2321 2322 2323
            index_req_swn_.Wait_For(std::chrono::seconds(WAIT_BUILD_INDEX_INTERVAL));
            GetFilesToBuildIndex(collection_id, file_types, files_holder);
            ++times;
G
groot 已提交
2324 2325 2326 2327
        }
    }

    // build index for partition
J
Jin Hai 已提交
2328
    std::vector<meta::CollectionSchema> partition_array;
G
groot 已提交
2329
    auto status = meta_ptr_->ShowPartitions(collection_id, partition_array);
G
groot 已提交
2330
    for (auto& schema : partition_array) {
2331 2332
        status = WaitCollectionIndexRecursively(schema.collection_id_, index);
        fiu_do_on("DBImpl.WaitCollectionIndexRecursively.fail_build_collection_Index_for_partition",
S
shengjh 已提交
2333
                  status = Status(DB_ERROR, ""));
G
groot 已提交
2334 2335 2336 2337 2338
        if (!status.ok()) {
            return status;
        }
    }

G
groot 已提交
2339
    // failed to build index for some files, return error
2340
    std::string err_msg;
2341 2342
    index_failed_checker_.GetErrMsgForCollection(collection_id, err_msg);
    fiu_do_on("DBImpl.WaitCollectionIndexRecursively.not_empty_err_msg", err_msg.append("fiu"));
2343 2344
    if (!err_msg.empty()) {
        return Status(DB_ERROR, err_msg);
G
groot 已提交
2345 2346
    }

G
groot 已提交
2347 2348
    LOG_ENGINE_DEBUG_ << "WaitCollectionIndexRecursively finished";

G
groot 已提交
2349 2350 2351 2352
    return Status::OK();
}

Status
2353
DBImpl::DropCollectionIndexRecursively(const std::string& collection_id) {
2354
    LOG_ENGINE_DEBUG_ << "Drop index for collection: " << collection_id;
2355 2356
    index_failed_checker_.CleanFailedIndexFileOfCollection(collection_id);
    auto status = meta_ptr_->DropCollectionIndex(collection_id);
G
groot 已提交
2357 2358 2359 2360 2361
    if (!status.ok()) {
        return status;
    }

    // drop partition index
J
Jin Hai 已提交
2362 2363
    std::vector<meta::CollectionSchema> partition_array;
    status = meta_ptr_->ShowPartitions(collection_id, partition_array);
G
groot 已提交
2364
    for (auto& schema : partition_array) {
2365 2366
        status = DropCollectionIndexRecursively(schema.collection_id_);
        fiu_do_on("DBImpl.DropCollectionIndexRecursively.fail_drop_collection_Index_for_partition",
S
shengjh 已提交
2367
                  status = Status(DB_ERROR, ""));
G
groot 已提交
2368 2369 2370 2371 2372 2373 2374 2375 2376
        if (!status.ok()) {
            return status;
        }
    }

    return Status::OK();
}

Status
2377
DBImpl::GetCollectionRowCountRecursively(const std::string& collection_id, uint64_t& row_count) {
G
groot 已提交
2378
    row_count = 0;
J
Jin Hai 已提交
2379
    auto status = meta_ptr_->Count(collection_id, row_count);
G
groot 已提交
2380 2381 2382 2383 2384
    if (!status.ok()) {
        return status;
    }

    // get partition row count
J
Jin Hai 已提交
2385 2386
    std::vector<meta::CollectionSchema> partition_array;
    status = meta_ptr_->ShowPartitions(collection_id, partition_array);
G
groot 已提交
2387
    for (auto& schema : partition_array) {
G
groot 已提交
2388
        uint64_t partition_row_count = 0;
2389 2390
        status = GetCollectionRowCountRecursively(schema.collection_id_, partition_row_count);
        fiu_do_on("DBImpl.GetCollectionRowCountRecursively.fail_get_collection_rowcount_for_partition",
S
shengjh 已提交
2391
                  status = Status(DB_ERROR, ""));
G
groot 已提交
2392 2393 2394 2395 2396 2397 2398 2399 2400 2401
        if (!status.ok()) {
            return status;
        }

        row_count += partition_row_count;
    }

    return Status::OK();
}

2402 2403 2404 2405
Status
DBImpl::ExecWalRecord(const wal::MXLogRecord& record) {
    fiu_return_on("DBImpl.ExexWalRecord.return", Status(););

2406 2407
    auto collections_flushed = [&](const std::set<std::string>& collection_ids) -> uint64_t {
        if (collection_ids.empty()) {
2408 2409 2410 2411 2412
            return 0;
        }

        uint64_t max_lsn = 0;
        if (options_.wal_enable_) {
2413
            for (auto& collection : collection_ids) {
2414
                uint64_t lsn = 0;
2415 2416
                meta_ptr_->GetCollectionFlushLSN(collection, lsn);
                wal_mgr_->CollectionFlushed(collection, lsn);
2417 2418 2419 2420 2421 2422 2423
                if (lsn > max_lsn) {
                    max_lsn = lsn;
                }
            }
        }

        std::lock_guard<std::mutex> lck(merge_result_mutex_);
2424 2425
        for (auto& collection : collection_ids) {
            merge_collection_ids_.insert(collection);
2426 2427 2428 2429 2430 2431 2432
        }
        return max_lsn;
    };

    Status status;

    switch (record.type) {
2433 2434 2435 2436
        case wal::MXLogType::Entity: {
            std::string target_collection_name;
            status = GetPartitionByTag(record.collection_id, record.partition_tag, target_collection_name);
            if (!status.ok()) {
2437
                LOG_WAL_ERROR_ << LogOut("[%s][%ld] ", "insert", 0) << "Get partition fail: " << status.message();
2438 2439 2440
                return status;
            }

2441
            std::set<std::string> flushed_collections;
2442 2443 2444
            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,
2445 2446
                                              record.attr_data, record.lsn, flushed_collections);
            collections_flushed(flushed_collections);
2447 2448 2449 2450

            milvus::server::CollectInsertMetrics metrics(record.length, status);
            break;
        }
2451
        case wal::MXLogType::InsertBinary: {
2452 2453
            std::string target_collection_name;
            status = GetPartitionByTag(record.collection_id, record.partition_tag, target_collection_name);
2454
            if (!status.ok()) {
2455
                LOG_WAL_ERROR_ << LogOut("[%s][%ld] ", "insert", 0) << "Get partition fail: " << status.message();
2456 2457 2458
                return status;
            }

2459 2460
            std::set<std::string> flushed_collections;
            status = mem_mgr_->InsertVectors(target_collection_name, record.length, record.ids,
2461
                                             (record.data_size / record.length / sizeof(uint8_t)),
2462
                                             (const u_int8_t*)record.data, record.lsn, flushed_collections);
2463
            // even though !status.ok, run
2464
            collections_flushed(flushed_collections);
2465 2466 2467 2468 2469 2470 2471

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

        case wal::MXLogType::InsertVector: {
2472 2473
            std::string target_collection_name;
            status = GetPartitionByTag(record.collection_id, record.partition_tag, target_collection_name);
2474
            if (!status.ok()) {
2475
                LOG_WAL_ERROR_ << LogOut("[%s][%ld] ", "insert", 0) << "Get partition fail: " << status.message();
2476 2477 2478
                return status;
            }

2479 2480
            std::set<std::string> flushed_collections;
            status = mem_mgr_->InsertVectors(target_collection_name, record.length, record.ids,
2481
                                             (record.data_size / record.length / sizeof(float)),
2482
                                             (const float*)record.data, record.lsn, flushed_collections);
2483
            // even though !status.ok, run
2484
            collections_flushed(flushed_collections);
2485 2486 2487 2488 2489 2490 2491

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

        case wal::MXLogType::Delete: {
J
Jin Hai 已提交
2492 2493
            std::vector<meta::CollectionSchema> partition_array;
            status = meta_ptr_->ShowPartitions(record.collection_id, partition_array);
2494 2495 2496 2497
            if (!status.ok()) {
                return status;
            }

2498
            std::vector<std::string> collection_ids{record.collection_id};
2499
            for (auto& partition : partition_array) {
2500 2501
                auto& partition_collection_id = partition.collection_id_;
                collection_ids.emplace_back(partition_collection_id);
2502 2503 2504
            }

            if (record.length == 1) {
2505
                for (auto& collection_id : collection_ids) {
J
Jin Hai 已提交
2506
                    status = mem_mgr_->DeleteVector(collection_id, *record.ids, record.lsn);
2507 2508 2509 2510 2511
                    if (!status.ok()) {
                        return status;
                    }
                }
            } else {
2512
                for (auto& collection_id : collection_ids) {
J
Jin Hai 已提交
2513
                    status = mem_mgr_->DeleteVectors(collection_id, record.length, record.ids, record.lsn);
2514 2515 2516 2517 2518 2519 2520 2521 2522
                    if (!status.ok()) {
                        return status;
                    }
                }
            }
            break;
        }

        case wal::MXLogType::Flush: {
J
Jin Hai 已提交
2523 2524 2525 2526
            if (!record.collection_id.empty()) {
                // flush one collection
                std::vector<meta::CollectionSchema> partition_array;
                status = meta_ptr_->ShowPartitions(record.collection_id, partition_array);
2527 2528 2529 2530
                if (!status.ok()) {
                    return status;
                }

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

2537 2538
                std::set<std::string> flushed_collections;
                for (auto& collection_id : collection_ids) {
2539
                    const std::lock_guard<std::mutex> lock(flush_merge_compact_mutex_);
J
Jin Hai 已提交
2540
                    status = mem_mgr_->Flush(collection_id);
2541 2542 2543
                    if (!status.ok()) {
                        break;
                    }
2544
                    flushed_collections.insert(collection_id);
2545 2546
                }

2547
                collections_flushed(flushed_collections);
2548 2549

            } else {
2550 2551
                // flush all collections
                std::set<std::string> collection_ids;
2552 2553
                {
                    const std::lock_guard<std::mutex> lock(flush_merge_compact_mutex_);
2554
                    status = mem_mgr_->Flush(collection_ids);
2555 2556
                }

2557
                uint64_t lsn = collections_flushed(collection_ids);
2558 2559 2560 2561 2562 2563 2564 2565 2566 2567 2568 2569
                if (options_.wal_enable_) {
                    wal_mgr_->RemoveOldFiles(lsn);
                }
            }
            break;
        }
    }

    return status;
}

void
G
groot 已提交
2570 2571 2572 2573 2574 2575 2576 2577 2578 2579 2580
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() {
2581
    SetThreadName("wal_thread");
2582 2583
    server::SystemInfo::GetInstance().Init();

2584
    std::chrono::system_clock::time_point next_auto_flush_time;
2585
    auto get_next_auto_flush_time = [&]() {
2586
        return std::chrono::system_clock::now() + std::chrono::seconds(options_.auto_flush_interval_);
2587
    };
2588 2589 2590
    if (options_.auto_flush_interval_ > 0) {
        next_auto_flush_time = get_next_auto_flush_time();
    }
2591 2592

    while (true) {
2593 2594
        if (options_.auto_flush_interval_ > 0) {
            if (std::chrono::system_clock::now() >= next_auto_flush_time) {
G
groot 已提交
2595
                InternalFlush();
2596 2597
                next_auto_flush_time = get_next_auto_flush_time();
            }
2598 2599
        }

G
groot 已提交
2600
        wal::MXLogRecord record;
2601 2602
        auto error_code = wal_mgr_->GetNextRecord(record);
        if (error_code != WAL_SUCCESS) {
2603
            LOG_ENGINE_ERROR_ << "WAL background GetNextRecord error";
2604 2605 2606 2607 2608 2609
            break;
        }

        if (record.type != wal::MXLogType::None) {
            ExecWalRecord(record);
            if (record.type == wal::MXLogType::Flush) {
G
groot 已提交
2610 2611
                // notify flush request to return
                flush_req_swn_.Notify();
2612 2613

                // if user flush all manually, update auto flush also
J
Jin Hai 已提交
2614
                if (record.collection_id.empty() && options_.auto_flush_interval_ > 0) {
2615 2616 2617 2618 2619 2620
                    next_auto_flush_time = get_next_auto_flush_time();
                }
            }

        } else {
            if (!initialized_.load(std::memory_order_acquire)) {
G
groot 已提交
2621 2622
                InternalFlush();
                flush_req_swn_.Notify();
2623 2624
                WaitMergeFileFinish();
                WaitBuildIndexFinish();
2625
                LOG_ENGINE_DEBUG_ << "WAL background thread exit";
2626 2627 2628
                break;
            }

2629
            if (options_.auto_flush_interval_ > 0) {
G
groot 已提交
2630
                swn_wal_.Wait_Until(next_auto_flush_time);
2631
            } else {
G
groot 已提交
2632
                swn_wal_.Wait();
2633
            }
2634 2635 2636 2637
        }
    }
}

G
groot 已提交
2638 2639
void
DBImpl::BackgroundFlushThread() {
2640
    SetThreadName("flush_thread");
G
groot 已提交
2641 2642 2643
    server::SystemInfo::GetInstance().Init();
    while (true) {
        if (!initialized_.load(std::memory_order_acquire)) {
2644
            LOG_ENGINE_DEBUG_ << "DB background flush thread exit";
G
groot 已提交
2645 2646 2647 2648 2649 2650 2651 2652 2653 2654 2655 2656 2657 2658 2659 2660 2661
            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)) {
2662
            LOG_ENGINE_DEBUG_ << "DB background metric thread exit";
G
groot 已提交
2663 2664 2665 2666 2667
            break;
        }

        swn_metric_.Wait_For(std::chrono::seconds(BACKGROUND_METRIC_INTERVAL));
        StartMetricTask();
G
groot 已提交
2668
        meta::FilesHolder::PrintInfo();
G
groot 已提交
2669 2670 2671
    }
}

2672 2673 2674 2675 2676
void
DBImpl::OnCacheInsertDataChanged(bool value) {
    options_.insert_cache_immediately_ = value;
}

2677 2678 2679 2680 2681
void
DBImpl::OnUseBlasThresholdChanged(int64_t threshold) {
    faiss::distance_compute_blas_threshold = threshold;
}

S
starlord 已提交
2682 2683
}  // namespace engine
}  // namespace milvus