HappenLee commented on code in PR #54276: URL: https://github.com/apache/doris/pull/54276#discussion_r2287990195
########## be/src/olap/rowset/segment_v2/ann_index/faiss_ann_index.cpp: ########## @@ -0,0 +1,450 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you 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 +// +// 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. + +#include "faiss_ann_index.h" + +#include <faiss/index_io.h> +#include <omp.h> + +#include <cmath> +#include <cstddef> +#include <cstdint> +#include <limits> +#include <memory> +#include <string> + +#include "CLucene/store/IndexInput.h" +#include "CLucene/store/IndexOutput.h" +#include "common/config.h" +#include "common/exception.h" +#include "common/logging.h" +#include "common/status.h" +#include "faiss/IndexHNSW.h" +#include "faiss/MetricType.h" +#include "faiss/impl/IDSelector.h" +#include "faiss/impl/io.h" +#include "olap/rowset/segment_v2/ann_index/ann_index.h" +#include "olap/rowset/segment_v2/ann_index/ann_index_files.h" +#include "olap/rowset/segment_v2/ann_index/ann_search_params.h" +#include "util/time.h" +#include "vec/core/types.h" + +namespace doris::segment_v2 { +#include "common/compile_check_begin.h" +std::unique_ptr<faiss::IDSelector> FaissVectorIndex::roaring_to_faiss_selector( + const roaring::Roaring& roaring) { + std::vector<faiss::idx_t> ids; + ids.reserve(roaring.cardinality()); + + for (roaring::Roaring::const_iterator it = roaring.begin(); it != roaring.end(); ++it) { + ids.push_back(static_cast<faiss::idx_t>(*it)); + } + // construct derived and wrap into base unique_ptr explicitly + return std::unique_ptr<faiss::IDSelector>(new faiss::IDSelectorBatch(ids.size(), ids.data())); +} + +void FaissVectorIndex::update_roaring(const faiss::idx_t* labels, const size_t n, + roaring::Roaring& roaring) { + // make sure roaring is empty before adding new elements + DCHECK(roaring.cardinality() == 0); + for (size_t i = 0; i < n; ++i) { + if (labels[i] >= 0) { + roaring.add(cast_set<vectorized::UInt32>(labels[i])); + } + } +} + +FaissVectorIndex::FaissVectorIndex() : _index(nullptr) {} + +struct FaissIndexWriter : faiss::IOWriter { +public: + FaissIndexWriter() = default; + FaissIndexWriter(lucene::store::IndexOutput* output) : _output(output) {} + ~FaissIndexWriter() override { + if (_output != nullptr) { + _output->close(); + delete _output; + } + } + + size_t operator()(const void* ptr, size_t size, size_t nitems) override { + size_t bytes = size * nitems; + if (bytes > 0) { + const auto* data = reinterpret_cast<const uint8_t*>(ptr); + // CLucene IndexOutput::writeBytes accepts at most Int32 bytes at a time. + const size_t kMaxChunk = + static_cast<size_t>(std::numeric_limits<vectorized::Int32>::max()); + size_t written = 0; + while (written < bytes) { + size_t to_write = bytes - written; + if (to_write > kMaxChunk) to_write = kMaxChunk; + try { + _output->writeBytes(data + written, cast_set<vectorized::Int32>(to_write)); + } catch (const std::exception& e) { + throw doris::Exception(doris::ErrorCode::IO_ERROR, + "Failed to write vector index {}", e.what()); + } + written += to_write; + } + } + return nitems; + }; + + lucene::store::IndexOutput* _output = nullptr; +}; + +struct FaissIndexReader : faiss::IOReader { +public: + FaissIndexReader() = default; + FaissIndexReader(lucene::store::IndexInput* input) : _input(input) {} + ~FaissIndexReader() override { + if (_input != nullptr) { + _input->close(); + delete _input; + } + } + size_t operator()(void* ptr, size_t size, size_t nitems) override { + size_t bytes = size * nitems; + if (bytes > 0) { + auto* data = reinterpret_cast<uint8_t*>(ptr); + const size_t kMaxChunk = + static_cast<size_t>(std::numeric_limits<vectorized::Int32>::max()); + size_t read = 0; + while (read < bytes) { + size_t to_read = bytes - read; + if (to_read > kMaxChunk) to_read = kMaxChunk; + try { + _input->readBytes(data + read, cast_set<vectorized::Int32>(to_read)); + } catch (const std::exception& e) { + throw doris::Exception(doris::ErrorCode::IO_ERROR, + "Failed to read vector index {}", e.what()); + } + read += to_read; + } + } + return nitems; + }; + + lucene::store::IndexInput* _input = nullptr; +}; + +/** Add n vectors of dimension d to the index. +* +* Vectors are implicitly assigned labels ntotal .. ntotal + n - 1 +* This function slices the input vectors in chunks smaller than +* blocksize_add and calls add_core. +* @param n number of vectors +* @param x input matrix, size n * d +*/ +doris::Status FaissVectorIndex::add(int n, const float* vec) { + DCHECK(vec != nullptr); + DCHECK(_index != nullptr); + omp_set_num_threads(config::omp_threads_limit); Review Comment: why we need each time set the config? should we set once time ? -- This is an automated message from the Apache Git Service. To respond to the message, please log on to GitHub and use the URL above to go to the specific comment. To unsubscribe, e-mail: [email protected] For queries about this service, please contact Infrastructure at: [email protected] --------------------------------------------------------------------- To unsubscribe, e-mail: [email protected] For additional commands, e-mail: [email protected]
