diff --git a/.gitmodules b/.gitmodules index 2a53cb71dad0ee..54c1a8a36366af 100644 --- a/.gitmodules +++ b/.gitmodules @@ -21,3 +21,11 @@ path = contrib/apache-orc url = https://github.com/apache/doris-thirdparty.git branch = orc +[submodule "doris-faiss"] + path = contrib/faiss + url = https://github.com/apache/doris-thirdparty.git + branch = faiss +[submodule "doris-openblas"] + path = contrib/openblas + url = https://github.com/apache/doris-thirdparty.git + branch = openblas diff --git a/be/CMakeLists.txt b/be/CMakeLists.txt index 48cd273ee593ef..9feb4b016ab6fd 100644 --- a/be/CMakeLists.txt +++ b/be/CMakeLists.txt @@ -727,7 +727,7 @@ endif () # use this to avoid some runtime tracker. reuse BE_TEST symbol, no need another. if (BUILD_BENCHMARK) add_definitions(-DBE_TEST) -# The separate BENCHMARK marker is introduced here because +# The separate BENCHMARK marker is introduced here because # some BE UTs mock certain functions, and BENCHMARK cannot find their definitions. add_definitions(-DBE_BENCHMARK) endif() @@ -767,6 +767,7 @@ function(pch_reuse target) endif() endfunction(pch_reuse target) + add_subdirectory(${SRC_DIR}/agent) add_subdirectory(${SRC_DIR}/common) add_subdirectory(${SRC_DIR}/exec) @@ -775,6 +776,7 @@ add_subdirectory(${SRC_DIR}/gen_cpp) add_subdirectory(${SRC_DIR}/geo) add_subdirectory(${SRC_DIR}/http) add_subdirectory(${SRC_DIR}/io) +add_subdirectory(${SRC_DIR}/olap/rowset/segment_v2/ann_index) add_subdirectory(${SRC_DIR}/olap) add_subdirectory(${SRC_DIR}/runtime) add_subdirectory(${SRC_DIR}/runtime_filter) diff --git a/be/src/cloud/cloud_internal_service.cpp b/be/src/cloud/cloud_internal_service.cpp index 68715c69560bd1..32aad81279ec35 100644 --- a/be/src/cloud/cloud_internal_service.cpp +++ b/be/src/cloud/cloud_internal_service.cpp @@ -357,9 +357,8 @@ void CloudInternalServiceImpl::warm_up_rowset(google::protobuf::RpcController* c // inverted index auto schema_ptr = rs_meta.tablet_schema(); auto idx_version = schema_ptr->get_inverted_index_storage_format(); - bool has_inverted_index = schema_ptr->has_inverted_index(); - if (has_inverted_index) { + if (schema_ptr->has_inverted_index() || schema_ptr->has_ann_index()) { if (idx_version == InvertedIndexStorageFormatPB::V1) { auto&& inverted_index_info = rs_meta.inverted_index_file_info(segment_id); std::unordered_map index_size_map; diff --git a/be/src/cloud/cloud_rowset_writer.cpp b/be/src/cloud/cloud_rowset_writer.cpp index fcf6115907e11e..5b77495f8c2f18 100644 --- a/be/src/cloud/cloud_rowset_writer.cpp +++ b/be/src/cloud/cloud_rowset_writer.cpp @@ -112,7 +112,7 @@ Status CloudRowsetWriter::build(RowsetSharedPtr& rowset) { } else { _rowset_meta->add_segments_file_size(seg_file_size.value()); } - if (_context.tablet_schema->has_inverted_index()) { + if (_context.tablet_schema->has_inverted_index() || _context.tablet_schema->has_ann_index()) { if (auto idx_files_info = _idx_files.inverted_index_file_info(_segment_start_id); !idx_files_info.has_value()) [[unlikely]] { LOG(ERROR) << "expected inverted index files info, but none presents: " diff --git a/be/src/cloud/cloud_snapshot_mgr.cpp b/be/src/cloud/cloud_snapshot_mgr.cpp index 6cb78b0ed55e14..7832969e231f37 100644 --- a/be/src/cloud/cloud_snapshot_mgr.cpp +++ b/be/src/cloud/cloud_snapshot_mgr.cpp @@ -235,7 +235,8 @@ Status CloudSnapshotMgr::_create_rowset_meta( file_mapping[src_index_file] = dst_index_file; } } else { - if (context.tablet_schema->has_inverted_index()) { + if (context.tablet_schema->has_inverted_index() || + context.tablet_schema->has_ann_index()) { std::string src_index_file = InvertedIndexDescriptor::get_index_file_path_v2( InvertedIndexDescriptor::get_index_file_path_prefix(src_segment_file)); std::string dst_index_file = InvertedIndexDescriptor::get_index_file_path_v2( diff --git a/be/src/cloud/cloud_tablet.cpp b/be/src/cloud/cloud_tablet.cpp index 88a3a66c2ab972..ad29b781d2b7a3 100644 --- a/be/src/cloud/cloud_tablet.cpp +++ b/be/src/cloud/cloud_tablet.cpp @@ -354,7 +354,7 @@ void CloudTablet::add_rowsets(std::vector to_add, bool version_ download_idx_file(idx_path, index_size_map[index->index_id()]); } } else { - if (schema_ptr->has_inverted_index()) { + if (schema_ptr->has_inverted_index() || schema_ptr->has_ann_index()) { auto&& inverted_index_info = rowset_meta->inverted_index_file_info(seg_id); int64_t idx_size = 0; diff --git a/be/src/cloud/cloud_warm_up_manager.cpp b/be/src/cloud/cloud_warm_up_manager.cpp index 17906672cdfbd2..1f67abd7f14866 100644 --- a/be/src/cloud/cloud_warm_up_manager.cpp +++ b/be/src/cloud/cloud_warm_up_manager.cpp @@ -277,7 +277,7 @@ void CloudWarmUpManager::handle_jobs() { expiration_time, wait, true); } } else { - if (schema_ptr->has_inverted_index()) { + if (schema_ptr->has_inverted_index() || schema_ptr->has_ann_index()) { auto idx_path = storage_resource.value()->remote_idx_v2_path(*rs, seg_id); file_size = idx_file_info.has_index_size() ? idx_file_info.index_size() @@ -556,14 +556,13 @@ void CloudWarmUpManager::warm_up_rowset(RowsetMeta& rs_meta, int64_t sync_wait_t // update metrics auto schema_ptr = rs_meta.tablet_schema(); - bool has_inverted_index = schema_ptr->has_inverted_index(); auto idx_version = schema_ptr->get_inverted_index_storage_format(); for (int64_t segment_id = 0; segment_id < rs_meta.num_segments(); segment_id++) { g_file_cache_event_driven_warm_up_requested_segment_num << 1; g_file_cache_event_driven_warm_up_requested_segment_size << rs_meta.segment_file_size(segment_id); - if (has_inverted_index) { + if (schema_ptr->has_inverted_index() || schema_ptr->has_ann_index()) { if (idx_version == InvertedIndexStorageFormatPB::V1) { auto&& inverted_index_info = rs_meta.inverted_index_file_info(segment_id); if (inverted_index_info.index_info().empty()) { diff --git a/be/src/common/cast_set.h b/be/src/common/cast_set.h index ae9f6ae9bed8fd..035f09e306b46f 100644 --- a/be/src/common/cast_set.h +++ b/be/src/common/cast_set.h @@ -31,24 +31,21 @@ template void check_cast_value(U b) { if constexpr (IsUnsignedV) { if (b > std::numeric_limits::max()) { - throw doris::Exception(ErrorCode::INVALID_ARGUMENT, - "value {} cast to type {} out of range [{},{}]", b, - typeid(T).name(), std::numeric_limits::min(), - std::numeric_limits::max()); + throw doris::Exception( + ErrorCode::INTERNAL_ERROR, "value {} cast to type {} out of range [{},{}]", b, + typeid(T).name(), std::numeric_limits::min(), std::numeric_limits::max()); } } else if constexpr (IsUnsignedV) { if (b < 0 || b > std::numeric_limits::max()) { - throw doris::Exception(ErrorCode::INVALID_ARGUMENT, - "value {} cast to type {} out of range [{},{}]", b, - typeid(T).name(), std::numeric_limits::min(), - std::numeric_limits::max()); + throw doris::Exception( + ErrorCode::INTERNAL_ERROR, "value {} cast to type {} out of range [{},{}]", b, + typeid(T).name(), std::numeric_limits::min(), std::numeric_limits::max()); } } else { if (b < std::numeric_limits::min() || b > std::numeric_limits::max()) { - throw doris::Exception(ErrorCode::INVALID_ARGUMENT, - "value {} cast to type {} out of range [{},{}]", b, - typeid(T).name(), std::numeric_limits::min(), - std::numeric_limits::max()); + throw doris::Exception( + ErrorCode::INTERNAL_ERROR, "value {} cast to type {} out of range [{},{}]", b, + typeid(T).name(), std::numeric_limits::min(), std::numeric_limits::max()); } } } diff --git a/be/src/common/config.cpp b/be/src/common/config.cpp index c7eb4136c52a7e..52cb73ce4918cd 100644 --- a/be/src/common/config.cpp +++ b/be/src/common/config.cpp @@ -1584,6 +1584,12 @@ DEFINE_mBool(enable_auto_clone_on_mow_publish_missing_version, "false"); // The maximum number of threads supported when executing LLMFunction DEFINE_mInt32(llm_max_concurrent_requests, "1"); +// Maximum number of openmp threads can be used by each doris threads. +// This configuration controls the parallelism level for OpenMP operations within Doris, +// helping to prevent resource contention and ensure stable performance when multiple +// Doris threads are executing OpenMP-accelerated operations simultaneously. +DEFINE_mInt32(omp_threads_limit, "8"); + // clang-format off #ifdef BE_TEST // test s3 diff --git a/be/src/common/config.h b/be/src/common/config.h index 5f5900165b02ea..d2e869c9193342 100644 --- a/be/src/common/config.h +++ b/be/src/common/config.h @@ -1636,6 +1636,9 @@ DECLARE_String(fuzzy_test_type); // The maximum number of threads supported when executing LLMFunction DECLARE_mInt32(llm_max_concurrent_requests); +// Maximum number of OpenMP threads that can be used by each Doris thread +DECLARE_Int32(omp_threads_limit); + #ifdef BE_TEST // test s3 DECLARE_String(test_s3_resource); diff --git a/be/src/olap/CMakeLists.txt b/be/src/olap/CMakeLists.txt index f74e53ea995333..ef9e44562d5344 100644 --- a/be/src/olap/CMakeLists.txt +++ b/be/src/olap/CMakeLists.txt @@ -22,7 +22,14 @@ set(LIBRARY_OUTPUT_PATH "${BUILD_DIR}/src/olap") set(EXECUTABLE_OUTPUT_PATH "${BUILD_DIR}/src/olap") file(GLOB_RECURSE SRC_FILES CONFIGURE_DEPENDS *.cpp) + +# Files in the ann_index directory use faiss header files +# Some of Doris's compilation check options fail on these header files, so exclude files in the ann_index directory +# They are compiled separately as a .a library and linked by Olap +list(FILTER SRC_FILES EXCLUDE REGEX ".*/olap/rowset/segment_v2/ann_index/.*\\.cpp$") + add_library(Olap STATIC ${SRC_FILES}) +target_link_libraries(Olap PRIVATE ann_index) if (OS_MACOSX) target_compile_options(Olap PRIVATE -Wno-unused-lambda-capture) diff --git a/be/src/olap/compaction.cpp b/be/src/olap/compaction.cpp index c0c394ae635fa3..cad5f5b4401f56 100644 --- a/be/src/olap/compaction.cpp +++ b/be/src/olap/compaction.cpp @@ -77,6 +77,7 @@ #include "runtime/memory/mem_tracker_limiter.h" #include "runtime/thread_context.h" #include "util/doris_metrics.h" +#include "util/pretty_printer.h" #include "util/time.h" #include "util/trace.h" #include "vec/common/schema_util.h" @@ -556,13 +557,19 @@ Status CompactionMixin::execute_compact_impl(int64_t permits) { LOG(INFO) << "succeed to do " << compaction_name() << " is_vertical=" << _is_vertical << ". tablet=" << _tablet->tablet_id() << ", output_version=" << _output_version << ", current_max_version=" << tablet()->max_version().second - << ", disk=" << tablet()->data_dir()->path() << ", segments=" << _input_num_segments - << ", input_rowsets_data_size=" << _input_rowsets_data_size - << ", input_rowsets_index_size=" << _input_rowsets_index_size - << ", input_rowsets_total_size=" << _input_rowsets_total_size - << ", output_rowset_data_size=" << _output_rowset->data_disk_size() - << ", output_rowset_index_size=" << _output_rowset->index_disk_size() - << ", output_rowset_total_size=" << _output_rowset->total_disk_size() + << ", disk=" << tablet()->data_dir()->path() + << ", input_segments=" << _input_num_segments << ", input_rowsets_data_size=" + << PrettyPrinter::print_bytes(_input_rowsets_data_size) + << ", input_rowsets_index_size=" + << PrettyPrinter::print_bytes(_input_rowsets_index_size) + << ", input_rowsets_total_size=" + << PrettyPrinter::print_bytes(_input_rowsets_total_size) + << ", output_rowset_data_size=" + << PrettyPrinter::print_bytes(_output_rowset->data_disk_size()) + << ", output_rowset_index_size=" + << PrettyPrinter::print_bytes(_output_rowset->index_disk_size()) + << ", output_rowset_total_size=" + << PrettyPrinter::print_bytes(_output_rowset->total_disk_size()) << ", input_row_num=" << _input_row_num << ", output_row_num=" << _output_rowset->num_rows() << ", filtered_row_num=" << _stats.filtered_rows @@ -769,8 +776,8 @@ Status Compaction::do_inverted_index_compaction() { // dest index files // format: rowsetId_segmentId - auto& inverted_index_file_writers = dynamic_cast(_output_rs_writer.get()) - ->inverted_index_file_writers(); + auto& inverted_index_file_writers = + dynamic_cast(_output_rs_writer.get())->index_file_writers(); DBUG_EXECUTE_IF( "Compaction::do_inverted_index_compaction_inverted_index_file_writers_size_error", { inverted_index_file_writers.clear(); }) diff --git a/be/src/olap/iterators.h b/be/src/olap/iterators.h index df05856f86afeb..69c1eae621367b 100644 --- a/be/src/olap/iterators.h +++ b/be/src/olap/iterators.h @@ -25,6 +25,7 @@ #include "olap/block_column_predicate.h" #include "olap/column_predicate.h" #include "olap/olap_common.h" +#include "olap/rowset/segment_v2/ann_index/ann_topn_runtime.h" #include "olap/rowset/segment_v2/row_ranges.h" #include "olap/tablet_schema.h" #include "runtime/runtime_state.h" @@ -122,6 +123,7 @@ class StorageReadOptions { size_t topn_limit = 0; std::map virtual_column_exprs; + std::shared_ptr ann_topn_runtime; std::map vir_cid_to_idx_in_block; std::map vir_col_idx_to_type; diff --git a/be/src/olap/olap_common.h b/be/src/olap/olap_common.h index 3f4b7afbd007ba..8db63359235658 100644 --- a/be/src/olap/olap_common.h +++ b/be/src/olap/olap_common.h @@ -382,6 +382,28 @@ struct OlapReaderStatistics { int64_t inverted_index_lookup_timer = 0; InvertedIndexStatistics inverted_index_stats; + int64_t ann_index_load_ns = 0; + int64_t ann_topn_search_ns = 0; + int64_t ann_index_topn_search_cnt = 0; + + // Detailed timing for ANN operations + int64_t ann_index_topn_engine_search_ns = 0; // time spent in engine for range search + int64_t ann_index_topn_result_process_ns = 0; // time spent processing TopN results + int64_t ann_index_topn_engine_convert_ns = 0; // time spent on FAISS-side conversions (TopN) + int64_t ann_index_topn_engine_prepare_ns = + 0; // time spent preparing before engine search (TopN) + int64_t rows_ann_index_topn_filtered = 0; + + int64_t ann_index_range_search_ns = 0; + int64_t ann_index_range_search_cnt = 0; + // Detailed timing for ANN Range search + int64_t ann_range_engine_search_ns = 0; // time spent in engine for range search + int64_t ann_range_pre_process_ns = 0; // time spent preparing before engine search + + int64_t ann_range_result_convert_ns = 0; // time spent processing range results + int64_t ann_range_engine_convert_ns = 0; // time spent on FAISS-side conversions (Range) + int64_t rows_ann_index_range_filtered = 0; + int64_t output_index_result_column_timer = 0; // number of segment filtered by column stat when creating seg iterator int64_t filtered_segment_number = 0; diff --git a/be/src/olap/rowset/beta_rowset.cpp b/be/src/olap/rowset/beta_rowset.cpp index d0e3aec380da64..78f6d304315f3a 100644 --- a/be/src/olap/rowset/beta_rowset.cpp +++ b/be/src/olap/rowset/beta_rowset.cpp @@ -255,7 +255,7 @@ Status BetaRowset::remove() { } } } else { - if (_schema->has_inverted_index()) { + if (_schema->has_inverted_index() || _schema->has_ann_index()) { std::string inverted_index_file = InvertedIndexDescriptor::get_index_file_path_v2( InvertedIndexDescriptor::get_index_file_path_prefix(seg_path)); st = fs->delete_file(inverted_index_file); @@ -367,7 +367,7 @@ Status BetaRowset::link_files_to(const std::string& dir, RowsetId new_rowset_id, } } } else { - if (_schema->has_inverted_index() && + if ((_schema->has_inverted_index() || _schema->has_ann_index()) && (without_index_uids == nullptr || without_index_uids->empty())) { std::string inverted_index_file_src = InvertedIndexDescriptor::get_index_file_path_v2( @@ -434,7 +434,7 @@ Status BetaRowset::copy_files_to(const std::string& dir, const RowsetId& new_row } } } else { - if (_schema->has_inverted_index()) { + if (_schema->has_inverted_index() || _schema->has_ann_index()) { std::string inverted_index_src_file = InvertedIndexDescriptor::get_index_file_path_v2( InvertedIndexDescriptor::get_index_file_path_prefix(src_path)); @@ -492,7 +492,7 @@ Status BetaRowset::upload_to(const StorageResource& dest_fs, const RowsetId& new } } } else { - if (_schema->has_inverted_index()) { + if (_schema->has_inverted_index() || _schema->has_ann_index()) { std::string remote_inverted_index_file = InvertedIndexDescriptor::get_index_file_path_v2( InvertedIndexDescriptor::get_index_file_path_prefix( @@ -646,7 +646,7 @@ Status BetaRowset::add_to_binlog() { linked_success_files.push_back(binlog_index_file); } } else { - if (_schema->has_inverted_index()) { + if (_schema->has_inverted_index() || _schema->has_ann_index()) { auto index_file = InvertedIndexDescriptor::get_index_file_path_v2( InvertedIndexDescriptor::get_index_file_path_prefix(seg_file)); auto binlog_index_file = (std::filesystem::path(binlog_dir) / @@ -693,7 +693,7 @@ Status BetaRowset::calc_file_crc(uint32_t* crc_value, int64_t* file_count) { } } } else { - if (_schema->has_inverted_index()) { + if (_schema->has_inverted_index() || _schema->has_ann_index()) { std::string inverted_index_file = InvertedIndexDescriptor::get_index_file_path_v2( InvertedIndexDescriptor::get_index_file_path_prefix(seg_path)); file_paths.emplace_back(std::move(inverted_index_file)); diff --git a/be/src/olap/rowset/beta_rowset_reader.cpp b/be/src/olap/rowset/beta_rowset_reader.cpp index d19665d618cfd3..0d2dccf50633b3 100644 --- a/be/src/olap/rowset/beta_rowset_reader.cpp +++ b/be/src/olap/rowset/beta_rowset_reader.cpp @@ -44,6 +44,7 @@ #include "olap/schema_cache.h" #include "olap/tablet_meta.h" #include "olap/tablet_schema.h" +#include "runtime/descriptors.h" #include "util/runtime_profile.h" #include "vec/core/block.h" #include "vec/olap/vgeneric_iterators.h" @@ -102,6 +103,7 @@ Status BetaRowsetReader::get_segment_iterators(RowsetReaderContext* read_context _read_options.remaining_conjunct_roots = _read_context->remaining_conjunct_roots; _read_options.common_expr_ctxs_push_down = _read_context->common_expr_ctxs_push_down; _read_options.virtual_column_exprs = _read_context->virtual_column_exprs; + _read_options.ann_topn_runtime = _read_context->ann_topn_runtime; _read_options.vir_cid_to_idx_in_block = _read_context->vir_cid_to_idx_in_block; _read_options.vir_col_idx_to_type = _read_context->vir_col_idx_to_type; _read_options.score_runtime = _read_context->score_runtime; diff --git a/be/src/olap/rowset/beta_rowset_writer.cpp b/be/src/olap/rowset/beta_rowset_writer.cpp index 4794f143460172..b189d27fc68619 100644 --- a/be/src/olap/rowset/beta_rowset_writer.cpp +++ b/be/src/olap/rowset/beta_rowset_writer.cpp @@ -52,6 +52,7 @@ #include "olap/tablet_schema.h" #include "runtime/thread_context.h" #include "util/debug_points.h" +#include "util/pretty_printer.h" #include "util/slice.h" #include "util/time.h" #include "vec/columns/column.h" @@ -550,7 +551,8 @@ Status BetaRowsetWriter::_rename_compacted_indices(int64_t begin, int64_t end, u if (_context.tablet_schema->get_inverted_index_storage_format() >= InvertedIndexStorageFormatPB::V2) { - if (_context.tablet_schema->has_inverted_index()) { + if (_context.tablet_schema->has_inverted_index() || + _context.tablet_schema->has_ann_index()) { auto src_idx_path = InvertedIndexDescriptor::get_index_file_path_v2(src_index_path_prefix); auto dst_idx_path = @@ -842,7 +844,8 @@ Status BetaRowsetWriter::build(RowsetSharedPtr& rowset) { _rowset_meta->set_tablet_schema(_context.tablet_schema); // If segment compaction occurs, the idx file info will become inaccurate. - if (_context.tablet_schema->has_inverted_index() && _num_segcompacted == 0) { + if ((_context.tablet_schema->has_inverted_index() || _context.tablet_schema->has_ann_index()) && + _num_segcompacted == 0) { if (auto idx_files_info = _idx_files.inverted_index_file_info(_segment_start_id); !idx_files_info.has_value()) [[unlikely]] { LOG(ERROR) << "expected inverted index files info, but none presents: " @@ -991,7 +994,7 @@ Status BetaRowsetWriter::create_segment_writer_for_segcompaction( RETURN_IF_ERROR(_create_file_writer(path, file_writer)); IndexFileWriterPtr index_file_writer; - if (_context.tablet_schema->has_inverted_index()) { + if (_context.tablet_schema->has_inverted_index() || _context.tablet_schema->has_ann_index()) { io::FileWriterPtr idx_file_writer; std::string prefix(InvertedIndexDescriptor::get_index_file_path_prefix(path)); if (_context.tablet_schema->get_inverted_index_storage_format() != @@ -1123,7 +1126,8 @@ Status BetaRowsetWriter::flush_segment_writer_for_segcompaction( _segid_statistics_map.emplace(segid, segstat); } VLOG_DEBUG << "_segid_statistics_map add new record. segid:" << segid << " row_num:" << row_num - << " data_size:" << segment_size << " index_size:" << index_size; + << " data_size:" << PrettyPrinter::print_bytes(segment_size) + << " index_size:" << PrettyPrinter::print_bytes(inverted_index_file_size); writer->reset(); diff --git a/be/src/olap/rowset/beta_rowset_writer.h b/be/src/olap/rowset/beta_rowset_writer.h index 8429e30f4c359b..e21bdf4009d0d2 100644 --- a/be/src/olap/rowset/beta_rowset_writer.h +++ b/be/src/olap/rowset/beta_rowset_writer.h @@ -189,7 +189,7 @@ class BaseBetaRowsetWriter : public RowsetWriter { return _seg_files.get_file_writers(); } - std::unordered_map& inverted_index_file_writers() { + std::unordered_map& index_file_writers() { return this->_idx_files.get_file_writers(); } diff --git a/be/src/olap/rowset/rowset_reader_context.h b/be/src/olap/rowset/rowset_reader_context.h index e56fba65e5968f..c6e8dc718c76c4 100644 --- a/be/src/olap/rowset/rowset_reader_context.h +++ b/be/src/olap/rowset/rowset_reader_context.h @@ -22,6 +22,7 @@ #include "olap/column_predicate.h" #include "olap/olap_common.h" #include "olap/rowid_conversion.h" +#include "olap/rowset/segment_v2/ann_index/ann_topn_runtime.h" #include "runtime/runtime_state.h" #include "vec/exprs/score_runtime.h" #include "vec/exprs/vexpr.h" @@ -91,6 +92,7 @@ struct RowsetReaderContext { std::shared_ptr score_runtime; CollectionStatisticsPtr collection_statistics; + std::shared_ptr ann_topn_runtime; }; } // namespace doris diff --git a/be/src/olap/rowset/segcompaction.cpp b/be/src/olap/rowset/segcompaction.cpp index 2b37f166568197..456d64d82b5c78 100644 --- a/be/src/olap/rowset/segcompaction.cpp +++ b/be/src/olap/rowset/segcompaction.cpp @@ -168,7 +168,7 @@ Status SegcompactionWorker::_delete_original_segments(uint32_t begin, uint32_t e // message when we encounter an error. RETURN_NOT_OK_STATUS_WITH_WARN(fs->delete_file(seg_path), absl::Substitute("Failed to delete file=$0", seg_path)); - if (schema->has_inverted_index() && + if ((schema->has_inverted_index() || schema->has_ann_index()) && schema->get_inverted_index_storage_format() >= InvertedIndexStorageFormatPB::V2) { auto idx_path = InvertedIndexDescriptor::get_index_file_path_v2( InvertedIndexDescriptor::get_index_file_path_prefix(seg_path)); diff --git a/be/src/olap/rowset/segment_creator.cpp b/be/src/olap/rowset/segment_creator.cpp index 08d331d9559c17..989ccf787f05ef 100644 --- a/be/src/olap/rowset/segment_creator.cpp +++ b/be/src/olap/rowset/segment_creator.cpp @@ -136,7 +136,7 @@ Status SegmentFlusher::_create_segment_writer(std::unique_ptrcreate(segment_id, segment_file_writer)); IndexFileWriterPtr index_file_writer; - if (_context.tablet_schema->has_inverted_index()) { + if (_context.tablet_schema->has_inverted_index() || _context.tablet_schema->has_ann_index()) { RETURN_IF_ERROR(_context.file_writer_creator->create(segment_id, &index_file_writer)); } @@ -154,7 +154,7 @@ Status SegmentFlusher::_create_segment_writer(std::unique_ptrhas_inverted_index()) { + if (_context.tablet_schema->has_inverted_index() || _context.tablet_schema->has_ann_index()) { RETURN_IF_ERROR(_idx_files.add(segment_id, std::move(index_file_writer))); } auto s = writer->init(); @@ -173,7 +173,7 @@ Status SegmentFlusher::_create_segment_writer( RETURN_IF_ERROR(_context.file_writer_creator->create(segment_id, segment_file_writer)); IndexFileWriterPtr index_file_writer; - if (_context.tablet_schema->has_inverted_index()) { + if (_context.tablet_schema->has_inverted_index() || _context.tablet_schema->has_ann_index()) { RETURN_IF_ERROR(_context.file_writer_creator->create(segment_id, &index_file_writer)); } @@ -190,7 +190,7 @@ Status SegmentFlusher::_create_segment_writer( segment_file_writer.get(), segment_id, _context.tablet_schema, _context.tablet, _context.data_dir, writer_options, index_file_writer.get()); RETURN_IF_ERROR(_seg_files.add(segment_id, std::move(segment_file_writer))); - if (_context.tablet_schema->has_inverted_index()) { + if (_context.tablet_schema->has_inverted_index() || _context.tablet_schema->has_ann_index()) { RETURN_IF_ERROR(_idx_files.add(segment_id, std::move(index_file_writer))); } auto s = writer->init(); @@ -248,7 +248,7 @@ Status SegmentFlusher::_flush_segment_writer( << ", flushing rowset_dir: " << _context.tablet_path << ", rowset_id:" << _context.rowset_id << ", data size:" << PrettyPrinter::print_bytes(segstat.data_size) - << ", index size:" << segstat.index_size; + << ", index size:" << PrettyPrinter::print_bytes(segstat.index_size); writer.reset(); diff --git a/be/src/olap/rowset/segment_v2/ann_index/CMakeLists.txt b/be/src/olap/rowset/segment_v2/ann_index/CMakeLists.txt new file mode 100644 index 00000000000000..7f9ad28ed8e42e --- /dev/null +++ b/be/src/olap/rowset/segment_v2/ann_index/CMakeLists.txt @@ -0,0 +1,29 @@ +# 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. + +add_subdirectory(cmake-protect) + +# Use all .cpp files in this directory as sources +file(GLOB VECTOR_LIB_SRC CONFIGURE_DEPENDS *.cpp) + +find_package(OpenMP REQUIRED) + +add_library(ann_index STATIC ${VECTOR_LIB_SRC}) +target_link_libraries(ann_index PUBLIC faiss OpenMP::OpenMP_CXX) + +# Some header files from faiss are used by doris, they will break compile check. +target_compile_options(ann_index PRIVATE -Wno-shadow-field) diff --git a/be/src/olap/rowset/segment_v2/ann_index/ann_index.cpp b/be/src/olap/rowset/segment_v2/ann_index/ann_index.cpp new file mode 100644 index 00000000000000..62bf13369e5e86 --- /dev/null +++ b/be/src/olap/rowset/segment_v2/ann_index/ann_index.cpp @@ -0,0 +1,64 @@ +// 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 "olap/rowset/segment_v2/ann_index/ann_index.h" + +#include "vec/functions/array/function_array_distance.h" + +namespace doris::segment_v2 { + +std::string metric_to_string(AnnIndexMetric metric) { + switch (metric) { + case AnnIndexMetric::L2: + return vectorized::L2Distance::name; + case AnnIndexMetric::IP: + return vectorized::InnerProduct::name; + default: + return "UNKNOWN"; + } +} + +AnnIndexMetric string_to_metric(const std::string& metric) { + if (metric == vectorized::L2Distance::name) { + return AnnIndexMetric::L2; + } else if (metric == vectorized::InnerProduct::name) { + return AnnIndexMetric::IP; + } else { + return AnnIndexMetric::UNKNOWN; + } +} + +std::string ann_index_type_to_string(AnnIndexType type) { + switch (type) { + case AnnIndexType::UNKNOWN: + return "unknown"; + case AnnIndexType::HNSW: + return "hnsw"; + default: + return "unknown"; + } +} + +AnnIndexType string_to_ann_index_type(const std::string& type) { + if (type == "hnsw") { + return AnnIndexType::HNSW; + } else { + return AnnIndexType::UNKNOWN; + } +} + +} // namespace doris::segment_v2 \ No newline at end of file diff --git a/be/src/olap/rowset/segment_v2/ann_index/ann_index.h b/be/src/olap/rowset/segment_v2/ann_index/ann_index.h new file mode 100644 index 00000000000000..2448aba0f489c4 --- /dev/null +++ b/be/src/olap/rowset/segment_v2/ann_index/ann_index.h @@ -0,0 +1,124 @@ +// 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. + +/** + * @file ann_index.h + * @brief Abstract interface for vector similarity search indexes in Doris. + * + * This file defines the abstract VectorIndex interface that provides a unified + * API for different vector index implementations (FAISS, etc.). The interface + * supports both approximate k-nearest neighbor search and range search operations. + * + * Key operations supported: + * - Adding vectors to the index during build phase + * - K-nearest neighbor search for Top-N queries + * - Range search for finding vectors within a distance threshold + * - Persistence to/from storage for index durability + * + * This abstraction allows Doris to support multiple vector index libraries + * through a consistent interface. + */ + +#pragma once + +#include + +#include "common/status.h" + +namespace lucene::store { +class Directory; +} + +#include "common/compile_check_begin.h" +namespace doris::segment_v2 { +struct IndexSearchParameters; +struct IndexSearchResult; + +enum class AnnIndexMetric { L2, IP, UNKNOWN }; + +std::string metric_to_string(AnnIndexMetric metric); + +AnnIndexMetric string_to_metric(const std::string& metric); + +enum class AnnIndexType { UNKNOWN, HNSW }; + +std::string ann_index_type_to_string(AnnIndexType type); + +AnnIndexType string_to_ann_index_type(const std::string& type); + +/** + * @brief Abstract base class for vector similarity search indexes. + * + * This class defines the interface that all vector index implementations + * must follow. It provides the core operations needed for vector similarity + * search in Doris, including index building, searching, and persistence. + * + * Implementations of this interface (like FaissVectorIndex) handle the + * specifics of different vector index libraries while providing a consistent + * API for the Doris query execution engine. + */ +class VectorIndex { +public: + virtual ~VectorIndex() = default; + + /** Add n vectors of dimension d vectors 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 + */ + virtual doris::Status add(int n, const float* x) = 0; + + /** Return approximate nearest neighbors of a query vector. + * The result is stored in the result object. + * @param query_vec input vector, size d + * @param k number of nearest neighbors to return + * @param params search parameters + * @param result output search result + * @return status of the operation + */ + virtual doris::Status ann_topn_search(const float* query_vec, int k, + const segment_v2::IndexSearchParameters& params, + segment_v2::IndexSearchResult& result) = 0; + /** + * Search for the nearest neighbors of a query vector within a given radius. + * @param query_vec input vector, size d + * @param radius search radius + * @param result output search result + * @return status of the operation + */ + virtual doris::Status range_search(const float* query_vec, const float& radius, + const segment_v2::IndexSearchParameters& params, + segment_v2::IndexSearchResult& result) = 0; + + virtual doris::Status save(lucene::store::Directory*) = 0; + + virtual doris::Status load(lucene::store::Directory*) = 0; + + size_t get_dimension() const { return _dimension; } + + void set_metric(AnnIndexMetric metric) { _metric = metric; } + +protected: + // When adding vectors to the index, use this variable to check the dimension of the vectors. + size_t _dimension = 0; + AnnIndexMetric _metric = AnnIndexMetric::L2; // Default metric is L2 distance +}; +#include "common/compile_check_end.h" +} // namespace doris::segment_v2 \ No newline at end of file diff --git a/be/src/olap/rowset/segment_v2/ann_index/ann_index_files.h b/be/src/olap/rowset/segment_v2/ann_index/ann_index_files.h new file mode 100644 index 00000000000000..c4bdd96dbbc6d3 --- /dev/null +++ b/be/src/olap/rowset/segment_v2/ann_index/ann_index_files.h @@ -0,0 +1,27 @@ +// 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. + +#pragma once + +// Centralized constants for ANN index filenames. +// Use a versioned, library-agnostic name to allow future upgrades (e.g., disk-ANN). + +namespace doris::segment_v2 { + +inline constexpr char faiss_index_fila_name[] = "ann.faiss"; + +} // namespace doris::segment_v2 diff --git a/be/src/olap/rowset/segment_v2/ann_index/ann_index_iterator.cpp b/be/src/olap/rowset/segment_v2/ann_index/ann_index_iterator.cpp new file mode 100644 index 00000000000000..b43547d001c75f --- /dev/null +++ b/be/src/olap/rowset/segment_v2/ann_index/ann_index_iterator.cpp @@ -0,0 +1,57 @@ +// 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 "ann_index_iterator.h" + +#include + +#include "olap/rowset/segment_v2/ann_index/ann_search_params.h" + +namespace doris::segment_v2 { +#include "common/compile_check_begin.h" +AnnIndexIterator::AnnIndexIterator(const IndexReaderPtr& reader) : IndexIterator() { + _ann_reader = std::dynamic_pointer_cast(reader); +} + +Status AnnIndexIterator::read_from_index(const IndexParam& param) { + auto* a_param = std::get(param); + if (a_param == nullptr) { + return Status::Error("a_param is null"); + } + if (_ann_reader == nullptr) { + return Status::Error("_ann_reader is null"); + } + + // _context may be unset in some test scenarios; pass nullptr IOContext in that case. + io::IOContext* io_ctx = (_context != nullptr) ? _context->io_ctx : nullptr; + return _ann_reader->query(io_ctx, a_param, a_param->stats.get()); +} + +Status AnnIndexIterator::range_search(const AnnRangeSearchParams& params, + const VectorSearchUserParams& custom_params, + segment_v2::AnnRangeSearchResult* result, + segment_v2::AnnIndexStats* stats) { + if (_ann_reader == nullptr) { + return Status::Error("_ann_reader is null"); + } + + // _context may be null when iterator is used in isolation (e.g., unit tests). + io::IOContext* io_ctx = (_context != nullptr) ? _context->io_ctx : nullptr; + return _ann_reader->range_search(params, custom_params, result, stats, io_ctx); +} +#include "common/compile_check_end.h" +} // namespace doris::segment_v2 \ No newline at end of file diff --git a/be/src/olap/rowset/segment_v2/ann_index/ann_index_iterator.h b/be/src/olap/rowset/segment_v2/ann_index/ann_index_iterator.h new file mode 100644 index 00000000000000..98eaecfaf067e5 --- /dev/null +++ b/be/src/olap/rowset/segment_v2/ann_index/ann_index_iterator.h @@ -0,0 +1,57 @@ +// 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. + +#pragma once + +#include +#include + +#include "olap/rowset/segment_v2/ann_index/ann_index_reader.h" +#include "olap/rowset/segment_v2/index_iterator.h" +#include "runtime/runtime_state.h" + +namespace doris::segment_v2 { +struct AnnRangeSearchParams; +struct AnnRangeSearchResult; +#include "common/compile_check_begin.h" +class AnnIndexIterator : public IndexIterator { +public: + AnnIndexIterator(const IndexReaderPtr& reader); + ~AnnIndexIterator() override = default; + + IndexReaderPtr get_reader(IndexReaderType reader_type) const override { + return std::static_pointer_cast(_ann_reader); + } + MOCK_FUNCTION Status read_from_index(const IndexParam& param) override; + + Status read_null_bitmap(InvertedIndexQueryCacheHandle* cache_handle) override { + return Status::OK(); + } + + Result has_null() override { return true; } + + MOCK_FUNCTION Status range_search(const AnnRangeSearchParams& params, + const VectorSearchUserParams& custom_params, + AnnRangeSearchResult* result, AnnIndexStats* stats); + +private: + std::shared_ptr _ann_reader; + + ENABLE_FACTORY_CREATOR(AnnIndexIterator); +}; +#include "common/compile_check_end.h" +} // namespace doris::segment_v2 \ No newline at end of file diff --git a/be/src/olap/rowset/segment_v2/ann_index/ann_index_reader.cpp b/be/src/olap/rowset/segment_v2/ann_index/ann_index_reader.cpp new file mode 100644 index 00000000000000..0f3f56b1aa4c98 --- /dev/null +++ b/be/src/olap/rowset/segment_v2/ann_index/ann_index_reader.cpp @@ -0,0 +1,220 @@ +// 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 "ann_index_reader.h" + +#include +#include + +#include "ann_index_iterator.h" +#include "common/config.h" +#include "io/io_common.h" +#include "olap/rowset/segment_v2/ann_index/ann_index.h" +#include "olap/rowset/segment_v2/ann_index/ann_search_params.h" +#include "olap/rowset/segment_v2/ann_index/faiss_ann_index.h" +#include "olap/rowset/segment_v2/index_file_reader.h" +#include "olap/rowset/segment_v2/inverted_index_compound_reader.h" +#include "runtime/runtime_state.h" +#include "util/doris_metrics.h" +#include "util/once.h" +#include "util/runtime_profile.h" + +namespace doris::segment_v2 { +#include "common/compile_check_begin.h" +void AnnIndexReader::update_result(const IndexSearchResult& search_result, + std::vector& distance, roaring::Roaring& roaring) { + DCHECK(search_result.distances != nullptr); + DCHECK(search_result.roaring != nullptr); + size_t limit = search_result.roaring->cardinality(); + // Use search result to update distance and row_id + distance.resize(limit); + for (size_t i = 0; i < limit; ++i) { + distance[i] = search_result.distances[i]; + } + roaring = *search_result.roaring; +} + +AnnIndexReader::AnnIndexReader(const TabletIndex* index_meta, + std::shared_ptr index_file_reader) + : _index_meta(*index_meta), _index_file_reader(index_file_reader) { + const auto index_properties = _index_meta.properties(); + auto it = index_properties.find("index_type"); + DCHECK(it != index_properties.end()); + _index_type = it->second; + it = index_properties.find("metric_type"); + DCHECK(it != index_properties.end()); + _metric_type = string_to_metric(it->second); +} + +Status AnnIndexReader::new_iterator(std::unique_ptr* iterator) { + *iterator = AnnIndexIterator::create_unique(shared_from_this()); + return Status::OK(); +} + +Status AnnIndexReader::load_index(io::IOContext* io_ctx) { + return _load_index_once.call([&]() { + DorisMetrics::instance()->ann_index_load_cnt->increment(1); + + try { + RETURN_IF_ERROR( + _index_file_reader->init(config::inverted_index_read_buffer_size, io_ctx)); + Result> compound_dir; + compound_dir = _index_file_reader->open(&_index_meta, io_ctx); + if (!compound_dir.has_value()) { + return Status::IOError("Failed to open index file: {}", + compound_dir.error().to_string()); + } + _vector_index = std::make_unique(); + _vector_index->set_metric(_metric_type); + RETURN_IF_ERROR(_vector_index->load(compound_dir->get())); + } catch (CLuceneError& err) { + return Status::Error( + "CLuceneError occur when open ann idx file, error msg: {}", err.what()); + } + return Status::OK(); + }); +} + +Status AnnIndexReader::query(io::IOContext* io_ctx, AnnTopNParam* param, AnnIndexStats* stats) { +#ifndef BE_TEST + { + SCOPED_TIMER(&(stats->load_index_costs_ns)); + RETURN_IF_ERROR(load_index(io_ctx)); + double load_costs_ms = static_cast(stats->load_index_costs_ns.value()) / 1000.0; + DorisMetrics::instance()->ann_index_load_costs_ms->increment( + static_cast(load_costs_ms)); + } +#endif + { + DorisMetrics::instance()->ann_index_search_cnt->increment(1); + SCOPED_TIMER(&(stats->search_costs_ns)); + DCHECK(_vector_index != nullptr); + const float* query_vec = param->query_value; + const int limit = static_cast(param->limit); + IndexSearchResult index_search_result; + if (string_to_ann_index_type(_index_type) == AnnIndexType::HNSW) { + HNSWSearchParameters hnsw_search_params; + hnsw_search_params.roaring = param->roaring; + hnsw_search_params.rows_of_segment = param->rows_of_segment; + hnsw_search_params.ef_search = param->_user_params.hnsw_ef_search; + hnsw_search_params.check_relative_distance = + param->_user_params.hnsw_check_relative_distance; + hnsw_search_params.bounded_queue = param->_user_params.hnsw_bounded_queue; + RETURN_IF_ERROR(_vector_index->ann_topn_search(query_vec, limit, hnsw_search_params, + index_search_result)); + // Accumulate detailed engine timings + stats->engine_search_ns.update(index_search_result.engine_search_ns); + stats->engine_convert_ns.update(index_search_result.engine_convert_ns); + stats->engine_prepare_ns.update(index_search_result.engine_prepare_ns); + } else { + throw Status::NotSupported("Unsupported index type: {}", _index_type); + } + + DCHECK(index_search_result.roaring != nullptr); + DCHECK(index_search_result.distances != nullptr); + DCHECK(index_search_result.row_ids != nullptr); + param->distance = std::make_unique>(); + { + SCOPED_TIMER(&(stats->result_process_costs_ns)); + update_result(index_search_result, *param->distance, *param->roaring); + } + param->row_ids = std::move(index_search_result.row_ids); + } + + double search_costs_ms = static_cast(stats->search_costs_ns.value()) / 1000.0; + DorisMetrics::instance()->ann_index_search_costs_ms->increment( + static_cast(search_costs_ms)); + return Status::OK(); +} + +Status AnnIndexReader::range_search(const AnnRangeSearchParams& params, + const VectorSearchUserParams& custom_params, + segment_v2::AnnRangeSearchResult* result, + segment_v2::AnnIndexStats* stats, io::IOContext* io_ctx) { + DCHECK(stats != nullptr); +#ifndef BE_TEST + { + SCOPED_TIMER(&(stats->load_index_costs_ns)); + RETURN_IF_ERROR(load_index(io_ctx)); + double load_costs_ms = static_cast(stats->load_index_costs_ns.value()) / 1000.0; + DorisMetrics::instance()->ann_index_load_costs_ms->increment( + static_cast(load_costs_ms)); + } +#endif + { + DorisMetrics::instance()->ann_index_search_cnt->increment(1); + SCOPED_TIMER(&(stats->search_costs_ns)); + DCHECK(_vector_index != nullptr); + segment_v2::IndexSearchResult search_result; + std::unique_ptr search_param = nullptr; + + if (string_to_ann_index_type(_index_type) == AnnIndexType::HNSW) { + auto hnsw_param = std::make_unique(); + hnsw_param->ef_search = custom_params.hnsw_ef_search; + hnsw_param->check_relative_distance = custom_params.hnsw_check_relative_distance; + hnsw_param->bounded_queue = custom_params.hnsw_bounded_queue; + search_param = std::move(hnsw_param); + } else { + throw Status::NotSupported("Unsupported index type: {}", _index_type); + } + + search_param->is_le_or_lt = params.is_le_or_lt; + search_param->roaring = params.roaring; + DCHECK(search_param->roaring != nullptr); + + RETURN_IF_ERROR(_vector_index->range_search(params.query_value, params.radius, + *search_param, search_result)); + // Accumulate detailed engine timings + stats->engine_prepare_ns.update(search_result.engine_prepare_ns); + stats->engine_search_ns.update(search_result.engine_search_ns); + stats->engine_convert_ns.update(search_result.engine_convert_ns); + + DCHECK(search_result.roaring != nullptr); + result->roaring = search_result.roaring; + + if (params.is_le_or_lt == false) { + DCHECK(search_result.distances == nullptr); + DCHECK(search_result.row_ids == nullptr); + } + + { + SCOPED_TIMER(&(stats->result_process_costs_ns)); + if (search_result.row_ids != nullptr) { + DCHECK(search_result.row_ids->size() == search_result.roaring->cardinality()) + << "Row ids size: " << search_result.row_ids->size() + << ", roaring size: " << search_result.roaring->cardinality(); + result->row_ids = std::move(search_result.row_ids); + } else { + result->row_ids = nullptr; + } + + if (search_result.distances != nullptr) { + result->distance = std::move(search_result.distances); + } else { + result->distance = nullptr; + } + } + } + + double search_costs_ms = static_cast(stats->search_costs_ns.value()) / 1000.0; + DorisMetrics::instance()->ann_index_search_costs_ms->increment( + static_cast(search_costs_ms)); + + return Status::OK(); +} + +} // namespace doris::segment_v2 \ No newline at end of file diff --git a/be/src/olap/rowset/segment_v2/ann_index/ann_index_reader.h b/be/src/olap/rowset/segment_v2/ann_index/ann_index_reader.h new file mode 100644 index 00000000000000..fdf655a20b8e0a --- /dev/null +++ b/be/src/olap/rowset/segment_v2/ann_index/ann_index_reader.h @@ -0,0 +1,75 @@ +// 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. + +#pragma once + +#include "olap/rowset/segment_v2/ann_index/ann_index.h" +#include "olap/rowset/segment_v2/ann_index/ann_search_params.h" +#include "olap/rowset/segment_v2/index_reader.h" +#include "olap/tablet_schema.h" +#include "runtime/runtime_state.h" +#include "util/once.h" +namespace doris::segment_v2 { +#include "common/compile_check_begin.h" + +struct AnnTopNParam; +struct AnnRangeSearchParams; +struct AnnRangeSearchResult; +struct IndexSearchResult; + +class IndexFileReader; +class IndexIterator; + +class AnnIndexReader : public IndexReader { +public: + AnnIndexReader(const TabletIndex* index_meta, + std::shared_ptr index_file_reader); + ~AnnIndexReader() override = default; + + static void update_result(const IndexSearchResult&, std::vector& distance, + roaring::Roaring& row_id); + + Status load_index(io::IOContext* io_ctx); + + Status query(io::IOContext* io_ctx, AnnTopNParam* param, AnnIndexStats* stats); + + Status range_search(const AnnRangeSearchParams& params, + const VectorSearchUserParams& custom_params, AnnRangeSearchResult* result, + AnnIndexStats* stats, io::IOContext* io_ctx = nullptr); + + IndexType index_type() override { return IndexType::ANN; } + + uint64_t get_index_id() const override { return _index_meta.index_id(); } + + Status new_iterator(std::unique_ptr* iterator) override; + + AnnIndexMetric get_metric_type() const { return _metric_type; } + +private: + TabletIndex _index_meta; + std::shared_ptr _index_file_reader; + std::unique_ptr _vector_index; + // TODO: Use integer. + std::string _index_type; + AnnIndexMetric _metric_type; + + DorisCallOnce _load_index_once; +}; + +using AnnIndexReaderPtr = std::shared_ptr; +#include "common/compile_check_end.h" +} // namespace doris::segment_v2 diff --git a/be/src/olap/rowset/segment_v2/ann_index/ann_index_writer.cpp b/be/src/olap/rowset/segment_v2/ann_index/ann_index_writer.cpp new file mode 100644 index 00000000000000..287e75dce91047 --- /dev/null +++ b/be/src/olap/rowset/segment_v2/ann_index/ann_index_writer.cpp @@ -0,0 +1,124 @@ +// 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 "olap/rowset/segment_v2/ann_index/ann_index_writer.h" + +#include +#include +#include + +#include "common/cast_set.h" +#include "olap/rowset/segment_v2/ann_index/faiss_ann_index.h" +#include "olap/rowset/segment_v2/inverted_index_fs_directory.h" + +namespace doris::segment_v2 { +#include "common/compile_check_begin.h" +static std::string get_or_default(const std::map& properties, + const std::string& key, const std::string& default_value) { + auto it = properties.find(key); + if (it != properties.end()) { + return it->second; + } + return default_value; +} + +AnnIndexColumnWriter::AnnIndexColumnWriter(IndexFileWriter* index_file_writer, + const TabletIndex* index_meta) + : _index_file_writer(index_file_writer), _index_meta(index_meta) {} + +AnnIndexColumnWriter::~AnnIndexColumnWriter() {} + +Status AnnIndexColumnWriter::init() { + Result> compound_dir = _index_file_writer->open(_index_meta); + + if (!compound_dir.has_value()) { + return Status::IOError("Failed to open index file: {}", compound_dir.error().to_string()); + } + + _dir = compound_dir.value(); + + _vector_index = nullptr; + const auto& properties = _index_meta->properties(); + const std::string index_type = get_or_default(properties, INDEX_TYPE, "hnsw"); + const std::string metric_type = get_or_default(properties, METRIC_TYPE, "l2_distance"); + FaissBuildParameter build_parameter; + std::shared_ptr faiss_index = std::make_shared(); + build_parameter.index_type = FaissBuildParameter::string_to_index_type(index_type); + build_parameter.dim = std::stoi(get_or_default(properties, DIM, "512")); + build_parameter.max_degree = std::stoi(get_or_default(properties, MAX_DEGREE, "32")); + build_parameter.metric_type = FaissBuildParameter::string_to_metric_type(metric_type); + + faiss_index->build(build_parameter); + + _vector_index = faiss_index; + LOG_INFO("Create a new faiss index, index_type {} dim {} metric_type {} max_degree {}", + index_type, build_parameter.dim, metric_type, build_parameter.max_degree); + return Status::OK(); +} + +Status AnnIndexColumnWriter::add_values(const std::string fn, const void* values, size_t count) { + return Status::OK(); +} + +void AnnIndexColumnWriter::close_on_error() {} + +Status AnnIndexColumnWriter::add_array_values(size_t field_size, const void* value_ptr, + const uint8_t* null_map, const uint8_t* offsets_ptr, + size_t num_rows) { + // TODO: Performance optimization + if (num_rows == 0) { + return Status::OK(); + } + + const auto* offsets = reinterpret_cast(offsets_ptr); + const size_t dim = _vector_index->get_dimension(); + for (size_t i = 0; i < num_rows; ++i) { + auto array_elem_size = offsets[i + 1] - offsets[i]; + if (array_elem_size != dim) { + return Status::InvalidArgument("Ann index expect array with {} dim, got {}.", dim, + array_elem_size); + } + } + + const float* p = reinterpret_cast(value_ptr); + RETURN_IF_ERROR(_vector_index->add(cast_set(num_rows), p)); + + return Status::OK(); +} + +Status AnnIndexColumnWriter::add_array_values(size_t field_size, const CollectionValue* values, + size_t count) { + return Status::InternalError("Ann index should not be used on nullable column"); +} + +Status AnnIndexColumnWriter::add_nulls(uint32_t count) { + return Status::InternalError("Ann index should not be used on nullable column"); +} + +Status AnnIndexColumnWriter::add_array_nulls(const uint8_t* null_map, size_t row_id) { + return Status::InternalError("Ann index should not be used on nullable column"); +} + +int64_t AnnIndexColumnWriter::size() const { + return 0; +} + +Status AnnIndexColumnWriter::finish() { + return _vector_index->save(_dir.get()); +} +#include "common/compile_check_end.h" +} // namespace doris::segment_v2 diff --git a/be/src/olap/rowset/segment_v2/ann_index/ann_index_writer.h b/be/src/olap/rowset/segment_v2/ann_index/ann_index_writer.h new file mode 100644 index 00000000000000..c8e87b0f122ed9 --- /dev/null +++ b/be/src/olap/rowset/segment_v2/ann_index/ann_index_writer.h @@ -0,0 +1,74 @@ +// 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. + +#pragma once + +#include // IWYU pragma: keep +#include +#include +#include + +#include +#include +#include +#include + +#include "common/config.h" +#include "olap/rowset/segment_v2/ann_index/ann_index.h" +#include "olap/rowset/segment_v2/index_file_writer.h" +#include "olap/rowset/segment_v2/index_writer.h" +#include "olap/rowset/segment_v2/inverted_index_fs_directory.h" +#include "olap/tablet_schema.h" +#include "runtime/collection_value.h" + +namespace doris::segment_v2 { +#include "common/compile_check_begin.h" +class AnnIndexColumnWriter : public IndexColumnWriter { +public: + static constexpr const char* INDEX_TYPE = "index_type"; + static constexpr const char* METRIC_TYPE = "metric_type"; + static constexpr const char* DIM = "dim"; + static constexpr const char* MAX_DEGREE = "max_degree"; + + explicit AnnIndexColumnWriter(IndexFileWriter* index_file_writer, + const TabletIndex* index_meta); + + ~AnnIndexColumnWriter() override; + + Status init() override; + void close_on_error() override; + Status add_nulls(uint32_t count) override; + Status add_array_nulls(const uint8_t* null_map, size_t num_rows) override; + Status add_values(const std::string fn, const void* values, size_t count) override; + Status add_array_values(size_t field_size, const void* value_ptr, const uint8_t* null_map, + const uint8_t* offsets_ptr, size_t count) override; + Status add_array_values(size_t field_size, const CollectionValue* values, + size_t count) override; + int64_t size() const override; + Status finish() override; + +private: + // VectorIndex shoule be managed by some cache. + // VectorIndex should be weak shared by AnnIndexWriter and VectorIndexReader + // This should be a weak_ptr + std::shared_ptr _vector_index; + IndexFileWriter* _index_file_writer; + const TabletIndex* _index_meta; + std::shared_ptr _dir; +}; +#include "common/compile_check_end.h" +} // namespace doris::segment_v2 diff --git a/be/src/olap/rowset/segment_v2/ann_index/ann_range_search_runtime.cpp b/be/src/olap/rowset/segment_v2/ann_index/ann_range_search_runtime.cpp new file mode 100644 index 00000000000000..416af7826d2a89 --- /dev/null +++ b/be/src/olap/rowset/segment_v2/ann_index/ann_range_search_runtime.cpp @@ -0,0 +1,63 @@ +// 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 "olap/rowset/segment_v2/ann_index/ann_range_search_runtime.h" + +#include + +#include "olap/rowset/segment_v2/ann_index/ann_search_params.h" + +namespace doris::segment_v2 { +#include "common/compile_check_begin.h" + +/** + * @brief Converts the runtime info to search parameters for execution. + * + * This method creates a AnnRangeSearchParams structure that can be passed + * to the underlying ANN index implementation for performing the actual + * range search operation. + * + * @return AnnRangeSearchParams configured with the runtime information + */ +AnnRangeSearchParams AnnRangeSearchRuntime::to_range_search_params() const { + AnnRangeSearchParams params; + params.query_value = query_value.get(); + params.radius = static_cast(radius); + params.roaring = nullptr; + params.is_le_or_lt = is_le_or_lt; + return params; +} + +/** + * @brief Generates a human-readable string representation for debugging. + * + * Creates a formatted string containing all the important runtime + * information including search parameters, column indices, and + * configuration flags. This is primarily used for logging and + * debugging purposes. + * + * @return Formatted string with runtime information + */ +std::string AnnRangeSearchRuntime::to_string() const { + return fmt::format( + "is_ann_range_search: {}, is_le_or_lt: {}, src_col_idx: {}, " + "dst_col_idx: {}, metric_type {}, radius: {}, user params: {}, query_vector is null: " + "{}", + is_ann_range_search, is_le_or_lt, src_col_idx, dst_col_idx, + metric_to_string(metric_type), radius, user_params.to_string(), query_value == nullptr); +} +} // namespace doris::segment_v2 \ No newline at end of file diff --git a/be/src/olap/rowset/segment_v2/ann_index/ann_range_search_runtime.h b/be/src/olap/rowset/segment_v2/ann_index/ann_range_search_runtime.h new file mode 100644 index 00000000000000..c3d112c9bf2412 --- /dev/null +++ b/be/src/olap/rowset/segment_v2/ann_index/ann_range_search_runtime.h @@ -0,0 +1,148 @@ +// 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. + +#pragma once + +#include + +#include + +#include "olap/rowset/segment_v2/ann_index/ann_index.h" +#include "vec/runtime/vector_search_user_params.h" + +namespace doris::segment_v2 { +struct AnnRangeSearchParams; +#include "common/compile_check_begin.h" + +/** + * @brief Runtime information structure for ANN (Approximate Nearest Neighbor) range search operations. + * + * This structure encapsulates all the necessary runtime parameters required for performing + * range search queries on ANN indexes. Range search finds all vectors within a specified + * distance radius from a query vector, which is different from traditional K-NN search + * that finds a fixed number of nearest neighbors. + * + * The structure supports: + * - L2 distance and inner product metrics + * - Configurable search radius for distance thresholds + * - Deep copy semantics for query vectors + * - Integration with Doris vectorized execution engine + */ +struct AnnRangeSearchRuntime { + /** + * @brief Default constructor initializing all fields to safe default values. + * + * Initializes the structure with: + * - Range search disabled by default + * - Less-than-or-equal comparison mode + * - Zero radius and invalid metric type + * - Null query vector pointer + */ + // DefaultConstructor + AnnRangeSearchRuntime() + : is_ann_range_search(false), + is_le_or_lt(true), + src_col_idx(0), + dst_col_idx(-1), + radius(0.0), + metric_type(AnnIndexMetric::UNKNOWN) { + query_value = nullptr; + } + + /** + * @brief Copy constructor with deep copy semantics for query vector. + * + * Performs deep copying of all fields including the query_value array. + * This is crucial for thread safety and preventing memory corruption + * when the runtime info is passed between different execution contexts. + * + * @param other The source RangeSearchRuntimeInfo to copy from + */ + // CopyConstructor + AnnRangeSearchRuntime(const AnnRangeSearchRuntime& other) + : is_ann_range_search(other.is_ann_range_search), + is_le_or_lt(other.is_le_or_lt), + src_col_idx(other.src_col_idx), + dim(other.dim), + dst_col_idx(other.dst_col_idx), + radius(other.radius), + metric_type(other.metric_type), + user_params(other.user_params) { + // Do deep copy to query_value. + if (other.query_value) { + query_value = std::make_unique(other.dim); + std::copy(other.query_value.get(), other.query_value.get() + other.dim, + query_value.get()); + } else { + query_value = nullptr; + } + } + + /** + * @brief Assignment operator with deep copy semantics. + * + * Ensures proper assignment of all fields with deep copying of the query vector. + * Maintains the same memory safety guarantees as the copy constructor. + * + * @param other The source RangeSearchRuntimeInfo to assign from + * @return Reference to this object for chaining + */ + AnnRangeSearchRuntime& operator=(const AnnRangeSearchRuntime& other) { + is_ann_range_search = other.is_ann_range_search; + is_le_or_lt = other.is_le_or_lt; + src_col_idx = other.src_col_idx; + dst_col_idx = other.dst_col_idx; + radius = other.radius; + metric_type = other.metric_type; + user_params = other.user_params; + dim = other.dim; + // Do deep copy to query_value. + if (other.query_value) { + query_value = std::make_unique(other.dim); + std::copy(other.query_value.get(), other.query_value.get() + other.dim, + query_value.get()); + } else { + query_value = nullptr; + } + return *this; + } + + /** + * @brief Converts the runtime info to AnnRangeSearchParams for actual search execution. + * @return AnnRangeSearchParams structure suitable for index operations + */ + AnnRangeSearchParams to_range_search_params() const; + + /** + * @brief Generates a string representation for debugging and logging. + * @return String containing all relevant runtime information + */ + std::string to_string() const; + + // Core search configuration + bool is_ann_range_search = false; ///< Flag indicating if ANN range search is enabled + bool is_le_or_lt = true; ///< Comparison mode: true for <=, false for < + size_t src_col_idx = 0; ///< Source column index in the schema + size_t dim = 0; ///< Dimensionality of the vector space + int64_t dst_col_idx = -1; ///< Destination column index (-1 if not applicable) + double radius = 0.0; ///< Search radius/distance threshold + AnnIndexMetric metric_type; ///< Distance metric (L2, Inner Product, etc.) + doris::VectorSearchUserParams user_params; ///< User-defined search parameters + std::unique_ptr query_value; ///< Query vector data (deep copied) +}; +#include "common/compile_check_end.h" +} // namespace doris::segment_v2 \ No newline at end of file diff --git a/be/src/olap/rowset/segment_v2/ann_index/ann_search_params.h b/be/src/olap/rowset/segment_v2/ann_index/ann_search_params.h new file mode 100644 index 00000000000000..c38c8d2138dbe3 --- /dev/null +++ b/be/src/olap/rowset/segment_v2/ann_index/ann_search_params.h @@ -0,0 +1,146 @@ +// 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. + +/** + * @file ann_search_params.h + * @brief Parameter structures and statistics for ANN (Approximate Nearest Neighbor) search operations. + * + * This file defines the core parameter structures used for configuring and executing + * ANN search operations in Doris. It includes both top-N search and range search + * parameter definitions, as well as statistics collection structures. + * + * The structures defined here serve as the interface between the query execution + * engine and the underlying vector index implementations (FAISS, etc.). + */ + +#pragma once + +#include +#include +#include + +#include +#include + +#include "util/runtime_profile.h" +#include "vec/runtime/vector_search_user_params.h" + +namespace doris::segment_v2 { +#include "common/compile_check_begin.h" + +struct AnnIndexStats { + AnnIndexStats() + : search_costs_ns(TUnit::TIME_NS, 0), + load_index_costs_ns(TUnit::TIME_NS, 0), + engine_search_ns(TUnit::TIME_NS, 0), + result_process_costs_ns(TUnit::TIME_NS, 0), + engine_convert_ns(TUnit::TIME_NS, 0), + engine_prepare_ns(TUnit::TIME_NS, 0) {} + + AnnIndexStats(const AnnIndexStats& other) + : search_costs_ns(TUnit::TIME_NS, other.search_costs_ns.value()), + load_index_costs_ns(TUnit::TIME_NS, other.load_index_costs_ns.value()), + engine_search_ns(TUnit::TIME_NS, other.engine_search_ns.value()), + result_process_costs_ns(TUnit::TIME_NS, other.result_process_costs_ns.value()), + engine_convert_ns(TUnit::TIME_NS, other.engine_convert_ns.value()), + engine_prepare_ns(TUnit::TIME_NS, other.engine_prepare_ns.value()) {} + + AnnIndexStats& operator=(const AnnIndexStats& other) { + if (this != &other) { + search_costs_ns.set(other.search_costs_ns.value()); + load_index_costs_ns.set(other.load_index_costs_ns.value()); + engine_search_ns.set(other.engine_search_ns.value()); + result_process_costs_ns.set(other.result_process_costs_ns.value()); + engine_convert_ns.set(other.engine_convert_ns.value()); + engine_prepare_ns.set(other.engine_prepare_ns.value()); + } + return *this; + } + + RuntimeProfile::Counter search_costs_ns; // total time cost of TopN search + RuntimeProfile::Counter load_index_costs_ns; // time cost of loading ANN index + RuntimeProfile::Counter engine_search_ns; // time cost of calling FAISS/search engine + RuntimeProfile::Counter result_process_costs_ns; // time cost of processing search results + RuntimeProfile::Counter engine_convert_ns; // time cost of engine-side conversions + RuntimeProfile::Counter + engine_prepare_ns; // time cost before engine search (allocations, setup) +}; + +struct AnnTopNParam { + const float* query_value; + const size_t query_value_size; + size_t limit; + doris::VectorSearchUserParams _user_params; + roaring::Roaring* roaring; + size_t rows_of_segment = 0; + std::unique_ptr> distance = nullptr; + std::unique_ptr> row_ids = nullptr; + std::unique_ptr stats = nullptr; +}; + +struct AnnRangeSearchParams { + bool is_le_or_lt = true; + float* query_value = nullptr; + float radius = -1; + roaring::Roaring* roaring; // roaring from segment_iterator + std::string to_string() const { + DCHECK(roaring != nullptr); + return fmt::format("is_le_or_lt: {}, radius: {}, input rows {}", is_le_or_lt, radius, + roaring->cardinality()); + } + virtual ~AnnRangeSearchParams() = default; +}; + +struct AnnRangeSearchResult { + std::shared_ptr roaring; + std::unique_ptr> row_ids; + std::unique_ptr distance; +}; + +/* +This struct is used to wrap the search result of a vector index. +roaring is a bitmap that contains the row ids that satisfy the search condition. +row_ids is a vector of row ids that are returned by the search, it could be used by virtual_column_iterator to do column filter. +distances is a vector of distances that are returned by the search. +For range search, is condition is not le_or_lt, the row_ids and distances will be nullptr. +*/ +struct IndexSearchResult { + IndexSearchResult() = default; + + std::unique_ptr distances = nullptr; + std::unique_ptr> row_ids = nullptr; + std::shared_ptr roaring = nullptr; + // Internal engine timings (ns) + int64_t engine_search_ns = 0; // time spent in the underlying index search call + int64_t engine_convert_ns = 0; // time spent building selectors/results inside the engine + int64_t engine_prepare_ns = 0; // time spent preparing buffers before engine search +}; + +struct IndexSearchParameters { + roaring::Roaring* roaring = nullptr; + bool is_le_or_lt = true; + size_t rows_of_segment = 0; + virtual ~IndexSearchParameters() = default; +}; + +struct HNSWSearchParameters : public IndexSearchParameters { + int ef_search = 16; + bool check_relative_distance = true; + bool bounded_queue = true; +}; +#include "common/compile_check_end.h" +} // namespace doris::segment_v2 \ No newline at end of file diff --git a/be/src/olap/rowset/segment_v2/ann_index/ann_topn_runtime.cpp b/be/src/olap/rowset/segment_v2/ann_index/ann_topn_runtime.cpp new file mode 100644 index 00000000000000..6076c3620c9702 --- /dev/null +++ b/be/src/olap/rowset/segment_v2/ann_index/ann_topn_runtime.cpp @@ -0,0 +1,191 @@ +// 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 "ann_topn_runtime.h" + +#include +#include +#include +#include + +#include "common/logging.h" +#include "olap/rowset/segment_v2/ann_index/ann_index_iterator.h" +#include "olap/rowset/segment_v2/ann_index/ann_search_params.h" +#include "runtime/runtime_state.h" +#include "vec/columns/column.h" +#include "vec/columns/column_array.h" +#include "vec/columns/column_const.h" +#include "vec/columns/column_nullable.h" +#include "vec/common/assert_cast.h" +#include "vec/exprs/varray_literal.h" +#include "vec/exprs/vexpr_context.h" +#include "vec/exprs/vexpr_fwd.h" +#include "vec/exprs/virtual_slot_ref.h" +#include "vec/exprs/vslot_ref.h" +#include "vec/functions/array/function_array_distance.h" + +namespace doris::segment_v2 { +#include "common/compile_check_begin.h" +Status AnnTopNRuntime::prepare(RuntimeState* state, const RowDescriptor& row_desc) { + RETURN_IF_ERROR(_order_by_expr_ctx->prepare(state, row_desc)); + RETURN_IF_ERROR(_order_by_expr_ctx->open(state)); + + // Check the structure of the _order_by_expr_ctx + + /* + vectorized::VirtualSlotRef + | + | + FuncationCall + |---------------- + | | + | | + CastToArray ArrayLiteral + | + | + SlotRef + */ + std::shared_ptr vir_slot_ref = + std::dynamic_pointer_cast(_order_by_expr_ctx->root()); + DCHECK(vir_slot_ref != nullptr); + if (vir_slot_ref == nullptr) { + return Status::InternalError( + "root of order by expr of ann topn must be a vectorized::VirtualSlotRef, got\n{}", + _order_by_expr_ctx->root()->debug_string()); + } + DCHECK(vir_slot_ref->column_id() >= 0); + _dest_column_idx = vir_slot_ref->column_id(); + auto vir_col_expr = vir_slot_ref->get_virtual_column_expr(); + std::shared_ptr distance_fn_call = + std::dynamic_pointer_cast(vir_col_expr); + + if (distance_fn_call == nullptr) { + return Status::InternalError("Ann topn expr expect FuncationCall, got\n{}", + vir_col_expr->debug_string()); + } + + std::shared_ptr cast_to_array_expr = + std::dynamic_pointer_cast(distance_fn_call->children()[0]); + + if (cast_to_array_expr == nullptr) { + return Status::InternalError("Ann topn expr expect cast_to_array_expr, got\n{}", + distance_fn_call->children()[0]->debug_string()); + } + + std::shared_ptr slot_ref = + std::dynamic_pointer_cast(cast_to_array_expr->children()[0]); + if (slot_ref == nullptr) { + return Status::InternalError("Ann topn expr expect SlotRef, got\n{}", + cast_to_array_expr->children()[0]->debug_string()); + } + + // slot_ref->column_id() is acutually the columnd idx in block. + _src_column_idx = slot_ref->column_id(); + + std::shared_ptr array_literal = + std::dynamic_pointer_cast(distance_fn_call->children()[1]); + if (array_literal == nullptr) { + return Status::InternalError("Ann topn expr expect ArrayLiteral, got\n{}", + distance_fn_call->children()[1]->debug_string()); + } + _query_array = array_literal->get_column_ptr(); + _user_params = state->get_vector_search_params(); + + std::set distance_func_names = {vectorized::L2DistanceApproximate::name, + vectorized::InnerProductApproximate::name}; + if (distance_func_names.contains(distance_fn_call->function_name()) == false) { + return Status::InternalError("Ann topn expr expect distance function, got {}", + distance_fn_call->function_name()); + } + std::string metric_name = distance_fn_call->function_name(); + // Strip the "_approximate" suffix + metric_name = metric_name.substr(0, metric_name.size() - 12); + + _metric_type = segment_v2::string_to_metric(metric_name); + + VLOG_DEBUG << "AnnTopNRuntime: {}" << this->debug_string(); + return Status::OK(); +} + +Status AnnTopNRuntime::evaluate_vector_ann_search(segment_v2::IndexIterator* ann_index_iterator, + roaring::Roaring* roaring, size_t rows_of_segment, + vectorized::IColumn::MutablePtr& result_column, + std::unique_ptr>& row_ids, + segment_v2::AnnIndexStats& ann_index_stats) { + DCHECK(ann_index_iterator != nullptr); + segment_v2::AnnIndexIterator* ann_index_iterator_casted = + dynamic_cast(ann_index_iterator); + DCHECK(ann_index_iterator_casted != nullptr); + DCHECK(_order_by_expr_ctx != nullptr); + DCHECK(_order_by_expr_ctx->root() != nullptr); + + const vectorized::ColumnConst* const_column = + assert_cast(_query_array.get()); + const vectorized::ColumnArray* column_array = + assert_cast(const_column->get_data_column_ptr().get()); + const vectorized::ColumnNullable* column_nullable = + assert_cast(column_array->get_data_ptr().get()); + const vectorized::ColumnFloat64* cf64 = assert_cast( + column_nullable->get_nested_column_ptr().get()); + + const double* query_value = cf64->get_data().data(); + const size_t query_value_size = cf64->get_data().size(); + + std::unique_ptr query_value_f32 = std::make_unique(query_value_size); + for (size_t i = 0; i < query_value_size; ++i) { + query_value_f32[i] = static_cast(query_value[i]); + } + + segment_v2::AnnTopNParam ann_query_params { + .query_value = query_value_f32.get(), + .query_value_size = query_value_size, + .limit = _limit, + ._user_params = _user_params, + .roaring = roaring, + .rows_of_segment = rows_of_segment, + .distance = nullptr, + .row_ids = nullptr, + .stats = std::make_unique()}; + + RETURN_IF_ERROR(ann_index_iterator->read_from_index(&ann_query_params)); + + DCHECK(ann_query_params.distance != nullptr); + DCHECK(ann_query_params.row_ids != nullptr); + + size_t num_results = ann_query_params.distance->size(); + auto result_column_double = vectorized::ColumnFloat64::create(num_results); + auto result_null_map = vectorized::ColumnUInt8::create(num_results, 0); + + for (size_t i = 0; i < num_results; ++i) { + result_column_double->get_data()[i] = (*ann_query_params.distance)[i]; + } + + result_column = vectorized::ColumnNullable::create(std::move(result_column_double), + std::move(result_null_map)); + row_ids = std::move(ann_query_params.row_ids); + ann_index_stats = *ann_query_params.stats; + return Status::OK(); +} + +std::string AnnTopNRuntime::debug_string() const { + return fmt::format( + "AnnTopNRuntime: limit={}, src_col_idx={}, dest_col_idx={}, asc={}, user_params={}, " + "metric_type={}, order_by_expr={}", + _limit, _src_column_idx, _dest_column_idx, _asc, _user_params.to_string(), + segment_v2::metric_to_string(_metric_type), _order_by_expr_ctx->root()->debug_string()); +} +} // namespace doris::segment_v2 \ No newline at end of file diff --git a/be/src/olap/rowset/segment_v2/ann_index/ann_topn_runtime.h b/be/src/olap/rowset/segment_v2/ann_index/ann_topn_runtime.h new file mode 100644 index 00000000000000..8fd4dcee8a69ca --- /dev/null +++ b/be/src/olap/rowset/segment_v2/ann_index/ann_topn_runtime.h @@ -0,0 +1,168 @@ +// 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. + +/** + * @file ann_topn_runtime.h + * @brief Runtime execution engine for ANN (Approximate Nearest Neighbor) Top-N queries in Doris. + * + * This file contains the runtime infrastructure for executing ANN Top-N queries efficiently + * using vector indexes. It provides the bridge between Doris's SQL execution engine and + * underlying vector similarity search libraries like FAISS. + * + * The main class AnnTopNRuntime handles: + * - SQL expression analysis to extract vector search parameters + * - Integration with segment-level ANN indexes + * - Result collection and sorting for Top-K nearest neighbor queries + * - Performance statistics and monitoring + * + * This is used internally by the query execution engine when processing SQL queries like: + * SELECT * FROM table ORDER BY l2_distance(vector_column, [1,2,3]) LIMIT 10; + */ + +#pragma once + +#include "runtime/runtime_state.h" +#include "vec/columns/column.h" +#include "vec/exprs/varray_literal.h" +#include "vec/exprs/vcast_expr.h" +#include "vec/exprs/vectorized_fn_call.h" +#include "vec/exprs/vexpr.h" +#include "vec/exprs/vexpr_context.h" +#include "vec/exprs/vexpr_fwd.h" +#include "vec/exprs/vslot_ref.h" + +namespace doris::segment_v2 { +#include "common/compile_check_begin.h" +struct AnnIndexStats; + +/** + * @brief Runtime execution engine for ANN (Approximate Nearest Neighbor) Top-N queries. + * + * This class implements the runtime execution logic for ANN Top-N queries, which find + * the K nearest neighbors to a given query vector. It integrates with Doris's vectorized + * execution framework and supports various distance metrics and search parameters. + * + * Key features: + * - Supports both ascending and descending order results + * - Configurable K (limit) parameter for top-N results + * - Integration with segment-level ANN indexes + * - Performance statistics collection + * - Thread-safe execution in parallel query contexts + * + * Typical usage in SQL: + * SELECT * FROM table ORDER BY l2_distance(vec_column, [1,2,3]) LIMIT 10; + */ +class AnnTopNRuntime { + ENABLE_FACTORY_CREATOR(AnnTopNRuntime); + +public: + /** + * @brief Constructs an AnnTopNRuntime instance. + * + * @param asc Sort order: true for ascending (smallest distances first), false for descending + * @param limit Maximum number of results to return (K in K-NN) + * @param order_by_expr_ctx Expression context for the distance function (e.g., l2_distance) + */ + AnnTopNRuntime(bool asc, size_t limit, vectorized::VExprContextSPtr order_by_expr_ctx) + : _asc(asc), _limit(limit), _order_by_expr_ctx(order_by_expr_ctx) {}; + + /** + * @brief Prepares the runtime for execution by analyzing the distance expression. + * + * This method analyzes the ORDER BY expression to extract: + * - Source column index (the vector column being searched) + * - Distance metric type (L2, Inner Product, etc.) + * - Query vector from the literal array + * - User-defined search parameters + * + * @param state Runtime state containing session and query context + * @param row_desc Row descriptor for the input schema + * @return Status indicating success or failure + */ + Status prepare(RuntimeState* state, const RowDescriptor& row_desc); + + vectorized::VExprContextSPtr get_order_by_expr_ctx() const { return _order_by_expr_ctx; } + + /** + * @brief Executes the ANN search on the given index iterator. + * + * This is the core method that performs the actual ANN search by: + * 1. Calling the underlying index search method (e.g., HNSW, IVF) + * 2. Filtering results based on the provided row bitmap + * 3. Collecting performance statistics + * 4. Returning the top-K results with their distances and row IDs + * + * @param ann_index_iterator Iterator for the ANN index on the segment + * @param row_bitmap Bitmap indicating which rows are valid for the search + * @param result_column Output column containing the computed distances + * @param row_ids Output vector containing the row IDs of matching results + * @param ann_index_stats Statistics collector for performance monitoring + * @return Status indicating success or failure + */ + Status evaluate_vector_ann_search(segment_v2::IndexIterator* ann_index_iterator, + roaring::Roaring* row_bitmap, size_t rows_of_segment, + vectorized::IColumn::MutablePtr& result_column, + std::unique_ptr>& row_ids, + segment_v2::AnnIndexStats& ann_index_stats); + + /** + * @brief Gets the distance metric type used by this runtime. + * @return The metric type (L2_DISTANCE, INNER_PRODUCT, etc.) + */ + AnnIndexMetric get_metric_type() const { return _metric_type; } + + /** + * @brief Returns a debug string representation of this runtime. + * @return String containing runtime configuration and state information + */ + std::string debug_string() const; + + /** + * @brief Gets the source column index (vector column being searched). + * @return Column index in the table schema + */ + size_t get_src_column_idx() const { return _src_column_idx; } + + /** + * @brief Gets the destination column index for distance results. + * @return Column index where distance values will be stored + */ + size_t get_dest_column_idx() const { return _dest_column_idx; } + + /** + * @brief Gets the sort order for results. + * @return true for ascending order (smallest distances first), false for descending + */ + bool is_asc() const { return _asc; } + +private: + // Core configuration + const bool _asc; ///< Sort order for results + const size_t _limit; ///< Maximum number of results (K in K-NN) + vectorized::VExprContextSPtr + _order_by_expr_ctx; ///< Expression context for distance calculation + + // Runtime metadata + std::string _name = "ann_topn_runtime"; ///< Runtime identifier for logging + size_t _src_column_idx = -1; ///< Source vector column index + size_t _dest_column_idx = -1; ///< Destination distance column index + segment_v2::AnnIndexMetric _metric_type; ///< Distance metric type + vectorized::IColumn::Ptr _query_array; ///< Query vector data + doris::VectorSearchUserParams _user_params; ///< User-defined search parameters +}; +#include "common/compile_check_end.h" +} // namespace doris::segment_v2 \ No newline at end of file diff --git a/be/src/olap/rowset/segment_v2/ann_index/cmake-protect/CMakeLists.txt b/be/src/olap/rowset/segment_v2/ann_index/cmake-protect/CMakeLists.txt new file mode 100644 index 00000000000000..8e5272711b855e --- /dev/null +++ b/be/src/olap/rowset/segment_v2/ann_index/cmake-protect/CMakeLists.txt @@ -0,0 +1,46 @@ +# 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. + +# Make sure compile check in doris will not break compilation of faiss and openblas +set(CMAKE_C_FLAGS "${CMAKE_C_FLAGS} -w") +set(CMAKE_CXX_FLAGS "${CMAKE_CXX_FLAGS} -w") + +set(BUILD_WITHOUT_LAPACK OFF CACHE BOOL "Disable LAPACK support in OpenBLAS") +set(NO_SHARED TRUE CACHE BOOL "Disable shared library in OpenBLAS") +set(C_LAPACK TRUE CACHE BOOL "Enable C interface for LAPACK in OpenBLAS") +set(USE_OPENMP TRUE CACHE BOOL "Enable OpenMP in OpenBLAS") +set(NOFORTRAN ON CACHE BOOL "Disable Fortran in OpenBLAS") +set(BUILD_STATIC_LIBS ON CACHE BOOL "Build static libraries in OpenBLAS") +set(BUILD_TESTING OFF CACHE BOOL "Build shared libraries in OpenBLAS") +set(BUILD_RELAPACK ON CACHE BOOL "Build relapack in OpenBLAS") +set(BUILD_BENCHMARKS OFF CACHE BOOL "Build benchmarks in OpenBLAS") +set(NO_LAPACK OFF CACHE BOOL "Disable LAPACK in OpenBLAS") +set(NO_CBLAS ON CACHE BOOL "Disable CBLAS in OpenBLAS") +set(NO_AVX512 ON CACHE BOOL "Disable AVX512 in OpenBLAS") + +# EXCLUDE_FROM_ALL so that binary in openblas is not installed. +add_subdirectory(${PROJECT_SOURCE_DIR}/../contrib/openblas ${PROJECT_BINARY_DIR}/openblas EXCLUDE_FROM_ALL) + +set(OPENBLAS_LIBRARY "${PROJECT_BINARY_DIR}/openblas/lib/libopenblas.a" CACHE PATH "Path to OpenBLAS build directory") +set(FAISS_ENABLE_MKL OFF CACHE BOOL "Disable MKL support in FAISS") +set(FAISS_ENABLE_GPU OFF CACHE BOOL "Disable GPU support in FAISS") +set(FAISS_ENABLE_PYTHON OFF CACHE BOOL "Disable Python support in FAISS") +set(FAISS_ENABLE_EXTRAS OFF CACHE BOOL "Disable FAISS extras") +set(BUILD_TESTING OFF CACHE BOOL "Disable FAISS testing") + +# EXCLUDE_FROM_ALL so that binary in faiss is not installed. +add_subdirectory(${PROJECT_SOURCE_DIR}/../contrib/faiss ${PROJECT_BINARY_DIR}/faiss EXCLUDE_FROM_ALL) \ No newline at end of file diff --git a/be/src/olap/rowset/segment_v2/ann_index/faiss_ann_index.cpp b/be/src/olap/rowset/segment_v2/ann_index/faiss_ann_index.cpp new file mode 100644 index 00000000000000..ce8866fff18f8d --- /dev/null +++ b/be/src/olap/rowset/segment_v2/ann_index/faiss_ann_index.cpp @@ -0,0 +1,452 @@ +// 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 +#include + +#include +#include +#include +#include +#include +#include + +#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 FaissVectorIndex::roaring_to_faiss_selector( + const roaring::Roaring& roaring) { + std::vector ids; + ids.resize(roaring.cardinality()); + + size_t i = 0; + for (roaring::Roaring::const_iterator it = roaring.begin(); it != roaring.end(); ++it, ++i) { + ids[i] = cast_set(*it); + } + // construct derived and wrap into base unique_ptr explicitly + return std::unique_ptr(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(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(ptr); + // CLucene IndexOutput::writeBytes accepts at most Int32 bytes at a time. + const size_t kMaxChunk = + static_cast(std::numeric_limits::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(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(ptr); + const size_t kMaxChunk = + static_cast(std::numeric_limits::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(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); + _index->add(n, vec); + return doris::Status::OK(); +} + +void FaissVectorIndex::build(const FaissBuildParameter& params) { + _dimension = params.dim; + switch (params.metric_type) { + case FaissBuildParameter::MetricType::L2: + _metric = AnnIndexMetric::L2; + break; + case FaissBuildParameter::MetricType::IP: + _metric = AnnIndexMetric::IP; + break; + default: + throw doris::Exception(doris::ErrorCode::INVALID_ARGUMENT, "Unsupported metric type: {}", + static_cast(params.metric_type)); + break; + } + + if (params.index_type == FaissBuildParameter::IndexType::HNSW) { + if (params.metric_type == FaissBuildParameter::MetricType::L2) { + _index = std::make_unique(params.dim, params.max_degree, + faiss::METRIC_L2); + } else if (params.metric_type == FaissBuildParameter::MetricType::IP) { + _index = std::make_unique(params.dim, params.max_degree, + faiss::METRIC_INNER_PRODUCT); + } else { + throw doris::Exception(doris::ErrorCode::INVALID_ARGUMENT, + "Unsupported metric type: {}", + static_cast(params.metric_type)); + } + } else { + throw doris::Exception(doris::ErrorCode::INVALID_ARGUMENT, "Unsupported index type: {}", + static_cast(params.index_type)); + } +} + +// TODO: Support batch search +doris::Status FaissVectorIndex::ann_topn_search(const float* query_vec, int k, + const segment_v2::IndexSearchParameters& params, + segment_v2::IndexSearchResult& result) { + std::unique_ptr distances_ptr; + std::unique_ptr> labels_ptr; + { + SCOPED_RAW_TIMER(&result.engine_prepare_ns); + distances_ptr = std::make_unique(k); + // Initialize labels with -1 + // Even if there are N vectors in the index, limit N search in faiss could return less than N(eg, HNSW) + // so we need to initialize labels with -1 to tell the end of the result ids. + labels_ptr = std::make_unique>(k, -1); + } + float* distances = distances_ptr.get(); + faiss::idx_t* labels = (*labels_ptr).data(); + DCHECK(params.roaring != nullptr) + << "Roaring should not be null for topN search, please set roaring in params"; + + faiss::SearchParametersHNSW param; + const HNSWSearchParameters* hnsw_params = dynamic_cast(¶ms); + if (hnsw_params == nullptr) { + return doris::Status::InvalidArgument( + "HNSW search parameters should not be null for HNSW index"); + } + param.efSearch = hnsw_params->ef_search; + param.check_relative_distance = hnsw_params->check_relative_distance; + param.bounded_queue = hnsw_params->bounded_queue; + param.sel = nullptr; + std::unique_ptr id_sel = nullptr; + // Costs of roaring to faiss selector is very high especially when the cardinality is very high. + if (params.roaring->cardinality() != params.rows_of_segment) { + LOG_INFO("Roaring to faiss selector, roaring {} rows, segment {} rows", + params.roaring->cardinality(), params.rows_of_segment); + { + SCOPED_RAW_TIMER(&result.engine_prepare_ns); + id_sel = roaring_to_faiss_selector(*params.roaring); + } + param.sel = id_sel.get(); + } + { + SCOPED_RAW_TIMER(&result.engine_search_ns); + _index->search(1, query_vec, k, distances, labels, ¶m); + } + { + SCOPED_RAW_TIMER(&result.engine_convert_ns); + result.roaring = std::make_shared(); + update_roaring(labels, k, *result.roaring); + size_t roaring_cardinality = result.roaring->cardinality(); + result.distances = std::make_unique(roaring_cardinality); + result.row_ids = std::make_unique>(); + result.row_ids->resize(roaring_cardinality); + + if (_metric == AnnIndexMetric::L2) { + // For l2_distance, we need to convert the distance to the actual distance. + // The distance returned by Faiss is actually the squared distance. + // So we need to take the square root of the squared distance. + for (size_t i = 0; i < roaring_cardinality; ++i) { + (*result.row_ids)[i] = labels[i]; + result.distances[i] = std::sqrt(distances[i]); + } + } else if (_metric == AnnIndexMetric::IP) { + // For inner product, we can use the distance directly. + for (size_t i = 0; i < roaring_cardinality; ++i) { + (*result.row_ids)[i] = labels[i]; + result.distances[i] = distances[i]; + } + } else { + throw doris::Exception(doris::ErrorCode::INVALID_ARGUMENT, + "Unsupported metric type: {}", static_cast(_metric)); + } + + DCHECK(result.row_ids->size() == result.roaring->cardinality()) + << "Row ids size: " << result.row_ids->size() + << ", roaring size: " << result.roaring->cardinality(); + } + // distance/row_ids conversion above already timed via SCOPED_RAW_TIMER + return doris::Status::OK(); +} + +// For l2 distance, range search radius is the squared distance. +// For inner product, range search radius is the actual distance. +// range search on inner product returns all vectors with inner product greater than or equal to the radius. +// For l2 distance, range search returns all vectors with squared distance less than or equal to the radius. +doris::Status FaissVectorIndex::range_search(const float* query_vec, const float& radius, + const segment_v2::IndexSearchParameters& params, + segment_v2::IndexSearchResult& result) { + DCHECK(_index != nullptr); + DCHECK(query_vec != nullptr); + DCHECK(params.roaring != nullptr) + << "Roaring should not be null for range search, please set roaring in params"; + std::unique_ptr sel; + { + // Engine prepare: convert roaring bitmap to FAISS selector + SCOPED_RAW_TIMER(&result.engine_prepare_ns); + sel = roaring_to_faiss_selector(*params.roaring); + } + faiss::RangeSearchResult native_search_result(1, true); + const HNSWSearchParameters* hnsw_params = dynamic_cast(¶ms); + // Currently only support HNSW index for range search. + DCHECK(hnsw_params != nullptr) << "HNSW search parameters should not be null for HNSW index"; + + faiss::SearchParametersHNSW param; + { + // Engine prepare: set search parameters and bind selector + SCOPED_RAW_TIMER(&result.engine_prepare_ns); + param.efSearch = hnsw_params->ef_search; + param.check_relative_distance = hnsw_params->check_relative_distance; + param.bounded_queue = hnsw_params->bounded_queue; + param.sel = sel.get(); + } + { + // Engine search: FAISS range_search + SCOPED_RAW_TIMER(&result.engine_search_ns); + if (_metric == AnnIndexMetric::L2) { + if (radius <= 0) { + _index->range_search(1, query_vec, 0.0f, &native_search_result, ¶m); + } else { + _index->range_search(1, query_vec, radius * radius, &native_search_result, ¶m); + } + } else if (_metric == AnnIndexMetric::IP) { + _index->range_search(1, query_vec, radius, &native_search_result, ¶m); + } + } + + size_t begin = native_search_result.lims[0]; + size_t end = native_search_result.lims[1]; + auto row_ids = std::make_unique>(); + row_ids->resize(end - begin); + if (params.is_le_or_lt) { + if (_metric == AnnIndexMetric::L2) { + std::unique_ptr distances_ptr; + float* distances = nullptr; + auto roaring = std::make_shared(); + { + // Engine convert: build roaring, row_ids, distances from FAISS result + SCOPED_RAW_TIMER(&result.engine_convert_ns); + distances_ptr = std::make_unique(end - begin); + distances = distances_ptr.get(); + // The distance returned by Faiss is actually the squared distance. + // So we need to take the square root of the squared distance. + for (size_t i = begin; i < end; ++i) { + (*row_ids)[i] = native_search_result.labels[i]; + roaring->add(cast_set(native_search_result.labels[i])); + distances[i - begin] = sqrt(native_search_result.distances[i]); + } + } + result.distances = std::move(distances_ptr); + result.row_ids = std::move(row_ids); + result.roaring = roaring; + + DCHECK(result.row_ids->size() == result.roaring->cardinality()) + << "row_ids size: " << result.row_ids->size() + << ", roaring size: " << result.roaring->cardinality(); + } else if (_metric == AnnIndexMetric::IP) { + // For IP, we can use the distance directly. + // range search on ip gets all vectors with inner product greater than or equal to the radius. + // so we need to do a convertion. + const roaring::Roaring& origin_row_ids = *params.roaring; + std::shared_ptr roaring = std::make_shared(); + { + // Engine convert: compute roaring difference + SCOPED_RAW_TIMER(&result.engine_convert_ns); + for (size_t i = begin; i < end; ++i) { + roaring->add(cast_set(native_search_result.labels[i])); + } + result.roaring = std::make_shared(); + // remove all rows that should not be included. + *(result.roaring) = origin_row_ids - *roaring; + // Just update the roaring. distance can not be used. + } + } else { + throw doris::Exception(doris::ErrorCode::INVALID_ARGUMENT, + "Unsupported metric type: {}", static_cast(_metric)); + } + } else { + if (_metric == AnnIndexMetric::L2) { + // Faiss can only return labels in the range of radius. + // If the precidate is not less than, we need to to a convertion. + const roaring::Roaring& origin_row_ids = *params.roaring; + std::shared_ptr roaring = std::make_shared(); + { + // Engine convert: compute roaring difference + SCOPED_RAW_TIMER(&result.engine_convert_ns); + for (size_t i = begin; i < end; ++i) { + roaring->add(cast_set(native_search_result.labels[i])); + } + result.roaring = std::make_shared(); + *(result.roaring) = origin_row_ids - *roaring; + result.distances = nullptr; + result.row_ids = nullptr; + } + } else if (_metric == AnnIndexMetric::IP) { + // For inner product, we can use the distance directly. + // range search on ip gets all vectors with inner product greater than or equal to the radius. + // when query condition is not le_or_lt, we can use the roaring and distance directly. + std::unique_ptr distances_ptr = std::make_unique(end - begin); + float* distances = distances_ptr.get(); + auto roaring = std::make_shared(); + // The distance returned by Faiss is actually the squared distance. + // So we need to take the square root of the squared distance. + for (size_t i = begin; i < end; ++i) { + (*row_ids)[i] = native_search_result.labels[i]; + roaring->add(cast_set(native_search_result.labels[i])); + distances[i - begin] = native_search_result.distances[i]; + } + result.distances = std::move(distances_ptr); + result.row_ids = std::move(row_ids); + result.roaring = roaring; + + DCHECK(result.row_ids->size() == result.roaring->cardinality()) + << "row_ids size: " << result.row_ids->size() + << ", roaring size: " << result.roaring->cardinality(); + + } else { + throw doris::Exception(doris::ErrorCode::INVALID_ARGUMENT, + "Unsupported metric type: {}", static_cast(_metric)); + } + } + + return Status::OK(); +} + +doris::Status FaissVectorIndex::save(lucene::store::Directory* dir) { + auto start_time = std::chrono::high_resolution_clock::now(); + + lucene::store::IndexOutput* idx_output = dir->createOutput(faiss_index_fila_name); + auto writer = std::make_unique(idx_output); + faiss::write_index(_index.get(), writer.get()); + + auto end_time = std::chrono::high_resolution_clock::now(); + auto duration = std::chrono::duration_cast(end_time - start_time); + LOG_INFO(fmt::format("Faiss index saved to {}, {}, rows {}, cost {} ms", dir->toString(), + faiss_index_fila_name, _index->ntotal, duration.count())); + return doris::Status::OK(); +} + +doris::Status FaissVectorIndex::load(lucene::store::Directory* dir) { + auto start_time = std::chrono::high_resolution_clock::now(); + lucene::store::IndexInput* idx_input = nullptr; + try { + idx_input = dir->openInput(faiss_index_fila_name); + } catch (const CLuceneError& e) { + return doris::Status::Error( + "Failed to open index file: {}, error: {}", faiss_index_fila_name, e.what()); + } + + auto reader = std::make_unique(idx_input); + faiss::Index* idx = faiss::read_index(reader.get()); + auto end_time = std::chrono::high_resolution_clock::now(); + auto duration = std::chrono::duration_cast(end_time - start_time); + VLOG_DEBUG << fmt::format("Load index from {} costs {} ms, rows {}", dir->getObjectName(), + duration.count(), idx->ntotal); + _index.reset(idx); + return doris::Status::OK(); +} + +} // namespace doris::segment_v2 \ No newline at end of file diff --git a/be/src/olap/rowset/segment_v2/ann_index/faiss_ann_index.h b/be/src/olap/rowset/segment_v2/ann_index/faiss_ann_index.h new file mode 100644 index 00000000000000..f0c1270f2dd6f5 --- /dev/null +++ b/be/src/olap/rowset/segment_v2/ann_index/faiss_ann_index.h @@ -0,0 +1,230 @@ +// 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. + +#pragma once + +#include +#include +#include +#include +#include + +#include + +#include "common/status.h" +#include "olap/rowset/segment_v2/ann_index/ann_index.h" +#include "vec/core/types.h" + +namespace doris::segment_v2 { +#include "common/compile_check_begin.h" +struct IndexSearchParameters; +struct IndexSearchResult; +/** + * @brief Build parameters for constructing FAISS-based vector indexes. + * + * This structure encapsulates all configuration parameters needed to build + * various types of FAISS indexes. It supports different index types and + * distance metrics commonly used in vector similarity search. + */ +struct FaissBuildParameter { + /** + * @brief Supported vector index types. + */ + enum class IndexType { + HNSW ///< Hierarchical Navigable Small World (HNSW) index for high performance + }; + + /** + * @brief Supported distance metrics for vector similarity. + */ + enum class MetricType { + L2, ///< Euclidean distance (L2 norm) + IP, ///< Inner product (cosine similarity when vectors are normalized) + }; + + /** + * @brief Converts string representation to IndexType enum. + * @param type String representation of index type (e.g., "hnsw") + * @return Corresponding IndexType enum value + * @throws doris::Exception for unsupported index types + */ + static IndexType string_to_index_type(const std::string& type) { + if (type == "hnsw") { + return IndexType::HNSW; + } else { + throw doris::Exception(doris::ErrorCode::INVALID_ARGUMENT, "Unsupported index type: {}", + type); + } + } + + /** + * @brief Converts string representation to MetricType enum. + * @param type String representation of metric type (e.g., "l2_distance", "inner_product") + * @return Corresponding MetricType enum value + * @throws doris::Exception for unsupported metric types + */ + static MetricType string_to_metric_type(const std::string& type) { + if (type == "l2_distance") { + return MetricType::L2; + } else if (type == "inner_product") { + return MetricType::IP; + } else { + throw doris::Exception(doris::ErrorCode::INVALID_ARGUMENT, + "Unsupported metric type: {}", type); + } + } + + // HNSW-specific parameters + int dim = 0; ///< Vector dimensionality (must match data vectors) + int max_degree = 0; ///< Maximum number of connections per node in HNSW graph + IndexType index_type = IndexType::HNSW; ///< Type of index to build + MetricType metric_type = MetricType::L2; ///< Distance metric to use +}; + +/** + * @brief FAISS-based implementation of vector index for approximate nearest neighbor search. + * + * This class provides a concrete implementation of the VectorIndex interface using + * the FAISS library. It supports various index types (currently HNSW) and distance + * metrics (L2, Inner Product) for efficient vector similarity search. + * + * Key features: + * - High-performance approximate nearest neighbor search + * - Support for both exact and range search queries + * - Integration with Doris storage and query execution + * - Persistence to/from Lucene directory storage + * - Bitmap-based result filtering + * + * Thread safety: This class is NOT thread-safe. Concurrent access should be + * synchronized externally. + */ +class FaissVectorIndex : public VectorIndex { +public: + /** + * @brief Converts a Roaring bitmap to a FAISS IDSelector for filtered search. + * + * This utility method creates a FAISS IDSelector that can be used to filter + * search results to only include vectors whose IDs are present in the bitmap. + * This is essential for supporting WHERE clause filtering in vector queries. + * + * @param bitmap Roaring bitmap containing valid vector IDs + * @return Unique pointer to a FAISS IDSelector for the given bitmap + */ + static std::unique_ptr roaring_to_faiss_selector( + const roaring::Roaring& bitmap); + + /** + * @brief Updates a Roaring bitmap with the given labels/IDs. + * + * This method is used to update result bitmaps with the vector IDs + * returned from FAISS search operations. + * + * @param labels Array of vector IDs returned from search + * @param n Number of labels in the array + * @param roaring Reference to the Roaring bitmap to update + */ + static void update_roaring(const faiss::idx_t* labels, const size_t n, + roaring::Roaring& roaring); + + /** + * @brief Default constructor. + */ + FaissVectorIndex(); + + /** + * @brief Adds vectors to the index for future searches. + * + * This method is used during index building to add vectors to the FAISS index. + * The vectors must have the same dimensionality as specified in build parameters. + * + * @param n Number of vectors to add + * @param vec Pointer to vector data (n * dim float values) + * @return Status indicating success or failure + */ + doris::Status add(int n, const float* vec) override; + + /** + * @brief Sets the build parameters for the index. + * + * This method must be called before adding vectors or performing searches. + * It configures the underlying FAISS index with the specified parameters. + * + * @param params Build parameters including index type, metric, and dimensions + */ + void build(const FaissBuildParameter& params); + + /** + * @brief Performs approximate k-nearest neighbor search. + * + * Finds the k most similar vectors to the query vector using the configured + * distance metric. Results are ordered by similarity (closest first for L2, + * highest score first for inner product). + * + * @param query_vec Query vector (must be same dimensionality as index) + * @param k Number of nearest neighbors to find + * @param params Search parameters including any filtering criteria + * @param result Output structure containing distances and vector IDs + * @return Status indicating success or failure + */ + doris::Status ann_topn_search(const float* query_vec, int k, + const segment_v2::IndexSearchParameters& params, + segment_v2::IndexSearchResult& result) override; + + /** + * @brief Performs range search to find all vectors within a distance threshold. + * + * Finds all vectors within the specified radius from the query vector. + * This is useful for similarity queries where you want all "similar enough" + * vectors rather than a fixed number of nearest neighbors. + * + * @param query_vec Query vector (must be same dimensionality as index) + * @param radius Maximum distance threshold for results + * @param params Search parameters including any filtering criteria + * @param result Output structure containing distances and vector IDs + * @return Status indicating success or failure + */ + doris::Status range_search(const float* query_vec, const float& radius, + const segment_v2::IndexSearchParameters& params, + segment_v2::IndexSearchResult& result) override; + + /** + * @brief Saves the index to persistent storage. + * + * Serializes the complete FAISS index to the provided Lucene directory + * for later loading. This enables index persistence across restarts. + * + * @param directory Lucene directory for writing index data + * @return Status indicating success or failure + */ + doris::Status save(lucene::store::Directory*) override; + + /** + * @brief Loads the index from persistent storage. + * + * Deserializes a previously saved FAISS index from the provided Lucene + * directory. The loaded index is ready for search operations. + * + * @param directory Lucene directory containing saved index data + * @return Status indicating success or failure + */ + doris::Status load(lucene::store::Directory*) override; + +private: + std::unique_ptr _index = nullptr; ///< Underlying FAISS index instance +}; +#include "common/compile_check_end.h" +} // namespace doris::segment_v2 \ No newline at end of file diff --git a/be/src/olap/rowset/segment_v2/column_reader.cpp b/be/src/olap/rowset/segment_v2/column_reader.cpp index 69ab951cd54056..f9350b4dd2b6a9 100644 --- a/be/src/olap/rowset/segment_v2/column_reader.cpp +++ b/be/src/olap/rowset/segment_v2/column_reader.cpp @@ -38,6 +38,7 @@ #include "olap/inverted_index_parser.h" #include "olap/iterators.h" #include "olap/olap_common.h" +#include "olap/rowset/segment_v2/ann_index/ann_index_reader.h" #include "olap/rowset/segment_v2/binary_dict_page.h" // for BinaryDictPageDecoder #include "olap/rowset/segment_v2/binary_plain_page.h" #include "olap/rowset/segment_v2/bitmap_index_reader.h" @@ -684,7 +685,14 @@ Status ColumnReader::_load_index(const std::shared_ptr& index_f type = _type_info->type(); } + if (index_meta->index_type() == IndexType::ANN) { + _index_readers[index_meta->index_id()] = + std::make_shared(index_meta, index_file_reader); + return Status::OK(); + } + IndexReaderPtr index_reader; + if (is_string_type(type)) { if (should_analyzer) { try { diff --git a/be/src/olap/rowset/segment_v2/column_reader.h b/be/src/olap/rowset/segment_v2/column_reader.h index 382d2bf0a8d98f..17117fbc428e82 100644 --- a/be/src/olap/rowset/segment_v2/column_reader.h +++ b/be/src/olap/rowset/segment_v2/column_reader.h @@ -245,6 +245,7 @@ class ColumnReader : public MetadataAdder { [[nodiscard]] Status _load_ordinal_index(bool use_page_cache, bool kept_in_memory, const ColumnIteratorOptions& iter_opts); [[nodiscard]] Status _load_bitmap_index(bool use_page_cache, bool kept_in_memory); + [[nodiscard]] Status _load_index(const std::shared_ptr& index_file_reader, const TabletIndex* index_meta); [[nodiscard]] Status _load_bloom_filter_index(bool use_page_cache, bool kept_in_memory, diff --git a/be/src/olap/rowset/segment_v2/column_writer.cpp b/be/src/olap/rowset/segment_v2/column_writer.cpp index edc5b457c1fed6..461a5046e71eee 100644 --- a/be/src/olap/rowset/segment_v2/column_writer.cpp +++ b/be/src/olap/rowset/segment_v2/column_writer.cpp @@ -924,6 +924,14 @@ Status ArrayColumnWriter::init() { _opts.inverted_indexes[0])); } } + if (_opts.need_ann_index) { + auto* writer = dynamic_cast(_item_writer.get()); + if (writer != nullptr) { + _ann_index_writer = std::make_unique(_opts.index_file_writer, + _opts.ann_index); + RETURN_IF_ERROR(_ann_index_writer->init()); + } + } return Status::OK(); } @@ -934,6 +942,13 @@ Status ArrayColumnWriter::write_inverted_index() { return Status::OK(); } +Status ArrayColumnWriter::write_ann_index() { + if (_opts.need_ann_index) { + return _ann_index_writer->finish(); + } + return Status::OK(); +} + // batch append data for array Status ArrayColumnWriter::append_data(const uint8_t** ptr, size_t num_rows) { // data_ptr contains @@ -960,6 +975,21 @@ Status ArrayColumnWriter::append_data(const uint8_t** ptr, size_t num_rows) { } } + if (_opts.need_ann_index) { + auto* writer = dynamic_cast(_item_writer.get()); + // now only support nested type is scala + if (writer != nullptr) { + //NOTE: use array field name as index field, but item_writer size should be used when moving item_data_ptr + RETURN_IF_ERROR(_ann_index_writer->add_array_values( + _item_writer->get_field()->size(), reinterpret_cast(data), + reinterpret_cast(nested_null_map), offsets_ptr, num_rows)); + } else { + return Status::NotSupported( + "Ann index can only be build on array with scalar type. but got {} as nested", + _item_writer->get_field()->type()); + } + } + RETURN_IF_ERROR(_offset_writer->append_data(&offsets_ptr, num_rows)); return Status::OK(); } diff --git a/be/src/olap/rowset/segment_v2/column_writer.h b/be/src/olap/rowset/segment_v2/column_writer.h index 1868b4bd448346..4f42d6bb7502d1 100644 --- a/be/src/olap/rowset/segment_v2/column_writer.h +++ b/be/src/olap/rowset/segment_v2/column_writer.h @@ -30,6 +30,7 @@ #include "common/status.h" // for Status #include "olap/field.h" // for Field +#include "olap/rowset/segment_v2/ann_index/ann_index_writer.h" #include "olap/rowset/segment_v2/bloom_filter.h" #include "olap/rowset/segment_v2/common.h" #include "olap/rowset/segment_v2/inverted_index_writer.h" @@ -64,6 +65,7 @@ struct ColumnWriterOptions { bool need_bloom_filter = false; bool is_ngram_bf_index = false; bool need_inverted_index = false; + bool need_ann_index = false; uint8_t gram_size; uint16_t gram_bf_size; BloomFilterOptions bf_options; @@ -76,6 +78,7 @@ struct ColumnWriterOptions { RowsetWriterContext* rowset_ctx = nullptr; // For collect segment statistics for compaction std::vector input_rs_readers; + const TabletIndex* ann_index = nullptr; std::string to_string() const { std::stringstream ss; ss << std::boolalpha << "meta=" << meta->DebugString() @@ -172,6 +175,8 @@ class ColumnWriter { virtual Status write_inverted_index() = 0; + virtual Status write_ann_index() { return Status::OK(); } + virtual Status write_bloom_filter_index() = 0; virtual ordinal_t get_next_rowid() const = 0; @@ -404,6 +409,7 @@ class ArrayColumnWriter final : public ColumnWriter { return Status::OK(); } Status write_inverted_index() override; + Status write_ann_index() override; Status write_bloom_filter_index() override { if (_opts.need_bloom_filter) { return Status::NotSupported("array not support bloom filter index"); @@ -421,6 +427,7 @@ class ArrayColumnWriter final : public ColumnWriter { std::unique_ptr _null_writer; std::unique_ptr _item_writer; std::unique_ptr _inverted_index_builder; + std::unique_ptr _ann_index_writer; ColumnWriterOptions _opts; }; diff --git a/be/src/olap/rowset/segment_v2/index_file_writer.cpp b/be/src/olap/rowset/segment_v2/index_file_writer.cpp index 703509faa8a0d5..0ec7d5f6859271 100644 --- a/be/src/olap/rowset/segment_v2/index_file_writer.cpp +++ b/be/src/olap/rowset/segment_v2/index_file_writer.cpp @@ -25,6 +25,7 @@ #include "common/status.h" #include "io/fs/s3_file_writer.h" #include "io/fs/stream_sink_file_writer.h" +#include "olap/rowset/segment_v2/ann_index/ann_index_files.h" #include "olap/rowset/segment_v2/index_file_reader.h" #include "olap/rowset/segment_v2/index_storage_format_v1.h" #include "olap/rowset/segment_v2/index_storage_format_v2.h" @@ -134,6 +135,11 @@ Status IndexFileWriter::add_into_searcher_cache() { for (const auto& entry : _indices_dirs) { auto index_meta = entry.first; auto dir = DORIS_TRY(index_file_reader->_open(index_meta.first, index_meta.second)); + std::vector file_names; + dir->list(&file_names); + if (file_names.size() == 1 && (file_names[0] == faiss_index_fila_name)) { + continue; + } auto index_file_key = InvertedIndexDescriptor::get_index_file_cache_key( _index_path_prefix, index_meta.first, index_meta.second); InvertedIndexSearcherCache::CacheKey searcher_cache_key(index_file_key); @@ -212,6 +218,8 @@ Status IndexFileWriter::close() { err.what()); } } + LOG_INFO("IndexFileWriter closing, enable_write_index_searcher_cache: {}", + config::enable_write_index_searcher_cache); if (config::enable_write_index_searcher_cache) { return add_into_searcher_cache(); } diff --git a/be/src/olap/rowset/segment_v2/index_file_writer.h b/be/src/olap/rowset/segment_v2/index_file_writer.h index b293883fc017f2..45fa71a540dc9f 100644 --- a/be/src/olap/rowset/segment_v2/index_file_writer.h +++ b/be/src/olap/rowset/segment_v2/index_file_writer.h @@ -25,6 +25,7 @@ #include #include +#include "common/be_mock_util.h" #include "io/fs/file_system.h" #include "io/fs/file_writer.h" #include "io/fs/local_file_system.h" @@ -52,7 +53,7 @@ class IndexFileWriter { io::FileWriterPtr file_writer = nullptr, bool can_use_ram_dir = true); virtual ~IndexFileWriter() = default; - Result> open(const TabletIndex* index_meta); + MOCK_FUNCTION Result> open(const TabletIndex* index_meta); Status delete_index(const TabletIndex* index_meta); Status initialize(InvertedIndexDirectoryMap& indices_dirs); Status add_into_searcher_cache(); diff --git a/be/src/olap/rowset/segment_v2/index_iterator.h b/be/src/olap/rowset/segment_v2/index_iterator.h index 96cd8f45914a79..b97069f4089c57 100644 --- a/be/src/olap/rowset/segment_v2/index_iterator.h +++ b/be/src/olap/rowset/segment_v2/index_iterator.h @@ -23,19 +23,28 @@ #include "common/exception.h" #include "common/factory_creator.h" +#include "olap/rowset/segment_v2/ann_index/ann_index_reader.h" #include "olap/rowset/segment_v2/index_query_context.h" #include "olap/rowset/segment_v2/index_reader.h" #include "olap/rowset/segment_v2/inverted_index_query_type.h" #include "runtime/runtime_state.h" +namespace doris::vectorized { +struct AnnTopNParam; +} + namespace doris::segment_v2 { class InvertedIndexQueryCacheHandle; struct InvertedIndexParam; -using IndexParam = std::variant; -using IndexReaderType = std::variant; +using IndexParam = std::variant; + +enum class AnnIndexReaderType { + ANN = 0, +}; +using IndexReaderType = std::variant; class IndexIterator { public: IndexIterator() = default; diff --git a/be/src/olap/rowset/segment_v2/index_reader.h b/be/src/olap/rowset/segment_v2/index_reader.h index 65ad9ea1a47333..cf0fc5f34106ee 100644 --- a/be/src/olap/rowset/segment_v2/index_reader.h +++ b/be/src/olap/rowset/segment_v2/index_reader.h @@ -33,9 +33,6 @@ class IndexIterator; class InvertedIndexReader; using InvertedIndexReaderPtr = std::shared_ptr; -class AnnIndexReader; -using AnnIndexReaderPtr = std::shared_ptr; - class IndexReader : public std::enable_shared_from_this, public MetadataAdder { public: diff --git a/be/src/olap/rowset/segment_v2/index_storage_format_v2.cpp b/be/src/olap/rowset/segment_v2/index_storage_format_v2.cpp index 2b4c8672710843..2342e15d34b53e 100644 --- a/be/src/olap/rowset/segment_v2/index_storage_format_v2.cpp +++ b/be/src/olap/rowset/segment_v2/index_storage_format_v2.cpp @@ -52,6 +52,7 @@ Status IndexStorageFormatV2::write() { auto result = create_output_stream(); out_dir = std::move(result.first); compound_file_output = std::move(result.second); + VLOG_DEBUG << fmt::format("Output compound index file to streams: {}", out_dir->toString()); // Write version and number of indices write_version_and_indices_count(compound_file_output.get()); @@ -190,7 +191,6 @@ IndexStorageFormatV2::create_output_stream() { DCHECK(_index_file_writer->_idx_v2_writer != nullptr) << "inverted index file writer v2 is nullptr"; auto compound_file_output = out_dir->createOutputV2(_index_file_writer->_idx_v2_writer.get()); - return {std::move(out_dir_ptr), std::move(compound_file_output)}; } diff --git a/be/src/olap/rowset/segment_v2/index_writer.cpp b/be/src/olap/rowset/segment_v2/index_writer.cpp index 6bba37eb1e49fb..d5cf7f11b0332c 100644 --- a/be/src/olap/rowset/segment_v2/index_writer.cpp +++ b/be/src/olap/rowset/segment_v2/index_writer.cpp @@ -39,6 +39,11 @@ bool IndexColumnWriter::check_support_inverted_index(const TabletColumn& column) return true; } +bool IndexColumnWriter::check_support_ann_index(const TabletColumn& column) { + // bellow types are not supported in inverted index for extracted columns + return column.is_array_type(); +} + Status IndexColumnWriter::create(const Field* field, std::unique_ptr* res, IndexFileWriter* index_file_writer, const TabletIndex* index_meta) { diff --git a/be/src/olap/rowset/segment_v2/index_writer.h b/be/src/olap/rowset/segment_v2/index_writer.h index f956a60d8d5cd9..c166cdd49ea8e7 100644 --- a/be/src/olap/rowset/segment_v2/index_writer.h +++ b/be/src/olap/rowset/segment_v2/index_writer.h @@ -75,6 +75,8 @@ class IndexColumnWriter { // check if the column is valid for inverted index, some columns // are generated from variant, but not all of them are supported static bool check_support_inverted_index(const TabletColumn& column); + + static bool check_support_ann_index(const TabletColumn& column); }; class TmpFileDirs { diff --git a/be/src/olap/rowset/segment_v2/segment_iterator.cpp b/be/src/olap/rowset/segment_v2/segment_iterator.cpp index 1ae3e7ad107ec5..5660eca4c8ba5b 100644 --- a/be/src/olap/rowset/segment_v2/segment_iterator.cpp +++ b/be/src/olap/rowset/segment_v2/segment_iterator.cpp @@ -49,6 +49,9 @@ #include "olap/like_column_predicate.h" #include "olap/olap_common.h" #include "olap/primary_key_index.h" +#include "olap/rowset/segment_v2/ann_index/ann_index.h" +#include "olap/rowset/segment_v2/ann_index/ann_index_reader.h" +#include "olap/rowset/segment_v2/ann_index/ann_topn_runtime.h" #include "olap/rowset/segment_v2/bitmap_index_reader.h" #include "olap/rowset/segment_v2/column_reader.h" #include "olap/rowset/segment_v2/index_file_reader.h" @@ -316,6 +319,7 @@ Status SegmentIterator::_init_impl(const StorageReadOptions& opts) { _virtual_column_exprs = _opts.virtual_column_exprs; _vir_cid_to_idx_in_block = _opts.vir_cid_to_idx_in_block; _score_runtime = _opts.score_runtime; + _ann_topn_runtime = _opts.ann_topn_runtime; RETURN_IF_ERROR(init_iterators()); @@ -331,7 +335,8 @@ Status SegmentIterator::_init_impl(const StorageReadOptions& opts) { auto storage_type = _segment->get_data_type_of( col->get_desc(), _opts.io_ctx.reader_type != ReaderType::READER_QUERY); if (storage_type == nullptr) { - storage_type = vectorized::DataTypeFactory::instance().create_data_type(*col); + storage_type = vectorized::DataTypeFactory::instance().create_data_type( + col->get_desc(), col->is_nullable()); } // Currently, when writing a lucene index, the field of the document is column_name, and the column name is // bound to the index field. Since version 1.2, the data file storage has been changed from column_name to @@ -412,6 +417,8 @@ Status SegmentIterator::_lazy_init() { _prepare_score_column_materialization(); + RETURN_IF_ERROR(_apply_ann_topn_predicate()); + if (_opts.read_orderby_key_reverse) { _range_iter.reset(new BackwardBitmapRangeIterator(_row_bitmap)); } else { @@ -601,6 +608,98 @@ Status SegmentIterator::_get_row_ranges_by_column_conditions() { return Status::OK(); } +Status SegmentIterator::_apply_ann_topn_predicate() { + if (_ann_topn_runtime == nullptr) { + return Status::OK(); + } + + VLOG_DEBUG << fmt::format("Try apply ann topn: {}", _ann_topn_runtime->debug_string()); + size_t src_col_idx = _ann_topn_runtime->get_src_column_idx(); + ColumnId src_cid = _schema->column_id(src_col_idx); + IndexIterator* ann_index_iterator = _index_iterators[src_cid].get(); + bool has_ann_index = ann_index_iterator != nullptr; + bool has_common_expr_push_down = !_common_expr_ctxs_push_down.empty(); + bool has_column_predicate = std::any_of(_is_pred_column.begin(), _is_pred_column.end(), + [](bool is_pred) { return is_pred; }); + if (!has_ann_index || has_common_expr_push_down || has_column_predicate) { + VLOG_DEBUG << fmt::format( + "Ann topn can not be evaluated by ann index, has_ann_index: {}, " + "has_common_expr_push_down: {}, has_column_predicate: {}", + has_ann_index, has_common_expr_push_down, has_column_predicate); + return Status::OK(); + } + + // Process asc & desc according to the type of metric + auto index_reader = ann_index_iterator->get_reader(AnnIndexReaderType::ANN); + auto ann_index_reader = dynamic_cast(index_reader.get()); + DCHECK(ann_index_reader != nullptr); + if (ann_index_reader->get_metric_type() == AnnIndexMetric::IP) { + if (_ann_topn_runtime->is_asc()) { + VLOG_DEBUG << fmt::format( + "Asc topn for inner product can not be evaluated by ann index"); + return Status::OK(); + } + } else { + if (!_ann_topn_runtime->is_asc()) { + VLOG_DEBUG << fmt::format("Desc topn for l2/cosine can not be evaluated by ann index"); + return Status::OK(); + } + } + + if (ann_index_reader->get_metric_type() != _ann_topn_runtime->get_metric_type()) { + VLOG_DEBUG << fmt::format( + "Ann topn metric type {} not match index metric type {}, can not be evaluated by " + "ann index", + metric_to_string(_ann_topn_runtime->get_metric_type()), + metric_to_string(ann_index_reader->get_metric_type())); + return Status::OK(); + } + + size_t pre_size = _row_bitmap.cardinality(); + size_t rows_of_segment = _segment->num_rows(); + if (static_cast(pre_size) < static_cast(rows_of_segment) * 0.3) { + VLOG_DEBUG << fmt::format( + "Ann topn predicate input rows {} < 30% of segment rows {}, will not use ann index " + "to " + "filter", + pre_size, rows_of_segment); + return Status::OK(); + } + vectorized::IColumn::MutablePtr result_column; + std::unique_ptr> result_row_ids; + segment_v2::AnnIndexStats ann_index_stats; + RETURN_IF_ERROR(_ann_topn_runtime->evaluate_vector_ann_search(ann_index_iterator, &_row_bitmap, + rows_of_segment, result_column, + result_row_ids, ann_index_stats)); + + VLOG_DEBUG << fmt::format("Ann topn filtered {} - {} = {} rows", pre_size, + _row_bitmap.cardinality(), pre_size - _row_bitmap.cardinality()); + + int64_t rows_filterd = pre_size - _row_bitmap.cardinality(); + _opts.stats->rows_ann_index_topn_filtered += rows_filterd; + _opts.stats->ann_index_load_ns += ann_index_stats.load_index_costs_ns.value(); + _opts.stats->ann_topn_search_ns += ann_index_stats.search_costs_ns.value(); + _opts.stats->ann_index_topn_engine_search_ns += ann_index_stats.engine_search_ns.value(); + _opts.stats->ann_index_topn_result_process_ns += + ann_index_stats.result_process_costs_ns.value(); + _opts.stats->ann_index_topn_engine_convert_ns += ann_index_stats.engine_convert_ns.value(); + _opts.stats->ann_index_topn_engine_prepare_ns += ann_index_stats.engine_prepare_ns.value(); + _opts.stats->ann_index_topn_search_cnt += 1; + const size_t dst_col_idx = _ann_topn_runtime->get_dest_column_idx(); + ColumnIterator* column_iter = _column_iterators[_schema->column_id(dst_col_idx)].get(); + DCHECK(column_iter != nullptr); + VirtualColumnIterator* virtual_column_iter = dynamic_cast(column_iter); + DCHECK(virtual_column_iter != nullptr); + VLOG_DEBUG << fmt::format( + "Virtual column iterator, column_idx {}, is materialized with {} rows", dst_col_idx, + result_row_ids->size()); + // reference count of result_column should be 1, so move will not issue any data copy. + virtual_column_iter->prepare_materialization(std::move(result_column), + std::move(result_row_ids)); + + return Status::OK(); +} + Status SegmentIterator::_get_row_ranges_from_conditions(RowRanges* condition_row_ranges) { std::set cids; for (auto& entry : _opts.col_id_to_predicates) { @@ -848,6 +947,7 @@ bool SegmentIterator::_check_apply_by_inverted_index(ColumnPredicate* pred) { return true; } +// TODO: optimization when all expr can not evaluate by inverted/ann index, Status SegmentIterator::_apply_index_expr() { for (const auto& expr_ctx : _common_expr_ctxs_push_down) { if (Status st = expr_ctx->evaluate_inverted_index(num_rows()); !st.ok()) { @@ -862,6 +962,33 @@ Status SegmentIterator::_apply_index_expr() { } } } + + // Apply ann range search + segment_v2::AnnIndexStats ann_index_stats; + for (const auto& expr_ctx : _common_expr_ctxs_push_down) { + size_t origin_rows = _row_bitmap.cardinality(); + RETURN_IF_ERROR(expr_ctx->evaluate_ann_range_search(_index_iterators, _schema->column_ids(), + _column_iterators, _row_bitmap, + ann_index_stats)); + _opts.stats->rows_ann_index_range_filtered += (origin_rows - _row_bitmap.cardinality()); + _opts.stats->ann_index_load_ns += ann_index_stats.load_index_costs_ns.value(); + _opts.stats->ann_index_range_search_ns += ann_index_stats.search_costs_ns.value(); + _opts.stats->ann_range_engine_search_ns += ann_index_stats.engine_search_ns.value(); + _opts.stats->ann_range_result_convert_ns += ann_index_stats.result_process_costs_ns.value(); + _opts.stats->ann_range_engine_convert_ns += ann_index_stats.engine_convert_ns.value(); + _opts.stats->ann_range_pre_process_ns += ann_index_stats.engine_prepare_ns.value(); + } + + for (auto it = _common_expr_ctxs_push_down.begin(); it != _common_expr_ctxs_push_down.end();) { + if ((*it)->root()->has_been_executed()) { + _opts.stats->ann_index_range_search_cnt++; + it = _common_expr_ctxs_push_down.erase(it); + } else { + ++it; + } + } + // TODO：Do we need to remove these expr root from _remaining_conjunct_roots? + return Status::OK(); } @@ -1201,6 +1328,22 @@ Status SegmentIterator::_init_index_iterators() { } } + // Ann index iterators + for (auto cid : _schema->column_ids()) { + if (_index_iterators[cid] == nullptr) { + const auto& column = _opts.tablet_schema->column(cid); + int32_t col_unique_id = + column.is_extracted_column() ? column.parent_unique_id() : column.unique_id(); + RETURN_IF_ERROR(_segment->new_index_iterator( + column, + _segment->_tablet_schema->ann_index(col_unique_id, column.suffix_path()), _opts, + &_index_iterators[cid])); + if (_index_iterators[cid] != nullptr) { + _index_iterators[cid]->set_context(_index_query_context); + } + } + } + return Status::OK(); } @@ -2218,8 +2361,8 @@ Status SegmentIterator::_next_batch_internal(vectorized::Block* block) { _is_char_type.resize(_schema->columns().size(), false); _vec_init_char_column_id(block); } - for (size_t i = 0; i < _schema->num_column_ids(); i++) { - auto cid = _schema->column_id(i); + for (size_t i = 0; i < _schema->column_ids().size(); i++) { + ColumnId cid = _schema->column_ids()[i]; auto column_desc = _schema->column(cid); if (_is_pred_column[cid]) { auto storage_column_type = _storage_name_and_type[cid].second; @@ -2227,6 +2370,8 @@ Status SegmentIterator::_next_batch_internal(vectorized::Block* block) { // both are DataTypeString, but DataTypeString only return FieldType::OLAP_FIELD_TYPE_STRING // in get_storage_field_type. RETURN_IF_CATCH_EXCEPTION( + // Here, cid will not go out of bounds + // because the size of _current_return_columns equals _schema->tablet_columns().size() _current_return_columns[cid] = Schema::get_predicate_column_ptr( _is_char_type[cid] ? FieldType::OLAP_FIELD_TYPE_CHAR : storage_column_type->get_storage_field_type(), @@ -2679,6 +2824,7 @@ Status SegmentIterator::_construct_compound_expr_context() { _common_expr_inverted_index_status); for (const auto& expr_ctx : _opts.common_expr_ctxs_push_down) { vectorized::VExprContextSPtr context; + // _ann_range_search_runtime will do deep copy. RETURN_IF_ERROR(expr_ctx->clone(_opts.runtime_state, context)); context->set_inverted_index_context(inverted_index_context); _common_expr_ctxs_push_down.emplace_back(context); diff --git a/be/src/olap/rowset/segment_v2/segment_iterator.h b/be/src/olap/rowset/segment_v2/segment_iterator.h index a99b538ac991b0..7709426669c0c9 100644 --- a/be/src/olap/rowset/segment_v2/segment_iterator.h +++ b/be/src/olap/rowset/segment_v2/segment_iterator.h @@ -40,6 +40,7 @@ #include "olap/olap_common.h" #include "olap/row_cursor.h" #include "olap/row_cursor_cell.h" +#include "olap/rowset/segment_v2/ann_index/ann_topn_runtime.h" #include "olap/rowset/segment_v2/common.h" #include "olap/rowset/segment_v2/index_iterator.h" #include "olap/rowset/segment_v2/segment.h" @@ -171,6 +172,8 @@ class SegmentIterator : public RowwiseIterator { [[nodiscard]] Status _init_return_column_iterators(); [[nodiscard]] Status _init_bitmap_index_iterators(); [[nodiscard]] Status _init_index_iterators(); + + Status _apply_ann_topn_predicate(); // calculate row ranges that fall into requested key ranges using short key index [[nodiscard]] Status _get_row_ranges_by_keys(); [[nodiscard]] Status _prepare_seek(const StorageReadOptions::KeyRange& key_range); @@ -486,6 +489,8 @@ class SegmentIterator : public RowwiseIterator { std::shared_ptr _score_runtime; + std::shared_ptr _ann_topn_runtime; + // cid to virtual column expr std::map _virtual_column_exprs; std::map _vir_cid_to_idx_in_block; diff --git a/be/src/olap/rowset/segment_v2/segment_writer.cpp b/be/src/olap/rowset/segment_v2/segment_writer.cpp index 136da1f7f94aba..1f77221d232d66 100644 --- a/be/src/olap/rowset/segment_v2/segment_writer.cpp +++ b/be/src/olap/rowset/segment_v2/segment_writer.cpp @@ -47,7 +47,7 @@ #include "olap/rowset/segment_creator.h" #include "olap/rowset/segment_v2/column_writer.h" // ColumnWriter #include "olap/rowset/segment_v2/index_file_writer.h" -#include "olap/rowset/segment_v2/inverted_index_writer.h" +#include "olap/rowset/segment_v2/index_writer.h" #include "olap/rowset/segment_v2/page_io.h" #include "olap/rowset/segment_v2/page_pointer.h" #include "olap/rowset/segment_v2/variant_stats_calculator.h" @@ -233,6 +233,13 @@ Status SegmentWriter::_create_column_writer(uint32_t cid, const TabletColumn& co DCHECK(_index_file_writer != nullptr); } } + // indexes for this column + if (const auto& index = schema->ann_index(column); index != nullptr) { + opts.ann_index = index; + opts.need_ann_index = true; + DCHECK(_index_file_writer != nullptr); + } + opts.index_file_writer = _index_file_writer; #define DISABLE_INDEX_IF_FIELD_TYPE(TYPE, type_name) \ @@ -960,6 +967,7 @@ Status SegmentWriter::finalize_columns_index(uint64_t* index_size) { RETURN_IF_ERROR(_write_zone_map()); RETURN_IF_ERROR(_write_bitmap_index()); RETURN_IF_ERROR(_write_inverted_index()); + RETURN_IF_ERROR(_write_ann_index()); RETURN_IF_ERROR(_write_bloom_filter_index()); *index_size = _file_writer->bytes_appended() - index_start; @@ -1090,6 +1098,13 @@ Status SegmentWriter::_write_inverted_index() { return Status::OK(); } +Status SegmentWriter::_write_ann_index() { + for (auto& column_writer : _column_writers) { + RETURN_IF_ERROR(column_writer->write_ann_index()); + } + return Status::OK(); +} + Status SegmentWriter::_write_bloom_filter_index() { for (auto& column_writer : _column_writers) { RETURN_IF_ERROR(column_writer->write_bloom_filter_index()); diff --git a/be/src/olap/rowset/segment_v2/segment_writer.h b/be/src/olap/rowset/segment_v2/segment_writer.h index cd6d838b9875c2..b5c99bbafa95d4 100644 --- a/be/src/olap/rowset/segment_v2/segment_writer.h +++ b/be/src/olap/rowset/segment_v2/segment_writer.h @@ -166,6 +166,7 @@ class SegmentWriter { Status _write_zone_map(); Status _write_bitmap_index(); Status _write_inverted_index(); + Status _write_ann_index(); Status _write_bloom_filter_index(); Status _write_short_key_index(); Status _write_primary_key_index(); diff --git a/be/src/olap/rowset/segment_v2/vertical_segment_writer.cpp b/be/src/olap/rowset/segment_v2/vertical_segment_writer.cpp index 45ebb660756fef..1f73fff305915a 100644 --- a/be/src/olap/rowset/segment_v2/vertical_segment_writer.cpp +++ b/be/src/olap/rowset/segment_v2/vertical_segment_writer.cpp @@ -235,6 +235,13 @@ Status VerticalSegmentWriter::_create_column_writer(uint32_t cid, const TabletCo } opts.index_file_writer = _index_file_writer; + if (const auto& index = tablet_schema->ann_index(column); index != nullptr) { + opts.ann_index = index; + opts.need_ann_index = true; + DCHECK(_index_file_writer != nullptr); + opts.index_file_writer = _index_file_writer; + } + #define DISABLE_INDEX_IF_FIELD_TYPE(TYPE, type_name) \ if (column.type() == FieldType::OLAP_FIELD_TYPE_##TYPE) { \ opts.need_zone_map = false; \ @@ -1194,6 +1201,7 @@ Status VerticalSegmentWriter::finalize_columns_index(uint64_t* index_size) { RETURN_IF_ERROR(_write_zone_map()); RETURN_IF_ERROR(_write_bitmap_index()); RETURN_IF_ERROR(_write_inverted_index()); + RETURN_IF_ERROR(_write_ann_index()); RETURN_IF_ERROR(_write_bloom_filter_index()); *index_size = _file_writer->bytes_appended() - index_start; @@ -1289,6 +1297,13 @@ Status VerticalSegmentWriter::_write_inverted_index() { return Status::OK(); } +Status VerticalSegmentWriter::_write_ann_index() { + for (auto& column_writer : _column_writers) { + RETURN_IF_ERROR(column_writer->write_ann_index()); + } + return Status::OK(); +} + Status VerticalSegmentWriter::_write_bloom_filter_index() { for (auto& column_writer : _column_writers) { RETURN_IF_ERROR(column_writer->write_bloom_filter_index()); diff --git a/be/src/olap/rowset/segment_v2/vertical_segment_writer.h b/be/src/olap/rowset/segment_v2/vertical_segment_writer.h index 398ac0bb583952..294756788c4828 100644 --- a/be/src/olap/rowset/segment_v2/vertical_segment_writer.h +++ b/be/src/olap/rowset/segment_v2/vertical_segment_writer.h @@ -139,6 +139,7 @@ class VerticalSegmentWriter { Status _write_zone_map(); Status _write_bitmap_index(); Status _write_inverted_index(); + Status _write_ann_index(); Status _write_bloom_filter_index(); Status _write_short_key_index(); Status _write_primary_key_index(); diff --git a/be/src/olap/rowset/vertical_beta_rowset_writer.cpp b/be/src/olap/rowset/vertical_beta_rowset_writer.cpp index 684b6c797056c3..462729623d1090 100644 --- a/be/src/olap/rowset/vertical_beta_rowset_writer.cpp +++ b/be/src/olap/rowset/vertical_beta_rowset_writer.cpp @@ -170,7 +170,7 @@ Status VerticalBetaRowsetWriter::_create_segment_writer( DCHECK(segment_file_writer != nullptr); IndexFileWriterPtr index_file_writer; - if (context.tablet_schema->has_inverted_index()) { + if (context.tablet_schema->has_inverted_index() || context.tablet_schema->has_ann_index()) { RETURN_IF_ERROR(RowsetWriter::create_index_file_writer(seg_id, &index_file_writer)); } @@ -185,7 +185,7 @@ Status VerticalBetaRowsetWriter::_create_segment_writer( context.data_dir, writer_options, index_file_writer.get()); RETURN_IF_ERROR(this->_seg_files.add(seg_id, std::move(segment_file_writer))); - if (context.tablet_schema->has_inverted_index()) { + if (context.tablet_schema->has_inverted_index() || context.tablet_schema->has_ann_index()) { RETURN_IF_ERROR(this->_idx_files.add(seg_id, std::move(index_file_writer))); } diff --git a/be/src/olap/snapshot_manager.cpp b/be/src/olap/snapshot_manager.cpp index 5a6ff80d36cf74..3a7634c62f21aa 100644 --- a/be/src/olap/snapshot_manager.cpp +++ b/be/src/olap/snapshot_manager.cpp @@ -753,7 +753,7 @@ Status SnapshotManager::_create_snapshot_files(const TabletSharedPtr& ref_tablet linked_success_files.push_back(snapshot_segment_index_file_path); } } else { - if (tablet_schema.has_inverted_index()) { + if (tablet_schema.has_inverted_index() || tablet_schema.has_ann_index()) { auto index_file = InvertedIndexDescriptor::get_index_file_path_v2( InvertedIndexDescriptor::get_index_file_path_prefix( segment_file_path)); diff --git a/be/src/olap/tablet_meta.cpp b/be/src/olap/tablet_meta.cpp index ced9ad54e7b57c..d6759fda55811a 100644 --- a/be/src/olap/tablet_meta.cpp +++ b/be/src/olap/tablet_meta.cpp @@ -294,6 +294,9 @@ TabletMeta::TabletMeta(int64_t table_id, int64_t partition_id, int64_t tablet_id case TIndexType::INVERTED: index_pb->set_index_type(IndexType::INVERTED); break; + case TIndexType::ANN: + index_pb->set_index_type(IndexType::ANN); + break; case TIndexType::BLOOMFILTER: index_pb->set_index_type(IndexType::BLOOMFILTER); break; diff --git a/be/src/olap/tablet_reader.cpp b/be/src/olap/tablet_reader.cpp index fea6efed0bd853..7eb53414ddbfd5 100644 --- a/be/src/olap/tablet_reader.cpp +++ b/be/src/olap/tablet_reader.cpp @@ -270,6 +270,7 @@ Status TabletReader::_capture_rs_readers(const ReaderParams& read_params) { _reader_context.virtual_column_exprs = read_params.virtual_column_exprs; _reader_context.vir_cid_to_idx_in_block = read_params.vir_cid_to_idx_in_block; _reader_context.vir_col_idx_to_type = read_params.vir_col_idx_to_type; + _reader_context.ann_topn_runtime = read_params.ann_topn_runtime; return Status::OK(); } diff --git a/be/src/olap/tablet_reader.h b/be/src/olap/tablet_reader.h index 820cc9f2f4b55f..81fb03ef7b2411 100644 --- a/be/src/olap/tablet_reader.h +++ b/be/src/olap/tablet_reader.h @@ -201,6 +201,7 @@ class TabletReader { std::shared_ptr score_runtime; CollectionStatisticsPtr collection_statistics; + std::shared_ptr ann_topn_runtime; }; TabletReader() = default; diff --git a/be/src/olap/tablet_schema.cpp b/be/src/olap/tablet_schema.cpp index 13d6b8776a0b69..3d94e982954d83 100644 --- a/be/src/olap/tablet_schema.cpp +++ b/be/src/olap/tablet_schema.cpp @@ -874,6 +874,9 @@ void TabletIndex::init_from_thrift(const TOlapTableIndex& index, case TIndexType::INVERTED: _index_type = IndexType::INVERTED; break; + case TIndexType::ANN: + _index_type = IndexType::ANN; + break; case TIndexType::BLOOMFILTER: _index_type = IndexType::BLOOMFILTER; break; @@ -901,6 +904,9 @@ void TabletIndex::init_from_thrift(const TOlapTableIndex& index, case TIndexType::INVERTED: _index_type = IndexType::INVERTED; break; + case TIndexType::ANN: + _index_type = IndexType::ANN; + break; case TIndexType::BLOOMFILTER: _index_type = IndexType::BLOOMFILTER; break; @@ -1614,6 +1620,7 @@ std::vector TabletSchema::inverted_indexs(const TabletColumn if (!segment_v2::IndexColumnWriter::check_support_inverted_index(col)) { return {}; } + // TODO use more efficient impl // Use parent id if unique not assigned, this could happend when accessing subcolumns of variants int32_t col_unique_id = col.is_extracted_column() ? col.parent_unique_id() : col.unique_id(); @@ -1656,6 +1663,32 @@ std::vector TabletSchema::inverted_indexs(const TabletColumn return result; } +const TabletIndex* TabletSchema::ann_index(int32_t col_unique_id, + const std::string& suffix_path) const { + for (size_t i = 0; i < _indexes.size(); i++) { + if (_indexes[i]->index_type() == IndexType::ANN) { + for (int32_t id : _indexes[i]->col_unique_ids()) { + if (id == col_unique_id && + _indexes[i]->get_index_suffix() == escape_for_path_name(suffix_path)) { + return _indexes[i].get(); + } + } + } + } + return nullptr; +} + +const TabletIndex* TabletSchema::ann_index(const TabletColumn& col) const { + // Some columns(Float, Double, JSONB ...) from the variant do not support inverted index + if (!segment_v2::IndexColumnWriter::check_support_ann_index(col)) { + return nullptr; + } + // TODO use more efficient impl + // Use parent id if unique not assigned, this could happend when accessing subcolumns of variants + int32_t col_unique_id = col.is_extracted_column() ? col.parent_unique_id() : col.unique_id(); + return ann_index(col_unique_id, escape_for_path_name(col.suffix_path())); +} + bool TabletSchema::has_ngram_bf_index(int32_t col_unique_id) const { IndexKey index_key(IndexType::NGRAM_BF, col_unique_id, ""); auto it = _col_id_suffix_to_index.find(index_key); diff --git a/be/src/olap/tablet_schema.h b/be/src/olap/tablet_schema.h index 82ec82ea99901b..9d3d740cedb2ee 100644 --- a/be/src/olap/tablet_schema.h +++ b/be/src/olap/tablet_schema.h @@ -76,6 +76,11 @@ class TabletColumn : public MetadataAdder { TabletColumn(FieldAggregationMethod agg, FieldType filed_type, bool is_nullable); TabletColumn(FieldAggregationMethod agg, FieldType filed_type, bool is_nullable, int32_t unique_id, size_t length); + +#ifdef BE_TEST + virtual ~TabletColumn() = default; +#endif + void init_from_pb(const ColumnPB& column); void init_from_thrift(const TColumn& column); void to_schema_pb(ColumnPB* column) const; @@ -88,7 +93,7 @@ class TabletColumn : public MetadataAdder { _col_name = col_name; _col_name_lower_case = to_lower(_col_name); } - FieldType type() const { return _type; } + MOCK_FUNCTION FieldType type() const { return _type; } void set_type(FieldType type) { _type = type; } bool is_key() const { return _is_key; } bool is_nullable() const { return _is_nullable; } @@ -154,7 +159,7 @@ class TabletColumn : public MetadataAdder { void add_sub_column(TabletColumn& sub_column); uint32_t get_subtype_count() const { return _sub_column_count; } - const TabletColumn& get_sub_column(uint64_t i) const { return *_sub_columns[i]; } + MOCK_FUNCTION const TabletColumn& get_sub_column(uint64_t i) const { return *_sub_columns[i]; } const std::vector& get_sub_columns() const { return _sub_columns; } friend bool operator==(const TabletColumn& a, const TabletColumn& b); @@ -290,9 +295,11 @@ class TabletIndex : public MetadataAdder { int64_t index_id() const { return _index_id; } const std::string& index_name() const { return _index_name; } - IndexType index_type() const { return _index_type; } + MOCK_FUNCTION IndexType index_type() const { return _index_type; } const std::vector& col_unique_ids() const { return _col_unique_ids; } - const std::map& properties() const { return _properties; } + MOCK_FUNCTION const std::map& properties() const { + return _properties; + } int32_t get_gram_size() const { if (_properties.contains("gram_size")) { return std::stoi(_properties.at("gram_size")); @@ -487,6 +494,18 @@ class TabletSchema : public MetadataAdder { } return false; } + + bool has_ann_index() const { + for (const auto& index : _indexes) { + if (index->index_type() == IndexType::ANN) { + if (!index->col_unique_ids().empty() && index->col_unique_ids()[0] >= 0) { + return true; + } + } + } + return false; + } + bool has_inverted_index_with_index_id(int64_t index_id) const; void update_index(const TabletColumn& column, const IndexType& index_type, @@ -496,6 +515,12 @@ class TabletSchema : public MetadataAdder { std::vector inverted_indexs(int32_t col_unique_id, const std::string& suffix_path = "") const; + const TabletIndex* ann_index(const TabletColumn& col) const; + + // Regardless of whether this column supports inverted index + // TabletIndex information will be returned as long as it exists. + const TabletIndex* ann_index(int32_t col_unique_id, const std::string& suffix_path = "") const; + std::vector inverted_index_by_field_pattern( int32_t col_unique_id, const std::string& field_pattern) const; diff --git a/be/src/olap/task/engine_storage_migration_task.cpp b/be/src/olap/task/engine_storage_migration_task.cpp index 435236108ea9b3..f3746be8755ae2 100644 --- a/be/src/olap/task/engine_storage_migration_task.cpp +++ b/be/src/olap/task/engine_storage_migration_task.cpp @@ -426,7 +426,7 @@ Status EngineStorageMigrationTask::_copy_index_and_data_files( return status; } } - } else if (tablet_schema.has_inverted_index()) { + } else if (tablet_schema.has_inverted_index() || tablet_schema.has_ann_index()) { auto index_file = InvertedIndexDescriptor::get_index_file_path_v2( InvertedIndexDescriptor::get_index_file_path_prefix(segment_file_path)); auto snapshot_segment_index_file_path = diff --git a/be/src/pipeline/exec/olap_scan_operator.cpp b/be/src/pipeline/exec/olap_scan_operator.cpp index 8503db7e0d33a8..dfb8386877b516 100644 --- a/be/src/pipeline/exec/olap_scan_operator.cpp +++ b/be/src/pipeline/exec/olap_scan_operator.cpp @@ -28,10 +28,12 @@ #include "cloud/cloud_tablet_hotspot.h" #include "cloud/config.h" #include "olap/parallel_scanner_builder.h" +#include "olap/rowset/segment_v2/ann_index/ann_topn_runtime.h" #include "olap/storage_engine.h" #include "olap/tablet_manager.h" #include "pipeline/exec/scan_operator.h" #include "pipeline/query_cache/query_cache.h" +#include "runtime/runtime_state.h" #include "runtime_filter/runtime_filter_consumer_helper.h" #include "service/backend_options.h" #include "util/runtime_profile.h" @@ -59,6 +61,22 @@ Status OlapScanLocalState::init(RuntimeState* state, LocalStateInfo& info) { _score_runtime = vectorized::ScoreRuntime::create_shared(ordering_expr_ctx, asc, limit); } + if (olap_scan_node.__isset.ann_sort_info || olap_scan_node.__isset.ann_sort_limit) { + DCHECK(olap_scan_node.__isset.ann_sort_info); + DCHECK(olap_scan_node.__isset.ann_sort_limit); + DCHECK(olap_scan_node.ann_sort_info.ordering_exprs.size() == 1); + const doris::TExpr& ordering_expr = olap_scan_node.ann_sort_info.ordering_exprs.front(); + DCHECK(ordering_expr.nodes[0].__isset.slot_ref); + DCHECK(ordering_expr.nodes[0].slot_ref.is_virtual_slot); + DCHECK(olap_scan_node.ann_sort_info.is_asc_order.size() == 1); + const bool asc = olap_scan_node.ann_sort_info.is_asc_order[0]; + const size_t limit = olap_scan_node.ann_sort_limit; + std::shared_ptr ordering_expr_ctx; + RETURN_IF_ERROR(vectorized::VExpr::create_expr_tree(ordering_expr, ordering_expr_ctx)); + _ann_topn_runtime = + segment_v2::AnnTopNRuntime::create_shared(asc, limit, ordering_expr_ctx); + } + RETURN_IF_ERROR(Base::init(state, info)); RETURN_IF_ERROR(_sync_cloud_tablets(state)); return Status::OK(); @@ -257,6 +275,55 @@ Status OlapScanLocalState::_init_profile() { _index_filter_profile = std::make_unique("IndexFilter"); _scanner_profile->add_child(_index_filter_profile.get(), true, nullptr); + /* + SegmentIterator: + - AnnIndexLoadCosts: 102.262us + - AnnIndexRangeSearchCosts: 0ns + - AnnIndexRangeSearchFiltered: 0 + - AnnIndexTopNCosts: 658.303ms + - AnnIndexTopNFiltered: 9.49791M (9497910) + - AnnIndexTopNSearchCnt: 209ns + */ + _ann_range_search_filter_counter = + ADD_COUNTER(_segment_profile, "AnnIndexRangeSearchFiltered", TUnit::UNIT); + _ann_topn_filter_counter = ADD_COUNTER(_segment_profile, "AnnIndexTopNFiltered", TUnit::UNIT); + + _ann_topn_search_costs = ADD_TIMER(_segment_profile, "AnnIndexTopNSearchCosts"); + _ann_topn_search_cnt = ADD_COUNTER(_segment_profile, "AnnIndexTopNSearchCnt", TUnit::UNIT); + _ann_range_search_costs = ADD_TIMER(_segment_profile, "AnnIndexRangeSearchCosts"); + _ann_range_search_cnt = ADD_COUNTER(_segment_profile, "AnnIndexRangeSearchCnt", TUnit::UNIT); + + // Detailed ANN timers (TopN) + // Create child timers under AnnIndexTopNSearchCosts for better readability + _ann_topn_engine_search_costs = ADD_CHILD_TIMER( + _segment_profile, "AnnIndexTopNEngineSearchCosts", "AnnIndexTopNSearchCosts"); + _ann_index_load_costs = ADD_TIMER(_segment_profile, "AnnIndexLoadCosts"); + _ann_topn_post_process_costs = ADD_CHILD_TIMER( + _segment_profile, "AnnIndexTopNResultPostProcessCosts", "AnnIndexTopNSearchCosts"); + _ann_topn_pre_process_costs = ADD_CHILD_TIMER( + _segment_profile, "AnnIndexTopNEnginePrepareCosts", "AnnIndexTopNSearchCosts"); + // Detailed ANN timers (Range) + // Create child timers under AnnIndexRangeSearchCosts to mirror TopN hierarchy + _ann_range_engine_search_costs = ADD_CHILD_TIMER( + _segment_profile, "AnnIndexRangeEngineSearchCosts", "AnnIndexRangeSearchCosts"); + _ann_range_post_process_costs = ADD_CHILD_TIMER( + _segment_profile, "AnnIndexRangeResultPostProcessCosts", "AnnIndexRangeSearchCosts"); + _ann_range_pre_process_costs = ADD_CHILD_TIMER( + _segment_profile, "AnnIndexRangeEnginePrepareCosts", "AnnIndexRangeSearchCosts"); + // Conversion inside FAISS wrappers (TopN): two separate sub counters under post process + _ann_topn_engine_convert_costs = + ADD_CHILD_TIMER(_segment_profile, "AnnIndexTopNEngineConvertCosts", + "AnnIndexTopNResultPostProcessCosts"); + _ann_range_engine_convert_costs = + ADD_CHILD_TIMER(_segment_profile, "AnnIndexRangeEngineConvertCosts", + "AnnIndexRangeResultPostProcessCosts"); + // Keep this as a child of post process to show the sum for Doris-side handling + _ann_topn_result_convert_costs = + ADD_CHILD_TIMER(_segment_profile, "AnnIndexTopNResultConvertCosts", + "AnnIndexTopNResultPostProcessCosts"); + _ann_range_result_convert_costs = + ADD_CHILD_TIMER(_segment_profile, "AnnIndexRangeResultConvertCosts", + "AnnIndexRangeResultPostProcessCosts"); return Status::OK(); } @@ -417,6 +484,7 @@ Status OlapScanLocalState::_init_scanners(std::list* sc auto* olap_scanner = assert_cast(scanner.get()); RETURN_IF_ERROR(olap_scanner->init(state(), _conjuncts)); } + return Status::OK(); } @@ -670,6 +738,10 @@ Status OlapScanLocalState::open(RuntimeState* state) { RETURN_IF_ERROR(_score_runtime->prepare(state, p.intermediate_row_desc())); } + if (_ann_topn_runtime) { + RETURN_IF_ERROR(_ann_topn_runtime->prepare(state, p.intermediate_row_desc())); + } + RETURN_IF_ERROR(ScanLocalState::open(state)); return Status::OK(); diff --git a/be/src/pipeline/exec/olap_scan_operator.h b/be/src/pipeline/exec/olap_scan_operator.h index 42251be75459d8..7cb48e88d7c8fa 100644 --- a/be/src/pipeline/exec/olap_scan_operator.h +++ b/be/src/pipeline/exec/olap_scan_operator.h @@ -26,6 +26,7 @@ #include "olap/tablet_reader.h" #include "operator.h" #include "pipeline/exec/scan_operator.h" +#include "util/runtime_profile.h" namespace doris::vectorized { class OlapScanner; @@ -209,6 +210,31 @@ class OlapScanLocalState final : public ScanLocalState { RuntimeProfile::Counter* _inverted_index_analyzer_timer = nullptr; RuntimeProfile::Counter* _inverted_index_lookup_timer = nullptr; + RuntimeProfile::Counter* _ann_topn_filter_counter = nullptr; + // topn_search_costs = index_load_costs + engine_search_costs + pre_process_costs + post_process_costs + RuntimeProfile::Counter* _ann_topn_search_costs = nullptr; + RuntimeProfile::Counter* _ann_topn_search_cnt = nullptr; + + RuntimeProfile::Counter* _ann_index_load_costs = nullptr; + RuntimeProfile::Counter* _ann_topn_pre_process_costs = nullptr; + RuntimeProfile::Counter* _ann_topn_engine_search_costs = nullptr; + RuntimeProfile::Counter* _ann_topn_post_process_costs = nullptr; + // post_process_costs = engine_convert_costs + result_convert_costs + RuntimeProfile::Counter* _ann_topn_engine_convert_costs = nullptr; + RuntimeProfile::Counter* _ann_topn_result_convert_costs = nullptr; + + RuntimeProfile::Counter* _ann_range_search_filter_counter = nullptr; + // range_Search_costs = index_load_costs + engine_search_costs + pre_process_costs + post_process_costs + RuntimeProfile::Counter* _ann_range_search_costs = nullptr; + RuntimeProfile::Counter* _ann_range_search_cnt = nullptr; + + RuntimeProfile::Counter* _ann_range_pre_process_costs = nullptr; + RuntimeProfile::Counter* _ann_range_engine_search_costs = nullptr; + RuntimeProfile::Counter* _ann_range_post_process_costs = nullptr; + + RuntimeProfile::Counter* _ann_range_engine_convert_costs = nullptr; + RuntimeProfile::Counter* _ann_range_result_convert_costs = nullptr; + RuntimeProfile::Counter* _output_index_result_column_timer = nullptr; // number of segment filtered by column stat when creating seg iterator diff --git a/be/src/pipeline/exec/operator.h b/be/src/pipeline/exec/operator.h index 1dea1253aeaf19..87e1feb4743bab 100644 --- a/be/src/pipeline/exec/operator.h +++ b/be/src/pipeline/exec/operator.h @@ -51,6 +51,7 @@ class TDataSink; namespace vectorized { class AsyncResultWriter; class ScoreRuntime; +class AnnTopNRuntime; } // namespace vectorized } // namespace doris @@ -275,6 +276,7 @@ class PipelineXLocalStateBase { vectorized::VExprContextSPtrs _conjuncts; vectorized::VExprContextSPtrs _projections; std::shared_ptr _score_runtime; + std::shared_ptr _ann_topn_runtime; // Used in common subexpression elimination to compute intermediate results. std::vector _intermediate_projections; diff --git a/be/src/pipeline/exec/scan_operator.h b/be/src/pipeline/exec/scan_operator.h index 82493176f95dae..51637ef5be8d90 100644 --- a/be/src/pipeline/exec/scan_operator.h +++ b/be/src/pipeline/exec/scan_operator.h @@ -133,7 +133,7 @@ class ScanLocalState : public ScanLocalStateBase { : ScanLocalStateBase(state, parent) {} ~ScanLocalState() override = default; - Status init(RuntimeState* state, LocalStateInfo& info) override; + virtual Status init(RuntimeState* state, LocalStateInfo& info) override; virtual Status open(RuntimeState* state) override; diff --git a/be/src/runtime/runtime_state.h b/be/src/runtime/runtime_state.h index cb5c8db291f4bb..1d27944da314bb 100644 --- a/be/src/runtime/runtime_state.h +++ b/be/src/runtime/runtime_state.h @@ -49,6 +49,7 @@ #include "runtime/workload_group/workload_group.h" #include "util/debug_util.h" #include "util/runtime_profile.h" +#include "vec/runtime/vector_search_user_params.h" namespace doris { class RuntimeFilter; @@ -670,6 +671,11 @@ class RuntimeState { std::shared_ptr& get_id_file_map() { return _id_file_map; } void set_id_file_map(); + VectorSearchUserParams get_vector_search_params() const { + return VectorSearchUserParams(_query_options.hnsw_ef_search, + _query_options.hnsw_check_relative_distance, + _query_options.hnsw_bounded_queue); + } private: Status create_error_log_file(); diff --git a/be/src/service/backend_service.cpp b/be/src/service/backend_service.cpp index 4ea4c4255ee83a..b7bd4de1a186cf 100644 --- a/be/src/service/backend_service.cpp +++ b/be/src/service/backend_service.cpp @@ -451,7 +451,7 @@ void _ingest_binlog(StorageEngine& engine, IngestBinlogArg* arg) { } } else { for (int64_t segment_index = 0; segment_index < num_segments; ++segment_index) { - if (tablet_schema->has_inverted_index()) { + if (tablet_schema->has_inverted_index() || tablet_schema->has_ann_index()) { auto get_segment_index_file_size_url = fmt::format( "{}?method={}&tablet_id={}&rowset_id={}&segment_index={}&segment_index_id={" "}", diff --git a/be/src/util/doris_metrics.cpp b/be/src/util/doris_metrics.cpp index c0c0660b4c08d2..60c16fbf9a96e9 100644 --- a/be/src/util/doris_metrics.cpp +++ b/be/src/util/doris_metrics.cpp @@ -240,6 +240,11 @@ DEFINE_GAUGE_CORE_METRIC_PROTOTYPE_2ARG(runtime_filter_consumer_timeout_num, Met DEFINE_COUNTER_METRIC_PROTOTYPE_2ARG(get_remote_tablet_slow_time_ms, MetricUnit::MILLISECONDS); DEFINE_COUNTER_METRIC_PROTOTYPE_2ARG(get_remote_tablet_slow_cnt, MetricUnit::NOUNIT); +DEFINE_COUNTER_METRIC_PROTOTYPE_2ARG(ann_index_load_costs_ms, MetricUnit::MILLISECONDS); +DEFINE_COUNTER_METRIC_PROTOTYPE_2ARG(ann_index_load_cnt, MetricUnit::NOUNIT); +DEFINE_COUNTER_METRIC_PROTOTYPE_2ARG(ann_index_search_costs_ms, MetricUnit::MILLISECONDS); +DEFINE_COUNTER_METRIC_PROTOTYPE_2ARG(ann_index_search_cnt, MetricUnit::NOUNIT); + const std::string DorisMetrics::_s_registry_name = "doris_be"; const std::string DorisMetrics::_s_hook_name = "doris_metrics"; @@ -397,6 +402,11 @@ DorisMetrics::DorisMetrics() : _metric_registry(_s_registry_name) { INT_COUNTER_METRIC_REGISTER(_server_metric_entity, get_remote_tablet_slow_cnt); INT_COUNTER_METRIC_REGISTER(_server_metric_entity, pipeline_task_queue_size); + + INT_COUNTER_METRIC_REGISTER(_server_metric_entity, ann_index_load_costs_ms); + INT_COUNTER_METRIC_REGISTER(_server_metric_entity, ann_index_load_cnt); + INT_COUNTER_METRIC_REGISTER(_server_metric_entity, ann_index_search_costs_ms); + INT_COUNTER_METRIC_REGISTER(_server_metric_entity, ann_index_search_cnt); } void DorisMetrics::initialize(bool init_system_metrics, const std::set& disk_devices, diff --git a/be/src/util/doris_metrics.h b/be/src/util/doris_metrics.h index 10f1ff0fca25f3..ddd79fc1ed2916 100644 --- a/be/src/util/doris_metrics.h +++ b/be/src/util/doris_metrics.h @@ -249,6 +249,10 @@ class DorisMetrics { IntCounter* scanner_cnt = nullptr; IntCounter* scanner_task_cnt = nullptr; IntCounter* pipeline_task_queue_size = nullptr; + IntCounter* ann_index_load_costs_ms = nullptr; + IntCounter* ann_index_load_cnt = nullptr; + IntCounter* ann_index_search_costs_ms = nullptr; + IntCounter* ann_index_search_cnt = nullptr; IntGauge* runtime_filter_consumer_num = nullptr; IntGauge* runtime_filter_consumer_ready_num = nullptr; diff --git a/be/src/vec/columns/column_dummy.h b/be/src/vec/columns/column_dummy.h index 2136b2d4046acf..8b7ce9e3d0a0c4 100644 --- a/be/src/vec/columns/column_dummy.h +++ b/be/src/vec/columns/column_dummy.h @@ -38,7 +38,7 @@ class IColumnDummy : public IColumn { public: virtual MutableColumnPtr clone_dummy(size_t s_) const = 0; - MutableColumnPtr clone_resized(size_t s) const override { return clone_dummy(s); } + MutableColumnPtr clone_resized(size_t size) const override { return clone_dummy(size); } size_t size() const override { return s; } void resize(size_t _s) override { s = _s; } void insert_default() override { ++s; } diff --git a/be/src/vec/exec/scan/olap_scanner.cpp b/be/src/vec/exec/scan/olap_scanner.cpp index e2ac9e902ebe20..55d1a776a35ddf 100644 --- a/be/src/vec/exec/scan/olap_scanner.cpp +++ b/be/src/vec/exec/scan/olap_scanner.cpp @@ -94,9 +94,11 @@ OlapScanner::OlapScanner(pipeline::ScanLocalStateBase* parent, OlapScanner::Para .vir_cid_to_idx_in_block {}, .vir_col_idx_to_type {}, .score_runtime {}, - .collection_statistics {}}) { + .collection_statistics {}, + .ann_topn_runtime {}}) { _tablet_reader_params.set_read_source(std::move(params.read_source)); _has_prepared = false; + _vector_search_params = params.state->get_vector_search_params(); } static std::string read_columns_to_string(TabletSchemaSPtr tablet_schema, @@ -137,6 +139,7 @@ Status OlapScanner::prepare() { VExprContextSPtr context; RETURN_IF_ERROR(ctx->clone(_state, context)); _common_expr_ctxs_push_down.emplace_back(context); + context->prepare_ann_range_search(_vector_search_params); } for (auto pair : local_state->_slot_id_to_virtual_column_expr) { @@ -150,6 +153,9 @@ Status OlapScanner::prepare() { _slot_id_to_col_type = local_state->_slot_id_to_col_type; _score_runtime = local_state->_score_runtime; + _score_runtime = local_state->_score_runtime; + _ann_topn_runtime = local_state->_ann_topn_runtime; + // set limit to reduce end of rowset and segment mem use _tablet_reader = std::make_unique(); // batch size is passed down to segment iterator, use _state->batch_size() @@ -319,6 +325,7 @@ Status OlapScanner::_init_tablet_reader_params( _tablet_reader_params.score_runtime = _score_runtime; _tablet_reader_params.output_columns = ((pipeline::OlapScanLocalState*)_local_state)->_maybe_read_column_ids; + _tablet_reader_params.ann_topn_runtime = _ann_topn_runtime; for (const auto& ele : ((pipeline::OlapScanLocalState*)_local_state)->_cast_types_for_variants) { _tablet_reader_params.target_cast_type_for_variants[ele.first] = ele.second; @@ -540,6 +547,7 @@ Status OlapScanner::_init_return_columns() { if (_return_columns.empty()) { return Status::InternalError("failed to build storage scanner, no materialized slot!"); } + return Status::OK(); } @@ -792,6 +800,47 @@ void OlapScanner::_collect_profile_before_close() { tablet->query_scan_bytes->increment(local_state->_read_uncompressed_counter->value()); tablet->query_scan_rows->increment(local_state->_scan_rows->value()); tablet->query_scan_count->increment(1); + + COUNTER_UPDATE(local_state->_ann_range_search_filter_counter, + stats.rows_ann_index_range_filtered); + COUNTER_UPDATE(local_state->_ann_topn_filter_counter, stats.rows_ann_index_topn_filtered); + COUNTER_UPDATE(local_state->_ann_index_load_costs, stats.ann_index_load_ns); + COUNTER_UPDATE(local_state->_ann_range_search_costs, stats.ann_index_range_search_ns); + COUNTER_UPDATE(local_state->_ann_range_search_cnt, stats.ann_index_range_search_cnt); + COUNTER_UPDATE(local_state->_ann_range_engine_search_costs, stats.ann_range_engine_search_ns); + // Engine prepare before search + COUNTER_UPDATE(local_state->_ann_range_pre_process_costs, stats.ann_range_pre_process_ns); + // Post process parent: Doris result process + engine convert + COUNTER_UPDATE(local_state->_ann_range_post_process_costs, + stats.ann_range_result_convert_ns + stats.ann_range_engine_convert_ns); + // Engine convert (child under post-process) + COUNTER_UPDATE(local_state->_ann_range_engine_convert_costs, stats.ann_range_engine_convert_ns); + // Doris-side result convert (child under post-process) + COUNTER_UPDATE(local_state->_ann_range_result_convert_costs, stats.ann_range_result_convert_ns); + + COUNTER_UPDATE(local_state->_ann_topn_search_costs, stats.ann_topn_search_ns); + COUNTER_UPDATE(local_state->_ann_topn_search_cnt, stats.ann_index_topn_search_cnt); + + // Detailed ANN timers + // ANN TopN timers with hierarchy + // Engine search time (FAISS) + COUNTER_UPDATE(local_state->_ann_topn_engine_search_costs, + stats.ann_index_topn_engine_search_ns); + // Engine prepare time (allocations/buffer setup before search) + COUNTER_UPDATE(local_state->_ann_topn_pre_process_costs, + stats.ann_index_topn_engine_prepare_ns); + // Post process parent includes Doris result processing + engine convert + COUNTER_UPDATE(local_state->_ann_topn_post_process_costs, + stats.ann_index_topn_result_process_ns + stats.ann_index_topn_engine_convert_ns); + // Engine-side conversion time inside FAISS wrappers (child under post-process) + COUNTER_UPDATE(local_state->_ann_topn_engine_convert_costs, + stats.ann_index_topn_engine_convert_ns); + + // Doris-side result convert costs (show separately as another child counter); use pure process time + COUNTER_UPDATE(local_state->_ann_topn_result_convert_costs, + stats.ann_index_topn_result_process_ns); + + // Overhead counter removed; precise instrumentation is reported via engine_prepare above. } } // namespace doris::vectorized diff --git a/be/src/vec/exec/scan/olap_scanner.h b/be/src/vec/exec/scan/olap_scanner.h index c72614d636f0e8..f12b3f37444036 100644 --- a/be/src/vec/exec/scan/olap_scanner.h +++ b/be/src/vec/exec/scan/olap_scanner.h @@ -118,6 +118,10 @@ class OlapScanner : public Scanner { // The idx of vir_col in block to its data type. std::map _vir_col_idx_to_type; std::shared_ptr _score_runtime; + + std::shared_ptr _ann_topn_runtime; + + VectorSearchUserParams _vector_search_params; }; } // namespace vectorized } // namespace doris diff --git a/be/src/vec/exprs/vectorized_fn_call.cpp b/be/src/vec/exprs/vectorized_fn_call.cpp index ef6a253b7b1038..941af8ca7a29b6 100644 --- a/be/src/vec/exprs/vectorized_fn_call.cpp +++ b/be/src/vec/exprs/vectorized_fn_call.cpp @@ -19,21 +19,40 @@ #include #include // IWYU pragma: keep +#include #include +#include #include #include "common/config.h" +#include "common/logging.h" #include "common/status.h" #include "common/utils.h" +#include "olap/rowset/segment_v2/ann_index/ann_index.h" +#include "olap/rowset/segment_v2/ann_index/ann_index_iterator.h" +#include "olap/rowset/segment_v2/ann_index/ann_search_params.h" +#include "olap/rowset/segment_v2/column_reader.h" +#include "olap/rowset/segment_v2/index_reader.h" +#include "olap/rowset/segment_v2/virtual_column_iterator.h" #include "pipeline/pipeline_task.h" #include "runtime/runtime_state.h" #include "udf/udf.h" #include "vec/columns/column.h" +#include "vec/columns/column_array.h" +#include "vec/columns/column_nullable.h" +#include "vec/columns/column_vector.h" #include "vec/core/block.h" +#include "vec/core/column_numbers.h" +#include "vec/core/types.h" #include "vec/data_types/data_type.h" #include "vec/data_types/data_type_agg_state.h" +#include "vec/exprs/varray_literal.h" +#include "vec/exprs/vcast_expr.h" #include "vec/exprs/vexpr_context.h" +#include "vec/exprs/virtual_slot_ref.h" +#include "vec/exprs/vliteral.h" +#include "vec/functions/array/function_array_distance.h" #include "vec/functions/function_agg_state.h" #include "vec/functions/function_fake.h" #include "vec/functions/function_java_udf.h" @@ -52,6 +71,12 @@ namespace doris::vectorized { const std::string AGG_STATE_SUFFIX = "_state"; +// Now left child is a function call, we need to check if it is a distance function +const static std::set DISTANCE_FUNCS = {L2DistanceApproximate::name, + InnerProductApproximate::name}; +const static std::set OPS_FOR_ANN_RANGE_SEARCH = { + TExprOpcode::GE, TExprOpcode::LE, TExprOpcode::LE, TExprOpcode::GT, TExprOpcode::LT}; + VectorizedFnCall::VectorizedFnCall(const TExprNode& node) : VExpr(node) {} Status VectorizedFnCall::prepare(RuntimeState* state, const RowDescriptor& desc, @@ -250,6 +275,10 @@ const std::string& VectorizedFnCall::expr_name() const { return _expr_name; } +std::string VectorizedFnCall::function_name() const { + return _function_name; +} + std::string VectorizedFnCall::debug_string() const { std::stringstream out; out << "VectorizedFn["; @@ -301,5 +330,253 @@ bool VectorizedFnCall::equals(const VExpr& other) { return true; } +/* + FuncationCall(LE/LT/GE/GT) + |---------------- + | | + | | + VirtualSlotRef Float64Literal + | + | + FuncationCall + |---------------- + | | + | | + CastToArray ArrayLiteral + | + | + SlotRef +*/ + +void VectorizedFnCall::prepare_ann_range_search( + const doris::VectorSearchUserParams& user_params, + segment_v2::AnnRangeSearchRuntime& range_search_runtime, bool& suitable_for_ann_index) { + if (!suitable_for_ann_index) { + return; + } + + if (OPS_FOR_ANN_RANGE_SEARCH.find(this->op()) == OPS_FOR_ANN_RANGE_SEARCH.end()) { + suitable_for_ann_index = false; + // Not a range search function. + return; + } + + range_search_runtime.is_le_or_lt = + (this->op() == TExprOpcode::LE || this->op() == TExprOpcode::LT); + + DCHECK(_children.size() == 2); + + auto left_child = get_child(0); + auto right_child = get_child(1); + + // Return type of L2Distance is always double. + auto right_literal = std::dynamic_pointer_cast(right_child); + if (right_literal == nullptr) { + suitable_for_ann_index = false; + // Right child is not a literal. + return; + } + + auto right_col = right_literal->get_column_ptr()->convert_to_full_column_if_const(); + auto right_type = right_literal->get_data_type(); + if (right_type->get_primitive_type() != PrimitiveType::TYPE_DOUBLE) { + suitable_for_ann_index = false; + // Right child is not a Float64Literal. + return; + } + + const ColumnFloat64* cf64_right = assert_cast(right_col.get()); + range_search_runtime.radius = cf64_right->get_data()[0]; + + std::shared_ptr function_call; + auto vir_slot_ref = std::dynamic_pointer_cast(left_child); + if (vir_slot_ref != nullptr) { + DCHECK(vir_slot_ref->get_virtual_column_expr() != nullptr); + function_call = std::dynamic_pointer_cast( + vir_slot_ref->get_virtual_column_expr()); + } else { + function_call = std::dynamic_pointer_cast(left_child); + } + + if (function_call == nullptr) { + suitable_for_ann_index = false; + // Left child is not a function call. + return; + } + + if (DISTANCE_FUNCS.find(function_call->_function_name) == DISTANCE_FUNCS.end()) { + // Left child is not a approximate distance function. Got function_call->_function_name + suitable_for_ann_index = false; + return; + } else { + // Strip the _approximate suffix. + std::string metric_name = function_call->_function_name; + metric_name = metric_name.substr(0, metric_name.size() - 12); + range_search_runtime.metric_type = segment_v2::string_to_metric(metric_name); + } + + UInt16 idx_of_cast_to_array = 0; + UInt16 idx_of_array_literal = 0; + for (UInt16 i = 0; i < function_call->get_num_children(); ++i) { + auto child = function_call->get_child(i); + if (std::dynamic_pointer_cast(child) != nullptr) { + idx_of_cast_to_array = i; + } else if (std::dynamic_pointer_cast(child) != nullptr) { + idx_of_array_literal = i; + } + } + + std::shared_ptr cast_to_array_expr = + std::dynamic_pointer_cast(function_call->get_child(idx_of_cast_to_array)); + std::shared_ptr array_literal = std::dynamic_pointer_cast( + function_call->get_child(idx_of_array_literal)); + + if (cast_to_array_expr == nullptr || array_literal == nullptr) { + suitable_for_ann_index = false; + // Cast to array expr or array literal is null. + return; + } + + // One of the children is a slot ref, and the other is an array literal, now begin to create search params. + std::shared_ptr slot_ref = + std::dynamic_pointer_cast(cast_to_array_expr->get_child(0)); + if (slot_ref == nullptr) { + suitable_for_ann_index = false; + // Cast to array expr's child is not a slot ref. + return; + } + + range_search_runtime.src_col_idx = slot_ref->column_id(); + range_search_runtime.dst_col_idx = vir_slot_ref == nullptr ? -1 : vir_slot_ref->column_id(); + auto col_const = array_literal->get_column_ptr(); + auto col_array = col_const->convert_to_full_column_if_const(); + const ColumnArray* array_col = assert_cast(col_array.get()); + DCHECK(array_col->size() == 1); + size_t dim = array_col->get_offsets()[0]; + range_search_runtime.dim = dim; + range_search_runtime.query_value = std::make_unique(dim); + + const ColumnNullable* cn = assert_cast(array_col->get_data_ptr().get()); + const ColumnFloat64* cf64 = + assert_cast(cn->get_nested_column_ptr().get()); + for (size_t i = 0; i < dim; ++i) { + range_search_runtime.query_value[i] = static_cast(cf64->get_data()[i]); + } + range_search_runtime.is_ann_range_search = true; + range_search_runtime.user_params = user_params; + VLOG_DEBUG << fmt::format("Ann range search params: {}", range_search_runtime.to_string()); + return; +} + +Status VectorizedFnCall::evaluate_ann_range_search( + const segment_v2::AnnRangeSearchRuntime& range_search_runtime, + const std::vector>& cid_to_index_iterators, + const std::vector& idx_to_cid, + const std::vector>& column_iterators, + roaring::Roaring& row_bitmap, segment_v2::AnnIndexStats& ann_index_stats) { + if (range_search_runtime.is_ann_range_search == false) { + return Status::OK(); + } + + VLOG_DEBUG << fmt::format("Try apply ann range search. Local search params: {}", + range_search_runtime.to_string()); + size_t origin_num = row_bitmap.cardinality(); + + int idx_in_block = static_cast(range_search_runtime.src_col_idx); + DCHECK(idx_in_block < idx_to_cid.size()) + << "idx_in_block: " << idx_in_block << ", idx_to_cid.size(): " << idx_to_cid.size(); + + ColumnId src_col_cid = idx_to_cid[idx_in_block]; + DCHECK(src_col_cid < cid_to_index_iterators.size()); + segment_v2::IndexIterator* index_iterator = cid_to_index_iterators[src_col_cid].get(); + if (index_iterator == nullptr) { + // No index iterator for column cid + return Status::OK(); + } + + segment_v2::AnnIndexIterator* ann_index_iterator = + dynamic_cast(index_iterator); + if (ann_index_iterator == nullptr) { + // No ann index iterator for column cid + return Status::OK(); + } + DCHECK(ann_index_iterator->get_reader(AnnIndexReaderType::ANN) != nullptr) + << "Ann index iterator should have reader. Column cid: " << src_col_cid; + std::shared_ptr ann_index_reader = std::dynamic_pointer_cast( + ann_index_iterator->get_reader(segment_v2::AnnIndexReaderType::ANN)); + DCHECK(ann_index_reader != nullptr) + << "Ann index reader should not be null. Column cid: " << src_col_cid; + // Check if metrics type is match. + if (ann_index_reader->get_metric_type() != range_search_runtime.metric_type) { + // Metric type not match, can not execute range search by index. + return Status::OK(); + } + + AnnRangeSearchParams params = range_search_runtime.to_range_search_params(); + + params.roaring = &row_bitmap; + DCHECK(params.roaring != nullptr); + DCHECK(params.query_value != nullptr); + segment_v2::AnnRangeSearchResult result; + auto stats = std::make_unique(); + RETURN_IF_ERROR(ann_index_iterator->range_search(params, range_search_runtime.user_params, + &result, stats.get())); + +#ifndef NDEBUG + if (range_search_runtime.is_le_or_lt == false) { + DCHECK(result.distance == nullptr) << "Should not have distance"; + } +#endif + + DCHECK(result.roaring != nullptr); + row_bitmap = *result.roaring; + + if (params.is_le_or_lt == false) { + DCHECK(result.distance == nullptr); + DCHECK(result.row_ids == nullptr); + } + + // Process virtual column + if (range_search_runtime.dst_col_idx >= 0) { + // Prepare materialization if we can use result from index. + // Typical situation: range search and operator is LE or LT. + if (result.distance != nullptr) { + DCHECK(result.row_ids != nullptr); + ColumnId dst_col_cid = idx_to_cid[range_search_runtime.dst_col_idx]; + DCHECK(dst_col_cid < column_iterators.size()); + DCHECK(column_iterators[dst_col_cid] != nullptr); + segment_v2::ColumnIterator* column_iterator = column_iterators[dst_col_cid].get(); + DCHECK(column_iterator != nullptr); + segment_v2::VirtualColumnIterator* virtual_column_iterator = + dynamic_cast(column_iterator); + DCHECK(virtual_column_iterator != nullptr); + // Now convert distance to column + size_t size = result.roaring->cardinality(); + auto distance_col = ColumnFloat64::create(size); + auto null_map = ColumnUInt8::create(size, 0); + // TODO: Return type of L2DistanceApproximate/InnerProductApproximate should be changed to float. + const float* src = reinterpret_cast(result.distance.get()); + double* dst = distance_col->get_data().data(); + for (size_t i = 0; i < size; ++i) { + dst[i] = static_cast(src[i]); + } + auto nullable_distance_col = + ColumnNullable::create(std::move(distance_col), std::move(null_map)); + virtual_column_iterator->prepare_materialization(std::move(nullable_distance_col), + std::move(result.row_ids)); + } else { + DCHECK(this->op() != TExprOpcode::LE && this->op() != TExprOpcode::LT) + << "Should not have distance"; + } + } + + _has_been_executed = true; + VLOG_DEBUG << fmt::format("Ann range search filtered {} rows, origin {} rows", + origin_num - row_bitmap.cardinality(), origin_num); + + ann_index_stats = *stats; + return Status::OK(); +} + #include "common/compile_check_end.h" } // namespace doris::vectorized diff --git a/be/src/vec/exprs/vectorized_fn_call.h b/be/src/vec/exprs/vectorized_fn_call.h index 6246698e86f215..9d105436459876 100644 --- a/be/src/vec/exprs/vectorized_fn_call.h +++ b/be/src/vec/exprs/vectorized_fn_call.h @@ -22,9 +22,12 @@ #include #include "common/status.h" +#include "olap/rowset/segment_v2/ann_index/ann_range_search_runtime.h" +#include "runtime/runtime_state.h" #include "udf/udf.h" #include "vec/core/column_numbers.h" #include "vec/exprs/vexpr.h" +#include "vec/exprs/vexpr_context.h" #include "vec/exprs/vliteral.h" #include "vec/exprs/vslot_ref.h" #include "vec/functions/function.h" @@ -59,6 +62,7 @@ class VectorizedFnCall : public VExpr { FunctionContext::FunctionStateScope scope) override; void close(VExprContext* context, FunctionContext::FunctionStateScope scope) override; const std::string& expr_name() const override; + std::string function_name() const; std::string debug_string() const override; bool is_constant() const override { if (!_function->is_use_default_implementation_for_constants() || @@ -75,6 +79,17 @@ class VectorizedFnCall : public VExpr { size_t estimate_memory(const size_t rows) override; + Status evaluate_ann_range_search( + const segment_v2::AnnRangeSearchRuntime& runtime, + const std::vector>& cid_to_index_iterators, + const std::vector& idx_to_cid, + const std::vector>& column_iterators, + roaring::Roaring& row_bitmap, segment_v2::AnnIndexStats& ann_index_stats) override; + + void prepare_ann_range_search(const doris::VectorSearchUserParams& params, + segment_v2::AnnRangeSearchRuntime& runtime, + bool& suitable_for_ann_index) override; + protected: FunctionBasePtr _function; std::string _expr_name; diff --git a/be/src/vec/exprs/vexpr.cpp b/be/src/vec/exprs/vexpr.cpp index dc6947893fa100..f404cbbd357ce3 100644 --- a/be/src/vec/exprs/vexpr.cpp +++ b/be/src/vec/exprs/vexpr.cpp @@ -33,6 +33,8 @@ #include "common/config.h" #include "common/exception.h" #include "common/status.h" +#include "olap/rowset/segment_v2/ann_index/ann_search_params.h" +#include "olap/rowset/segment_v2/ann_index/ann_topn_runtime.h" #include "pipeline/pipeline_task.h" #include "runtime/define_primitive_type.h" #include "vec/columns/column_vector.h" @@ -397,7 +399,9 @@ Status VExpr::prepare(RuntimeState* state, const RowDescriptor& row_desc, VExprC RETURN_IF_ERROR(i->prepare(state, row_desc, context)); } --context->_depth_num; +#ifndef BE_TEST _enable_inverted_index_query = state->query_options().enable_inverted_index_query; +#endif return Status::OK(); } @@ -963,5 +967,32 @@ bool VExpr::equals(const VExpr& other) { return false; } +Status VExpr::evaluate_ann_range_search( + const segment_v2::AnnRangeSearchRuntime& runtime, + const std::vector>& index_iterators, + const std::vector& idx_to_cid, + const std::vector>& column_iterators, + roaring::Roaring& row_bitmap, AnnIndexStats& ann_index_stats) { + return Status::OK(); +} + +void VExpr::prepare_ann_range_search(const doris::VectorSearchUserParams& params, + segment_v2::AnnRangeSearchRuntime& range_search_runtime, + bool& suitable_for_ann_index) { + if (!suitable_for_ann_index) { + return; + } + for (auto& child : _children) { + child->prepare_ann_range_search(params, range_search_runtime, suitable_for_ann_index); + if (!suitable_for_ann_index) { + return; + } + } +} + +bool VExpr::has_been_executed() { + return _has_been_executed; +} + #include "common/compile_check_end.h" } // namespace doris::vectorized diff --git a/be/src/vec/exprs/vexpr.h b/be/src/vec/exprs/vexpr.h index 96c1a2097f1369..6b6d8881685053 100644 --- a/be/src/vec/exprs/vexpr.h +++ b/be/src/vec/exprs/vexpr.h @@ -31,6 +31,9 @@ #include #include "common/status.h" +#include "olap/rowset/segment_v2/ann_index/ann_search_params.h" +#include "olap/rowset/segment_v2/column_reader.h" +#include "olap/rowset/segment_v2/index_reader.h" #include "olap/rowset/segment_v2/inverted_index_reader.h" #include "runtime/define_primitive_type.h" #include "runtime/large_int_value.h" @@ -45,6 +48,7 @@ #include "vec/core/types.h" #include "vec/data_types/data_type.h" #include "vec/data_types/data_type_ipv6.h" +#include "vec/exprs/vexpr_context.h" #include "vec/exprs/vexpr_fwd.h" #include "vec/functions/function.h" @@ -56,9 +60,14 @@ class ObjectPool; class RowDescriptor; class RuntimeState; +namespace segment_v2 { +class IndexIterator; +class ColumnIterator; +struct AnnRangeSearchRuntime; +}; // namespace segment_v2 + namespace vectorized { #include "common/compile_check_begin.h" - #define RETURN_IF_ERROR_OR_PREPARED(stmt) \ if (_prepared) { \ return Status::OK(); \ @@ -268,6 +277,26 @@ class VExpr { } } +#ifdef BE_TEST + void set_node_type(TExprNodeType::type node_type) { _node_type = node_type; } +#endif + virtual Status evaluate_ann_range_search( + const segment_v2::AnnRangeSearchRuntime& runtime, + const std::vector>& cid_to_index_iterators, + const std::vector& idx_to_cid, + const std::vector>& column_iterators, + roaring::Roaring& row_bitmap, segment_v2::AnnIndexStats& ann_index_stats); + + // Prepare the runtime for ANN range search. + // AnnRangeSearchRuntime is used to store the runtime information of ann range search. + // suitable_for_ann_index is used to indicate whether the current expr can be used for ANN range search. + // If suitable_for_ann_index is false, the we will do exhausted search. + virtual void prepare_ann_range_search(const doris::VectorSearchUserParams& params, + segment_v2::AnnRangeSearchRuntime& range_search_runtime, + bool& suitable_for_ann_index); + + bool has_been_executed(); + protected: /// Simple debug string that provides no expr subclass-specific information std::string debug_string(const std::string& expr_name) const { @@ -336,6 +365,8 @@ class VExpr { // ensuring uniqueness during index traversal uint32_t _index_unique_id = 0; bool _enable_inverted_index_query = true; + + bool _has_been_executed = false; }; } // namespace vectorized diff --git a/be/src/vec/exprs/vexpr_context.cpp b/be/src/vec/exprs/vexpr_context.cpp index bbad95b6967b91..597d1a928887e0 100644 --- a/be/src/vec/exprs/vexpr_context.cpp +++ b/be/src/vec/exprs/vexpr_context.cpp @@ -116,6 +116,9 @@ Status VExprContext::clone(RuntimeState* state, VExprContextSPtr& new_ctx) { new_ctx->_is_clone = true; new_ctx->_prepared = true; new_ctx->_opened = true; + // segment_v2::AnnRangeSearchRuntime should be cloned as well. + // The object of segment_v2::AnnRangeSearchRuntime is not shared by threads. + new_ctx->_ann_range_search_runtime = this->_ann_range_search_runtime; return _root->open(state, new_ctx.get(), FunctionContext::THREAD_LOCAL); } @@ -434,6 +437,31 @@ void VExprContext::_reset_memory_usage(const VExprContextSPtrs& contexts) { [](auto&& context) { context->_memory_usage = 0; }); } +void VExprContext::prepare_ann_range_search(const doris::VectorSearchUserParams& params) { + if (_root == nullptr) { + return; + } + + _root->prepare_ann_range_search(params, _ann_range_search_runtime, _suitable_for_ann_index); + VLOG_DEBUG << fmt::format("Prepare ann range search result {}, _suitable_for_ann_index {}", + this->_ann_range_search_runtime.to_string(), + this->_suitable_for_ann_index); + return; +} + +Status VExprContext::evaluate_ann_range_search( + const std::vector>& cid_to_index_iterators, + const std::vector& idx_to_cid, + const std::vector>& column_iterators, + roaring::Roaring& row_bitmap, segment_v2::AnnIndexStats& ann_index_stats) { + if (_root != nullptr) { + return _root->evaluate_ann_range_search(_ann_range_search_runtime, cid_to_index_iterators, + idx_to_cid, column_iterators, row_bitmap, + ann_index_stats); + } + return Status::OK(); +} + #include "common/compile_check_end.h" } // namespace doris::vectorized diff --git a/be/src/vec/exprs/vexpr_context.h b/be/src/vec/exprs/vexpr_context.h index e9984b703576e4..3571f95f851c07 100644 --- a/be/src/vec/exprs/vexpr_context.h +++ b/be/src/vec/exprs/vexpr_context.h @@ -27,7 +27,11 @@ #include "common/factory_creator.h" #include "common/status.h" +#include "olap/rowset/segment_v2/ann_index/ann_range_search_runtime.h" +#include "olap/rowset/segment_v2/ann_index/ann_search_params.h" +#include "olap/rowset/segment_v2/column_reader.h" #include "olap/rowset/segment_v2/inverted_index_reader.h" +#include "runtime/runtime_state.h" #include "runtime/types.h" #include "udf/udf.h" #include "vec/core/block.h" @@ -38,6 +42,10 @@ class RowDescriptor; class RuntimeState; } // namespace doris +namespace doris::segment_v2 { +class ColumnIterator; +} // namespace doris::segment_v2 + namespace doris::vectorized { class InvertedIndexContext { @@ -278,6 +286,14 @@ class VExprContext { [[nodiscard]] size_t get_memory_usage() const { return _memory_usage; } + void prepare_ann_range_search(const doris::VectorSearchUserParams& params); + + Status evaluate_ann_range_search( + const std::vector>& cid_to_index_iterators, + const std::vector& idx_to_cid, + const std::vector>& column_iterators, + roaring::Roaring& row_bitmap, segment_v2::AnnIndexStats& ann_index_stats); + private: // Close method is called in vexpr context dector, not need call expicility void close(); @@ -311,5 +327,8 @@ class VExprContext { std::shared_ptr _inverted_index_context; size_t _memory_usage = 0; + + segment_v2::AnnRangeSearchRuntime _ann_range_search_runtime; + bool _suitable_for_ann_index = true; }; } // namespace doris::vectorized diff --git a/be/src/vec/exprs/virtual_slot_ref.cpp b/be/src/vec/exprs/virtual_slot_ref.cpp index d3db6197a758ca..8a83dca3b6f32d 100644 --- a/be/src/vec/exprs/virtual_slot_ref.cpp +++ b/be/src/vec/exprs/virtual_slot_ref.cpp @@ -22,6 +22,7 @@ #include #include +#include #include "common/exception.h" #include "common/logging.h" @@ -35,9 +36,8 @@ #include "vec/exprs/vectorized_fn_call.h" #include "vec/exprs/vexpr_context.h" #include "vec/exprs/vexpr_fwd.h" - namespace doris::vectorized { - +#include "common/compile_check_begin.h" VirtualSlotRef::VirtualSlotRef(const doris::TExprNode& node) : VExpr(node), _column_id(-1), @@ -218,4 +218,33 @@ bool VirtualSlotRef::equals(const VExpr& other) { return true; } +/** + * @brief Implements ANN range search evaluation for virtual slot references. + * + * This method handles the case where a virtual slot reference wraps a distance + * function call that can be optimized using ANN index range search. Instead of + * computing distances for all rows, it delegates to the underlying virtual + * expression to perform the optimized search. + * + * @param range_search_runtime Runtime parameters for the range search + * @param cid_to_index_iterators Index iterators for each column + * @param idx_to_cid Column ID mapping + * @param column_iterators Data column iterators + * @param row_bitmap Result bitmap to be updated with matching rows + * @param ann_index_stats Performance statistics collector + * @return Status::OK() if successful, error status otherwise + */ +Status VirtualSlotRef::evaluate_ann_range_search( + const segment_v2::AnnRangeSearchRuntime& range_search_runtime, + const std::vector>& cid_to_index_iterators, + const std::vector& idx_to_cid, + const std::vector>& column_iterators, + roaring::Roaring& row_bitmap, segment_v2::AnnIndexStats& ann_index_stats) { + return _virtual_column_expr->evaluate_ann_range_search( + range_search_runtime, cid_to_index_iterators, idx_to_cid, column_iterators, row_bitmap, + ann_index_stats); + + return Status::OK(); +} +#include "common/compile_check_end.h" } // namespace doris::vectorized diff --git a/be/src/vec/exprs/virtual_slot_ref.h b/be/src/vec/exprs/virtual_slot_ref.h index c4e326501f80d0..85ca13d1ae43b2 100644 --- a/be/src/vec/exprs/virtual_slot_ref.h +++ b/be/src/vec/exprs/virtual_slot_ref.h @@ -20,7 +20,7 @@ #include "vec/exprs/vexpr.h" namespace doris::vectorized { - +#include "common/compile_check_begin.h" class VirtualSlotRef MOCK_REMOVE(final) : public VExpr { ENABLE_FACTORY_CREATOR(VirtualSlotRef); @@ -49,13 +49,84 @@ class VirtualSlotRef MOCK_REMOVE(final) : public VExpr { return _virtual_column_expr->evaluate_inverted_index(context, segment_num_rows); } + /* + @brief SQL expression tree patterns for ANN range search optimization. + + Pattern 1 (should not happen): + SELECT * FROM tbl WHERE distance_function(columnA, ArrayLiteral) > 100 + VirtualSlotRef + | + BINARY_PRED + |---------------------------------------| + | | + FUNCTION_CALL(l2_distance_approximate) IntLiteral + | + |-----------------------| + | | + SlotRef ArrayLiteral + + Pattern 2 (optimizable case): + SELECT distance_function(columnA, ArrayLiteral) AS dis FROM tbl WHERE dis > 100 + BINARY_PRED + | + |---------------------------------------| + | | + VIRTUAL_SLOT_REF IntLiteral + | + FUNCTION_CALL(l2_distance_approximate) + | + |-----------------------| + | | + SlotRef ArrayLiteral + */ + + /** + * @brief Evaluates ANN range search using index-based optimization. + * + * This method implements the core logic for ANN range search optimization. + * Instead of computing distances for all rows and then filtering, it uses + * the ANN index to efficiently find only the rows within the specified range. + * + * The method: + * 1. Extracts query parameters from the range search runtime info + * 2. Calls the ANN index to perform range search + * 3. Updates the row bitmap with matching results + * 4. Collects performance statistics + * + * @param range_search_runtime Runtime info containing query vector, radius, and metrics + * @param cid_to_index_iterators Vector of index iterators for each column + * @param idx_to_cid Mapping from index position to column ID + * @param column_iterators Vector of column iterators for data access + * @param row_bitmap Output bitmap updated with matching row IDs + * @param ann_index_stats Statistics collector for performance monitoring + * @return Status indicating success or failure of the search operation + */ + Status evaluate_ann_range_search( + const segment_v2::AnnRangeSearchRuntime& range_search_runtime, + const std::vector>& cid_to_index_iterators, + const std::vector& idx_to_cid, + const std::vector>& column_iterators, + roaring::Roaring& row_bitmap, segment_v2::AnnIndexStats& ann_index_stats) override; + +#ifdef BE_TEST + // Test-only setter methods for unit testing + void set_column_id(int column_id) { _column_id = column_id; } + void set_column_name(const std::string* column_name) { _column_name = column_name; } + void set_column_data_type(DataTypePtr column_data_type) { + _column_data_type = std::move(column_data_type); + } + void set_virtual_column_expr(std::shared_ptr virtual_column_expr) { + _virtual_column_expr = virtual_column_expr; + } +#endif + private: - int _column_id; - int _slot_id; - const std::string* _column_name; - const std::string _column_label; - std::shared_ptr _virtual_column_expr; - DataTypePtr _column_data_type; + int _column_id; ///< Column ID in the table schema + int _slot_id; ///< Slot ID in the expression context + const std::string* _column_name; ///< Column name for debugging/logging + const std::string _column_label; ///< Column label for display purposes + std::shared_ptr _virtual_column_expr; ///< Underlying virtual expression + DataTypePtr _column_data_type; ///< Data type of the column }; - +#include "common/compile_check_end.h" } // namespace doris::vectorized \ No newline at end of file diff --git a/be/src/vec/exprs/vslot_ref.h b/be/src/vec/exprs/vslot_ref.h index 08e09e32770cbf..3145476afed1ed 100644 --- a/be/src/vec/exprs/vslot_ref.h +++ b/be/src/vec/exprs/vslot_ref.h @@ -40,6 +40,8 @@ class VSlotRef MOCK_REMOVE(final) : public VExpr { VSlotRef(const SlotDescriptor* desc); #ifdef BE_TEST VSlotRef() = default; + void set_column_id(int column_id) { _column_id = column_id; } + void set_slot_id(int slot_id) { _slot_id = slot_id; } #endif Status prepare(RuntimeState* state, const RowDescriptor& desc, VExprContext* context) override; Status open(RuntimeState* state, VExprContext* context, diff --git a/be/src/vec/functions/array/function_array_distance.cpp b/be/src/vec/functions/array/function_array_distance.cpp index fc7ba9a0367a64..e89ce103e14613 100644 --- a/be/src/vec/functions/array/function_array_distance.cpp +++ b/be/src/vec/functions/array/function_array_distance.cpp @@ -26,6 +26,8 @@ void register_function_array_distance(SimpleFunctionFactory& factory) { factory.register_function >(); factory.register_function >(); factory.register_function >(); + factory.register_function >(); + factory.register_function >(); } } // namespace doris::vectorized diff --git a/be/src/vec/functions/array/function_array_distance.h b/be/src/vec/functions/array/function_array_distance.h index 8536f7e5b17f20..cb0661c4ccbf0a 100644 --- a/be/src/vec/functions/array/function_array_distance.h +++ b/be/src/vec/functions/array/function_array_distance.h @@ -79,6 +79,26 @@ class CosineDistance { } }; +class L2DistanceApproximate { +public: + static constexpr auto name = "l2_distance_approximate"; + struct State { + double sum = 0; + }; + static void accumulate(State& state, double x, double y) { state.sum += (x - y) * (x - y); } + static double finalize(const State& state) { return sqrt(state.sum); } +}; + +class InnerProductApproximate { +public: + static constexpr auto name = "inner_product_approximate"; + struct State { + double sum = 0; + }; + static void accumulate(State& state, double x, double y) { state.sum += x * y; } + static double finalize(const State& state) { return state.sum; } +}; + template class FunctionArrayDistance : public IFunction { public: diff --git a/be/src/vec/functions/array/function_array_element.h b/be/src/vec/functions/array/function_array_element.h index e9da6b9c2fda4a..5f42c2d1e38af8 100644 --- a/be/src/vec/functions/array/function_array_element.h +++ b/be/src/vec/functions/array/function_array_element.h @@ -279,7 +279,6 @@ class FunctionArrayElement : public IFunction { const auto& offsets = map_column.get_offsets(); const size_t rows = offsets.size(); - if (rows <= 0) { return nullptr; } diff --git a/be/src/vec/olap/block_reader.cpp b/be/src/vec/olap/block_reader.cpp index 87c6cf28e6527f..f40cb45652c749 100644 --- a/be/src/vec/olap/block_reader.cpp +++ b/be/src/vec/olap/block_reader.cpp @@ -147,7 +147,6 @@ Status BlockReader::_init_collect_iter(const ReaderParams& read_params) { } } } - { SCOPED_RAW_TIMER(&_stats.block_reader_build_heap_init_timer_ns); RETURN_IF_ERROR(_vcollect_iter.build_heap(valid_rs_readers)); @@ -210,6 +209,7 @@ Status BlockReader::init(const ReaderParams& read_params) { _return_columns_loc.resize(read_params.return_columns.size()); for (int i = 0; i < return_column_size; ++i) { auto cid = read_params.origin_return_columns->at(i); + // For each original cid, find the index in return_columns for (int j = 0; j < read_params.return_columns.size(); ++j) { if (read_params.return_columns[j] == cid) { if (j < _tablet->num_key_columns() || _tablet->keys_type() != AGG_KEYS) { diff --git a/be/src/vec/runtime/vector_search_user_params.cpp b/be/src/vec/runtime/vector_search_user_params.cpp new file mode 100644 index 00000000000000..19fb4e1f0829fe --- /dev/null +++ b/be/src/vec/runtime/vector_search_user_params.cpp @@ -0,0 +1,36 @@ +// 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 "vec/runtime/vector_search_user_params.h" + +#include + +namespace doris { +#include "common/compile_check_begin.h" +bool VectorSearchUserParams::operator==(const VectorSearchUserParams& other) const { + return hnsw_ef_search == other.hnsw_ef_search && + hnsw_check_relative_distance == other.hnsw_check_relative_distance && + hnsw_bounded_queue == other.hnsw_bounded_queue; +} + +std::string VectorSearchUserParams::to_string() const { + return fmt::format( + "hnsw_ef_search: {}, hnsw_check_relative_distance: {}, " + "hnsw_bounded_queue: {}", + hnsw_ef_search, hnsw_check_relative_distance, hnsw_bounded_queue); +} +} // namespace doris \ No newline at end of file diff --git a/be/src/vec/runtime/vector_search_user_params.h b/be/src/vec/runtime/vector_search_user_params.h new file mode 100644 index 00000000000000..600716651c51c9 --- /dev/null +++ b/be/src/vec/runtime/vector_search_user_params.h @@ -0,0 +1,35 @@ +// 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. + +#pragma once + +#include + +namespace doris { +#include "common/compile_check_begin.h" +// Constructed from session variables. +struct VectorSearchUserParams { + int hnsw_ef_search = 32; + bool hnsw_check_relative_distance = true; + bool hnsw_bounded_queue = true; + + bool operator==(const VectorSearchUserParams& other) const; + + std::string to_string() const; +}; +#include "common/compile_check_end.h" +} // namespace doris \ No newline at end of file diff --git a/be/src/vec/utils/util.hpp b/be/src/vec/utils/util.hpp index 9c22119cd8db1b..2297c411832be2 100644 --- a/be/src/vec/utils/util.hpp +++ b/be/src/vec/utils/util.hpp @@ -285,6 +285,17 @@ inline size_t calculate_false_number(ColumnPtr column) { } } +template +T read_from_json(const std::string& json_str) { + auto memBufferIn = std::make_shared( + reinterpret_cast(const_cast(json_str.data())), + static_cast(json_str.size())); + auto jsonProtocolIn = std::make_shared(memBufferIn); + T params; + params.read(jsonProtocolIn.get()); + return params; +} + } // namespace doris::vectorized namespace apache::thrift { diff --git a/be/test/CMakeLists.txt b/be/test/CMakeLists.txt index 1d15b4d231a6e2..fda1c82b718bf8 100644 --- a/be/test/CMakeLists.txt +++ b/be/test/CMakeLists.txt @@ -23,6 +23,13 @@ set(EXECUTABLE_OUTPUT_PATH "${BUILD_DIR}/test") file(GLOB_RECURSE UT_FILES CONFIGURE_DEPENDS *.cpp) +# Remove all cpp files from vector_search subdirectory +# Since cpp files in vector search subdirs use header files from faiss. +# The compile check used by doris can not be applied to faiss headers. +# So vector_search cpp files are compiled in a separate library using different compile options. +file(GLOB_RECURSE VECTOR_FILES CONFIGURE_DEPENDS olap/vector_search/*.cpp) +list(REMOVE_ITEM UT_FILES ${VECTOR_FILES}) + if(NOT DEFINED DORIS_WITH_LZO) list(REMOVE_ITEM UT_FILES ${CMAKE_CURRENT_SOURCE_DIR}/exec/plain_text_line_reader_lzop_test.cpp) endif() @@ -98,9 +105,13 @@ include_directories( ${CMAKE_BINARY_DIR}/apache-orc/c++/include ) +# Removed test files are added back using separate compile arguments. +add_subdirectory(olap/vector_search) + add_executable(doris_be_test ${UT_FILES}) -target_link_libraries(doris_be_test ${TEST_LINK_LIBS}) +target_link_libraries(doris_be_test ${TEST_LINK_LIBS} + -Wl,--whole-archive vector_search_test -Wl,--no-whole-archive) set_target_properties(doris_be_test PROPERTIES COMPILE_FLAGS "-fno-access-control") if (OS_MACOSX AND ARCH_ARM) diff --git a/be/test/olap/rowset/segment_v2/inverted_index/compaction/util/index_compaction_utils.cpp b/be/test/olap/rowset/segment_v2/inverted_index/compaction/util/index_compaction_utils.cpp index 392f4d6b6f2fb0..385e7c46ffdc58 100644 --- a/be/test/olap/rowset/segment_v2/inverted_index/compaction/util/index_compaction_utils.cpp +++ b/be/test/olap/rowset/segment_v2/inverted_index/compaction/util/index_compaction_utils.cpp @@ -482,7 +482,7 @@ class IndexCompactionUtils { } static void check_idx_file_writer_closed(BaseBetaRowsetWriter* writer, bool closed) { - for (const auto& [seg_id, idx_file_writer] : writer->inverted_index_file_writers()) { + for (const auto& [seg_id, idx_file_writer] : writer->index_file_writers()) { EXPECT_EQ(idx_file_writer->_closed, closed); } } diff --git a/be/test/olap/vector_search/CMakeLists.txt b/be/test/olap/vector_search/CMakeLists.txt new file mode 100644 index 00000000000000..e0618afbb6b581 --- /dev/null +++ b/be/test/olap/vector_search/CMakeLists.txt @@ -0,0 +1,39 @@ +# 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. + +# Collect all source files in this directory +file(GLOB_RECURSE VECTOR_SEARCH_TEST_FILES CONFIGURE_DEPENDS *.cpp) + +# Create a static library for vector search tests +add_library(vector_search_test STATIC ${VECTOR_SEARCH_TEST_FILES}) + +# Suppress warnings for this library to avoid compilation failures from third-party dependencies +target_compile_options(vector_search_test PRIVATE + -Wno-everything # Suppress all warnings + -fno-access-control # Disable access control error +) + +# Link necessary libraries +target_link_libraries(vector_search_test + PUBLIC + ann_index # Vector library that contains faiss integration + gtest # Google Test library + gmock # Google Mock library +) + +# Add dependencies to ensure proper build order +add_dependencies(vector_search_test ann_index faiss) \ No newline at end of file diff --git a/be/test/olap/vector_search/ann_index_edge_case_test.cpp b/be/test/olap/vector_search/ann_index_edge_case_test.cpp new file mode 100644 index 00000000000000..3ab22a63c37178 --- /dev/null +++ b/be/test/olap/vector_search/ann_index_edge_case_test.cpp @@ -0,0 +1,259 @@ +// 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 +#include + +#include +#include +#include + +#include "olap/rowset/segment_v2/ann_index/ann_index_iterator.h" +#include "olap/rowset/segment_v2/ann_index/ann_index_reader.h" +#include "olap/rowset/segment_v2/ann_index/ann_index_writer.h" +#include "olap/rowset/segment_v2/ann_index/faiss_ann_index.h" +#include "vector_search_utils.h" + +using namespace doris::vector_search_utils; + +namespace doris::vectorized { + +TEST_F(VectorSearchTest, TestAnnIndexStatsInitialization) { + doris::segment_v2::AnnIndexStats stats; + + // Test initial values + EXPECT_EQ(stats.search_costs_ns.value(), 0); + EXPECT_EQ(stats.load_index_costs_ns.value(), 0); + + // Test setting values + stats.search_costs_ns.set(1000L); + stats.load_index_costs_ns.set(2000L); + + EXPECT_EQ(stats.search_costs_ns.value(), 1000); + EXPECT_EQ(stats.load_index_costs_ns.value(), 2000); +} + +TEST_F(VectorSearchTest, TestAnnIndexStatsCopyConstructor) { + doris::segment_v2::AnnIndexStats original; + original.search_costs_ns.set(1500L); + original.load_index_costs_ns.set(2500L); + + doris::segment_v2::AnnIndexStats copied(original); + + EXPECT_EQ(copied.search_costs_ns.value(), 1500); + EXPECT_EQ(copied.load_index_costs_ns.value(), 2500); +} + +TEST_F(VectorSearchTest, TestAnnRangeSearchParamsToString) { + doris::segment_v2::AnnRangeSearchParams params; + params.is_le_or_lt = true; + params.radius = 5.5f; + + auto roaring = std::make_shared(); + roaring->add(1); + roaring->add(2); + roaring->add(3); + params.roaring = roaring.get(); + + std::string result = params.to_string(); + + EXPECT_TRUE(result.find("is_le_or_lt: true") != std::string::npos); + EXPECT_TRUE(result.find("radius: 5.5") != std::string::npos); + EXPECT_TRUE(result.find("input rows 3") != std::string::npos); +} + +TEST_F(VectorSearchTest, TestAnnRangeSearchParamsWithNullRoaring) { + auto roaring = std::make_unique(); + doris::segment_v2::AnnRangeSearchParams params; + params.is_le_or_lt = false; + params.radius = 10.0f; + params.roaring = roaring.get(); // Assigning a null pointer + + std::string result = params.to_string(); + + EXPECT_TRUE(result.find("is_le_or_lt: false") != std::string::npos); + EXPECT_TRUE(result.find("radius: 10") != std::string::npos); + EXPECT_TRUE(result.find("input rows 0") != std::string::npos); +} + +TEST_F(VectorSearchTest, TestAnnTopNParamValidation) { + // Test with zero limit + const float query_data[] = {1.0f, 2.0f, 3.0f, 4.0f}; + roaring::Roaring bitmap; + bitmap.add(1); + + doris::segment_v2::AnnTopNParam param = { + .query_value = query_data, + .query_value_size = 4, + .limit = 0, // Zero limit + ._user_params = doris::VectorSearchUserParams {}, + .roaring = &bitmap, + .distance = nullptr, + .row_ids = nullptr, + .stats = std::make_unique()}; + + // The parameter should be valid even with zero limit + EXPECT_EQ(param.limit, 0); + EXPECT_EQ(param.query_value_size, 4); + EXPECT_NE(param.query_value, nullptr); +} + +TEST_F(VectorSearchTest, TestVectorSearchUserParamsDefaultValues) { + doris::VectorSearchUserParams params; + + // Test default values + EXPECT_EQ(params.hnsw_ef_search, 32); + EXPECT_EQ(params.hnsw_check_relative_distance, true); + EXPECT_EQ(params.hnsw_bounded_queue, true); +} + +TEST_F(VectorSearchTest, TestVectorSearchUserParamsEquality) { + doris::VectorSearchUserParams params1; + params1.hnsw_ef_search = 100; + params1.hnsw_check_relative_distance = false; + params1.hnsw_bounded_queue = false; + + doris::VectorSearchUserParams params2; + params2.hnsw_ef_search = 100; + params2.hnsw_check_relative_distance = false; + params2.hnsw_bounded_queue = false; + + EXPECT_EQ(params1, params2); + + // Test inequality + params2.hnsw_ef_search = 50; + EXPECT_NE(params1, params2); +} + +TEST_F(VectorSearchTest, TestIndexSearchResultInitialization) { + doris::segment_v2::IndexSearchResult result; + + // Test initial state + EXPECT_EQ(result.roaring, nullptr); + EXPECT_EQ(result.distances, nullptr); + EXPECT_EQ(result.row_ids, nullptr); +} + +TEST_F(VectorSearchTest, TestAnnRangeSearchResultInitialization) { + doris::segment_v2::AnnRangeSearchResult result; + + // Test initial state + EXPECT_EQ(result.roaring, nullptr); + EXPECT_EQ(result.distance, nullptr); + EXPECT_EQ(result.row_ids, nullptr); +} + +TEST_F(VectorSearchTest, TestAnnIndexWriterWithEmptyProperties) { + // Test writer with empty properties (should use defaults) + std::map empty_properties; + + auto tablet_index = std::make_unique(); + tablet_index->_properties = empty_properties; + tablet_index->_index_id = 1; + + auto mock_file_writer = + std::make_unique(doris::io::global_local_filesystem()); + auto writer = std::make_unique(mock_file_writer.get(), + tablet_index.get()); + + auto fs_dir = std::make_shared(); + EXPECT_CALL(*mock_file_writer, open(testing::_)).WillOnce(testing::Return(fs_dir)); + + // Should not crash and should use default values + Status status = writer->init(); + EXPECT_TRUE(status.ok()); +} + +TEST_F(VectorSearchTest, TestLargeVectorDimensions) { + // Test with large vector dimensions + std::map properties; + properties["index_type"] = "hnsw"; + properties["metric_type"] = "l2_distance"; + properties["dim"] = "1024"; // Large dimension + properties["max_degree"] = "64"; + + auto tablet_index = std::make_unique(); + tablet_index->_properties = properties; + tablet_index->_index_id = 1; + + auto mock_file_writer = + std::make_unique(doris::io::global_local_filesystem()); + auto writer = std::make_unique(mock_file_writer.get(), + tablet_index.get()); + + auto fs_dir = std::make_shared(); + EXPECT_CALL(*mock_file_writer, open(testing::_)).WillOnce(testing::Return(fs_dir)); + + ASSERT_TRUE(writer->init().ok()); + + // Test adding vectors with correct large dimension + const size_t dim = 1024; + const size_t num_rows = 2; + std::vector vectors(num_rows * dim); + + // Fill with test data + for (size_t i = 0; i < vectors.size(); ++i) { + vectors[i] = static_cast(i % 100) / 100.0f; + } + + std::vector offsets = {0, dim, 2 * dim}; + + Status status = + writer->add_array_values(sizeof(float), vectors.data(), nullptr, + reinterpret_cast(offsets.data()), num_rows); + EXPECT_TRUE(status.ok()); +} + +TEST_F(VectorSearchTest, TestEmptyRoaringBitmap) { + // Test with empty roaring bitmap + doris::segment_v2::AnnRangeSearchParams params; + params.is_le_or_lt = true; + params.radius = 5.0f; + + float query_data[] = {1.0f, 2.0f, 3.0f, 4.0f}; + params.query_value = query_data; + + roaring::Roaring empty_bitmap; // Empty bitmap + params.roaring = &empty_bitmap; + + std::string result = params.to_string(); + + EXPECT_TRUE(result.find("input rows 0") != std::string::npos); +} + +TEST_F(VectorSearchTest, TestLargeRoaringBitmap) { + // Test with large roaring bitmap + doris::segment_v2::AnnRangeSearchParams params; + params.is_le_or_lt = false; + params.radius = 10.0f; + + float query_data[] = {1.0f, 2.0f, 3.0f, 4.0f}; + params.query_value = query_data; + + roaring::Roaring large_bitmap; + // Add many elements + for (uint32_t i = 0; i < 100000; ++i) { + large_bitmap.add(i); + } + params.roaring = &large_bitmap; + + std::string result = params.to_string(); + + EXPECT_TRUE(result.find("input rows 100000") != std::string::npos); +} + +} // namespace doris::vectorized diff --git a/be/test/olap/vector_search/ann_index_iterator_test.cpp b/be/test/olap/vector_search/ann_index_iterator_test.cpp new file mode 100644 index 00000000000000..c5c953da574947 --- /dev/null +++ b/be/test/olap/vector_search/ann_index_iterator_test.cpp @@ -0,0 +1,341 @@ +// 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 "olap/rowset/segment_v2/ann_index/ann_index_iterator.h" + +#include +#include + +#include +#include + +#include "olap/rowset/segment_v2/ann_index/ann_index_reader.h" +#include "olap/rowset/segment_v2/ann_index/ann_search_params.h" +#include "olap/rowset/segment_v2/ann_index/faiss_ann_index.h" +#include "vector_search_utils.h" + +using namespace doris::vector_search_utils; + +namespace doris::segment_v2 { + +class AnnIndexIteratorTest : public doris::vectorized::VectorSearchTest { +protected: + void SetUp() override { + doris::vectorized::VectorSearchTest::SetUp(); + + // Create test index properties + _properties["index_type"] = "hnsw"; + _properties["metric_type"] = "l2_distance"; + _properties["dim"] = "4"; + _properties["max_degree"] = "16"; + + // Create tablet index + _tablet_index = std::make_unique(); + _tablet_index->_properties = _properties; + _tablet_index->_index_id = 1; + _tablet_index->_index_name = "test_ann_index"; + + // Create mock index file reader + _mock_index_file_reader = std::make_shared(); + + // Create ann index reader + _ann_reader = + std::make_shared(_tablet_index.get(), _mock_index_file_reader); + } + + void TearDown() override { doris::vectorized::VectorSearchTest::TearDown(); } + + std::map _properties; + std::unique_ptr _tablet_index; + std::shared_ptr _mock_index_file_reader; + std::shared_ptr _ann_reader; +}; + +TEST_F(AnnIndexIteratorTest, TestConstructor) { + auto iterator = std::make_unique(_ann_reader); + EXPECT_NE(iterator, nullptr); +} + +TEST_F(AnnIndexIteratorTest, TestConstructorWithNullReader) { + auto iterator = std::make_unique(nullptr); + EXPECT_NE(iterator, nullptr); +} + +TEST_F(AnnIndexIteratorTest, TestReadFromIndexWithNullParam) { + auto iterator = std::make_unique(_ann_reader); + + // Test with null parameter - this should trigger the null check + IndexParam param = static_cast(nullptr); + auto status = iterator->read_from_index(param); + + EXPECT_FALSE(status.ok()); + EXPECT_TRUE(status.is()); + EXPECT_TRUE(status.msg().find("a_param is null") != std::string::npos); +} + +TEST_F(AnnIndexIteratorTest, TestReadFromIndexWithValidParam) { + auto iterator = std::make_unique(_ann_reader); + + // Set up the reader's _vector_index + auto doris_faiss_vector_index = std::make_unique(); + doris_faiss_vector_index->set_metric(AnnIndexMetric::L2); + + FaissBuildParameter build_params; + build_params.dim = 4; + build_params.max_degree = 16; + build_params.index_type = FaissBuildParameter::IndexType::HNSW; + build_params.metric_type = FaissBuildParameter::MetricType::L2; + doris_faiss_vector_index->build(build_params); + + _ann_reader->_vector_index = std::move(doris_faiss_vector_index); + + // Create valid AnnTopNParam + const float query_data[] = {1.0f, 2.0f, 3.0f, 4.0f}; + roaring::Roaring bitmap; + bitmap.add(1); + bitmap.add(2); + + doris::segment_v2::AnnTopNParam ann_param = { + .query_value = query_data, + .query_value_size = 4, + .limit = 10, + ._user_params = doris::VectorSearchUserParams {}, + .roaring = &bitmap, + .distance = nullptr, + .row_ids = nullptr, + .stats = std::make_unique()}; + + IndexParam param = &ann_param; + + auto status = iterator->read_from_index(param); + + // The query might succeed or fail depending on the internal index state, + // but it should not crash and should handle the parameter correctly + if (status.ok()) { + EXPECT_NE(ann_param.distance, nullptr); + EXPECT_NE(ann_param.row_ids, nullptr); + } +} + +TEST_F(AnnIndexIteratorTest, TestRangeSearchWithNullReader) { + auto iterator = std::make_unique(nullptr); + + doris::segment_v2::AnnRangeSearchParams params; + doris::VectorSearchUserParams user_params; + doris::segment_v2::AnnRangeSearchResult result; + auto stats = std::make_unique(); + + auto status = iterator->range_search(params, user_params, &result, stats.get()); + + EXPECT_FALSE(status.ok()); + EXPECT_TRUE(status.is()); + EXPECT_TRUE(status.msg().find("_ann_reader is null") != std::string::npos); +} + +TEST_F(AnnIndexIteratorTest, TestRangeSearchWithValidReader) { + auto iterator = std::make_unique(_ann_reader); + + // Set up the reader's _vector_index + auto doris_faiss_vector_index = std::make_unique(); + doris_faiss_vector_index->set_metric(AnnIndexMetric::L2); + + FaissBuildParameter build_params; + build_params.dim = 4; + build_params.max_degree = 16; + build_params.index_type = FaissBuildParameter::IndexType::HNSW; + build_params.metric_type = FaissBuildParameter::MetricType::L2; + doris_faiss_vector_index->build(build_params); + + _ann_reader->_vector_index = std::move(doris_faiss_vector_index); + + // Create range search parameters + float query_data[] = {1.0f, 2.0f, 3.0f, 4.0f}; + roaring::Roaring bitmap; + bitmap.add(1); + bitmap.add(2); + + doris::segment_v2::AnnRangeSearchParams params; + params.is_le_or_lt = true; + params.radius = 5.0f; + params.query_value = query_data; + params.roaring = &bitmap; + + doris::VectorSearchUserParams user_params; + user_params.hnsw_ef_search = 50; + user_params.hnsw_check_relative_distance = true; + user_params.hnsw_bounded_queue = true; + + doris::segment_v2::AnnRangeSearchResult result; + auto stats = std::make_unique(); + + auto status = iterator->range_search(params, user_params, &result, stats.get()); + + // The range search might succeed or fail depending on the internal index state, + // but it should not crash + if (status.ok()) { + EXPECT_NE(result.roaring, nullptr); + } +} + +TEST_F(AnnIndexIteratorTest, TestRangeSearchWithDifferentParameters) { + auto iterator = std::make_unique(_ann_reader); + + // Set up the reader's _vector_index + auto doris_faiss_vector_index = std::make_unique(); + doris_faiss_vector_index->set_metric(AnnIndexMetric::L2); + + FaissBuildParameter build_params; + build_params.dim = 4; + build_params.max_degree = 16; + build_params.index_type = FaissBuildParameter::IndexType::HNSW; + build_params.metric_type = FaissBuildParameter::MetricType::L2; + doris_faiss_vector_index->build(build_params); + + _ann_reader->_vector_index = std::move(doris_faiss_vector_index); + + // Test different parameter combinations + std::vector> test_cases = { + {true, 1.0f, 10}, // is_le_or_lt=true, small radius, small ef_search + {false, 5.0f, 50}, // is_le_or_lt=false, medium radius, medium ef_search + {true, 10.0f, 100}, // is_le_or_lt=true, large radius, large ef_search + }; + + for (const auto& [is_le_or_lt, radius, ef_search] : test_cases) { + float query_data[] = {1.0f, 2.0f, 3.0f, 4.0f}; + roaring::Roaring bitmap; + bitmap.add(1); + bitmap.add(2); + bitmap.add(3); + + doris::segment_v2::AnnRangeSearchParams params; + params.is_le_or_lt = is_le_or_lt; + params.radius = radius; + params.query_value = query_data; + params.roaring = &bitmap; + + doris::VectorSearchUserParams user_params; + user_params.hnsw_ef_search = ef_search; + user_params.hnsw_check_relative_distance = false; + user_params.hnsw_bounded_queue = false; + + doris::segment_v2::AnnRangeSearchResult result; + auto stats = std::make_unique(); + + auto status = iterator->range_search(params, user_params, &result, stats.get()); + + // Should not crash regardless of success/failure + if (status.ok()) { + EXPECT_NE(result.roaring, nullptr); + } + } +} + +TEST_F(AnnIndexIteratorTest, TestWithInnerProductMetric) { + // Test with inner product metric + auto properties = _properties; + properties["metric_type"] = "inner_product"; + + auto tablet_index = std::make_unique(); + tablet_index->_properties = properties; + tablet_index->_index_id = 1; + + auto ann_reader = std::make_shared(tablet_index.get(), _mock_index_file_reader); + auto iterator = std::make_unique(ann_reader); + + // Set up the reader's _vector_index with IP metric + auto doris_faiss_vector_index = std::make_unique(); + doris_faiss_vector_index->set_metric(AnnIndexMetric::IP); + + FaissBuildParameter build_params; + build_params.dim = 4; + build_params.max_degree = 16; + build_params.index_type = FaissBuildParameter::IndexType::HNSW; + build_params.metric_type = FaissBuildParameter::MetricType::IP; + doris_faiss_vector_index->build(build_params); + + ann_reader->_vector_index = std::move(doris_faiss_vector_index); + + // Test read_from_index with IP metric + const float query_data[] = {1.0f, 2.0f, 3.0f, 4.0f}; + roaring::Roaring bitmap; + bitmap.add(1); + bitmap.add(2); + + doris::segment_v2::AnnTopNParam ann_param = { + .query_value = query_data, + .query_value_size = 4, + .limit = 10, + ._user_params = doris::VectorSearchUserParams {}, + .roaring = &bitmap, + .distance = nullptr, + .row_ids = nullptr, + .stats = std::make_unique()}; + + IndexParam param = &ann_param; + + auto status = iterator->read_from_index(param); + + // Should not crash regardless of success/failure + if (status.ok()) { + EXPECT_NE(ann_param.distance, nullptr); + EXPECT_NE(ann_param.row_ids, nullptr); + } +} + +TEST_F(AnnIndexIteratorTest, TestSuccessfulWorkflow) { + // Test a complete successful workflow with mock + auto mock_iterator = std::make_unique(); + + // Create test data + const float query_data[] = {1.0f, 2.0f, 3.0f, 4.0f}; + roaring::Roaring bitmap; + bitmap.add(1); + bitmap.add(2); + + doris::segment_v2::AnnTopNParam ann_param = { + .query_value = query_data, + .query_value_size = 4, + .limit = 10, + ._user_params = doris::VectorSearchUserParams {}, + .roaring = &bitmap, + .distance = nullptr, + .row_ids = nullptr, + .stats = std::make_unique()}; + + IndexParam param = &ann_param; + + // Mock successful read_from_index + EXPECT_CALL(*mock_iterator, read_from_index(testing::_)) + .WillOnce(testing::Return(doris::Status::OK())); + + auto status = mock_iterator->read_from_index(param); + EXPECT_TRUE(status.ok()); + + // Mock successful range_search + doris::segment_v2::AnnRangeSearchParams range_params; + doris::VectorSearchUserParams user_params; + doris::segment_v2::AnnRangeSearchResult result; + auto stats = std::make_unique(); + + EXPECT_CALL(*mock_iterator, range_search(testing::_, testing::_, testing::_, testing::_)) + .WillOnce(testing::Return(doris::Status::OK())); + + status = mock_iterator->range_search(range_params, user_params, &result, stats.get()); + EXPECT_TRUE(status.ok()); +} + +} // namespace doris::segment_v2 diff --git a/be/test/olap/vector_search/ann_index_reader_test.cpp b/be/test/olap/vector_search/ann_index_reader_test.cpp new file mode 100644 index 00000000000000..7f914038982bbd --- /dev/null +++ b/be/test/olap/vector_search/ann_index_reader_test.cpp @@ -0,0 +1,495 @@ +// 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 "olap/rowset/segment_v2/ann_index/ann_index_reader.h" + +#include +#include +#include + +#include +#include +#include + +#include "olap/rowset/segment_v2/ann_index/ann_index_iterator.h" +#include "olap/rowset/segment_v2/ann_index/ann_search_params.h" +#include "olap/rowset/segment_v2/ann_index/faiss_ann_index.h" +#include "olap/tablet_schema.h" +#include "vector_search_utils.h" + +using namespace doris::vector_search_utils; + +namespace doris::vectorized { + +class AnnIndexReaderTest : public VectorSearchTest { +protected: + void SetUp() override { + VectorSearchTest::SetUp(); + + // Create test index properties + _properties["index_type"] = "hnsw"; + _properties["metric_type"] = "l2_distance"; + _properties["dim"] = "128"; + _properties["max_degree"] = "16"; + + // Create tablet index + _tablet_index = std::make_unique(); + _tablet_index->_properties = _properties; + _tablet_index->_index_id = 1; + _tablet_index->_index_name = "test_ann_index"; + + // Create mock index file reader + _mock_index_file_reader = std::make_shared(); + } + + void TearDown() override { VectorSearchTest::TearDown(); } + + std::map _properties; + std::unique_ptr _tablet_index; + std::shared_ptr _mock_index_file_reader; +}; + +TEST_F(AnnIndexReaderTest, TestConstructor) { + auto reader = std::make_unique(_tablet_index.get(), + _mock_index_file_reader); + + EXPECT_NE(reader, nullptr); + EXPECT_EQ(reader->get_index_id(), 1); + EXPECT_EQ(reader->index_type(), IndexType::ANN); + EXPECT_EQ(reader->get_metric_type(), segment_v2::AnnIndexMetric::L2); +} + +TEST_F(AnnIndexReaderTest, TestConstructorWithDifferentMetrics) { + // Test with inner product metric + auto properties = _properties; + properties["metric_type"] = "inner_product"; + + auto tablet_index = std::make_unique(); + tablet_index->_properties = properties; + tablet_index->_index_id = 2; + + auto reader = std::make_unique(tablet_index.get(), + _mock_index_file_reader); + + EXPECT_EQ(reader->get_metric_type(), segment_v2::AnnIndexMetric::IP); + EXPECT_EQ(reader->get_index_id(), 2); +} + +TEST_F(AnnIndexReaderTest, TestNewIterator) { + // TODO: Fix if we using unique_ptr here. + auto reader = std::make_shared(_tablet_index.get(), + _mock_index_file_reader); + + std::unique_ptr iterator; + Status status = reader->new_iterator(&iterator); + + EXPECT_TRUE(status.ok()); + EXPECT_NE(iterator, nullptr); + + // Verify it's an AnnIndexIterator + auto ann_iterator = dynamic_cast(iterator.get()); + EXPECT_NE(ann_iterator, nullptr); +} + +TEST_F(AnnIndexReaderTest, TestLoadIndexSuccess) { + auto reader = std::make_unique(_tablet_index.get(), + _mock_index_file_reader); + + // Mock successful index file operations + EXPECT_CALL(*_mock_index_file_reader, init(testing::_, testing::_)) + .WillOnce(testing::Return(Status::OK())); + + // For the open method that returns Result>, we need to use a different approach + // since gmock has issues with non-copyable return types + ON_CALL(*_mock_index_file_reader, open(testing::_, testing::_)) + .WillByDefault(testing::Invoke( + [](const doris::TabletIndex*, const doris::io::IOContext*) + -> doris::Result> { + return doris::ResultError(doris::Status::IOError("Mock not implemented")); + })); + + io::IOContext io_ctx; + Status status = reader->load_index(&io_ctx); + // We expect this to fail since we're not fully implementing the mock + // but it should not crash due to the copy constructor issue + EXPECT_FALSE(status.ok()); +} + +TEST_F(AnnIndexReaderTest, TestLoadIndexFailureInit) { + auto reader = std::make_unique(_tablet_index.get(), + _mock_index_file_reader); + + // Mock failed init + EXPECT_CALL(*_mock_index_file_reader, init(testing::_, testing::_)) + .WillOnce(testing::Return(Status::IOError("Init failed"))); + + io::IOContext io_ctx; + Status status = reader->load_index(&io_ctx); + EXPECT_FALSE(status.ok()); +} + +TEST_F(AnnIndexReaderTest, TestLoadIndexFailureOpen) { + auto reader = std::make_unique(_tablet_index.get(), + _mock_index_file_reader); + + // Mock successful init but failed open + EXPECT_CALL(*_mock_index_file_reader, init(testing::_, testing::_)) + .WillOnce(testing::Return(Status::OK())); + + ON_CALL(*_mock_index_file_reader, open(testing::_, testing::_)) + .WillByDefault(testing::Invoke( + [](const doris::TabletIndex*, const doris::io::IOContext*) + -> doris::Result> { + return doris::ResultError(doris::Status::IOError("Open failed")); + })); + + io::IOContext io_ctx; + Status status = reader->load_index(&io_ctx); + EXPECT_FALSE(status.ok()); +} + +TEST_F(AnnIndexReaderTest, TestQueryWithoutLoadIndex) { + auto reader = std::make_unique(_tablet_index.get(), + _mock_index_file_reader); + + // Set up _vector_index manually to bypass load_index for testing + auto doris_faiss_vector_index = std::make_unique(); + doris_faiss_vector_index->set_metric(doris::segment_v2::AnnIndexMetric::L2); + + doris::segment_v2::FaissBuildParameter build_params; + build_params.dim = 4; + build_params.max_degree = 16; + build_params.index_type = doris::segment_v2::FaissBuildParameter::IndexType::HNSW; + build_params.metric_type = doris::segment_v2::FaissBuildParameter::MetricType::L2; + doris_faiss_vector_index->build(build_params); + + reader->_vector_index = std::move(doris_faiss_vector_index); + + // Create query parameters + const float query_data[] = {1.0f, 2.0f, 3.0f, 4.0f}; + roaring::Roaring bitmap; + bitmap.add(1); + bitmap.add(2); + + segment_v2::AnnTopNParam param { + .query_value = query_data, + .query_value_size = 4, + .limit = 5, + ._user_params = VectorSearchUserParams {.hnsw_ef_search = 100, + .hnsw_check_relative_distance = false, + .hnsw_bounded_queue = false}, + .roaring = &bitmap}; + + segment_v2::AnnIndexStats stats; + io::IOContext io_ctx; + + Status status = reader->query(&io_ctx, ¶m, &stats); + + // The query might succeed or fail depending on the internal index state, + // but it should not crash and should properly initialize distance and row_ids + if (status.ok()) { + EXPECT_NE(param.distance, nullptr); + EXPECT_NE(param.row_ids, nullptr); + } +} + +TEST_F(AnnIndexReaderTest, TestRangeSearchWithoutLoadIndex) { + auto reader = std::make_unique(_tablet_index.get(), + _mock_index_file_reader); + + // Set up _vector_index manually to bypass load_index for testing + auto doris_faiss_vector_index = std::make_unique(); + doris_faiss_vector_index->set_metric(doris::segment_v2::AnnIndexMetric::L2); + + doris::segment_v2::FaissBuildParameter build_params; + build_params.dim = 4; + build_params.max_degree = 16; + build_params.index_type = doris::segment_v2::FaissBuildParameter::IndexType::HNSW; + build_params.metric_type = doris::segment_v2::FaissBuildParameter::MetricType::L2; + doris_faiss_vector_index->build(build_params); + + reader->_vector_index = std::move(doris_faiss_vector_index); + + // Create range search parameters + float query_data[] = {1.0f, 2.0f, 3.0f, 4.0f}; + roaring::Roaring bitmap; + bitmap.add(1); + bitmap.add(2); + + segment_v2::AnnRangeSearchParams params; + params.is_le_or_lt = true; + params.radius = 5.0f; + params.query_value = query_data; + params.roaring = &bitmap; + + VectorSearchUserParams user_params; + user_params.hnsw_ef_search = 50; + user_params.hnsw_check_relative_distance = true; + user_params.hnsw_bounded_queue = true; + + segment_v2::AnnRangeSearchResult result; + segment_v2::AnnIndexStats stats; + + Status status = reader->range_search(params, user_params, &result, &stats); + + // The range search might succeed or fail depending on the internal index state, + // but it should not crash + if (status.ok()) { + EXPECT_NE(result.roaring, nullptr); + } +} + +TEST_F(AnnIndexReaderTest, TestUpdateResultStatic) { + // Test the static update_result method + segment_v2::IndexSearchResult search_result; + + // Set up test data + auto roaring = std::make_shared(); + roaring->add(10); + roaring->add(20); + roaring->add(30); + + size_t num_results = 3; + auto distances = std::make_unique(num_results); + distances[0] = 1.5f; + distances[1] = 2.3f; + distances[2] = 3.1f; + + search_result.roaring = roaring; + search_result.distances = std::move(distances); + + // Call update_result + std::vector distance_vec; + roaring::Roaring result_roaring; + + segment_v2::AnnIndexReader::update_result(search_result, distance_vec, result_roaring); + + // Verify results + EXPECT_EQ(distance_vec.size(), num_results); + EXPECT_FLOAT_EQ(distance_vec[0], 1.5f); + EXPECT_FLOAT_EQ(distance_vec[1], 2.3f); + EXPECT_FLOAT_EQ(distance_vec[2], 3.1f); + EXPECT_EQ(result_roaring.cardinality(), num_results); + EXPECT_TRUE(result_roaring.contains(10)); + EXPECT_TRUE(result_roaring.contains(20)); + EXPECT_TRUE(result_roaring.contains(30)); +} + +TEST_F(AnnIndexReaderTest, TestRangeSearchWithDifferentParameters) { + auto reader = std::make_unique(_tablet_index.get(), + _mock_index_file_reader); + + // Set up _vector_index manually + auto doris_faiss_vector_index = std::make_unique(); + doris_faiss_vector_index->set_metric(doris::segment_v2::AnnIndexMetric::L2); + + doris::segment_v2::FaissBuildParameter build_params; + build_params.dim = 4; + build_params.max_degree = 16; + build_params.index_type = doris::segment_v2::FaissBuildParameter::IndexType::HNSW; + build_params.metric_type = doris::segment_v2::FaissBuildParameter::MetricType::L2; + doris_faiss_vector_index->build(build_params); + + reader->_vector_index = std::move(doris_faiss_vector_index); + + // Test case 1: is_le_or_lt = false + { + float query_data[] = {1.0f, 2.0f, 3.0f, 4.0f}; + roaring::Roaring bitmap; + bitmap.add(1); + + segment_v2::AnnRangeSearchParams params; + params.is_le_or_lt = false; // This should result in no distances/row_ids + params.radius = 5.0f; + params.query_value = query_data; + params.roaring = &bitmap; + + VectorSearchUserParams user_params; + user_params.hnsw_ef_search = 50; + + segment_v2::AnnRangeSearchResult result; + segment_v2::AnnIndexStats stats; + + Status status = reader->range_search(params, user_params, &result, &stats); + + if (status.ok()) { + // When is_le_or_lt = false, we expect no distance/row_ids + if (result.row_ids == nullptr) { + EXPECT_EQ(result.row_ids, nullptr); + } + if (result.distance == nullptr) { + EXPECT_EQ(result.distance, nullptr); + } + } + } + + // Test case 2: is_le_or_lt = true + { + float query_data[] = {1.0f, 2.0f, 3.0f, 4.0f}; + roaring::Roaring bitmap; + bitmap.add(1); + + segment_v2::AnnRangeSearchParams params; + params.is_le_or_lt = true; + params.radius = 5.0f; + params.query_value = query_data; + params.roaring = &bitmap; + + VectorSearchUserParams user_params; + user_params.hnsw_ef_search = 50; + user_params.hnsw_check_relative_distance = false; + user_params.hnsw_bounded_queue = false; + + segment_v2::AnnRangeSearchResult result; + segment_v2::AnnIndexStats stats; + + Status status = reader->range_search(params, user_params, &result, &stats); + + // This should not crash regardless of success/failure + if (status.ok()) { + EXPECT_NE(result.roaring, nullptr); + } + } +} + +TEST_F(AnnIndexReaderTest, TestWithInnerProductMetric) { + // Test with inner product metric type + auto properties = _properties; + properties["metric_type"] = "inner_product"; + + auto tablet_index = std::make_unique(); + tablet_index->_properties = properties; + tablet_index->_index_id = 1; + + auto reader = std::make_unique(tablet_index.get(), + _mock_index_file_reader); + + EXPECT_EQ(reader->get_metric_type(), segment_v2::AnnIndexMetric::IP); + + // Set up _vector_index with IP metric + auto doris_faiss_vector_index = std::make_unique(); + doris_faiss_vector_index->set_metric(doris::segment_v2::AnnIndexMetric::IP); + + doris::segment_v2::FaissBuildParameter build_params; + build_params.dim = 4; + build_params.max_degree = 16; + build_params.index_type = doris::segment_v2::FaissBuildParameter::IndexType::HNSW; + build_params.metric_type = doris::segment_v2::FaissBuildParameter::MetricType::IP; + doris_faiss_vector_index->build(build_params); + + reader->_vector_index = std::move(doris_faiss_vector_index); + + // Test query with IP metric + const float query_data[] = {1.0f, 2.0f, 3.0f, 4.0f}; + roaring::Roaring bitmap; + bitmap.add(1); + + segment_v2::AnnTopNParam param {.query_value = query_data, + .query_value_size = 4, + .limit = 5, + ._user_params = VectorSearchUserParams {}, + .roaring = &bitmap}; + + segment_v2::AnnIndexStats stats; + io::IOContext io_ctx; + + Status status = reader->query(&io_ctx, ¶m, &stats); + + // Should not crash regardless of success/failure + if (status.ok()) { + EXPECT_NE(param.distance, nullptr); + EXPECT_NE(param.row_ids, nullptr); + } +} + +TEST_F(AnnIndexReaderTest, AnnIndexReaderRangeSearch) { + size_t iterations = 5; + for (size_t i = 0; i < iterations; ++i) { + std::map index_properties; + index_properties["index_type"] = "hnsw"; + index_properties["metric_type"] = "l2"; + std::unique_ptr index_meta = std::make_unique(); + index_meta->_properties = index_properties; + auto mock_index_file_reader = std::make_shared(); + auto ann_index_reader = std::make_unique( + index_meta.get(), mock_index_file_reader); + doris::vector_search_utils::IndexType index_type = + doris::vector_search_utils::IndexType::HNSW; + const size_t dim = 128; + const size_t m = 16; + auto doris_faiss_index = doris::vector_search_utils::create_doris_index(index_type, dim, m); + auto native_faiss_index = + doris::vector_search_utils::create_native_index(index_type, dim, m); + const size_t num_vectors = 1000; + auto vectors = doris::vector_search_utils::generate_test_vectors_matrix(num_vectors, dim); + doris::vector_search_utils::add_vectors_to_indexes_serial_mode( + doris_faiss_index.get(), native_faiss_index.get(), vectors); + std::ignore = doris_faiss_index->save(this->_ram_dir.get()); + std::vector query_value = vectors[0]; + const float radius = doris::vector_search_utils::get_radius_from_matrix(query_value.data(), + dim, vectors, 0.3); + + // Make sure all rows are in the roaring + auto roaring = std::make_unique(); + for (size_t i = 0; i < num_vectors; ++i) { + roaring->add(i); + } + + doris::segment_v2::AnnRangeSearchParams params; + params.radius = radius; + params.query_value = query_value.data(); + params.roaring = roaring.get(); + doris::VectorSearchUserParams custom_params; + custom_params.hnsw_ef_search = 16; + doris::segment_v2::AnnRangeSearchResult result; + auto stats = std::make_unique(); + auto doris_faiss_vector_index = std::make_unique(); + std::ignore = doris_faiss_vector_index->load(this->_ram_dir.get()); + ann_index_reader->_vector_index = std::move(doris_faiss_vector_index); + std::ignore = ann_index_reader->range_search(params, custom_params, &result, stats.get()); + + ASSERT_TRUE(result.roaring != nullptr); + ASSERT_TRUE(result.distance != nullptr); + ASSERT_TRUE(result.row_ids != nullptr); + std::vector> doris_search_result_order_by_lables; + for (size_t i = 0; i < result.roaring->cardinality(); ++i) { + doris_search_result_order_by_lables.push_back( + {result.row_ids->at(i), result.distance[i]}); + } + + std::sort(doris_search_result_order_by_lables.begin(), + doris_search_result_order_by_lables.end(), + [](const auto& a, const auto& b) { return a.first < b.first; }); + + std::vector> native_search_result_order_by_lables = + doris::vector_search_utils::perform_native_index_range_search( + native_faiss_index.get(), query_value.data(), radius); + + ASSERT_EQ(result.roaring->cardinality(), native_search_result_order_by_lables.size()); + + for (size_t i = 0; i < native_search_result_order_by_lables.size(); ++i) { + ASSERT_EQ(doris_search_result_order_by_lables[i].first, + native_search_result_order_by_lables[i].first); + ASSERT_FLOAT_EQ(doris_search_result_order_by_lables[i].second, + native_search_result_order_by_lables[i].second); + } + } +} + +} // namespace doris::vectorized \ No newline at end of file diff --git a/be/test/olap/vector_search/ann_index_smoke_test.cpp b/be/test/olap/vector_search/ann_index_smoke_test.cpp new file mode 100644 index 00000000000000..08b08626e88fb7 --- /dev/null +++ b/be/test/olap/vector_search/ann_index_smoke_test.cpp @@ -0,0 +1,159 @@ +// 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 +#include + +// Add CLucene RAM Directory header +#include + +#include +#include + +#include "olap/olap_common.h" +#include "olap/rowset/segment_v2/ann_index/ann_index.h" +#include "olap/rowset/segment_v2/ann_index/ann_index_writer.h" +#include "olap/rowset/segment_v2/ann_index/ann_search_params.h" +#include "olap/rowset/segment_v2/ann_index/faiss_ann_index.h" +#include "olap/rowset/segment_v2/index_file_writer.h" +#include "vector_search_utils.h" + +using namespace doris::vector_search_utils; + +namespace doris { + +class AnnIndexTest : public testing::Test { +protected: + void SetUp() override { + // Create a tmp_file_dirs, this will be used by + const std::string testDir = "./ut_dir/AnnIndexTest"; + ASSERT_TRUE(io::global_local_filesystem()->delete_directory(testDir).ok()); + ASSERT_TRUE(io::global_local_filesystem()->create_directory(testDir).ok()); + std::vector paths; + paths.emplace_back(testDir, 1024); + auto tmp_file_dirs = std::make_unique(paths); + ASSERT_TRUE(tmp_file_dirs->init()); + ExecEnv::GetInstance()->set_tmp_file_dir(std::move(tmp_file_dirs)); + + _ram_dir = std::make_shared(); + auto fs = io::global_local_filesystem(); + _index_file_writer = std::make_unique(fs); + _index_meta = std::make_unique(); + _tablet_column_array = std::make_unique(); + _tablet_column_float = std::make_unique(); + + EXPECT_CALL(*_tablet_column_array, type()) + .WillRepeatedly(testing::Return(FieldType::OLAP_FIELD_TYPE_ARRAY)); + EXPECT_CALL(*_tablet_column_array, get_sub_column(0)) + .WillOnce(testing::ReturnRef(*_tablet_column_float)); + EXPECT_CALL(*_tablet_column_float, type()) + .WillOnce(testing::Return(FieldType::OLAP_FIELD_TYPE_FLOAT)); + + Field field(*_tablet_column_array); + + EXPECT_CALL(*_index_file_writer, open(_index_meta.get())) + .WillOnce(testing::Return(_ram_dir)); + + std::map properties = { + {segment_v2::AnnIndexColumnWriter::INDEX_TYPE, "hnsw"}, + {segment_v2::AnnIndexColumnWriter::DIM, "10"}, + {segment_v2::AnnIndexColumnWriter::MAX_DEGREE, "32"}}; + + EXPECT_CALL(*_index_meta, properties()).WillOnce(testing::ReturnRef(properties)); + _ann_index_col_writer = std::make_unique( + _index_file_writer.get(), _index_meta.get()); + EXPECT_TRUE(_ann_index_col_writer->init().ok()); + } + + void TearDown() override {} + + std::unique_ptr _ann_index_col_writer; + std::shared_ptr _ram_dir; + std::unique_ptr _tablet_column_array; + std::unique_ptr _tablet_column_float; + std::unique_ptr _index_file_writer; + std::unique_ptr _index_meta; +}; + +TEST_F(AnnIndexTest, SmokeTest) { + segment_v2::AnnIndexColumnWriter* ann_index_writer = + dynamic_cast(_ann_index_col_writer.get()); + ASSERT_NE(ann_index_writer, nullptr); + + // Add some dummy data + /* + [0,1,2,3,4,5,6,7,8,9] + [10,11,12,13,14,15,16,17,18,19] + [20,21,22,23,24,25,26,27,28,29] + [30,31,32,33,34,35,36,37,38,39] + [40,41,42,43,44,45,46,47,48,49] + [50,...] + */ + std::unique_ptr data(new float[10 * 10]); + for (int i = 0; i < 10 * 10; ++i) { + data[i] = static_cast(i); + } + + // Create offsets_data with size num_rows + 1; last entry is total elements + std::unique_ptr offsets_data(new size_t[11]); + size_t* offsets = offsets_data.get(); + + for (int i = 0; i < 10; ++i) { + offsets[i] = i * 10; // start offset of each row + } + offsets[10] = 10 * 10; // terminal offset + + auto st = ann_index_writer->add_array_values( + 0, data.get(), nullptr, reinterpret_cast(offsets_data.get()), 10); + EXPECT_TRUE(st.ok()) << st.to_string(); + + ASSERT_TRUE(ann_index_writer->finish().ok()); + + // Read the index file + auto index2 = std::make_unique(); + // Step 6: Load the index + auto load_status = index2->load(_ram_dir.get()); + ASSERT_TRUE(load_status.ok()) << load_status.to_string(); + // [0,1,2,3,4,5,6,7,8,9] + std::unique_ptr query_vec(new float[10]); + for (int i = 0; i < 10; ++i) { + query_vec[i] = static_cast(i); + } + + // Use HNSW search parameters for FAISS HNSW top-N search + segment_v2::HNSWSearchParameters params; + params.ef_search = 64; // reasonable default for test + // TopN search requires candidate roaring and rows_of_segment + auto all_rows = std::make_unique(); + for (int i = 0; i < 10; ++i) all_rows->add(i); + params.roaring = all_rows.get(); + params.rows_of_segment = 10; + segment_v2::IndexSearchResult result; + ASSERT_TRUE(index2->ann_topn_search(query_vec.get(), 1, params, result).ok()); + EXPECT_TRUE(result.roaring->cardinality() == 1); + EXPECT_TRUE(result.roaring->contains(0)); + + std::shared_ptr index_file_reader = + std::make_shared(); + // EXPECT_CALL(*index_file_reader, init(_, _)); + // EXPECT_CALL(*index_file_reader, open(_, _)) + // auto ann_index_reader = + // std::make_unique(_index_meta.get(), index_file_reader); + +} // namespace doris + +} // namespace doris \ No newline at end of file diff --git a/be/test/olap/vector_search/ann_index_writer_test.cpp b/be/test/olap/vector_search/ann_index_writer_test.cpp new file mode 100644 index 00000000000000..a62414181e19da --- /dev/null +++ b/be/test/olap/vector_search/ann_index_writer_test.cpp @@ -0,0 +1,387 @@ +// 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 "olap/rowset/segment_v2/ann_index/ann_index_writer.h" + +#include +#include +#include +#include + +#include +#include +#include + +#include "olap/rowset/segment_v2/index_file_writer.h" +#include "olap/rowset/segment_v2/inverted_index_fs_directory.h" +#include "olap/tablet_schema.h" +#include "runtime/collection_value.h" +#include "vector_search_utils.h" + +using namespace doris::vector_search_utils; + +namespace doris::segment_v2 { + +class AnnIndexWriterTest : public ::testing::Test { +protected: + void SetUp() override { + // Ensure ExecEnv has a valid tmp dir for IndexFileWriter (prevents nullptr deref) + if (ExecEnv::GetInstance()->get_tmp_file_dirs() == nullptr) { + const std::string tmp_dir = "./ut_dir/tmp_vector_search"; + (void)doris::io::global_local_filesystem()->delete_directory(tmp_dir); + (void)doris::io::global_local_filesystem()->create_directory(tmp_dir); + std::vector paths; + paths.emplace_back(tmp_dir, -1); + auto tmp_file_dirs = std::make_unique(paths); + ASSERT_TRUE(tmp_file_dirs->init().ok()); + ExecEnv::GetInstance()->set_tmp_file_dir(std::move(tmp_file_dirs)); + } + + // Create RAM directory for testing + _ram_dir = std::make_shared(); + + // Create test index properties + _properties["index_type"] = "hnsw"; + _properties["metric_type"] = "l2_distance"; + _properties["dim"] = "4"; + _properties["max_degree"] = "16"; + + // Create tablet index + _tablet_index = std::make_unique(); + _tablet_index->_properties = _properties; + _tablet_index->_index_id = 1; + _tablet_index->_index_name = "test_ann_index"; + + // Create mock index file writer + _index_file_writer = + std::make_unique(doris::io::global_local_filesystem()); + } + + void TearDown() override {} + + std::shared_ptr _ram_dir; + std::map _properties; + std::unique_ptr _tablet_index; + std::unique_ptr _index_file_writer; +}; + +TEST_F(AnnIndexWriterTest, TestConstructorAndDestructor) { + // Test constructor + auto writer = + std::make_unique(_index_file_writer.get(), _tablet_index.get()); + EXPECT_NE(writer, nullptr); + + // Destructor should be called automatically when writer goes out of scope +} + +TEST_F(AnnIndexWriterTest, TestInitWithDifferentProperties) { + // Test with different index types and parameters + std::vector> test_properties = { + {{"index_type", "hnsw"}, + {"metric_type", "inner_product"}, + {"dim", "8"}, + {"max_degree", "32"}}, + {{"index_type", "hnsw"}, + {"metric_type", "l2_distance"}, + {"dim", "128"}, + {"max_degree", "64"}}, + // Test with default values (missing properties) + {{"index_type", "hnsw"}}, + {}}; + + for (const auto& props : test_properties) { + auto tablet_index = std::make_unique(); + tablet_index->_properties = props; + tablet_index->_index_id = 1; + + auto writer = std::make_unique(_index_file_writer.get(), + tablet_index.get()); + + auto fs_dir = std::make_shared(); + EXPECT_CALL(*_index_file_writer, open(testing::_)).WillOnce(testing::Return(fs_dir)); + + Status status = writer->init(); + EXPECT_TRUE(status.ok()); + } +} + +TEST_F(AnnIndexWriterTest, TestAddArrayValuesSuccess) { + auto writer = + std::make_unique(_index_file_writer.get(), _tablet_index.get()); + + auto fs_dir = std::make_shared(); + fs_dir->init(doris::io::global_local_filesystem(), "./ut_dir/tmp_vector_search", nullptr); + EXPECT_CALL(*_index_file_writer, open(testing::_)).WillOnce(testing::Return(fs_dir)); + + ASSERT_TRUE(writer->init().ok()); + + // Prepare test data + const size_t dim = 4; + const size_t num_rows = 3; + std::vector vectors = { + 1.0f, 2.0f, 3.0f, 4.0f, // Row 0 + 5.0f, 6.0f, 7.0f, 8.0f, // Row 1 + 9.0f, 10.0f, 11.0f, 12.0f // Row 2 + }; + + std::vector offsets = {0, 4, 8, 12}; // Each row has 4 elements + + Status status = + writer->add_array_values(sizeof(float), vectors.data(), nullptr, + reinterpret_cast(offsets.data()), num_rows); + EXPECT_TRUE(status.ok()); +} + +TEST_F(AnnIndexWriterTest, TestAddArrayValuesEmptyRows) { + auto writer = + std::make_unique(_index_file_writer.get(), _tablet_index.get()); + + auto fs_dir = std::make_shared(); + fs_dir->init(doris::io::global_local_filesystem(), "./ut_dir/tmp_vector_search", nullptr); + EXPECT_CALL(*_index_file_writer, open(testing::_)).WillOnce(testing::Return(fs_dir)); + + ASSERT_TRUE(writer->init().ok()); + + // Test with zero rows + Status status = writer->add_array_values(sizeof(float), nullptr, nullptr, nullptr, 0); + EXPECT_TRUE(status.ok()); +} + +TEST_F(AnnIndexWriterTest, TestAddArrayValuesWrongDimension) { + auto writer = + std::make_unique(_index_file_writer.get(), _tablet_index.get()); + + auto fs_dir = std::make_shared(); + fs_dir->init(doris::io::global_local_filesystem(), "./ut_dir/tmp_vector_search", nullptr); + EXPECT_CALL(*_index_file_writer, open(testing::_)).WillOnce(testing::Return(fs_dir)); + + ASSERT_TRUE(writer->init().ok()); + + // Prepare test data with wrong dimension (expected 4, providing 3) + const size_t num_rows = 2; + std::vector vectors = { + 1.0f, 2.0f, 3.0f, // Row 0: 3 elements (wrong) + 4.0f, 5.0f, 6.0f // Row 1: 3 elements (wrong) + }; + + std::vector offsets = {0, 3, 6}; // Each row has 3 elements instead of 4 + + Status status = + writer->add_array_values(sizeof(float), vectors.data(), nullptr, + reinterpret_cast(offsets.data()), num_rows); + EXPECT_FALSE(status.ok()); + EXPECT_TRUE(status.is()); +} + +TEST_F(AnnIndexWriterTest, TestAddArrayValuesWithCollectionValue) { + auto writer = + std::make_unique(_index_file_writer.get(), _tablet_index.get()); + + auto fs_dir = std::make_shared(); + fs_dir->init(doris::io::global_local_filesystem(), "./ut_dir/tmp_vector_search", nullptr); + EXPECT_CALL(*_index_file_writer, open(testing::_)).WillOnce(testing::Return(fs_dir)); + + ASSERT_TRUE(writer->init().ok()); + + // This should return an error as ANN index doesn't support nullable columns + Status status = writer->add_array_values(sizeof(float), nullptr, 1); + EXPECT_FALSE(status.ok()); + EXPECT_TRUE(status.is()); +} + +TEST_F(AnnIndexWriterTest, TestAddValues) { + auto writer = + std::make_unique(_index_file_writer.get(), _tablet_index.get()); + + auto fs_dir = std::make_shared(); + fs_dir->init(doris::io::global_local_filesystem(), "./ut_dir/tmp_vector_search", nullptr); + EXPECT_CALL(*_index_file_writer, open(testing::_)).WillOnce(testing::Return(fs_dir)); + + ASSERT_TRUE(writer->init().ok()); + + // This method currently returns OK without doing anything + Status status = writer->add_values("test", nullptr, 0); + EXPECT_TRUE(status.ok()); +} + +TEST_F(AnnIndexWriterTest, TestAddNulls) { + auto writer = + std::make_unique(_index_file_writer.get(), _tablet_index.get()); + + auto fs_dir = std::make_shared(); + fs_dir->init(doris::io::global_local_filesystem(), "./ut_dir/tmp_vector_search", nullptr); + EXPECT_CALL(*_index_file_writer, open(testing::_)).WillOnce(testing::Return(fs_dir)); + + ASSERT_TRUE(writer->init().ok()); + + // This should return an error as ANN index doesn't support nullable columns + Status status = writer->add_nulls(10); + EXPECT_FALSE(status.ok()); + EXPECT_TRUE(status.is()); +} + +TEST_F(AnnIndexWriterTest, TestAddArrayNulls) { + auto writer = + std::make_unique(_index_file_writer.get(), _tablet_index.get()); + + auto fs_dir = std::make_shared(); + fs_dir->init(doris::io::global_local_filesystem(), "./ut_dir/tmp_vector_search", nullptr); + EXPECT_CALL(*_index_file_writer, open(testing::_)).WillOnce(testing::Return(fs_dir)); + + ASSERT_TRUE(writer->init().ok()); + + // This should return an error as ANN index doesn't support nullable columns + Status status = writer->add_array_nulls(nullptr, 10); + EXPECT_FALSE(status.ok()); + EXPECT_TRUE(status.is()); +} + +TEST_F(AnnIndexWriterTest, TestSize) { + auto writer = + std::make_unique(_index_file_writer.get(), _tablet_index.get()); + + auto fs_dir = std::make_shared(); + fs_dir->init(doris::io::global_local_filesystem(), "./ut_dir/tmp_vector_search", nullptr); + EXPECT_CALL(*_index_file_writer, open(testing::_)).WillOnce(testing::Return(fs_dir)); + + ASSERT_TRUE(writer->init().ok()); + + // Size method currently returns 0 + EXPECT_EQ(writer->size(), 0); +} + +TEST_F(AnnIndexWriterTest, TestCloseOnError) { + auto writer = + std::make_unique(_index_file_writer.get(), _tablet_index.get()); + + // close_on_error should not crash + writer->close_on_error(); +} + +TEST_F(AnnIndexWriterTest, TestFinish) { + auto writer = + std::make_unique(_index_file_writer.get(), _tablet_index.get()); + + auto fs_dir = std::make_shared(); + fs_dir->init(doris::io::global_local_filesystem(), "./ut_dir/tmp_vector_search", nullptr); + EXPECT_CALL(*_index_file_writer, open(testing::_)).WillOnce(testing::Return(fs_dir)); + + ASSERT_TRUE(writer->init().ok()); + + // Add some test data before finishing + const size_t dim = 4; + const size_t num_rows = 2; + std::vector vectors = { + 1.0f, 2.0f, 3.0f, 4.0f, // Row 0 + 5.0f, 6.0f, 7.0f, 8.0f // Row 1 + }; + + std::vector offsets = {0, 4, 8}; + + ASSERT_TRUE(writer->add_array_values(sizeof(float), vectors.data(), nullptr, + reinterpret_cast(offsets.data()), num_rows) + .ok()); + + // Finish should save the index + Status status = writer->finish(); + EXPECT_TRUE(status.ok()); +} + +TEST_F(AnnIndexWriterTest, TestFullWorkflow) { + // Test a complete workflow: init -> add_data -> finish + auto writer = + std::make_unique(_index_file_writer.get(), _tablet_index.get()); + + auto fs_dir = std::make_shared(); + fs_dir->init(doris::io::global_local_filesystem(), "./ut_dir/tmp_vector_search", nullptr); + EXPECT_CALL(*_index_file_writer, open(testing::_)).WillOnce(testing::Return(fs_dir)); + + // 1. Initialize + ASSERT_TRUE(writer->init().ok()); + + // 2. Add multiple batches of data + const size_t dim = 4; + + // Batch 1 + { + const size_t num_rows = 2; + std::vector vectors = {1.0f, 2.0f, 3.0f, 4.0f, 5.0f, 6.0f, 7.0f, 8.0f}; + std::vector offsets = {0, 4, 8}; + + ASSERT_TRUE(writer->add_array_values(sizeof(float), vectors.data(), nullptr, + reinterpret_cast(offsets.data()), + num_rows) + .ok()); + } + + // Batch 2 + { + const size_t num_rows = 3; + std::vector vectors = {9.0f, 10.0f, 11.0f, 12.0f, 13.0f, 14.0f, + 15.0f, 16.0f, 17.0f, 18.0f, 19.0f, 20.0f}; + std::vector offsets = {0, 4, 8, 12}; + + ASSERT_TRUE(writer->add_array_values(sizeof(float), vectors.data(), nullptr, + reinterpret_cast(offsets.data()), + num_rows) + .ok()); + } + + // 3. Finish + ASSERT_TRUE(writer->finish().ok()); +} + +TEST_F(AnnIndexWriterTest, TestInvalidIndexType) { + // Test with invalid index type + auto properties = _properties; + properties["index_type"] = "invalid_type"; + + auto tablet_index = std::make_unique(); + tablet_index->_properties = properties; + tablet_index->_index_id = 1; + + auto writer = + std::make_unique(_index_file_writer.get(), tablet_index.get()); + + auto fs_dir = std::make_shared(); + fs_dir->init(doris::io::global_local_filesystem(), "./ut_dir/tmp_vector_search", nullptr); + EXPECT_CALL(*_index_file_writer, open(testing::_)).WillOnce(testing::Return(fs_dir)); + + // This should throw an exception due to invalid index type + EXPECT_THROW(writer->init(), doris::Exception); +} + +TEST_F(AnnIndexWriterTest, TestInvalidMetricType) { + // Test with invalid metric type + auto properties = _properties; + properties["metric_type"] = "invalid_metric"; + + auto tablet_index = std::make_unique(); + tablet_index->_properties = properties; + tablet_index->_index_id = 1; + + auto writer = + std::make_unique(_index_file_writer.get(), tablet_index.get()); + + auto fs_dir = std::make_shared(); + EXPECT_CALL(*_index_file_writer, open(testing::_)).WillOnce(testing::Return(fs_dir)); + + // This should throw an exception due to invalid metric type + EXPECT_THROW(writer->init(), doris::Exception); +} + +} // namespace doris::segment_v2 diff --git a/be/test/olap/vector_search/ann_range_search_test.cpp b/be/test/olap/vector_search/ann_range_search_test.cpp new file mode 100644 index 00000000000000..119ca6a54cce63 --- /dev/null +++ b/be/test/olap/vector_search/ann_range_search_test.cpp @@ -0,0 +1,842 @@ +// 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 +#include +#include +#include + +#include +#include +#include +#include +#include + +#include "common/object_pool.h" +#include "olap/rowset/segment_v2/ann_index/ann_index_iterator.h" +#include "olap/rowset/segment_v2/ann_index/ann_index_reader.h" +#include "olap/rowset/segment_v2/ann_index/ann_range_search_runtime.h" +#include "olap/rowset/segment_v2/ann_index/ann_search_params.h" +#include "olap/rowset/segment_v2/ann_index/faiss_ann_index.h" +#include "olap/rowset/segment_v2/column_reader.h" +#include "olap/rowset/segment_v2/virtual_column_iterator.h" +#include "olap/vector_search/vector_search_utils.h" +#include "runtime/descriptors.h" +#include "runtime/runtime_state.h" +#include "vec/columns/column.h" +#include "vec/columns/column_nothing.h" +#include "vec/columns/column_nullable.h" +#include "vec/exprs/vexpr_context.h" +#include "vec/exprs/vexpr_fwd.h" +#include "vec/functions/functions_comparison.h" + +namespace doris::vectorized { +const std::string ann_range_search_thrift = + R"xxx({"1":{"lst":["rec",3,{"1":{"i32":2},"2":{"rec":{"1":{"lst":["rec",1,{"1":{"i32":0},"2":{"rec":{"1":{"i32":2}}}}]},"3":{"i64":-1}}},"3":{"i32":14},"4":{"i32":2},"20":{"i32":-1},"26":{"rec":{"1":{"rec":{"2":{"str":"ge"}}},"2":{"i32":0},"3":{"lst":["rec",2,{"1":{"lst":["rec",1,{"1":{"i32":0},"2":{"rec":{"1":{"i32":8}}}}]},"3":{"i64":-1}},{"1":{"lst":["rec",1,{"1":{"i32":0},"2":{"rec":{"1":{"i32":8}}}}]},"3":{"i64":-1}}]},"4":{"rec":{"1":{"lst":["rec",1,{"1":{"i32":0},"2":{"rec":{"1":{"i32":2}}}}]},"3":{"i64":-1}}},"5":{"tf":0},"7":{"str":"ge(double, double)"},"11":{"i64":0},"13":{"tf":1},"14":{"tf":0},"15":{"tf":0},"16":{"i64":360}}},"28":{"i32":8},"29":{"tf":1}},{"1":{"i32":16},"2":{"rec":{"1":{"lst":["rec",1,{"1":{"i32":0},"2":{"rec":{"1":{"i32":8}}}}]},"3":{"i64":-1}}},"4":{"i32":0},"15":{"rec":{"1":{"i32":3},"2":{"i32":0},"3":{"i32":-1},"4":{"tf":1}}},"20":{"i32":-1},"29":{"tf":1},"36":{"str":"dis"}},{"1":{"i32":8},"2":{"rec":{"1":{"lst":["rec",1,{"1":{"i32":0},"2":{"rec":{"1":{"i32":8}}}}]},"3":{"i64":-1}}},"4":{"i32":0},"9":{"rec":{"1":{"dbl":10}}},"20":{"i32":-1},"29":{"tf":0}}]}})xxx"; + +const std::string thrift_table_desc = + R"xxx({"1":{"lst":["rec",8,{"1":{"i32":0},"2":{"i32":0},"3":{"rec":{"1":{"lst":["rec",1,{"1":{"i32":0},"2":{"rec":{"1":{"i32":5}}}}]},"3":{"i64":-1}}},"4":{"i32":-1},"5":{"i32":-1},"6":{"i32":0},"7":{"i32":0},"8":{"str":"siteid"},"9":{"i32":0},"10":{"tf":1},"11":{"i32":0},"12":{"tf":1},"13":{"tf":1},"14":{"tf":0},"16":{"str":"10"},"17":{"i32":5}},{"1":{"i32":1},"2":{"i32":0},"3":{"rec":{"1":{"lst":["rec",2,{"1":{"i32":1},"4":{"tf":1},"5":{"lst":["tf",1,1]}},{"1":{"i32":0},"2":{"rec":{"1":{"i32":7}}}}]},"3":{"i64":-1}}},"4":{"i32":-1},"5":{"i32":-1},"6":{"i32":0},"7":{"i32":-1},"8":{"str":"embedding"},"9":{"i32":3},"10":{"tf":1},"11":{"i32":1},"12":{"tf":0},"13":{"tf":1},"14":{"tf":0},"17":{"i32":20}},{"1":{"i32":2},"2":{"i32":0},"3":{"rec":{"1":{"lst":["rec",1,{"1":{"i32":0},"2":{"rec":{"1":{"i32":23},"2":{"i32":2147483643}}}}]},"3":{"i64":-1}}},"4":{"i32":-1},"5":{"i32":-1},"6":{"i32":0},"7":{"i32":0},"8":{"str":"comment"},"9":{"i32":2},"10":{"tf":1},"11":{"i32":2},"12":{"tf":0},"13":{"tf":1},"14":{"tf":0},"17":{"i32":23}},{"1":{"i32":3},"2":{"i32":0},"3":{"rec":{"1":{"lst":["rec",1,{"1":{"i32":0},"2":{"rec":{"1":{"i32":8}}}}]},"3":{"i64":-1}}},"4":{"i32":-1},"5":{"i32":-1},"6":{"i32":0},"7":{"i32":0},"8":{"str":""},"9":{"i32":1},"10":{"tf":1},"11":{"i32":-1},"12":{"tf":0},"13":{"tf":1},"14":{"tf":0},"17":{"i32":0},"18":{"rec":{"1":{"lst":["rec",12,{"1":{"i32":20},"2":{"rec":{"1":{"lst":["rec",1,{"1":{"i32":0},"2":{"rec":{"1":{"i32":8}}}}]},"3":{"i64":-1}}},"4":{"i32":2},"20":{"i32":-1},"26":{"rec":{"1":{"rec":{"2":{"str":"l2_distance_approximate"}}},"2":{"i32":0},"3":{"lst":["rec",2,{"1":{"lst":["rec",2,{"1":{"i32":1},"4":{"tf":1},"5":{"lst":["tf",1,1]}},{"1":{"i32":0},"2":{"rec":{"1":{"i32":8}}}}]},"3":{"i64":-1}},{"1":{"lst":["rec",2,{"1":{"i32":1},"4":{"tf":1},"5":{"lst":["tf",1,1]}},{"1":{"i32":0},"2":{"rec":{"1":{"i32":8}}}}]},"3":{"i64":-1}}]},"4":{"rec":{"1":{"lst":["rec",1,{"1":{"i32":0},"2":{"rec":{"1":{"i32":8}}}}]},"3":{"i64":-1}}},"5":{"tf":0},"7":{"str":"l2_distance_approximate(array, array)"},"9":{"rec":{"1":{"str":""}}},"11":{"i64":0},"13":{"tf":1},"14":{"tf":0},"15":{"tf":0},"16":{"i64":360}}},"29":{"tf":1}},{"1":{"i32":5},"2":{"rec":{"1":{"lst":["rec",2,{"1":{"i32":1},"4":{"tf":1},"5":{"lst":["tf",1,1]}},{"1":{"i32":0},"2":{"rec":{"1":{"i32":8}}}}]},"3":{"i64":-1}}},"3":{"i32":4},"4":{"i32":1},"20":{"i32":-1},"26":{"rec":{"1":{"rec":{"2":{"str":"casttoarray"}}},"2":{"i32":0},"3":{"lst":["rec",1,{"1":{"lst":["rec",2,{"1":{"i32":1},"4":{"tf":1},"5":{"lst":["tf",1,1]}},{"1":{"i32":0},"2":{"rec":{"1":{"i32":7}}}}]},"3":{"i64":-1}}]},"4":{"rec":{"1":{"lst":["rec",2,{"1":{"i32":1},"4":{"tf":1},"5":{"lst":["tf",1,1]}},{"1":{"i32":0},"2":{"rec":{"1":{"i32":8}}}}]},"3":{"i64":-1}}},"5":{"tf":0},"7":{"str":"casttoarray(array)"},"9":{"rec":{"1":{"str":""}}},"11":{"i64":0},"13":{"tf":1},"14":{"tf":0},"15":{"tf":0},"16":{"i64":360}}},"29":{"tf":0}},{"1":{"i32":16},"2":{"rec":{"1":{"lst":["rec",2,{"1":{"i32":1},"4":{"tf":1},"5":{"lst":["tf",1,1]}},{"1":{"i32":0},"2":{"rec":{"1":{"i32":7}}}}]},"3":{"i64":-1}}},"4":{"i32":0},"15":{"rec":{"1":{"i32":1},"2":{"i32":0},"3":{"i32":1},"4":{"tf":0}}},"20":{"i32":-1},"29":{"tf":0},"36":{"str":"embedding"}},{"1":{"i32":21},"2":{"rec":{"1":{"lst":["rec",2,{"1":{"i32":1},"4":{"tf":1},"5":{"lst":["tf",1,1]}},{"1":{"i32":0},"2":{"rec":{"1":{"i32":8}}}}]},"3":{"i64":-1}}},"4":{"i32":8},"20":{"i32":-1},"28":{"i32":8},"29":{"tf":0}},{"1":{"i32":8},"2":{"rec":{"1":{"lst":["rec",1,{"1":{"i32":0},"2":{"rec":{"1":{"i32":8}}}}]},"3":{"i64":-1}}},"4":{"i32":0},"9":{"rec":{"1":{"dbl":1}}},"20":{"i32":-1},"29":{"tf":0}},{"1":{"i32":8},"2":{"rec":{"1":{"lst":["rec",1,{"1":{"i32":0},"2":{"rec":{"1":{"i32":8}}}}]},"3":{"i64":-1}}},"4":{"i32":0},"9":{"rec":{"1":{"dbl":2}}},"20":{"i32":-1},"29":{"tf":0}},{"1":{"i32":8},"2":{"rec":{"1":{"lst":["rec",1,{"1":{"i32":0},"2":{"rec":{"1":{"i32":8}}}}]},"3":{"i64":-1}}},"4":{"i32":0},"9":{"rec":{"1":{"dbl":3}}},"20":{"i32":-1},"29":{"tf":0}},{"1":{"i32":8},"2":{"rec":{"1":{"lst":["rec",1,{"1":{"i32":0},"2":{"rec":{"1":{"i32":8}}}}]},"3":{"i64":-1}}},"4":{"i32":0},"9":{"rec":{"1":{"dbl":4}}},"20":{"i32":-1},"29":{"tf":0}},{"1":{"i32":8},"2":{"rec":{"1":{"lst":["rec",1,{"1":{"i32":0},"2":{"rec":{"1":{"i32":8}}}}]},"3":{"i64":-1}}},"4":{"i32":0},"9":{"rec":{"1":{"dbl":5}}},"20":{"i32":-1},"29":{"tf":0}},{"1":{"i32":8},"2":{"rec":{"1":{"lst":["rec",1,{"1":{"i32":0},"2":{"rec":{"1":{"i32":8}}}}]},"3":{"i64":-1}}},"4":{"i32":0},"9":{"rec":{"1":{"dbl":6}}},"20":{"i32":-1},"29":{"tf":0}},{"1":{"i32":8},"2":{"rec":{"1":{"lst":["rec",1,{"1":{"i32":0},"2":{"rec":{"1":{"i32":8}}}}]},"3":{"i64":-1}}},"4":{"i32":0},"9":{"rec":{"1":{"dbl":7}}},"20":{"i32":-1},"29":{"tf":0}},{"1":{"i32":8},"2":{"rec":{"1":{"lst":["rec",1,{"1":{"i32":0},"2":{"rec":{"1":{"i32":8}}}}]},"3":{"i64":-1}}},"4":{"i32":0},"9":{"rec":{"1":{"dbl":20}}},"20":{"i32":-1},"29":{"tf":0}}]}}}},{"1":{"i32":4},"2":{"i32":1},"3":{"rec":{"1":{"lst":["rec",1,{"1":{"i32":0},"2":{"rec":{"1":{"i32":5}}}}]},"3":{"i64":-1}}},"4":{"i32":-1},"5":{"i32":-1},"6":{"i32":0},"7":{"i32":0},"8":{"str":"siteid"},"9":{"i32":0},"10":{"tf":1},"11":{"i32":0},"12":{"tf":1},"13":{"tf":1},"14":{"tf":0},"16":{"str":"10"},"17":{"i32":5}},{"1":{"i32":5},"2":{"i32":1},"3":{"rec":{"1":{"lst":["rec",2,{"1":{"i32":1},"4":{"tf":1},"5":{"lst":["tf",1,1]}},{"1":{"i32":0},"2":{"rec":{"1":{"i32":7}}}}]},"3":{"i64":-1}}},"4":{"i32":-1},"5":{"i32":-1},"6":{"i32":0},"7":{"i32":-1},"8":{"str":"embedding"},"9":{"i32":3},"10":{"tf":1},"11":{"i32":1},"12":{"tf":0},"13":{"tf":1},"14":{"tf":0},"17":{"i32":20}},{"1":{"i32":6},"2":{"i32":1},"3":{"rec":{"1":{"lst":["rec",1,{"1":{"i32":0},"2":{"rec":{"1":{"i32":23},"2":{"i32":2147483643}}}}]},"3":{"i64":-1}}},"4":{"i32":-1},"5":{"i32":-1},"6":{"i32":0},"7":{"i32":0},"8":{"str":"comment"},"9":{"i32":2},"10":{"tf":1},"11":{"i32":2},"12":{"tf":0},"13":{"tf":1},"14":{"tf":0},"17":{"i32":23}},{"1":{"i32":7},"2":{"i32":1},"3":{"rec":{"1":{"lst":["rec",1,{"1":{"i32":0},"2":{"rec":{"1":{"i32":8}}}}]},"3":{"i64":-1}}},"4":{"i32":-1},"5":{"i32":-1},"6":{"i32":0},"7":{"i32":0},"8":{"str":""},"9":{"i32":1},"10":{"tf":1},"11":{"i32":-1},"12":{"tf":0},"13":{"tf":1},"14":{"tf":0},"17":{"i32":0}}]},"2":{"lst":["rec",2,{"1":{"i32":0},"2":{"i32":0},"3":{"i32":0},"4":{"i64":-1273902531}},{"1":{"i32":1},"2":{"i32":0},"3":{"i32":0}}]},"3":{"lst":["rec",1,{"1":{"i64":1746777786941},"2":{"i32":1},"3":{"i32":3},"4":{"i32":0},"7":{"str":"vector_table_3"},"8":{"str":""},"11":{"rec":{"1":{"str":"vector_table_3"}}}}]}})xxx"; + +TEST_F(VectorSearchTest, TestPrepareAnnRangeSearch) { + TExpr texpr = read_from_json(ann_range_search_thrift); + // std::cout << "range_search thrift:\n" << apache::thrift::ThriftDebugString(texpr) << std::endl; + // TExpr distance_function_call = + // read_from_json(distance_function_call_thrift); + TDescriptorTable table1 = read_from_json(thrift_table_desc); + // std::cout << "table thrift:\n" << apache::thrift::ThriftDebugString(table1) << std::endl; + std::unique_ptr pool = std::make_unique(); + auto desc_tbl = std::make_unique(); + DescriptorTbl* desc_tbl_ptr = desc_tbl.get(); + ASSERT_TRUE(DescriptorTbl::create(pool.get(), table1, &(desc_tbl_ptr)).ok()); + RowDescriptor row_desc = RowDescriptor(*desc_tbl_ptr, {0}, {false}); + std::unique_ptr state = std::make_unique(); + state->set_desc_tbl(desc_tbl_ptr); + + VExprContextSPtr range_search_ctx; + doris::VectorSearchUserParams user_params; + ASSERT_TRUE(vectorized::VExpr::create_expr_tree(texpr, range_search_ctx).ok()); + ASSERT_TRUE(range_search_ctx->prepare(state.get(), row_desc).ok()); + ASSERT_TRUE(range_search_ctx->open(state.get()).ok()); + range_search_ctx->prepare_ann_range_search(user_params); + ASSERT_TRUE(range_search_ctx->_ann_range_search_runtime.is_ann_range_search == true); + ASSERT_EQ(range_search_ctx->_ann_range_search_runtime.is_le_or_lt, false); + ASSERT_EQ(range_search_ctx->_ann_range_search_runtime.dst_col_idx, 3); + ASSERT_EQ(range_search_ctx->_ann_range_search_runtime.src_col_idx, 1); + ASSERT_EQ(range_search_ctx->_ann_range_search_runtime.radius, 10); + + doris::segment_v2::AnnRangeSearchParams ann_range_search_runtime = + range_search_ctx->_ann_range_search_runtime.to_range_search_params(); + EXPECT_EQ(ann_range_search_runtime.radius, 10.0f); + std::vector query_array_groud_truth = {1, 2, 3, 4, 5, 6, 7, 20}; + std::vector query_array_f32; + for (int i = 0; i < query_array_groud_truth.size(); ++i) { + query_array_f32.push_back(static_cast(ann_range_search_runtime.query_value[i])); + } + for (int i = 0; i < query_array_f32.size(); ++i) { + EXPECT_EQ(query_array_f32[i], query_array_groud_truth[i]); + } +} + +TEST_F(VectorSearchTest, TestEvaluateAnnRangeSearch) { + TExpr texpr = read_from_json(ann_range_search_thrift); + TDescriptorTable table1 = read_from_json(thrift_table_desc); + std::unique_ptr pool = std::make_unique(); + auto desc_tbl = std::make_unique(); + DescriptorTbl* desc_tbl_ptr = desc_tbl.get(); + ASSERT_TRUE(DescriptorTbl::create(pool.get(), table1, &(desc_tbl_ptr)).ok()); + RowDescriptor row_desc = RowDescriptor(*desc_tbl_ptr, {0}, {false}); + std::unique_ptr state = std::make_unique(); + state->set_desc_tbl(desc_tbl_ptr); + + VExprContextSPtr range_search_ctx; + ASSERT_TRUE(vectorized::VExpr::create_expr_tree(texpr, range_search_ctx).ok()); + ASSERT_TRUE(range_search_ctx->prepare(state.get(), row_desc).ok()); + ASSERT_TRUE(range_search_ctx->open(state.get()).ok()); + doris::VectorSearchUserParams user_params; + range_search_ctx->prepare_ann_range_search(user_params); + ASSERT_EQ(range_search_ctx->_ann_range_search_runtime.user_params, user_params); + ASSERT_TRUE(range_search_ctx->_ann_range_search_runtime.is_ann_range_search == true); + ASSERT_EQ(range_search_ctx->_ann_range_search_runtime.is_le_or_lt, false); + ASSERT_EQ(range_search_ctx->_ann_range_search_runtime.src_col_idx, 1); + ASSERT_EQ(range_search_ctx->_ann_range_search_runtime.dst_col_idx, 3); + ASSERT_EQ(range_search_ctx->_ann_range_search_runtime.radius, 10); + + std::vector idx_to_cid; + idx_to_cid.resize(4); + idx_to_cid[0] = 0; + idx_to_cid[1] = 1; + idx_to_cid[2] = 2; + idx_to_cid[3] = 3; + std::vector> cid_to_index_iterators; + cid_to_index_iterators.resize(4); + cid_to_index_iterators[1] = + std::make_unique(); + std::vector> column_iterators; + column_iterators.resize(4); + column_iterators[3] = std::make_unique(); + + roaring::Roaring row_bitmap; + doris::vector_search_utils::MockAnnIndexIterator* mock_ann_index_iter = + dynamic_cast( + cid_to_index_iterators[1].get()); + + std::map properties; + properties["index_type"] = "hnsw"; + properties["metric_type"] = "l2_distance"; + auto pair = vector_search_utils::create_tmp_ann_index_reader(properties); + mock_ann_index_iter->_ann_reader = pair.second; + + // Explain: + // 1. predicate is dist >= 10, so it is not a within range search + // 2. return 10 results + EXPECT_CALL( + *mock_ann_index_iter, + range_search(testing::Truly([](const doris::segment_v2::AnnRangeSearchParams& params) { + return params.is_le_or_lt == false && params.radius == 10.0f; + }), + testing::_, testing::_, testing::_)) + .WillOnce(testing::Invoke([](const doris::segment_v2::AnnRangeSearchParams& params, + const doris::VectorSearchUserParams& custom_params, + doris::segment_v2::AnnRangeSearchResult* result, + doris::segment_v2::AnnIndexStats* stats) { + result->roaring = std::make_shared(); + result->row_ids = nullptr; + result->distance = nullptr; + return Status::OK(); + })); + + segment_v2::AnnIndexStats stats; + ASSERT_TRUE(range_search_ctx + ->evaluate_ann_range_search(cid_to_index_iterators, idx_to_cid, + column_iterators, row_bitmap, stats) + .ok()); + + doris::segment_v2::VirtualColumnIterator* virtual_column_iter = + dynamic_cast(column_iterators[3].get()); + vectorized::IColumn::Ptr column = virtual_column_iter->get_materialized_column(); + const vectorized::ColumnFloat32* float_column = + check_and_get_column(column.get()); + const vectorized::ColumnNothing* nothing_column = + check_and_get_column(column.get()); + ASSERT_EQ(float_column, nullptr); + ASSERT_NE(nothing_column, nullptr); + EXPECT_EQ(column->size(), 0); + + const auto& get_row_id_to_idx = virtual_column_iter->get_row_id_to_idx(); + EXPECT_EQ(get_row_id_to_idx.size(), 0); +} + +TEST_F(VectorSearchTest, TestEvaluateAnnRangeSearch2) { + // Modify expr from dis >= 10 to dis < 10 + TExpr texpr = read_from_json(ann_range_search_thrift); + TExprNode& opNode = texpr.nodes[0]; + opNode.opcode = TExprOpcode::LT; + opNode.fn.name.function_name = doris::vectorized::NameLess::name; + TDescriptorTable table1 = read_from_json(thrift_table_desc); + std::unique_ptr pool = std::make_unique(); + auto desc_tbl = std::make_unique(); + DescriptorTbl* desc_tbl_ptr = desc_tbl.get(); + ASSERT_TRUE(DescriptorTbl::create(pool.get(), table1, &(desc_tbl_ptr)).ok()); + RowDescriptor row_desc = RowDescriptor(*desc_tbl_ptr, {0}, {false}); + std::unique_ptr state = std::make_unique(); + state->set_desc_tbl(desc_tbl_ptr); + + VExprContextSPtr range_search_ctx; + ASSERT_TRUE(vectorized::VExpr::create_expr_tree(texpr, range_search_ctx).ok()); + ASSERT_TRUE(range_search_ctx->prepare(state.get(), row_desc).ok()); + ASSERT_TRUE(range_search_ctx->open(state.get()).ok()); + doris::VectorSearchUserParams user_params; + range_search_ctx->prepare_ann_range_search(user_params); + ASSERT_TRUE(range_search_ctx->_ann_range_search_runtime.is_ann_range_search == true); + ASSERT_EQ(range_search_ctx->_ann_range_search_runtime.is_le_or_lt, true); + ASSERT_EQ(range_search_ctx->_ann_range_search_runtime.src_col_idx, 1); + ASSERT_EQ(range_search_ctx->_ann_range_search_runtime.dst_col_idx, 3); + ASSERT_EQ(range_search_ctx->_ann_range_search_runtime.radius, 10); + + std::vector idx_to_cid; + idx_to_cid.resize(4); + idx_to_cid[0] = 0; + idx_to_cid[1] = 1; + idx_to_cid[2] = 2; + idx_to_cid[3] = 3; + std::vector> cid_to_index_iterators; + cid_to_index_iterators.resize(4); + cid_to_index_iterators[1] = + std::make_unique(); + std::vector> column_iterators; + column_iterators.resize(4); + column_iterators[3] = std::make_unique(); + roaring::Roaring row_bitmap; + doris::vector_search_utils::MockAnnIndexIterator* mock_ann_index_iter = + dynamic_cast( + cid_to_index_iterators[1].get()); + std::map properties; + properties["index_type"] = "hnsw"; + properties["metric_type"] = "l2_distance"; + auto pair = vector_search_utils::create_tmp_ann_index_reader(properties); + mock_ann_index_iter->_ann_reader = pair.second; + + // Explain: + // 1. predicate is dist >= 10, so it is not a within range search + // 2. return 10 results + EXPECT_CALL( + *mock_ann_index_iter, + range_search(testing::Truly([](const doris::segment_v2::AnnRangeSearchParams& params) { + return params.is_le_or_lt == true && params.radius == 10.0f; + }), + testing::_, testing::_, testing::_)) + .WillOnce(testing::Invoke([](const doris::segment_v2::AnnRangeSearchParams& params, + const doris::VectorSearchUserParams& custom_params, + doris::segment_v2::AnnRangeSearchResult* result, + doris::segment_v2::AnnIndexStats* stats) { + size_t num_results = 10; + result->roaring = std::make_shared(); + result->row_ids = std::make_unique>(); + for (size_t i = 0; i < num_results; ++i) { + result->roaring->add(i * 10); + result->row_ids->push_back(i * 10); + } + result->distance = std::make_unique(10); + return Status::OK(); + })); + + segment_v2::AnnIndexStats stats; + ASSERT_TRUE(range_search_ctx + ->evaluate_ann_range_search(cid_to_index_iterators, idx_to_cid, + column_iterators, row_bitmap, stats) + .ok()); + + doris::segment_v2::VirtualColumnIterator* virtual_column_iter = + dynamic_cast(column_iterators[3].get()); + + vectorized::IColumn::Ptr column = virtual_column_iter->get_materialized_column(); + const vectorized::ColumnNullable* nullable_column = + check_and_get_column(column.get()); + const vectorized::ColumnNothing* nothing_column = + check_and_get_column(column.get()); + ASSERT_NE(nullable_column, nullptr); + ASSERT_EQ(nothing_column, nullptr); + EXPECT_EQ(nullable_column->size(), 10); + EXPECT_EQ(row_bitmap.cardinality(), 10); + + const auto& get_row_id_to_idx = virtual_column_iter->get_row_id_to_idx(); + EXPECT_EQ(get_row_id_to_idx.size(), 10); +} + +TEST_F(VectorSearchTest, TestRangeSearchRuntimeInfoToString) { + // Test default constructor + doris::segment_v2::AnnRangeSearchRuntime runtime_info; + std::string result = runtime_info.to_string(); + + // Check that default values are included in the string + EXPECT_TRUE(result.find("is_ann_range_search: false") != std::string::npos); + EXPECT_TRUE(result.find("is_le_or_lt: true") != std::string::npos); + EXPECT_TRUE(result.find("src_col_idx: 0") != std::string::npos); + EXPECT_TRUE(result.find("dst_col_idx: -1") != std::string::npos); + EXPECT_TRUE(result.find("radius: 0") != std::string::npos); + EXPECT_TRUE(result.find("query_vector is null: true") != std::string::npos); + EXPECT_TRUE(result.find("metric_type UNKNOWN") != std::string::npos); + + // Test with configured values + doris::segment_v2::AnnRangeSearchRuntime runtime_info2; + runtime_info2.is_ann_range_search = true; + runtime_info2.is_le_or_lt = false; + runtime_info2.src_col_idx = 5; + runtime_info2.dst_col_idx = 3; + runtime_info2.radius = 15.5; + runtime_info2.metric_type = doris::segment_v2::AnnIndexMetric::L2; + runtime_info2.dim = 4; + runtime_info2.query_value = std::make_unique(4); + runtime_info2.query_value[0] = 1.0f; + runtime_info2.query_value[1] = 2.0f; + runtime_info2.query_value[2] = 3.0f; + runtime_info2.query_value[3] = 4.0f; + + doris::VectorSearchUserParams user_params; + user_params.hnsw_ef_search = 100; + user_params.hnsw_check_relative_distance = false; + user_params.hnsw_bounded_queue = false; + runtime_info2.user_params = user_params; + + std::string result2 = runtime_info2.to_string(); + + // Check that configured values are included in the string + EXPECT_TRUE(result2.find("is_ann_range_search: true") != std::string::npos); + EXPECT_TRUE(result2.find("is_le_or_lt: false") != std::string::npos); + EXPECT_TRUE(result2.find("src_col_idx: 5") != std::string::npos); + EXPECT_TRUE(result2.find("dst_col_idx: 3") != std::string::npos); + EXPECT_TRUE(result2.find("radius: 15.5") != std::string::npos); + EXPECT_TRUE(result2.find("query_vector is null: false") != std::string::npos); + EXPECT_TRUE(result2.find("metric_type l2_distance") != std::string::npos); + + // Test copy constructor preserves to_string output + doris::segment_v2::AnnRangeSearchRuntime runtime_info3(runtime_info2); + std::string result3 = runtime_info3.to_string(); + EXPECT_EQ(result2, result3); + + // Test assignment operator preserves to_string output + doris::segment_v2::AnnRangeSearchRuntime runtime_info4; + runtime_info4 = runtime_info2; + std::string result4 = runtime_info4.to_string(); + EXPECT_EQ(result2, result4); + + // Test with different metric types + doris::segment_v2::AnnRangeSearchRuntime runtime_info5; + runtime_info5.metric_type = doris::segment_v2::AnnIndexMetric::IP; + std::string result5 = runtime_info5.to_string(); + EXPECT_TRUE(result5.find("metric_type inner_product") != std::string::npos); + + // Test with null query_value + doris::segment_v2::AnnRangeSearchRuntime runtime_info6; + runtime_info6.query_value = nullptr; + std::string result6 = runtime_info6.to_string(); + EXPECT_TRUE(result6.find("query_vector is null: true") != std::string::npos); +} + +TEST_F(VectorSearchTest, TestAnnIndexIteratorErrorCases) { + // Test AnnIndexIterator::read_from_index with null param + // Create a mock reader first + std::map properties; + properties["index_type"] = "hnsw"; + properties["metric_type"] = "l2_distance"; + auto pair = vector_search_utils::create_tmp_ann_index_reader(properties); + + doris::segment_v2::AnnIndexIterator ann_iterator(pair.second); + + // Create a mock IndexParam with null AnnTopNParam + doris::segment_v2::IndexParam param; + param = static_cast(nullptr); + + auto status = ann_iterator.read_from_index(param); + EXPECT_FALSE(status.ok()); + EXPECT_TRUE(status.to_string().find("a_param is null") != std::string::npos); + + // Test AnnIndexIterator::range_search with null _ann_reader + // Create iterator with null reader + doris::segment_v2::AnnIndexIterator ann_iterator_null(nullptr); + doris::segment_v2::AnnRangeSearchParams range_params; + doris::VectorSearchUserParams user_params; + doris::segment_v2::AnnRangeSearchResult result; + doris::segment_v2::AnnIndexStats stats; + + status = ann_iterator_null.range_search(range_params, user_params, &result, &stats); + EXPECT_FALSE(status.ok()); + EXPECT_TRUE(status.to_string().find("_ann_reader is null") != std::string::npos); +} + +TEST_F(VectorSearchTest, TestAnnIndexIteratorSuccessCases) { + // Test successful cases to cover the remaining lines + std::map properties; + properties["index_type"] = "hnsw"; + properties["metric_type"] = "l2_distance"; + auto pair = vector_search_utils::create_tmp_ann_index_reader(properties); + + // Create iterator with valid reader + auto mock_iterator = std::make_unique(); + mock_iterator->_ann_reader = pair.second; + + // Test read_from_index with valid param + const float query_data[] = {1.0f, 2.0f, 3.0f, 4.0f}; + roaring::Roaring bitmap; + doris::segment_v2::AnnTopNParam ann_param = { + .query_value = query_data, + .query_value_size = 4, + .limit = 10, + ._user_params = doris::VectorSearchUserParams {}, + .roaring = &bitmap, + .distance = nullptr, + .row_ids = nullptr, + .stats = std::make_unique()}; + doris::segment_v2::IndexParam param = &ann_param; + + // Mock the query method to return OK + EXPECT_CALL(*mock_iterator, read_from_index(testing::_)) + .WillOnce(testing::Return(Status::OK())); + + auto status = mock_iterator->read_from_index(param); + EXPECT_TRUE(status.ok()); + + // Test range_search with valid parameters + doris::segment_v2::AnnRangeSearchParams range_params; + doris::VectorSearchUserParams user_params; + doris::segment_v2::AnnRangeSearchResult result; + doris::segment_v2::AnnIndexStats stats; + + EXPECT_CALL(*mock_iterator, range_search(testing::_, testing::_, testing::_, testing::_)) + .WillOnce(testing::Return(Status::OK())); + + status = mock_iterator->range_search(range_params, user_params, &result, &stats); + EXPECT_TRUE(status.ok()); +} + +TEST_F(VectorSearchTest, TestAnnIndexReaderUpdateResult) { + // Test AnnIndexReader::update_result method + std::map properties; + properties["index_type"] = "hnsw"; + properties["metric_type"] = "l2_distance"; + auto pair = vector_search_utils::create_tmp_ann_index_reader(properties); + auto reader = pair.second; + + // Create mock IndexSearchResult + doris::segment_v2::IndexSearchResult search_result; + search_result.roaring = std::make_shared(); + search_result.roaring->add(1); + search_result.roaring->add(5); + search_result.roaring->add(10); + + // Create distance array + size_t num_results = 3; + search_result.distances = std::make_unique(num_results); + search_result.distances[0] = 1.5f; + search_result.distances[1] = 2.3f; + search_result.distances[2] = 4.1f; + + // Call update_result + std::vector distance_vector; + roaring::Roaring result_roaring; + reader->update_result(search_result, distance_vector, result_roaring); + + // Verify results + EXPECT_EQ(distance_vector.size(), 3); + EXPECT_FLOAT_EQ(distance_vector[0], 1.5f); + EXPECT_FLOAT_EQ(distance_vector[1], 2.3f); + EXPECT_FLOAT_EQ(distance_vector[2], 4.1f); + EXPECT_EQ(result_roaring.cardinality(), 3); + EXPECT_TRUE(result_roaring.contains(1)); + EXPECT_TRUE(result_roaring.contains(5)); + EXPECT_TRUE(result_roaring.contains(10)); +} + +TEST_F(VectorSearchTest, TestAnnIndexReaderNewIterator) { + // Test AnnIndexReader::new_iterator method + std::map properties; + properties["index_type"] = "hnsw"; + properties["metric_type"] = "l2_distance"; + auto pair = vector_search_utils::create_tmp_ann_index_reader(properties); + auto reader = pair.second; + + std::unique_ptr iterator; + auto status = reader->new_iterator(&iterator); + + EXPECT_TRUE(status.ok()); + EXPECT_NE(iterator, nullptr); + + // Verify it's an AnnIndexIterator + auto ann_iterator = dynamic_cast(iterator.get()); + EXPECT_NE(ann_iterator, nullptr); +} + +TEST_F(VectorSearchTest, TestAnnIndexReaderQueryMethod) { + // Test AnnIndexReader::query method (coverage for lines that weren't covered) + std::map properties; + properties["index_type"] = "hnsw"; + properties["metric_type"] = "l2_distance"; + auto pair = vector_search_utils::create_tmp_ann_index_reader(properties); + auto reader = pair.second; + + // Set up _vector_index to avoid nullptr check failure + auto doris_faiss_vector_index = std::make_unique(); + doris_faiss_vector_index->set_metric(doris::segment_v2::AnnIndexMetric::L2); + + // Set up build parameters to initialize the internal _index + doris::segment_v2::FaissBuildParameter build_params; + build_params.dim = 4; + build_params.max_degree = 16; + build_params.index_type = doris::segment_v2::FaissBuildParameter::IndexType::HNSW; + build_params.metric_type = doris::segment_v2::FaissBuildParameter::MetricType::L2; + doris_faiss_vector_index->build(build_params); + + reader->_vector_index = std::move(doris_faiss_vector_index); + + // Create mock IO context + doris::io::IOContext io_ctx; + + // Create AnnTopNParam with proper initialization + const float query_data[] = {1.0f, 2.0f, 3.0f, 4.0f}; + roaring::Roaring bitmap; + bitmap.add(1); + bitmap.add(2); + bitmap.add(3); + + doris::segment_v2::AnnTopNParam param { + .query_value = query_data, + .query_value_size = 4, + .limit = 5, + ._user_params = doris::VectorSearchUserParams {.hnsw_ef_search = 100, + .hnsw_check_relative_distance = false, + .hnsw_bounded_queue = false}, + .roaring = &bitmap}; + + doris::segment_v2::AnnIndexStats stats; + + // This should cover the query method lines + auto status = reader->query(&io_ctx, ¶m, &stats); + // Note: This might fail in test environment since we don't have actual index file + // But it will cover the code paths we want to test + + // Verify that the distance and row_ids are properly initialized + if (status.ok()) { + EXPECT_NE(param.distance, nullptr); + EXPECT_NE(param.row_ids, nullptr); + } +} + +TEST_F(VectorSearchTest, TestAnnIndexReaderRangeSearchEdgeCases) { + // Test edge cases in range_search method to improve coverage + std::map properties; + properties["index_type"] = "hnsw"; + properties["metric_type"] = "l2_distance"; + auto pair = vector_search_utils::create_tmp_ann_index_reader(properties); + auto reader = pair.second; + + // Set up _vector_index to avoid nullptr check failure + auto doris_faiss_vector_index = std::make_unique(); + doris_faiss_vector_index->set_metric(doris::segment_v2::AnnIndexMetric::L2); + + // Set up build parameters to initialize the internal _index + doris::segment_v2::FaissBuildParameter build_params; + build_params.dim = 4; + build_params.max_degree = 16; + build_params.index_type = doris::segment_v2::FaissBuildParameter::IndexType::HNSW; + build_params.metric_type = doris::segment_v2::FaissBuildParameter::MetricType::L2; + doris_faiss_vector_index->build(build_params); + + reader->_vector_index = std::move(doris_faiss_vector_index); + + doris::io::IOContext io_ctx; + + // Test case 1: is_le_or_lt = false (covers lines 172-175) + { + doris::segment_v2::AnnRangeSearchParams params; + params.is_le_or_lt = false; // This should result in no distances/row_ids + params.radius = 5.0f; + float query_data[] = {1.0f, 2.0f, 3.0f, 4.0f}; + params.query_value = query_data; + + roaring::Roaring bitmap; + bitmap.add(1); + params.roaring = &bitmap; + + doris::VectorSearchUserParams user_params; + user_params.hnsw_ef_search = 50; + user_params.hnsw_check_relative_distance = true; + user_params.hnsw_bounded_queue = true; + + doris::segment_v2::AnnRangeSearchResult result; + doris::segment_v2::AnnIndexStats stats; + + auto status = reader->range_search(params, user_params, &result, &stats); + + if (status.ok()) { + // When is_le_or_lt = false, we expect no distance/row_ids + // This covers lines 183, 189 + if (result.row_ids == nullptr) { + EXPECT_EQ(result.row_ids, nullptr); + } + if (result.distance == nullptr) { + EXPECT_EQ(result.distance, nullptr); + } + } + } + + // Test case 2: Unsupported index type (covers lines 159-160) + { + // Create reader with unsupported index type + std::map unsupported_properties; + unsupported_properties["index_type"] = "ivf"; // Unsupported type + unsupported_properties["metric_type"] = "l2_distance"; + + // Since we can't easily create an AnnIndexReader with invalid type, + // we'll test this via a different approach or note it for manual testing + // The code path for line 159-160 requires actual index loading + } +} + +TEST_F(VectorSearchTest, TestAnnIndexReaderConstructor) { + // Test constructor and property parsing + std::map properties; + properties["index_type"] = "hnsw"; + properties["metric_type"] = "l2_distance"; + auto pair = vector_search_utils::create_tmp_ann_index_reader(properties); + auto reader = pair.second; + + // Verify that the constructor properly parsed the properties + EXPECT_EQ(reader->get_metric_type(), doris::segment_v2::AnnIndexMetric::L2); + + // Test with different metric type + std::map ip_properties; + ip_properties["index_type"] = "hnsw"; + ip_properties["metric_type"] = "inner_product"; + auto ip_pair = vector_search_utils::create_tmp_ann_index_reader(ip_properties); + auto ip_reader = ip_pair.second; + + EXPECT_EQ(ip_reader->get_metric_type(), doris::segment_v2::AnnIndexMetric::IP); +} + +TEST_F(VectorSearchTest, TestAnnIndexReader_UpdateResult) { + // Test AnnIndexReader::update_result method + std::map properties; + properties["index_type"] = "hnsw"; + properties["metric_type"] = "l2_distance"; + auto pair = vector_search_utils::create_tmp_ann_index_reader(properties); + auto reader = pair.second; + + // Create a search result to test update_result + doris::segment_v2::IndexSearchResult search_result; + + // Set up test data + size_t num_results = 3; + auto roaring = std::make_shared(); + roaring->add(10); + roaring->add(20); + roaring->add(30); + + auto distances = std::make_unique(num_results); + distances[0] = 1.5f; + distances[1] = 2.3f; + distances[2] = 3.1f; + + search_result.roaring = roaring; + search_result.distances = std::move(distances); + + // Test update_result method + std::vector distance_vec; + roaring::Roaring result_roaring; + + reader->update_result(search_result, distance_vec, result_roaring); + + // Verify results + EXPECT_EQ(distance_vec.size(), num_results); + EXPECT_FLOAT_EQ(distance_vec[0], 1.5f); + EXPECT_FLOAT_EQ(distance_vec[1], 2.3f); + EXPECT_FLOAT_EQ(distance_vec[2], 3.1f); + EXPECT_EQ(result_roaring.cardinality(), num_results); + EXPECT_TRUE(result_roaring.contains(10)); + EXPECT_TRUE(result_roaring.contains(20)); + EXPECT_TRUE(result_roaring.contains(30)); +} + +TEST_F(VectorSearchTest, TestAnnIndexReader_NewIterator) { + // Test new_iterator method + std::map properties; + properties["index_type"] = "hnsw"; + properties["metric_type"] = "l2_distance"; + auto pair = vector_search_utils::create_tmp_ann_index_reader(properties); + auto reader = pair.second; + + std::unique_ptr iterator; + auto status = reader->new_iterator(&iterator); + + EXPECT_TRUE(status.ok()); + EXPECT_NE(iterator, nullptr); + + // Verify that the iterator is actually an AnnIndexIterator + auto* ann_iterator = dynamic_cast(iterator.get()); + EXPECT_NE(ann_iterator, nullptr); +} + +TEST_F(VectorSearchTest, TestAnnIndexIterator_ReadFromIndex_NullParam) { + // Test AnnIndexIterator::read_from_index with null parameter + std::map properties; + properties["index_type"] = "hnsw"; + properties["metric_type"] = "l2_distance"; + auto pair = vector_search_utils::create_tmp_ann_index_reader(properties); + auto reader = pair.second; + + doris::segment_v2::AnnIndexIterator iterator(reader); + + // Test with null parameter - this should trigger the null check + doris::segment_v2::IndexParam param = static_cast(nullptr); + auto status = iterator.read_from_index(param); + + EXPECT_FALSE(status.ok()); + EXPECT_TRUE(status.is()); + EXPECT_TRUE(status.msg().find("a_param is null") != std::string::npos); +} + +TEST_F(VectorSearchTest, TestAnnIndexIterator_RangeSearch_NullReader) { + // Test AnnIndexIterator::range_search with null reader + doris::segment_v2::AnnIndexIterator iterator(nullptr); + + doris::segment_v2::AnnRangeSearchParams params; + doris::VectorSearchUserParams user_params; + doris::segment_v2::AnnRangeSearchResult result; + auto stats = std::make_unique(); + + auto status = iterator.range_search(params, user_params, &result, stats.get()); + + EXPECT_FALSE(status.ok()); + EXPECT_TRUE(status.is()); + EXPECT_TRUE(status.msg().find("_ann_reader is null") != std::string::npos); +} + +TEST_F(VectorSearchTest, TestAnnIndexStats_CopyConstructor) { + // Test AnnIndexStats copy constructor + doris::segment_v2::AnnIndexStats original; + original.search_costs_ns.set(static_cast(1000)); + original.load_index_costs_ns.set(static_cast(2000)); + + doris::segment_v2::AnnIndexStats copied(original); + + EXPECT_EQ(copied.search_costs_ns.value(), 1000); + EXPECT_EQ(copied.load_index_costs_ns.value(), 2000); +} + +TEST_F(VectorSearchTest, TestAnnRangeSearchParams_ToString) { + // Test AnnRangeSearchParams::to_string method + doris::segment_v2::AnnRangeSearchParams params; + params.is_le_or_lt = true; + params.radius = 5.5f; + + auto roaring = std::make_shared(); + roaring->add(1); + roaring->add(2); + roaring->add(3); + params.roaring = roaring.get(); + + std::string result = params.to_string(); + + EXPECT_TRUE(result.find("is_le_or_lt: true") != std::string::npos); + EXPECT_TRUE(result.find("radius: 5.5") != std::string::npos); + EXPECT_TRUE(result.find("input rows 3") != std::string::npos); +} + +TEST_F(VectorSearchTest, TestPrepareAnnRangeSearch_EarlyReturn_WrongOpcode) { + // Case 2: VectorizedFnCall with unsupported opcode (not GE, LE, GT, LT) + TExpr texpr = read_from_json(ann_range_search_thrift); + TExprNode& opNode = texpr.nodes[0]; + opNode.opcode = TExprOpcode::ADD; // Change to unsupported operation + opNode.fn.name.function_name = "add"; + opNode.fn.signature = "add(double, double)"; // Fix the signature for add operation + TDescriptorTable table1 = read_from_json(thrift_table_desc); + std::unique_ptr pool = std::make_unique(); + auto desc_tbl = std::make_unique(); + DescriptorTbl* desc_tbl_ptr = desc_tbl.get(); + ASSERT_TRUE(DescriptorTbl::create(pool.get(), table1, &(desc_tbl_ptr)).ok()); + RowDescriptor row_desc = RowDescriptor(*desc_tbl_ptr, {0}, {false}); + std::unique_ptr state = std::make_unique(); + state->set_desc_tbl(desc_tbl_ptr); + + VExprContextSPtr ctx; + // The VExpr::create_expr_tree might fail for invalid expressions + auto status = vectorized::VExpr::create_expr_tree(texpr, ctx); + if (!status.ok()) { + // If create_expr_tree fails, we can't test prepare_ann_range_search + // This is still a valid test case - invalid expressions should fail early + EXPECT_FALSE(status.ok()); + return; + } + + status = ctx->prepare(state.get(), row_desc); + if (!status.ok()) { + // If prepare fails, we can't test prepare_ann_range_search + // This is still a valid test case - invalid expressions should fail early + EXPECT_FALSE(status.ok()); + return; + } + + ASSERT_TRUE(ctx->open(state.get()).ok()); + + doris::VectorSearchUserParams user_params; + ctx->prepare_ann_range_search(user_params); + EXPECT_FALSE(ctx->_ann_range_search_runtime.is_ann_range_search); +} + +TEST_F(VectorSearchTest, TestPrepareAnnRangeSearch_EarlyReturn_NonLiteralRight) { + // Case 3: Right child is not a literal - we'll use a valid GE expression but modify the right operand to be non-literal + TExpr texpr = read_from_json(ann_range_search_thrift); + // Change the right operand (index 2) from literal to slot reference + TExprNode& rightNode = texpr.nodes[2]; + rightNode.node_type = TExprNodeType::SLOT_REF; + rightNode.__set_slot_ref(TSlotRef()); + rightNode.slot_ref.slot_id = 0; // Reference to a different slot + rightNode.slot_ref.tuple_id = 0; + rightNode.slot_ref.__set_is_virtual_slot(false); + rightNode.__isset.slot_ref = true; + + TDescriptorTable table1 = read_from_json(thrift_table_desc); + std::unique_ptr pool = std::make_unique(); + auto desc_tbl = std::make_unique(); + DescriptorTbl* desc_tbl_ptr = desc_tbl.get(); + ASSERT_TRUE(DescriptorTbl::create(pool.get(), table1, &(desc_tbl_ptr)).ok()); + RowDescriptor row_desc = RowDescriptor(*desc_tbl_ptr, {0}, {false}); + std::unique_ptr state = std::make_unique(); + state->set_desc_tbl(desc_tbl_ptr); + + VExprContextSPtr ctx; + auto status = vectorized::VExpr::create_expr_tree(texpr, ctx); + if (!status.ok()) { + // If create_expr_tree fails, that's still a valid test case + EXPECT_FALSE(status.ok()); + return; + } + + status = ctx->prepare(state.get(), row_desc); + if (!status.ok()) { + // If prepare fails, that's still a valid test case + EXPECT_FALSE(status.ok()); + return; + } + + ASSERT_TRUE(ctx->open(state.get()).ok()); + + doris::VectorSearchUserParams user_params; + ctx->prepare_ann_range_search(user_params); + EXPECT_FALSE(ctx->_ann_range_search_runtime.is_ann_range_search); +} + +} // namespace doris::vectorized \ No newline at end of file diff --git a/be/test/olap/vector_search/ann_topn_descriptor_test.cpp b/be/test/olap/vector_search/ann_topn_descriptor_test.cpp new file mode 100644 index 00000000000000..1dac2f07e7063a --- /dev/null +++ b/be/test/olap/vector_search/ann_topn_descriptor_test.cpp @@ -0,0 +1,188 @@ +// 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 +#include +#include +#include +#include +#include +#include + +#include +#include + +#include "olap/rowset/segment_v2/ann_index/ann_search_params.h" +#include "olap/rowset/segment_v2/ann_index/ann_topn_runtime.h" +#include "runtime/primitive_type.h" +#include "vec/columns/column_nullable.h" +#include "vec/exprs/virtual_slot_ref.h" +#include "vector_search_utils.h" + +using ::testing::_; +using ::testing::SetArgPointee; +using ::testing::Return; + +namespace doris::vectorized { + +TEST_F(VectorSearchTest, AnnTopNRuntimeConstructor) { + int limit = 10; + std::shared_ptr distanc_calcu_fn_call_ctx; + auto distance_function_call_thrift = read_from_json(_distance_function_call_thrift); + ASSERT_TRUE(distance_function_call_thrift.nodes.empty() != true); + auto st1 = vectorized::VExpr::create_expr_tree(distance_function_call_thrift, + distanc_calcu_fn_call_ctx); + ASSERT_TRUE(st1.ok()) << fmt::format( + "st: {}, expr {}", st1.to_string(), + apache::thrift::ThriftDebugString(distance_function_call_thrift)); + ASSERT_TRUE(distanc_calcu_fn_call_ctx != nullptr) << "create expr tree failed"; + ASSERT_TRUE(distanc_calcu_fn_call_ctx->root() != nullptr); + + std::shared_ptr virtual_slot_expr_ctx; + ASSERT_TRUE(vectorized::VExpr::create_expr_tree(_virtual_slot_ref_expr, virtual_slot_expr_ctx) + .ok()); + + ASSERT_TRUE(virtual_slot_expr_ctx != nullptr) << "create expr tree failed"; + ASSERT_TRUE(virtual_slot_expr_ctx->root() != nullptr); + + std::shared_ptr v = + std::dynamic_pointer_cast(virtual_slot_expr_ctx->root()); + if (v == nullptr) { + LOG(FATAL) << "VAnnTopNRuntime::SetUp() failed"; + } + + v->set_virtual_column_expr(distanc_calcu_fn_call_ctx->root()); + + std::shared_ptr predicate; + predicate = segment_v2::AnnTopNRuntime::create_shared(true, limit, virtual_slot_expr_ctx); + ASSERT_TRUE(predicate != nullptr) << "AnnTopNRuntime::create_shared(true,) failed"; +} + +TEST_F(VectorSearchTest, AnnTopNRuntimePrepare) { + int limit = 10; + std::shared_ptr distanc_calcu_fn_call_ctx; + auto distance_function_call_thrift = read_from_json(_distance_function_call_thrift); + Status st = vectorized::VExpr::create_expr_tree(distance_function_call_thrift, + distanc_calcu_fn_call_ctx); + + std::shared_ptr virtual_slot_expr_ctx; + st = vectorized::VExpr::create_expr_tree(_virtual_slot_ref_expr, virtual_slot_expr_ctx); + std::shared_ptr v = + std::dynamic_pointer_cast(virtual_slot_expr_ctx->root()); + if (v == nullptr) { + LOG(FATAL) << "VAnnTopNRuntime::SetUp() failed"; + } + + v->set_virtual_column_expr(distanc_calcu_fn_call_ctx->root()); + std::shared_ptr predicate; + predicate = segment_v2::AnnTopNRuntime::create_shared(true, limit, virtual_slot_expr_ctx); + st = predicate->prepare(&_runtime_state, _row_desc); + ASSERT_TRUE(st.ok()) << fmt::format("st: {}, expr {}", st.to_string(), + predicate->get_order_by_expr_ctx()->root()->debug_string()); + + std::cout << "predicate: " << predicate->debug_string() << std::endl; +} + +TEST_F(VectorSearchTest, AnnTopNRuntimeEvaluateTopN) { + int limit = 10; + std::shared_ptr distanc_calcu_fn_call_ctx; + auto distance_function_call_thrift = read_from_json(_distance_function_call_thrift); + Status st = vectorized::VExpr::create_expr_tree(distance_function_call_thrift, + distanc_calcu_fn_call_ctx); + + std::shared_ptr virtual_slot_expr_ctx; + st = vectorized::VExpr::create_expr_tree(_virtual_slot_ref_expr, virtual_slot_expr_ctx); + std::shared_ptr v = + std::dynamic_pointer_cast(virtual_slot_expr_ctx->root()); + if (v == nullptr) { + LOG(FATAL) << "VAnnTopNRuntime::SetUp() failed"; + } + + v->set_virtual_column_expr(distanc_calcu_fn_call_ctx->root()); + std::shared_ptr predicate; + predicate = segment_v2::AnnTopNRuntime::create_shared(true, limit, virtual_slot_expr_ctx); + st = predicate->prepare(&_runtime_state, _row_desc); + ASSERT_TRUE(st.ok()) << fmt::format("st: {}, expr {}", st.to_string(), + predicate->get_order_by_expr_ctx()->root()->debug_string()); + + const ColumnConst* const_column = + assert_cast(predicate->_query_array.get()); + const ColumnArray* column_array = + assert_cast(const_column->get_data_column_ptr().get()); + const ColumnNullable* column_nullable = + assert_cast(column_array->get_data_ptr().get()); + const ColumnFloat64* cf64 = + assert_cast(column_nullable->get_nested_column_ptr().get()); + + const double* query_value = cf64->get_data().data(); + const size_t query_value_size = cf64->get_data().size(); + ASSERT_EQ(query_value_size, 8); + std::vector query_value_f32; + for (size_t i = 0; i < query_value_size; ++i) { + query_value_f32.push_back(static_cast(query_value[i])); + } + ASSERT_FLOAT_EQ(query_value_f32[0], 1.0f) << "query_value_f32[0] = " << query_value_f32[0]; + ASSERT_FLOAT_EQ(query_value_f32[1], 2.0f) << "query_value_f32[1] = " << query_value_f32[1]; + ASSERT_FLOAT_EQ(query_value_f32[2], 3.0f) << "query_value_f32[2] = " << query_value_f32[2]; + ASSERT_FLOAT_EQ(query_value_f32[3], 4.0f) << "query_value_f32[3] = " << query_value_f32[3]; + ASSERT_FLOAT_EQ(query_value_f32[4], 5.0f) << "query_value_f32[4] = " << query_value_f32[4]; + ASSERT_FLOAT_EQ(query_value_f32[5], 6.0f) << "query_value_f32[5] = " << query_value_f32[5]; + ASSERT_FLOAT_EQ(query_value_f32[6], 7.0f) << "query_value_f32[6] = " << query_value_f32[6]; + ASSERT_FLOAT_EQ(query_value_f32[7], 20.0f) << "query_value_f32[7] = " << query_value_f32[7]; + + std::shared_ptr> query_vector = + std::make_shared>(10, 0.0); + for (size_t i = 0; i < 10; ++i) { + (*query_vector)[i] = static_cast(i); + } + + std::cout << "query_vector: " << fmt::format("[{}]", fmt::join(*query_vector, ",")) + << std::endl; + + EXPECT_CALL(*_ann_index_iterator, read_from_index(testing::_)) + .Times(1) + .WillOnce(testing::Invoke([](const segment_v2::IndexParam& value) { + auto* ann_param = std::get(value); + ann_param->distance = std::make_unique>(); + ann_param->row_ids = std::make_unique>(); + for (size_t i = 0; i < 10; ++i) { + ann_param->distance->push_back(static_cast(i)); + ann_param->row_ids->push_back(i); + } + return Status::OK(); + })); + + _result_column = ColumnNullable::create(ColumnFloat64::create(0, 0), ColumnUInt8::create(0, 0)); + std::unique_ptr> row_ids = std::make_unique>(); + + roaring::Roaring roaring; + doris::segment_v2::AnnIndexStats ann_index_stats; + // rows_of_segment is mocked as 10 to align with mocked iterator outputs + size_t rows_of_segment = 10; + st = predicate->evaluate_vector_ann_search(_ann_index_iterator.get(), &roaring, rows_of_segment, + _result_column, row_ids, ann_index_stats); + ColumnNullable* result_column_null = assert_cast(_result_column.get()); + ColumnFloat64* result_column_float = + assert_cast(result_column_null->get_nested_column_ptr().get()); + for (size_t i = 0; i < query_vector->size(); ++i) { + EXPECT_EQ(result_column_float->get_data()[i], (*query_vector)[i]); + } + ASSERT_TRUE(st.ok()); + ASSERT_EQ(row_ids->size(), 10); +} + +} // namespace doris::vectorized \ No newline at end of file diff --git a/be/test/olap/vector_search/faiss_vector_index_test.cpp b/be/test/olap/vector_search/faiss_vector_index_test.cpp new file mode 100644 index 00000000000000..af2c4e97596ed8 --- /dev/null +++ b/be/test/olap/vector_search/faiss_vector_index_test.cpp @@ -0,0 +1,1026 @@ +// 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 +#include +#include + +#include +#include +#include +#include +#include +#include + +#include "olap/rowset/segment_v2/ann_index/ann_index.h" +#include "olap/rowset/segment_v2/ann_index/ann_search_params.h" +#include "olap/rowset/segment_v2/ann_index/faiss_ann_index.h" +// metrics.h not used directly here +#include "vector_search_utils.h" + +using namespace doris::segment_v2; + +namespace doris::vectorized { + +// Test saving and loading an index +TEST_F(VectorSearchTest, TestSaveAndLoad) { + // Step 1: Create first index instance + auto index1 = std::make_unique(); + + // Step 2: Set build parameters + FaissBuildParameter params; + params.dim = 128; // Vector dimension + params.max_degree = 16; // HNSW max connections + params.index_type = FaissBuildParameter::IndexType::HNSW; + index1->build(params); + + // Step 3: Add vectors to the index + const int num_vectors = 100; + std::vector vectors; + for (int i = 0; i < num_vectors; i++) { + auto tmp = vector_search_utils::generate_random_vector(params.dim); + vectors.insert(vectors.end(), tmp.begin(), tmp.end()); + } + + std::ignore = index1->add(num_vectors, vectors.data()); + + // Step 4: Save the index + auto save_status = index1->save(_ram_dir.get()); + ASSERT_TRUE(save_status.ok()) << "Failed to save index: " << save_status.to_string(); + + // Step 5: Create a new index instance + auto index2 = std::make_unique(); + + // Step 6: Load the index + auto load_status = index2->load(_ram_dir.get()); + ASSERT_TRUE(load_status.ok()) << "Failed to load index: " << load_status.to_string(); + + // Step 7: Verify the loaded index works by searching + auto query_vec = vector_search_utils::generate_random_vector(params.dim); + const int top_k = 10; + + // TopN search requires a candidate roaring and rows_of_segment now. + HNSWSearchParameters topn_params; + auto topn_roaring = std::make_unique(); + for (int i = 0; i < num_vectors; ++i) topn_roaring->add(i); + topn_params.roaring = topn_roaring.get(); + topn_params.rows_of_segment = num_vectors; + + IndexSearchResult search_result1; + IndexSearchResult search_result2; + + std::ignore = index1->ann_topn_search(query_vec.data(), top_k, topn_params, search_result1); + + std::ignore = index2->ann_topn_search(query_vec.data(), top_k, topn_params, search_result2); + + // Compare the results + EXPECT_EQ(search_result1.roaring->cardinality(), search_result2.roaring->cardinality()) + << "Row ID cardinality mismatch"; + for (size_t i = 0; i < search_result1.roaring->cardinality(); ++i) { + EXPECT_EQ(search_result1.distances[i], search_result2.distances[i]) + << "Distance mismatch at index " << i; + } + + HNSWSearchParameters hnsw_params; + auto roaring_bitmap = std::make_unique(); + hnsw_params.roaring = roaring_bitmap.get(); + for (size_t i = 0; i < num_vectors; ++i) { + hnsw_params.roaring->add(i); + } + IndexSearchResult range_search_result1; + std::ignore = index1->range_search(vectors.data(), 10, hnsw_params, range_search_result1); + IndexSearchResult range_search_result2; + std::ignore = index2->range_search(vectors.data(), 10, hnsw_params, range_search_result2); + EXPECT_EQ(range_search_result1.roaring->cardinality(), + range_search_result2.roaring->cardinality()) + << "Row ID cardinality mismatch"; + for (size_t i = 0; i < range_search_result1.roaring->cardinality(); ++i) { + EXPECT_EQ(range_search_result1.distances[i], range_search_result2.distances[i]) + << "Distance mismatch at index " << i; + } +} + +TEST_F(VectorSearchTest, UpdateRoaring) { + // Create a roaring bitmap + roaring::Roaring roaring_bitmap; + // Create some dummy labels + const size_t n = 5; + faiss::idx_t labels[n] = {1, 2, 3, 4, 5}; + + // Call the update_roaring function + FaissVectorIndex::update_roaring(labels, n, roaring_bitmap); + + EXPECT_EQ(roaring_bitmap.cardinality(), n) << "Roaring bitmap size mismatch"; + + for (size_t i = 0; i < n; ++i) { + EXPECT_EQ(roaring_bitmap.contains(labels[i]), true) + << "Label " << labels[i] << " not found"; + } +} + +TEST_F(VectorSearchTest, CompareResultWithNativeFaiss1) { + const size_t iterations = 3; + // Create random number generator + std::random_device rd; + std::mt19937 gen(rd()); + // Define fixed parameter sets to choose from + const std::vector dimensions = {32, 64, 128, 256}; + const std::vector max_connections = {8, 16, 32, 64}; + const std::vector vector_counts = {100, 200, 500, 1000}; + + for (size_t iter = 0; iter < iterations; ++iter) { + // Randomly select parameters from the fixed sets + const int dimension = + dimensions[std::uniform_int_distribution<>(0, dimensions.size() - 1)(gen)]; + const int max_connection = max_connections[std::uniform_int_distribution<>( + 0, max_connections.size() - 1)(gen)]; + const int num_vectors = + vector_counts[std::uniform_int_distribution<>(0, vector_counts.size() - 1)(gen)]; + + // Step 1: Create indexes + auto doris_index = doris::vector_search_utils::create_doris_index( + doris::vector_search_utils::IndexType::HNSW, dimension, max_connection); + auto native_index = doris::vector_search_utils::create_native_index( + doris::vector_search_utils::IndexType::HNSW, dimension, max_connection); + + // Step 2: Generate vectors and add to indexes + auto vectors = + doris::vector_search_utils::generate_test_vectors_matrix(num_vectors, dimension); + doris::vector_search_utils::add_vectors_to_indexes_serial_mode(doris_index.get(), + native_index.get(), vectors); + + // Step 3: Search + int query_idx = num_vectors / 2; + const float* query_vec = vectors[query_idx].data(); + const int top_k = 10; + + // Search in Doris index + HNSWSearchParameters search_params; + auto roaring = std::make_unique(); + for (int i = 0; i < num_vectors; ++i) roaring->add(i); + search_params.roaring = roaring.get(); + search_params.rows_of_segment = num_vectors; + IndexSearchResult doris_results; + auto search_status = + doris_index->ann_topn_search(query_vec, top_k, search_params, doris_results); + ASSERT_EQ(search_status.ok(), true) + << "Search failed with dimension=" << dimension + << ", max_connections=" << max_connection << ", num_vectors=" << num_vectors; + + // Search in native Faiss index + std::vector native_distances(top_k); + std::vector native_indices(top_k); + native_index->search(1, query_vec, top_k, native_distances.data(), native_indices.data()); + size_t cnt = std::count_if(native_indices.begin(), native_indices.end(), + [](faiss::idx_t idx) { return idx != -1; }); + for (size_t i = 0; i < cnt; ++i) { + native_distances[i] = std::sqrt(native_distances[i]); + } + // Step 4: Compare results + vector_search_utils::compare_search_results(doris_results, native_distances, + native_indices); + } +} + +TEST_F(VectorSearchTest, CompareResultWithNativeFaiss2) { + const size_t iterations = 2; + // Create random number generator + std::random_device rd; + std::mt19937 gen(rd()); + // Define fixed parameter sets to choose from + const std::vector dimensions = {32, 64, 128, 256}; + const std::vector max_connections = {8, 16, 32, 64}; + const std::vector vector_counts = {100, 200, 500, 1000}; + + for (size_t i = 0; i < iterations; ++i) { + // Randomly select parameters from the fixed sets + const int dimension = + dimensions[std::uniform_int_distribution<>(0, dimensions.size() - 1)(gen)]; + const int max_connection = max_connections[std::uniform_int_distribution<>( + 0, max_connections.size() - 1)(gen)]; + const int num_vectors = + vector_counts[std::uniform_int_distribution<>(0, vector_counts.size() - 1)(gen)]; + + // Step 1: Create indexes + auto doris_index = doris::vector_search_utils::create_doris_index( + doris::vector_search_utils::IndexType::HNSW, dimension, max_connection); + auto native_index = doris::vector_search_utils::create_native_index( + doris::vector_search_utils::IndexType::HNSW, dimension, max_connection); + + // Step 2: Generate vectors and add to indexes + std::vector> vectors = + doris::vector_search_utils::generate_test_vectors_matrix(num_vectors, dimension); + doris::vector_search_utils::add_vectors_to_indexes_serial_mode(doris_index.get(), + native_index.get(), vectors); + + // Step 3: Search + int query_idx = num_vectors / 2; + const float* query_vec = vectors[query_idx].data(); + const int top_k = num_vectors; + HNSWSearchParameters search_params; + auto roaring = std::make_unique(); + for (int i = 0; i < num_vectors; ++i) roaring->add(i); + search_params.roaring = roaring.get(); + search_params.rows_of_segment = num_vectors; + IndexSearchResult doris_results; + std::ignore = doris_index->ann_topn_search(query_vec, top_k, search_params, doris_results); + + // Search in native Faiss index + std::vector native_distances(top_k, -1); + std::vector native_indices(top_k, -1); + native_index->search(1, query_vec, top_k, native_distances.data(), native_indices.data()); + size_t cnt = std::count_if(native_indices.begin(), native_indices.end(), + [](faiss::idx_t idx) { return idx != -1; }); + for (size_t i = 0; i < cnt; ++i) { + native_distances[i] = std::sqrt(native_distances[i]); + } + // Step 4: Compare results + doris::vector_search_utils::compare_search_results(doris_results, native_distances, + native_indices); + } +} + +TEST_F(VectorSearchTest, SearchAllVectors) { + size_t iterations = 5; + for (size_t i = 0; i < iterations; ++i) { + // Step 1: Create and build index + auto index1 = std::make_unique(); + + FaissBuildParameter params; + params.dim = 64; + params.max_degree = 32; + params.index_type = FaissBuildParameter::IndexType::HNSW; + index1->build(params); + + // Add 500 vectors + const int num_vectors = 500; + std::vector vectors; + for (int i = 0; i < num_vectors; i++) { + auto vec = doris::vector_search_utils::generate_random_vector(params.dim); + vectors.insert(vectors.end(), vec.begin(), vec.end()); + } + + ASSERT_EQ(index1->add(500, vectors.data()).ok(), true); + + // Save index + ASSERT_TRUE(index1->save(_ram_dir.get()).ok()); + + // Step 2: Load index + auto index2 = std::make_unique(); + ASSERT_TRUE(index2->load(_ram_dir.get()).ok()); + + // Step 3: Search all vectors + HNSWSearchParameters search_params; + auto roaring = std::make_unique(); + for (int i = 0; i < num_vectors; ++i) roaring->add(i); + search_params.roaring = roaring.get(); + search_params.rows_of_segment = num_vectors; + IndexSearchResult search_result; + + // Search for all vectors - use a vector we know is in the index + std::vector query_vec { + vectors.begin(), vectors.begin() + params.dim}; // Use the first vector we added + const int top_k = num_vectors; // Get all vectors + + ASSERT_EQ( + index2->ann_topn_search(query_vec.data(), top_k, search_params, search_result).ok(), + true); + // Step 4: Verify we got all vectors back + // Note: In practical ANN search with approximate algorithms like HNSW, + // we might not get exactly all vectors due to the nature of approximate search. + // So we verify we got a reasonable number back. + EXPECT_GE(search_result.roaring->cardinality(), num_vectors * 0.60) + << "Expected to find at least 60% of all vectors"; + + // Also verify the first result is the query vector itself (it should be an exact match) + ASSERT_EQ(search_result.roaring->isEmpty(), false) << "Search result should not be empty"; + size_t first = search_result.roaring->getIndex(0); + std::vector first_result_vec(vectors.begin() + first * params.dim, + vectors.begin() + (first + 1) * params.dim); + std::string query_vec_str = fmt::format("[{}]", fmt::join(query_vec, ",")); + std::string first_result_vec_str = fmt::format("[{}]", fmt::join(first_result_vec, ",")); + EXPECT_EQ(first_result_vec, query_vec) << "First result should be the query vector itself"; + } +} + +TEST_F(VectorSearchTest, CompRangeSearch) { + size_t iterations = 5; + std::vector metrics = {faiss::METRIC_L2, faiss::METRIC_INNER_PRODUCT}; + for (size_t i = 0; i < iterations; ++i) { + for (auto metric : metrics) { + // Random parameters for each test iteration + std::random_device rd; + std::mt19937 gen(rd()); + size_t random_d = std::uniform_int_distribution<>(1, 512)(gen); + size_t random_m = 4 << std::uniform_int_distribution<>(1, 4)(gen); + size_t random_n = std::uniform_int_distribution<>(10, 200)(gen); + + // Step 1: Create and build index + auto doris_index = std::make_unique(); + FaissBuildParameter params; + params.dim = random_d; + params.max_degree = random_m; + params.index_type = FaissBuildParameter::IndexType::HNSW; + if (metric == faiss::METRIC_L2) { + params.metric_type = FaissBuildParameter::MetricType::L2; + } else if (metric == faiss::METRIC_INNER_PRODUCT) { + params.metric_type = FaissBuildParameter::MetricType::IP; + } else { + throw std::runtime_error(fmt::format("Unsupported metric type: {}", metric)); + } + doris_index->build(params); + + const int num_vectors = random_n; + std::vector> vectors; + for (int i = 0; i < num_vectors; i++) { + auto vec = vector_search_utils::generate_random_vector(params.dim); + vectors.push_back(vec); + } + + std::unique_ptr native_index; + if (metric == faiss::METRIC_L2) { + native_index = std::make_unique(params.dim, params.max_degree, + faiss::METRIC_L2); + } else if (metric == faiss::METRIC_INNER_PRODUCT) { + native_index = std::make_unique(params.dim, params.max_degree, + faiss::METRIC_INNER_PRODUCT); + } else { + throw std::runtime_error(fmt::format("Unsupported metric type: {}", metric)); + } + + doris::vector_search_utils::add_vectors_to_indexes_serial_mode( + doris_index.get(), native_index.get(), vectors); + + std::vector query_vec = vectors.front(); + float radius = 0; + radius = doris::vector_search_utils::get_radius_from_matrix( + query_vec.data(), params.dim, vectors, 0.4f, metric); + + HNSWSearchParameters hnsw_params; + hnsw_params.ef_search = 16; + // Search on all rows; + auto roaring = std::make_unique(); + hnsw_params.roaring = roaring.get(); + for (size_t i = 0; i < vectors.size(); i++) { + hnsw_params.roaring->add(i); + } + hnsw_params.is_le_or_lt = metric == faiss::METRIC_L2; + IndexSearchResult doris_result; + std::ignore = + doris_index->range_search(query_vec.data(), radius, hnsw_params, doris_result); + + faiss::SearchParametersHNSW search_params_native; + search_params_native.efSearch = hnsw_params.ef_search; + faiss::RangeSearchResult search_result_native(1, true); + // 对于L2，radius要平方；对于IP，直接用 + float faiss_radius = (metric == faiss::METRIC_L2) ? radius * radius : radius; + native_index->range_search(1, query_vec.data(), faiss_radius, &search_result_native, + &search_params_native); + + std::vector> native_results; + size_t begin = search_result_native.lims[0]; + size_t end = search_result_native.lims[1]; + for (size_t i = begin; i < end; i++) { + native_results.push_back( + {search_result_native.labels[i], search_result_native.distances[i]}); + } + + // Make sure result is same + ASSERT_NEAR(doris_result.roaring->cardinality(), native_results.size(), 1) + << fmt::format("\nd: {}, m: {}, n: {}, metric: {}", random_d, random_m, + random_n, metric); + ASSERT_EQ(doris_result.distances != nullptr, true); + if (doris_result.roaring->cardinality() == native_results.size()) { + for (size_t i = 0; i < native_results.size(); i++) { + const size_t rowid = native_results[i].first; + const float dis = native_results[i].second; + ASSERT_EQ(doris_result.roaring->contains(rowid), true) + << "Row ID mismatch at rank " << i; + if (metric == faiss::METRIC_L2) { + ASSERT_FLOAT_EQ(doris_result.distances[i], sqrt(dis)) + << "Distance mismatch at rank " << i; + } else { + ASSERT_FLOAT_EQ(doris_result.distances[i], dis) + << "Distance mismatch at rank " << i; + } + } + } + } + } +} + +TEST_F(VectorSearchTest, RangeSearchAllRowsAsCandidates) { + size_t iterations = 5; + // Random parameters for each test iteration + + for (size_t i = 0; i < iterations; ++i) { + std::random_device rd; + std::mt19937 gen(rd()); + size_t random_d = + std::uniform_int_distribution<>(1, 512)(gen); // Random dimension from 32 to 256 + size_t random_m = + 4 << std::uniform_int_distribution<>(1, 4)(gen); // Random M (4, 8, 16, 32, 64) + size_t random_n = std::uniform_int_distribution<>(10, 200)(gen); // Random number of vectors + // Step 1: Create and build index + auto index1 = std::make_unique(); + + FaissBuildParameter params; + params.dim = random_d; + params.max_degree = random_m; + params.index_type = FaissBuildParameter::IndexType::HNSW; + index1->build(params); + + const int num_vectors = random_n; + std::vector> vectors; + for (int i = 0; i < num_vectors; i++) { + auto vec = vector_search_utils::generate_random_vector(params.dim); + vectors.push_back(vec); + } + std::unique_ptr native_index = + std::make_unique(params.dim, params.max_degree); + doris::vector_search_utils::add_vectors_to_indexes_serial_mode(index1.get(), + native_index.get(), vectors); + + std::vector query_vec = vectors.front(); + + std::vector> distances(num_vectors); + for (int i = 0; i < num_vectors; i++) { + double sum = 0; + auto& vec = vectors[i]; + for (int j = 0; j < params.dim; j++) { + accumulate(vec[j], query_vec[j], sum); + } + distances[i] = std::make_pair(i, finalize(sum)); + } + std::sort(distances.begin(), distances.end(), + [](const auto& a, const auto& b) { return a.second < b.second; }); + + float radius = distances[num_vectors / 4].second; + // Save index + ASSERT_TRUE(index1->save(_ram_dir.get()).ok()); + + // Step 2: Load index + auto index2 = std::make_unique(); + ASSERT_TRUE(index2->load(_ram_dir.get()).ok()); + + // Step 3: Range search + // Use a vector we know is in the index + + faiss::SearchParametersHNSW search_params; + std::unique_ptr all_rows = std::make_unique(); + for (size_t i = 0; i < num_vectors; ++i) { + all_rows->add(i); + } + auto sel = FaissVectorIndex::roaring_to_faiss_selector(*all_rows); + search_params.sel = sel.get(); + search_params.efSearch = 16; // Set efSearch for better accuracy + faiss::RangeSearchResult native_search_result(1, true); + native_index->range_search(1, query_vec.data(), radius * radius, &native_search_result, + &search_params); + + std::vector> native_results; + size_t begin = native_search_result.lims[0]; + size_t end = native_search_result.lims[1]; + for (size_t i = begin; i < end; i++) { + native_results.push_back( + {native_search_result.labels[i], native_search_result.distances[i]}); + } + + HNSWSearchParameters doris_search_params; + doris_search_params.ef_search = 16; // Set efSearch for better accuracy + doris_search_params.roaring = all_rows.get(); + IndexSearchResult search_result1; + IndexSearchResult search_result2; + + ASSERT_EQ( + index1->range_search(query_vec.data(), radius, doris_search_params, search_result1) + .ok(), + true); + ASSERT_EQ( + index2->range_search(query_vec.data(), radius, doris_search_params, search_result2) + .ok(), + true); + + ASSERT_EQ(search_result1.roaring->cardinality(), search_result2.roaring->cardinality()); + for (size_t i = 0; i < search_result1.roaring->cardinality(); i++) { + ASSERT_EQ(search_result1.distances[i], search_result2.distances[i]) + << "Distance mismatch at rank " << i; + } + + ASSERT_EQ(search_result2.roaring->cardinality(), native_results.size()); + + ASSERT_EQ(search_result2.distances != nullptr, true); + for (size_t i = 0; i < native_results.size(); i++) { + const size_t rowid = native_results[i].first; + const float dis = native_results[i].second; + ASSERT_EQ(search_result2.roaring->contains(rowid), true) + << "Row ID mismatch at rank " << i; + ASSERT_FLOAT_EQ(search_result2.distances[i], sqrt(dis)) + << "Distance mismatch at rank " << i; + } + + doris_search_params.is_le_or_lt = false; + doris_search_params.roaring = all_rows.get(); + for (size_t i = 0; i < num_vectors; ++i) { + doris_search_params.roaring->add(i); + } + IndexSearchResult search_result3; + std::ignore = + index1->range_search(query_vec.data(), radius, doris_search_params, search_result3); + roaring::Roaring ge_rows; + ASSERT_EQ(search_result3.distances == nullptr, true); + for (size_t i = 0; i < native_results.size(); ++i) { + ge_rows.add(native_results[i].first); + } + roaring::Roaring and_row_id = ge_rows & *search_result3.roaring; + roaring::Roaring or_row_id = ge_rows | *search_result3.roaring; + ASSERT_EQ(and_row_id.cardinality(), 0); + ASSERT_EQ(or_row_id.cardinality(), num_vectors); + } +} + +TEST_F(VectorSearchTest, RangeSearchWithSelector1) { + size_t iterations = 2; + for (size_t i = 0; i < iterations; ++i) { + // Step 1: Create and build index + auto index1 = std::make_unique(); + + FaissBuildParameter params; + params.dim = 100; + params.max_degree = 32; + params.index_type = FaissBuildParameter::IndexType::HNSW; + index1->build(params); + + const int num_vectors = 100; + std::vector> vectors; + for (int i = 0; i < num_vectors; i++) { + auto vec = vector_search_utils::generate_random_vector(params.dim); + vectors.push_back(vec); + } + + // Use a vector we know is in the index + std::vector query_vec = vectors.front(); + std::vector> distances(num_vectors); + for (int i = 0; i < num_vectors; i++) { + double sum = 0; + auto& vec = vectors[i]; + for (int j = 0; j < params.dim; j++) { + accumulate(vec[j], query_vec[j], sum); + } + distances[i] = std::make_pair(i, finalize(sum)); + } + std::sort(distances.begin(), distances.end(), + [](const auto& a, const auto& b) { return a.second < b.second; }); + // Use the median distance as the radius + float radius = distances[num_vectors / 2].second; + + std::unique_ptr native_index = std::make_unique( + params.dim, params.max_degree, faiss::METRIC_L2); + doris::vector_search_utils::add_vectors_to_indexes_serial_mode(index1.get(), + native_index.get(), vectors); + + std::unique_ptr all_rows = std::make_unique(); + std::unique_ptr sel_rows = std::make_unique(); + for (size_t i = 0; i < num_vectors; ++i) { + all_rows->add(i); + if (i % 2 == 0) { + sel_rows->add(i); + } + } + + // Step 3: Range search + faiss::SearchParametersHNSW search_params; + search_params.efSearch = 16; // Set efSearch for better accuracy + auto faiss_selector = segment_v2::FaissVectorIndex::roaring_to_faiss_selector(*sel_rows); + search_params.sel = faiss_selector.get(); + faiss::RangeSearchResult native_search_result(1, true); + native_index->range_search(1, query_vec.data(), radius * radius, &native_search_result, + &search_params); + // labels and distance + std::vector> native_results; + size_t begin = native_search_result.lims[0]; + size_t end = native_search_result.lims[1]; + for (size_t i = begin; i < end; i++) { + native_results.push_back( + {native_search_result.labels[i], native_search_result.distances[i]}); + } + + HNSWSearchParameters doris_search_params; + doris_search_params.ef_search = search_params.efSearch; + doris_search_params.is_le_or_lt = true; + doris_search_params.roaring = sel_rows.get(); + IndexSearchResult doris_search_result; + + ASSERT_EQ(index1->range_search(query_vec.data(), radius, doris_search_params, + doris_search_result) + .ok(), + true); + + ASSERT_EQ(native_results.size(), doris_search_result.roaring->cardinality()); + + ASSERT_EQ(doris_search_result.distances != nullptr, true); + for (size_t i = 0; i < native_results.size(); i++) { + const size_t rowid = native_results[i].first; + const float dis = native_results[i].second; + ASSERT_EQ(doris_search_result.roaring->contains(rowid), true) + << "Row ID mismatch at rank " << i; + ASSERT_FLOAT_EQ(doris_search_result.distances[i], sqrt(dis)) + << "Distance mismatch at rank " << i; + } + + doris_search_params.is_le_or_lt = false; + IndexSearchResult doris_search_result2; + ASSERT_EQ(index1->range_search(query_vec.data(), radius, doris_search_params, + doris_search_result2) + .ok(), + true); + roaring::Roaring ge_rows = *doris_search_result2.roaring; + roaring::Roaring less_rows; + for (size_t i = 0; i < native_results.size(); ++i) { + less_rows.add(native_results[i].first); + } + // result2 contains all rows that not included by result1 + roaring::Roaring and_row_id = ge_rows & less_rows; + roaring::Roaring or_row_id = ge_rows | less_rows; + ASSERT_NEAR(and_row_id.cardinality(), 0, 1); + ASSERT_EQ(or_row_id.cardinality(), sel_rows->cardinality()); + ASSERT_EQ(or_row_id, *sel_rows); + } +} + +TEST_F(VectorSearchTest, InnerProductTopKSearch) { + const size_t iterations = 1; + const std::vector dimensions = {32, 64}; + const std::vector vector_counts = {100, 500}; + const std::vector k_values = {5, 10, 20}; + + for (size_t iter = 0; iter < iterations; ++iter) { + for (int dim : dimensions) { + for (int n : vector_counts) { + for (int k : k_values) { + if (k > n) continue; + + // Create Doris index + auto doris_index = std::make_unique(); + FaissBuildParameter params; + params.dim = dim; + params.max_degree = 32; + params.index_type = FaissBuildParameter::IndexType::HNSW; + params.metric_type = FaissBuildParameter::MetricType::IP; + doris_index->build(params); + + // Generate normalized vectors (important for inner product) + std::vector flat_vectors; + for (int i = 0; i < n; ++i) { + auto vec = doris::vector_search_utils::generate_random_vector(dim); + // Normalize the vector + float norm = 0.0f; + for (float val : vec) { + norm += val * val; + } + norm = std::sqrt(norm); + if (norm > 0) { + for (float& val : vec) { + val /= norm; + } + } + flat_vectors.insert(flat_vectors.end(), vec.begin(), vec.end()); + } + + // Add vectors to index + doris_index->add(n, flat_vectors.data()); + + // Create query vector (also normalized) + auto query_vec = doris::vector_search_utils::generate_random_vector(dim); + float norm = 0.0f; + for (float val : query_vec) { + norm += val * val; + } + norm = std::sqrt(norm); + if (norm > 0) { + for (float& val : query_vec) { + val /= norm; + } + } + + // Perform search (top-N requires a candidate roaring and rows_of_segment) + HNSWSearchParameters search_params; + auto roaring = std::make_unique(); + for (int i = 0; i < n; ++i) roaring->add(i); + search_params.roaring = roaring.get(); + search_params.rows_of_segment = n; + IndexSearchResult search_result; + auto status = doris_index->ann_topn_search(query_vec.data(), k, search_params, + search_result); + + ASSERT_TRUE(status.ok()) << "Inner product search failed"; + ASSERT_GE(search_result.roaring->cardinality(), static_cast(k * 0.7)) + << "Expected to find at least 70% of requested results"; + + // Verify distances are in descending order (higher inner product is better) + for (size_t i = 1; i < search_result.roaring->cardinality(); ++i) { + ASSERT_GE(search_result.distances[i - 1], search_result.distances[i]) + << "Inner product distances should be in descending order"; + } + + // Verify all distances are valid (between -1 and 1 for normalized vectors) + for (size_t i = 0; i < search_result.roaring->cardinality(); ++i) { + ASSERT_GE(search_result.distances[i], -1.1f) + << "Inner product distance should be >= -1"; + ASSERT_LE(search_result.distances[i], 1.1f) + << "Inner product distance should be <= 1"; + } + } + } + } + } +} + +TEST_F(VectorSearchTest, InnerProductRangeSearchBasic) { + const size_t iterations = 3; + + for (size_t iter = 0; iter < iterations; ++iter) { + const int dim = 64; + const int n = 500; + const int m = 32; + + // Create Doris index + auto doris_index = std::make_unique(); + FaissBuildParameter params; + params.dim = dim; + params.max_degree = m; + params.index_type = FaissBuildParameter::IndexType::HNSW; + params.metric_type = FaissBuildParameter::MetricType::IP; + doris_index->build(params); + + // Create native index for comparison + faiss::IndexFlatIP native_index(dim); + + // Generate vectors + std::vector> vectors; + std::vector flat_vectors; + for (int i = 0; i < n; ++i) { + auto vec = doris::vector_search_utils::generate_random_vector(dim); + vectors.push_back(vec); + flat_vectors.insert(flat_vectors.end(), vec.begin(), vec.end()); + } + + // Add vectors to both indexes + doris_index->add(n, flat_vectors.data()); + native_index.add(n, flat_vectors.data()); + + // Use first vector as query + std::vector query_vec = vectors[0]; + + // Calculate radius based on inner product distribution + float radius = doris::vector_search_utils::get_radius_from_matrix( + query_vec.data(), dim, vectors, 0.5f, faiss::METRIC_INNER_PRODUCT); + + // Perform Doris range search + HNSWSearchParameters doris_params; + doris_params.ef_search = 100; + doris_params.is_le_or_lt = false; // For inner product, we want values >= radius + auto roaring = std::make_unique(); + for (int i = 0; i < n; ++i) { + roaring->add(i); + } + doris_params.roaring = roaring.get(); + + IndexSearchResult doris_result; + auto status = + doris_index->range_search(query_vec.data(), radius, doris_params, doris_result); + ASSERT_TRUE(status.ok()) << "Doris range search failed"; + + // Perform native range search + faiss::RangeSearchResult native_result(1, true); + native_index.range_search(1, query_vec.data(), radius, &native_result); + + // Compare results + size_t native_count = native_result.lims[1] - native_result.lims[0]; + ASSERT_NEAR(doris_result.roaring->cardinality(), native_count, 1) + << "Result count mismatch for inner product range search"; + + // Verify all returned distances are >= radius + for (size_t i = 0; i < doris_result.roaring->cardinality(); ++i) { + ASSERT_GE(doris_result.distances[i], radius - 1e-6) + << "Distance should be >= radius for inner product range search"; + } + } +} + +TEST_F(VectorSearchTest, InnerProductVsL2Comparison) { + const int dim = 32; + const int n = 100; + const int k = 10; + + // Generate the same set of vectors + std::vector flat_vectors; + for (int i = 0; i < n; ++i) { + auto vec = doris::vector_search_utils::generate_random_vector(dim); + flat_vectors.insert(flat_vectors.end(), vec.begin(), vec.end()); + } + + // Create L2 index + auto l2_index = std::make_unique(); + FaissBuildParameter l2_params; + l2_params.dim = dim; + l2_params.max_degree = 32; + l2_params.index_type = FaissBuildParameter::IndexType::HNSW; + l2_params.metric_type = FaissBuildParameter::MetricType::L2; + l2_index->build(l2_params); + l2_index->add(n, flat_vectors.data()); + + // Create Inner Product index + auto ip_index = std::make_unique(); + FaissBuildParameter ip_params; + ip_params.dim = dim; + ip_params.max_degree = 32; + ip_params.index_type = FaissBuildParameter::IndexType::HNSW; + ip_params.metric_type = FaissBuildParameter::MetricType::IP; + ip_index->build(ip_params); + ip_index->add(n, flat_vectors.data()); + + // Use first vector as query + std::vector query_vec(flat_vectors.begin(), flat_vectors.begin() + dim); + + // Search with L2 + HNSWSearchParameters search_params; + auto roaring = std::make_unique(); + for (int i = 0; i < n; ++i) roaring->add(i); + search_params.roaring = roaring.get(); + search_params.rows_of_segment = n; + IndexSearchResult l2_result; + auto l2_status = l2_index->ann_topn_search(query_vec.data(), k, search_params, l2_result); + ASSERT_EQ(l2_status.ok(), true) << "L2 search failed"; + + // Search with Inner Product + IndexSearchResult ip_result; + auto ip_status = ip_index->ann_topn_search(query_vec.data(), k, search_params, ip_result); + ASSERT_EQ(ip_status.ok(), true) << "Inner Product search failed"; + + // Both should find results + ASSERT_GT(l2_result.roaring->cardinality(), 0) << "L2 search should find results"; + ASSERT_GT(ip_result.roaring->cardinality(), 0) << "Inner Product search should find results"; + + // Results should be different (different metrics lead to different rankings) + // We'll check that at least some results are different + std::set l2_ids, ip_ids; + for (size_t i = 0; i < l2_result.roaring->cardinality(); ++i) { + l2_ids.insert(l2_result.roaring->getIndex(i)); + } + for (size_t i = 0; i < ip_result.roaring->cardinality(); ++i) { + ip_ids.insert(ip_result.roaring->getIndex(i)); + } + + // At least verify that both metrics return valid results + ASSERT_GT(l2_ids.size(), 0) << "L2 search should return valid IDs"; + ASSERT_GT(ip_ids.size(), 0) << "Inner Product search should return valid IDs"; + + // Verify distance ranges make sense + // L2 distances should be positive + for (size_t i = 0; i < l2_result.roaring->cardinality(); ++i) { + ASSERT_GE(l2_result.distances[i], 0.0f) << "L2 distance should be non-negative"; + } + + // Inner product distances can be negative or positive + bool has_valid_ip_distance = false; + for (size_t i = 0; i < ip_result.roaring->cardinality(); ++i) { + if (std::isfinite(ip_result.distances[i])) { + has_valid_ip_distance = true; + break; + } + } + ASSERT_EQ(has_valid_ip_distance, true) << "Inner Product should return valid distances"; +} + +TEST_F(VectorSearchTest, TestIdSelectorWithEmptyRoaring) { + auto roaring = std::make_unique(); + auto sel = FaissVectorIndex::roaring_to_faiss_selector(*roaring); + for (size_t i = 0; i < 10000; ++i) { + ASSERT_EQ(sel->is_member(i), false) << "Selector should be empty"; + } +} + +// New tests: radius == 0 or < 0 +TEST_F(VectorSearchTest, L2RangeSearchZeroAndNegativeRadius) { + const int dim = 32; + const int m = 32; + const int n = 200; + + auto index = std::make_unique(); + FaissBuildParameter params; + params.dim = dim; + params.max_degree = m; + params.index_type = FaissBuildParameter::IndexType::HNSW; + params.metric_type = FaissBuildParameter::MetricType::L2; + index->build(params); + + // Generate data + std::vector flat_vectors; + flat_vectors.reserve(static_cast(n) * dim); + for (int i = 0; i < n; ++i) { + auto v = doris::vector_search_utils::generate_random_vector(dim); + flat_vectors.insert(flat_vectors.end(), v.begin(), v.end()); + } + ASSERT_EQ(index->add(n, flat_vectors.data()).ok(), true); + + // Query uses the first vector -> exact match at distance 0 + std::vector query(flat_vectors.begin(), flat_vectors.begin() + dim); + + HNSWSearchParameters sp; + sp.ef_search = 64; + sp.is_le_or_lt = false; // Only test the ">=" branch (complement for L2) + auto all_rows = std::make_unique(); + for (int i = 0; i < n; ++i) all_rows->add(i); + sp.roaring = all_rows.get(); + + // radius == 0, ">=" for L2 means all rows except those with distance <= 0. + // With approximate HNSW range_search, the exact 0-distance self may or may not be found, + // so the complement size could be n (if none found) or n-1 (if self found). + IndexSearchResult res_ge0; + ASSERT_EQ(index->range_search(query.data(), 0.0f, sp, res_ge0).ok(), true); + ASSERT_EQ(res_ge0.distances, nullptr); + ASSERT_EQ(res_ge0.row_ids, nullptr); + ASSERT_TRUE(res_ge0.roaring->cardinality() == static_cast(n) || + res_ge0.roaring->cardinality() == static_cast(n - 1)); + + // radius < 0 (e.g., -1.0f) -> no vector has distance <= -1, so ">=" branch should return all rows + IndexSearchResult res_ge_neg; + ASSERT_EQ(index->range_search(query.data(), -1.0f, sp, res_ge_neg).ok(), true); + ASSERT_EQ(res_ge_neg.distances, nullptr); + ASSERT_EQ(res_ge_neg.roaring->cardinality(), static_cast(n)); +} + +TEST_F(VectorSearchTest, InnerProductRangeSearchZeroAndNegativeRadius) { + const int dim = 32; + const int m = 32; + const int n = 200; + + auto index = std::make_unique(); + FaissBuildParameter params; + params.dim = dim; + params.max_degree = m; + params.index_type = FaissBuildParameter::IndexType::HNSW; + params.metric_type = FaissBuildParameter::MetricType::IP; + index->build(params); + + // Generate normalized vectors to keep IP in [-1, 1] + std::vector flat_vectors; + flat_vectors.reserve(static_cast(n) * dim); + for (int i = 0; i < n; ++i) { + auto v = doris::vector_search_utils::generate_random_vector(dim); + float norm = 0.0f; + for (float x : v) norm += x * x; + norm = std::sqrt(norm); + if (norm > 0) + for (float& x : v) x /= norm; + flat_vectors.insert(flat_vectors.end(), v.begin(), v.end()); + } + ASSERT_EQ(index->add(n, flat_vectors.data()).ok(), true); + + std::vector query(flat_vectors.begin(), flat_vectors.begin() + dim); + // normalize query as well + float qn = 0.0f; + for (float x : query) qn += x * x; + qn = std::sqrt(qn); + if (qn > 0) + for (float& x : query) x /= qn; + + HNSWSearchParameters sp; + sp.ef_search = 100; + auto origin = std::make_unique(); + for (int i = 0; i < n; ++i) origin->add(i); + sp.roaring = origin.get(); + + // radius == 0, is_le_or_lt = false -> inner product >= 0 + sp.is_le_or_lt = false; + IndexSearchResult res_ge0; + ASSERT_EQ(index->range_search(query.data(), 0.0f, sp, res_ge0).ok(), true); + ASSERT_NE(res_ge0.distances, nullptr); + ASSERT_GT(res_ge0.roaring->cardinality(), 0u); + for (size_t i = 0; i < res_ge0.roaring->cardinality(); ++i) { + ASSERT_GE(res_ge0.distances[i], -1e-6f); + } + + // radius < 0, e.g., -1.0f -> almost all should satisfy IP >= -1 + sp.is_le_or_lt = false; + IndexSearchResult res_gen; + ASSERT_EQ(index->range_search(query.data(), -1.0f, sp, res_gen).ok(), true); + ASSERT_GE(res_gen.roaring->cardinality(), static_cast(n * 0.9)); +} + +} // namespace doris::vectorized \ No newline at end of file diff --git a/be/test/olap/vector_search/vector_search_utils.cpp b/be/test/olap/vector_search/vector_search_utils.cpp new file mode 100644 index 00000000000000..506d0e4ea8d5cc --- /dev/null +++ b/be/test/olap/vector_search/vector_search_utils.cpp @@ -0,0 +1,275 @@ +// 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 "vector_search_utils.h" + +#include + +#include +#include + +#include "olap/rowset/segment_v2/ann_index/ann_index.h" +#include "olap/rowset/segment_v2/ann_index/ann_search_params.h" +#include "olap/rowset/segment_v2/ann_index/faiss_ann_index.h" + +namespace doris::vector_search_utils { +static void accumulate(double x, double y, double& sum) { + sum += (x - y) * (x - y); +} + +static double finalize(double sum) { + return sqrt(sum); +} + +// Generate random vectors for testing +std::vector generate_random_vector(int dim) { + std::vector vector(dim); + std::random_device rd; + std::mt19937 gen(rd()); + std::uniform_real_distribution dis(-1.0f, 1.0f); + + for (int i = 0; i < dim; i++) { + vector[i] = dis(gen); + } + return vector; +} + +// Helper function to create and configure a Doris Vector index +std::unique_ptr create_doris_index(IndexType index_type, + int dimension, int m) { + auto index = std::make_unique(); + segment_v2::FaissBuildParameter params; + params.dim = dimension; + params.max_degree = m; + switch (index_type) { + case IndexType::HNSW: + params.index_type = segment_v2::FaissBuildParameter::IndexType::HNSW; + break; + default: + throw std::invalid_argument("Unsupported index type"); + } + index->build(params); + return std::move(index); +} + +// Helper function to create a native Faiss index +std::unique_ptr create_native_index(IndexType type, int dimension, int m) { + std::unique_ptr index; + + switch (type) { + case IndexType::FLAT_L2: + index = std::make_unique(dimension); + break; + case IndexType::HNSW: + index = std::make_unique(dimension, m, faiss::METRIC_L2); + break; + default: + throw std::invalid_argument("Unsupported index type"); + } + + return index; +} + +// Removed: create_native_hnsw_index_with_metric (not needed) + +// Helper function to generate a batch of random vectors +std::vector> generate_test_vectors_matrix(int num_vectors, int dimension) { + std::vector> vectors; + vectors.reserve(num_vectors); + + for (int i = 0; i < num_vectors; i++) { + vectors.push_back(generate_random_vector(dimension)); + } + + return vectors; +} + +std::vector generate_test_vectors_flatten(int num_vectors, int dimension) { + std::vector vectors; + vectors.reserve(num_vectors * dimension); + + for (int i = 0; i < num_vectors; i++) { + auto tmp = generate_random_vector(dimension); + vectors.insert(vectors.end(), tmp.begin(), tmp.end()); + } + + return vectors; +} + +// Helper function to add vectors to both Doris and native indexes +void add_vectors_to_indexes_serial_mode(segment_v2::VectorIndex* doris_index, + faiss::Index* native_index, + const std::vector>& vectors) { + for (size_t i = 0; i < vectors.size(); i++) { + if (doris_index) { + auto status = doris_index->add(1, vectors[i].data()); + ASSERT_TRUE(status.ok()) + << "Failed to add vector to Doris index: " << status.to_string(); + } + if (native_index) { + // Add vector to native Faiss index + native_index->add(1, vectors[i].data()); + } + } +} + +void add_vectors_to_indexes_batch_mode(segment_v2::VectorIndex* doris_index, + faiss::Index* native_index, size_t num_vectors, + const std::vector& flatten_vectors) { + if (doris_index) { + auto status = doris_index->add(num_vectors, flatten_vectors.data()); + ASSERT_TRUE(status.ok()) << "Failed to add vectors to Doris index: " << status.to_string(); + } + + if (native_index) { + // Add vectors to native Faiss index + native_index->add(num_vectors, flatten_vectors.data()); + } +} + +// Helper function to print search results for comparison +void print_search_results(const segment_v2::IndexSearchResult& doris_results, + const std::vector& native_distances, + const std::vector& native_indices, int query_idx) { + std::cout << "Query vector index: " << query_idx << std::endl; + + std::cout << "Doris Index Results:" << std::endl; + for (int i = 0; i < doris_results.roaring->cardinality(); i++) { + std::cout << "ID: " << doris_results.roaring->getIndex(i) + << ", Distance: " << doris_results.distances[i] << std::endl; + } + + std::cout << "Native Faiss Results:" << std::endl; + for (size_t i = 0; i < native_indices.size(); i++) { + if (native_indices[i] == -1) continue; + std::cout << "ID: " << native_indices[i] << ", Distance: " << native_distances[i] + << std::endl; + } +} + +// Helper function to compare search results between Doris and native Faiss +void compare_search_results(const segment_v2::IndexSearchResult& doris_results, + const std::vector& native_distances, + const std::vector& native_indices, float abs_error) { + EXPECT_EQ(doris_results.roaring->cardinality(), + std::count_if(native_indices.begin(), native_indices.end(), + [](faiss::idx_t id) { return id != -1; })); + + for (size_t i = 0; i < native_indices.size(); i++) { + if (native_indices[i] == -1) continue; + + EXPECT_TRUE(doris_results.roaring->contains(native_indices[i])) + << "ID mismatch at rank " << i; + EXPECT_NEAR(doris_results.distances[i], native_distances[i], abs_error) + << "Distance mismatch at rank " << i; + } +} + +// result is a vector of pairs, where each pair contains the labels and distance +// result is sorted by labels +std::vector> perform_native_index_range_search(faiss::Index* index, + const float* query_vector, + float radius) { + std::vector> results; + faiss::RangeSearchResult result(1); + index->range_search(1, query_vector, radius * radius, &result); + size_t begin = result.lims[0]; + size_t end = result.lims[1]; + results.reserve(end - begin); + for (size_t j = begin; j < end; ++j) { + results.push_back({result.labels[j], sqrt(result.distances[j])}); + } + std::sort(results.begin(), results.end(), + [](const auto& a, const auto& b) { return a.first < b.first; }); + return results; +} + +std::unique_ptr perform_doris_index_range_search( + segment_v2::VectorIndex* index, const float* query_vector, float radius, + const segment_v2::IndexSearchParameters& params) { + auto result = std::make_unique(); + std::ignore = index->range_search(query_vector, radius, params, *result); + return result; +} + +float get_radius_from_flatten(const float* vector, int dim, + const std::vector& flatten_vectors, float percentile) { + size_t n = flatten_vectors.size() / dim; + std::vector> distances(n); + for (int i = 0; i < n; i++) { + double sum = 0; + for (int j = 0; j < dim; j++) { + accumulate(flatten_vectors[i * dim + j], flatten_vectors[j], sum); + } + distances[i] = std::make_pair(i, finalize(sum)); + } + std::sort(distances.begin(), distances.end(), + [](const auto& a, const auto& b) { return a.second < b.second; }); + // Use the median distance as the radius + size_t percentile_index = static_cast(n * percentile); + float radius = distances[percentile_index].second; + + return radius; +} + +float get_radius_from_matrix(const float* vector, int dim, + const std::vector>& matrix_vectors, + float percentile, + faiss::MetricType metric_type /* = faiss::METRIC_L2 */) { + size_t n = matrix_vectors.size(); + std::vector> distances(n); + for (size_t i = 0; i < n; i++) { + double sum = 0; + if (metric_type == faiss::METRIC_L2) { + for (int j = 0; j < dim; j++) { + accumulate(matrix_vectors[i][j], vector[j], sum); + } + distances[i] = std::make_pair(i, finalize(sum)); + } else if (metric_type == faiss::METRIC_INNER_PRODUCT) { + for (int j = 0; j < dim; j++) { + sum += matrix_vectors[i][j] * vector[j]; + } + distances[i] = std::make_pair(i, static_cast(sum)); + } else { + throw std::invalid_argument("Unsupported metric type in get_radius_from_matrix"); + } + } + if (metric_type == faiss::METRIC_L2) { + std::sort(distances.begin(), distances.end(), + [](const auto& a, const auto& b) { return a.second < b.second; }); + } else if (metric_type == faiss::METRIC_INNER_PRODUCT) { + std::sort(distances.begin(), distances.end(), + [](const auto& a, const auto& b) { return a.second > b.second; }); + } + // Use the percentile distance as the radius + size_t percentile_index = static_cast(n * percentile); + if (percentile_index >= n) percentile_index = n - 1; + float radius = distances[percentile_index].second; + + return radius; +} + +std::pair, std::shared_ptr> +create_tmp_ann_index_reader(std::map properties) { + auto mock_tablet_index = std::make_unique(); + mock_tablet_index->_properties = properties; + auto mock_index_file_reader = std::make_shared(); + auto ann_reader = std::make_shared(mock_tablet_index.get(), + mock_index_file_reader); + return std::make_pair(std::move(mock_tablet_index), ann_reader); +} +} // namespace doris::vector_search_utils \ No newline at end of file diff --git a/be/test/olap/vector_search/vector_search_utils.h b/be/test/olap/vector_search/vector_search_utils.h new file mode 100644 index 00000000000000..a4ec4dbe73f693 --- /dev/null +++ b/be/test/olap/vector_search/vector_search_utils.h @@ -0,0 +1,352 @@ +// 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. + +#pragma once + +#include +#include +#include +#include +#include +#include +#include +#include +#include + +#include +#include +#include + +#include "common/object_pool.h" +#include "io/fs/local_file_system.h" +#include "olap/rowset/segment_v2/ann_index/ann_index.h" +#include "olap/rowset/segment_v2/ann_index/ann_index_iterator.h" +#include "olap/rowset/segment_v2/index_file_reader.h" +#include "olap/rowset/segment_v2/index_writer.h" +#include "olap/rowset/segment_v2/inverted_index_common.h" +#include "olap/rowset/segment_v2/inverted_index_compound_reader.h" +#include "olap/tablet_schema.h" +#include "runtime/descriptors.h" +#include "runtime/exec_env.h" +#include "vec/exprs/vexpr_context.h" +#include "vec/utils/util.hpp" +// Add CLucene RAM Directory header +#include +#include + +using doris::segment_v2::DorisCompoundReader; + +namespace faiss { +struct Index; +struct IndexHNSWFlat; +} // namespace faiss + +namespace doris::segment_v2 { +class FaissVectorIndex; +} + +namespace doris::vector_search_utils { + +// Generate random vectors for testing +std::vector generate_random_vector(int dim); +// Generate random vectors in matrix form +std::vector> generate_test_vectors_matrix(int num_vectors, int dimension); +// Generate random vectors as a flatten vector +std::vector generate_test_vectors_flatten(int num_vectors, int dimension); + +// Enum for different index types +enum class IndexType { + FLAT_L2, + HNSW, + // Add more index types as needed +}; +std::unique_ptr create_native_index(IndexType type, int dimension, int m); + +std::unique_ptr create_doris_index(IndexType index_type, + int dimension, int m); + +// Helper function to add vectors to both Doris and native indexes +void add_vectors_to_indexes_serial_mode(segment_v2::VectorIndex* doris_index, + faiss::Index* native_index, + const std::vector>& vectors); + +void add_vectors_to_indexes_batch_mode(segment_v2::VectorIndex* doris_index, + faiss::Index* native_index, size_t num_vectors, + const std::vector& flatten_vectors); + +void print_search_results(const segment_v2::IndexSearchResult& doris_results, + const std::vector& native_distances, + const std::vector& native_indices, int query_idx); + +float get_radius_from_flatten(const float* vector, int dim, + const std::vector& flatten_vectors, float percentile); +float get_radius_from_matrix(const float* vector, int dim, + const std::vector>& matrix_vectors, + float percentile, faiss::MetricType metric_type = faiss::METRIC_L2); +// Helper function to compare search results between Doris and native Faiss +void compare_search_results(const segment_v2::IndexSearchResult& doris_results, + const std::vector& native_distances, + const std::vector& native_indices, + float abs_error = 1e-5); + +// result is a vector of pairs, where each pair contains the labels and distance +// result is sorted by labels +std::vector> perform_native_index_range_search(faiss::Index* index, + const float* query_vector, + float radius); + +std::unique_ptr perform_doris_index_range_search( + segment_v2::VectorIndex* index, const float* query_vector, float radius, + const segment_v2::IndexSearchParameters& params); + +class MockIndexFileReader : public ::doris::segment_v2::IndexFileReader { +public: + MockIndexFileReader() + : IndexFileReader(doris::io::global_local_filesystem(), "", + doris::InvertedIndexStorageFormatPB::V2) {} + + MOCK_METHOD2(init, doris::Status(int, const doris::io::IOContext* io_ctx)); + + MOCK_CONST_METHOD2( + open, Result>( + const doris::TabletIndex*, const doris::io::IOContext*)); +}; + +class MockTabletSchema : public doris::TabletIndex {}; + +class MockTabletColumn : public doris::TabletColumn { + MOCK_METHOD(doris::FieldType, type, (), (const)); + // Match base class signature exactly to ensure override is used + MOCK_METHOD((const TabletColumn&), get_sub_column, (uint64_t), (const)); +}; + +class MockTabletIndex : public doris::TabletIndex { + MOCK_METHOD(doris::IndexType, index_type, (), (const)); + MOCK_METHOD((const std::map&), properties, (), (const)); +}; + +class MockIndexFileWriter : public doris::segment_v2::IndexFileWriter { +public: + MockIndexFileWriter(doris::io::FileSystemSPtr fs) + : doris::segment_v2::IndexFileWriter(fs, "test_index", "rowset_id", 1, + doris::InvertedIndexStorageFormatPB::V2) {} + + MOCK_METHOD(doris::Result>, open, + (const doris::TabletIndex* index_meta), (override)); +}; + +class MockAnnIndexIterator : public doris::segment_v2::AnnIndexIterator { +public: + MockAnnIndexIterator() : doris::segment_v2::AnnIndexIterator(nullptr) {} + + ~MockAnnIndexIterator() override = default; + + MOCK_METHOD(Status, read_from_index, (const doris::segment_v2::IndexParam& param), (override)); + MOCK_METHOD(Status, range_search, + (const segment_v2::AnnRangeSearchParams& params, + const VectorSearchUserParams& custom_params, + segment_v2::AnnRangeSearchResult* result, segment_v2::AnnIndexStats* stats), + (override)); + +private: + io::IOContext _io_ctx_mock; +}; + +class MockAnnIndexReader : public doris::segment_v2::AnnIndexReader {}; + +std::pair, std::shared_ptr> +create_tmp_ann_index_reader(std::map properties); + +} // namespace doris::vector_search_utils + +namespace doris::vectorized { + +class VectorSearchTest : public ::testing::Test { +public: + static void accumulate(double x, double y, double& sum) { sum += (x - y) * (x - y); } + static double finalize(double sum) { return sqrt(sum); } + +protected: + void SetUp() override { + // Ensure ExecEnv has a valid tmp dir for IndexFileWriter (prevents nullptr deref) + { + // Only set if not configured by other tests + if (ExecEnv::GetInstance()->get_tmp_file_dirs() == nullptr) { + const std::string tmp_dir = "./ut_dir/tmp_vector_search"; + (void)io::global_local_filesystem()->delete_directory(tmp_dir); + (void)io::global_local_filesystem()->create_directory(tmp_dir); + std::vector paths; + paths.emplace_back(tmp_dir, -1); + auto tmp_file_dirs = std::make_unique(paths); + ASSERT_TRUE(tmp_file_dirs->init().ok()); + ExecEnv::GetInstance()->set_tmp_file_dir(std::move(tmp_file_dirs)); + } + } + + TDescriptorTable thrift_tbl; + TTableDescriptor thrift_table_desc; + thrift_table_desc.id = 0; + thrift_tbl.tableDescriptors.push_back(thrift_table_desc); + TTupleDescriptor tuple_desc; + tuple_desc.__isset.tableId = true; + tuple_desc.id = 0; + tuple_desc.tableId = 0; + thrift_tbl.tupleDescriptors.push_back(tuple_desc); + TSlotDescriptor slot_desc; + slot_desc.id = 0; + slot_desc.parent = 0; + slot_desc.isMaterialized = true; + slot_desc.need_materialize = true; + slot_desc.__isset.need_materialize = true; + TTypeNode type_node; + type_node.type = TTypeNodeType::type::SCALAR; + TScalarType scalar_type; + scalar_type.__set_type(TPrimitiveType::DOUBLE); + type_node.__set_scalar_type(scalar_type); + slot_desc.slotType.types.push_back(type_node); + slot_desc.virtual_column_expr = read_from_json(_distance_function_call_thrift); + slot_desc.__isset.virtual_column_expr = true; + thrift_tbl.slotDescriptors.push_back(slot_desc); + slot_desc.id = 1; + slot_desc.__isset.virtual_column_expr = false; + thrift_tbl.slotDescriptors.push_back(slot_desc); + thrift_tbl.__isset.slotDescriptors = true; + // std::cout << "+++++++++++++++++++ thrift table descriptor:\n" + // << apache::thrift::ThriftDebugString(thrift_tbl) << std::endl; + // std::cout << "+++++++++++++++++++ thrift table descriptor end\n"; + ASSERT_TRUE(DescriptorTbl::create(&obj_pool, thrift_tbl, &_desc_tbl).ok()); + // std::cout << "+++++++++++++++++++ desc tbl\n" << _desc_tbl->debug_string() << std::endl; + // std::cout << "+++++++++++++++++++ desc tbl end\n"; + _runtime_state.set_desc_tbl(_desc_tbl); + + config::max_depth_of_expr_tree = 1000; + doris::TSlotRef slot_ref; + slot_ref.slot_id = 0; + slot_ref.__isset.is_virtual_slot = true; + slot_ref.is_virtual_slot = true; + + doris::TExprNode virtual_slot_ref_node; + virtual_slot_ref_node.slot_ref = slot_ref; + virtual_slot_ref_node.label = "virtual_slot_ref"; + virtual_slot_ref_node.node_type = TExprNodeType::SLOT_REF; + virtual_slot_ref_node.type = TTypeDesc(); + type_node.type = TTypeNodeType::type::SCALAR; + type_node.scalar_type.type = TPrimitiveType::DOUBLE; + type_node.__isset.scalar_type = true; + + virtual_slot_ref_node.type.types.push_back(type_node); + virtual_slot_ref_node.__isset.slot_ref = true; + virtual_slot_ref_node.__isset.label = true; + virtual_slot_ref_node.__isset.opcode = false; + + _virtual_slot_ref_expr.nodes.push_back(virtual_slot_ref_node); + _ann_index_iterator = std::make_unique(); + + _row_desc = RowDescriptor(*_desc_tbl, {0}, {false}); + + // Create CLucene RAM directory instead of mock + _ram_dir = std::make_shared(); + + // Optional: Create test file to simulate index presence + auto output = _ram_dir->createOutput("index_file"); + // Write some dummy data + const char* dummy_data = "dummy data"; + output->writeBytes((const uint8_t*)dummy_data, strlen(dummy_data)); + output->close(); + delete output; // CLucene requires manual delete + } + + void TearDown() override {} + +private: + doris::ObjectPool obj_pool; + RowDescriptor _row_desc; + std::unique_ptr _ann_index_iterator; + vectorized::IColumn::MutablePtr _result_column; + doris::TExpr _virtual_slot_ref_expr; + DescriptorTbl* _desc_tbl; + doris::RuntimeState _runtime_state; + std::shared_ptr _ram_dir; + + /* + [0] TExprNode { + num_children = 2 + fn = TFunctionName { + name = "l2_distance_approximate" + } + }, + [1] TExprNode { + num_children = 1 + fn = TFunctionName { + name = "casttoarray" + } + }, + [2] TExprNode { + num_children = 0, + slot_ref = TSlotRef { + slot_id = 1 + } + }, + [3] TExprNode { + node_type = ARRAY_LITERAL, + num_children = 8, + }, + [4] TExprNode { + float_literal = TFloatLiteral{ + value = 1 + } + }, + [5] TExprNode { + float_literal = TFloatLiteral{ + value = 2 + } + }, + [6] TExprNode { + float_literal = TFloatLiteral{ + value = 3 + } + }, + [7] TExprNode { + float_literal = TFloatLiteral{ + value = 4 + } + }, + [8] TExprNode { + float_literal = TFloatLiteral{ + value = 5 + } + }, + [9] TExprNode { + float_literal = TFloatLiteral{ + value = 6 + } + }, + [10] TExprNode { + float_literal = TFloatLiteral{ + value = 7 + } + }, + [11] TExprNode { + float_literal = TFloatLiteral{ + value = 20 + } + }, + */ + const std::string _distance_function_call_thrift = + R"xxx({"1":{"lst":["rec",12,{"1":{"i32":20},"2":{"rec":{"1":{"lst":["rec",1,{"1":{"i32":0},"2":{"rec":{"1":{"i32":8}}}}]},"3":{"i64":-1}}},"4":{"i32":2},"20":{"i32":-1},"26":{"rec":{"1":{"rec":{"2":{"str":"l2_distance_approximate"}}},"2":{"i32":0},"3":{"lst":["rec",2,{"1":{"lst":["rec",2,{"1":{"i32":1},"4":{"tf":1},"5":{"lst":["tf",1,1]}},{"1":{"i32":0},"2":{"rec":{"1":{"i32":8}}}}]},"3":{"i64":-1}},{"1":{"lst":["rec",2,{"1":{"i32":1},"4":{"tf":1},"5":{"lst":["tf",1,1]}},{"1":{"i32":0},"2":{"rec":{"1":{"i32":8}}}}]},"3":{"i64":-1}}]},"4":{"rec":{"1":{"lst":["rec",1,{"1":{"i32":0},"2":{"rec":{"1":{"i32":8}}}}]},"3":{"i64":-1}}},"5":{"tf":0},"7":{"str":"l2_distance_approximate(array, array)"},"9":{"rec":{"1":{"str":""}}},"11":{"i64":0},"13":{"tf":1},"14":{"tf":0},"15":{"tf":0},"16":{"i64":360}}},"29":{"tf":1}},{"1":{"i32":5},"2":{"rec":{"1":{"lst":["rec",2,{"1":{"i32":1},"4":{"tf":1},"5":{"lst":["tf",1,1]}},{"1":{"i32":0},"2":{"rec":{"1":{"i32":8}}}}]},"3":{"i64":-1}}},"3":{"i32":4},"4":{"i32":1},"20":{"i32":-1},"26":{"rec":{"1":{"rec":{"2":{"str":"casttoarray"}}},"2":{"i32":0},"3":{"lst":["rec",1,{"1":{"lst":["rec",2,{"1":{"i32":1},"4":{"tf":1},"5":{"lst":["tf",1,1]}},{"1":{"i32":0},"2":{"rec":{"1":{"i32":7}}}}]},"3":{"i64":-1}}]},"4":{"rec":{"1":{"lst":["rec",2,{"1":{"i32":1},"4":{"tf":1},"5":{"lst":["tf",1,1]}},{"1":{"i32":0},"2":{"rec":{"1":{"i32":8}}}}]},"3":{"i64":-1}}},"5":{"tf":0},"7":{"str":"casttoarray(array)"},"9":{"rec":{"1":{"str":""}}},"11":{"i64":0},"13":{"tf":1},"14":{"tf":0},"15":{"tf":0},"16":{"i64":360}}},"29":{"tf":0}},{"1":{"i32":16},"2":{"rec":{"1":{"lst":["rec",2,{"1":{"i32":1},"4":{"tf":1},"5":{"lst":["tf",1,1]}},{"1":{"i32":0},"2":{"rec":{"1":{"i32":7}}}}]},"3":{"i64":-1}}},"4":{"i32":0},"15":{"rec":{"1":{"i32":1},"2":{"i32":0},"3":{"i32":1}}},"20":{"i32":-1},"29":{"tf":0},"36":{"str":"embedding"}},{"1":{"i32":21},"2":{"rec":{"1":{"lst":["rec",2,{"1":{"i32":1},"4":{"tf":1},"5":{"lst":["tf",1,1]}},{"1":{"i32":0},"2":{"rec":{"1":{"i32":8}}}}]},"3":{"i64":-1}}},"4":{"i32":8},"20":{"i32":-1},"28":{"i32":8},"29":{"tf":0}},{"1":{"i32":8},"2":{"rec":{"1":{"lst":["rec",1,{"1":{"i32":0},"2":{"rec":{"1":{"i32":8}}}}]},"3":{"i64":-1}}},"4":{"i32":0},"9":{"rec":{"1":{"dbl":1}}},"20":{"i32":-1},"29":{"tf":0}},{"1":{"i32":8},"2":{"rec":{"1":{"lst":["rec",1,{"1":{"i32":0},"2":{"rec":{"1":{"i32":8}}}}]},"3":{"i64":-1}}},"4":{"i32":0},"9":{"rec":{"1":{"dbl":2}}},"20":{"i32":-1},"29":{"tf":0}},{"1":{"i32":8},"2":{"rec":{"1":{"lst":["rec",1,{"1":{"i32":0},"2":{"rec":{"1":{"i32":8}}}}]},"3":{"i64":-1}}},"4":{"i32":0},"9":{"rec":{"1":{"dbl":3}}},"20":{"i32":-1},"29":{"tf":0}},{"1":{"i32":8},"2":{"rec":{"1":{"lst":["rec",1,{"1":{"i32":0},"2":{"rec":{"1":{"i32":8}}}}]},"3":{"i64":-1}}},"4":{"i32":0},"9":{"rec":{"1":{"dbl":4}}},"20":{"i32":-1},"29":{"tf":0}},{"1":{"i32":8},"2":{"rec":{"1":{"lst":["rec",1,{"1":{"i32":0},"2":{"rec":{"1":{"i32":8}}}}]},"3":{"i64":-1}}},"4":{"i32":0},"9":{"rec":{"1":{"dbl":5}}},"20":{"i32":-1},"29":{"tf":0}},{"1":{"i32":8},"2":{"rec":{"1":{"lst":["rec",1,{"1":{"i32":0},"2":{"rec":{"1":{"i32":8}}}}]},"3":{"i64":-1}}},"4":{"i32":0},"9":{"rec":{"1":{"dbl":6}}},"20":{"i32":-1},"29":{"tf":0}},{"1":{"i32":8},"2":{"rec":{"1":{"lst":["rec",1,{"1":{"i32":0},"2":{"rec":{"1":{"i32":8}}}}]},"3":{"i64":-1}}},"4":{"i32":0},"9":{"rec":{"1":{"dbl":7}}},"20":{"i32":-1},"29":{"tf":0}},{"1":{"i32":8},"2":{"rec":{"1":{"lst":["rec",1,{"1":{"i32":0},"2":{"rec":{"1":{"i32":8}}}}]},"3":{"i64":-1}}},"4":{"i32":0},"9":{"rec":{"1":{"dbl":20}}},"20":{"i32":-1},"29":{"tf":0}}]}})xxx"; +}; +} // namespace doris::vectorized \ No newline at end of file diff --git a/contrib/faiss b/contrib/faiss new file mode 160000 index 00000000000000..b2b482fccec9f4 --- /dev/null +++ b/contrib/faiss @@ -0,0 +1 @@ +Subproject commit b2b482fccec9f488a0f5d90e415dec41a6ee6469 diff --git a/contrib/openblas b/contrib/openblas new file mode 160000 index 00000000000000..77986e49425532 --- /dev/null +++ b/contrib/openblas @@ -0,0 +1 @@ +Subproject commit 77986e49425532bf8f651db74cbe1579bcb4a5bf diff --git a/fe/fe-core/src/main/antlr4/org/apache/doris/nereids/DorisLexer.g4 b/fe/fe-core/src/main/antlr4/org/apache/doris/nereids/DorisLexer.g4 index e54d43dcdf7e22..db1198b8eb7e34 100644 --- a/fe/fe-core/src/main/antlr4/org/apache/doris/nereids/DorisLexer.g4 +++ b/fe/fe-core/src/main/antlr4/org/apache/doris/nereids/DorisLexer.g4 @@ -367,6 +367,7 @@ NEGATIVE: 'NEGATIVE'; NEVER: 'NEVER'; NEXT: 'NEXT'; NGRAM_BF: 'NGRAM_BF'; +ANN: 'ANN'; NO: 'NO'; NO_USE_MV: 'NO_USE_MV'; NON_NULLABLE: 'NON_NULLABLE'; diff --git a/fe/fe-core/src/main/antlr4/org/apache/doris/nereids/DorisParser.g4 b/fe/fe-core/src/main/antlr4/org/apache/doris/nereids/DorisParser.g4 index 91c98a6599d88a..0b32733b2aa3c6 100644 --- a/fe/fe-core/src/main/antlr4/org/apache/doris/nereids/DorisParser.g4 +++ b/fe/fe-core/src/main/antlr4/org/apache/doris/nereids/DorisParser.g4 @@ -199,7 +199,7 @@ supportedCreateStatement partitionSpec? #buildIndex | CREATE INDEX (IF NOT EXISTS)? name=identifier ON tableName=multipartIdentifier identifierList - (USING (BITMAP | NGRAM_BF | INVERTED))? + (USING (BITMAP | NGRAM_BF | INVERTED | ANN))? properties=propertyClause? (COMMENT STRING_LITERAL)? #createIndex | CREATE WORKLOAD POLICY (IF NOT EXISTS)? name=identifierOrText (CONDITIONS LEFT_PAREN workloadPolicyConditions RIGHT_PAREN)? @@ -1426,7 +1426,7 @@ indexDefs ; indexDef - : INDEX (ifNotExists=IF NOT EXISTS)? indexName=identifier cols=identifierList (USING indexType=(BITMAP | INVERTED | NGRAM_BF))? (PROPERTIES LEFT_PAREN properties=propertyItemList RIGHT_PAREN)? (COMMENT comment=STRING_LITERAL)? + : INDEX (ifNotExists=IF NOT EXISTS)? indexName=identifier cols=identifierList (USING indexType=(BITMAP | INVERTED | NGRAM_BF | ANN ))? (PROPERTIES LEFT_PAREN properties=propertyItemList RIGHT_PAREN)? (COMMENT comment=STRING_LITERAL)? ; partitionsDef @@ -1853,6 +1853,7 @@ nonReserved | ALIAS | ALWAYS | ANALYZED + | ANN | ARRAY | AT | AUTHORS diff --git a/fe/fe-core/src/main/java/org/apache/doris/alter/SchemaChangeHandler.java b/fe/fe-core/src/main/java/org/apache/doris/alter/SchemaChangeHandler.java index 7d274ef5c754e2..3ddb207e951a6b 100644 --- a/fe/fe-core/src/main/java/org/apache/doris/alter/SchemaChangeHandler.java +++ b/fe/fe-core/src/main/java/org/apache/doris/alter/SchemaChangeHandler.java @@ -20,6 +20,7 @@ import org.apache.doris.analysis.AddColumnClause; import org.apache.doris.analysis.AddColumnsClause; import org.apache.doris.analysis.AlterClause; +import org.apache.doris.analysis.AnnIndexPropertiesChecker; import org.apache.doris.analysis.BuildIndexClause; import org.apache.doris.analysis.ColumnPosition; import org.apache.doris.analysis.CreateIndexClause; @@ -2160,7 +2161,10 @@ public int getAsInt() { } } // only inverted index with local mode can do light drop index change - if (found != null && found.getIndexType() == IndexDef.IndexType.INVERTED + // not sure whether the logic here is correct for ann index + // just reuse the logic for inverted index + if (found != null && (found.getIndexType() == IndexDef.IndexType.INVERTED + || found.getIndexType() == IndexDef.IndexType.ANN) && Config.isNotCloudMode()) { alterIndexes.add(found); isDropIndex = true; @@ -2738,6 +2742,13 @@ private boolean processAddIndex(CreateIndexClause alterClause, OlapTable olapTab alterIndex.setIndexId(Env.getCurrentEnv().getNextId()); } + if (indexDef.isAnnIndex()) { + if (olapTable.getKeysType() != KeysType.DUP_KEYS) { + throw new AnalysisException("ANN index can only be built on table with DUP_KEYS"); + } + AnnIndexPropertiesChecker.checkProperties(indexDef.getProperties()); + } + for (String col : indexDef.getColumns()) { Column column = olapTable.getColumn(col); if (column != null) { diff --git a/fe/fe-core/src/main/java/org/apache/doris/analysis/AnnIndexPropertiesChecker.java b/fe/fe-core/src/main/java/org/apache/doris/analysis/AnnIndexPropertiesChecker.java new file mode 100644 index 00000000000000..cafe7e9248891f --- /dev/null +++ b/fe/fe-core/src/main/java/org/apache/doris/analysis/AnnIndexPropertiesChecker.java @@ -0,0 +1,88 @@ +// 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. + +package org.apache.doris.analysis; + +import org.apache.doris.nereids.exceptions.AnalysisException; + +import java.util.Map; + +public class AnnIndexPropertiesChecker { + public static void checkProperties(Map properties) { + // // ANN index is not supported in cloud mode + // if (Config.isCloudMode()) { + // throw new AnalysisException("ANN index is not supported in cloud mode"); + // } + + String type = null; + String metric = null; + String dim = null; + for (String key : properties.keySet()) { + switch (key) { + case "index_type": + type = properties.get(key); + if (!type.equals("hnsw")) { + throw new AnalysisException("only support ann index with type hnsw, got: " + type); + } + break; + case "metric_type": + metric = properties.get(key); + if (!metric.equals("l2_distance") && !metric.equals("inner_product")) { + throw new AnalysisException( + "only support ann index with metric l2_distance or inner_product, got: " + metric); + } + break; + case "dim": + dim = properties.get(key); + try { + int dimension = Integer.parseInt(dim); + if (dimension <= 0) { + throw new AnalysisException("dim of ann index must be a positive integer, got: " + dim); + } + } catch (NumberFormatException e) { + throw new AnalysisException("dim of ann index must be a positive integer, got: " + dim); + } + break; + case "max_degree": + String maxDegree = properties.get(key); + try { + int degree = Integer.parseInt(maxDegree); + if (degree <= 0) { + throw new AnalysisException( + "max_degree of ann index must be a positive integer, got: " + maxDegree); + } + } catch (NumberFormatException e) { + throw new AnalysisException( + "max_degree of ann index must be a positive integer, got: " + maxDegree); + } + break; + default: + throw new AnalysisException("unknown ann index property: " + key); + } + } + + if (type == null) { + throw new AnalysisException("index_type of ann index be specified."); + } + if (metric == null) { + throw new AnalysisException("metric_type of ann index must be specified."); + } + if (dim == null) { + throw new AnalysisException("dim of ann index must be specified"); + } + } +} diff --git a/fe/fe-core/src/main/java/org/apache/doris/analysis/IndexDef.java b/fe/fe-core/src/main/java/org/apache/doris/analysis/IndexDef.java index 701f3971bed383..2f342081e6e365 100644 --- a/fe/fe-core/src/main/java/org/apache/doris/analysis/IndexDef.java +++ b/fe/fe-core/src/main/java/org/apache/doris/analysis/IndexDef.java @@ -92,7 +92,7 @@ public IndexDef(String indexName, PartitionNames partitionNames, IndexType index } public void analyze() throws AnalysisException { - if (isBuildDeferred && indexType == IndexDef.IndexType.INVERTED) { + if (isBuildDeferred && (indexType == IndexDef.IndexType.INVERTED || indexType == IndexDef.IndexType.ANN)) { if (Strings.isNullOrEmpty(indexName)) { throw new AnalysisException("index name cannot be blank."); } @@ -103,7 +103,8 @@ public void analyze() throws AnalysisException { } if (indexType == IndexDef.IndexType.BITMAP - || indexType == IndexDef.IndexType.INVERTED) { + || indexType == IndexDef.IndexType.INVERTED + || indexType == IndexDef.IndexType.ANN) { if (columns == null || columns.size() != 1) { throw new AnalysisException(indexType.toString() + " index can only apply to a single column."); } @@ -210,13 +211,18 @@ public enum IndexType { BITMAP, INVERTED, BLOOMFILTER, - NGRAM_BF + NGRAM_BF, + ANN } public boolean isInvertedIndex() { return (this.indexType == IndexType.INVERTED); } + public boolean isAnnIndex() { + return (this.indexType == IndexType.ANN); + } + // Check if the column type is supported for inverted index public boolean isSupportIdxType(Type colType) { if (colType.isArrayType()) { @@ -235,6 +241,28 @@ public boolean isSupportIdxType(Type colType) { public void checkColumn(Column column, KeysType keysType, boolean enableUniqueKeyMergeOnWrite, TInvertedIndexFileStorageFormat invertedIndexFileStorageFormat) throws AnalysisException { + if (indexType == IndexType.ANN) { + if (column.isAllowNull()) { + throw new AnalysisException("ANN index must be built on a column that is not nullable"); + } + + String indexColName = column.getName(); + caseSensitivityColumns.add(indexColName); + PrimitiveType primitiveType = column.getDataType(); + if (!primitiveType.isArrayType()) { + throw new AnalysisException("ANN index column must be array type"); + } + Type columnType = column.getType(); + Type itemType = ((ArrayType) columnType).getItemType(); + if (!itemType.isFloatingPointType()) { + throw new AnalysisException("ANN index column item type must be float type"); + } + if (keysType != KeysType.DUP_KEYS) { + throw new AnalysisException("ANN index can only be used in DUP_KEYS table"); + } + return; + } + if (indexType == IndexType.BITMAP || indexType == IndexType.INVERTED || indexType == IndexType.BLOOMFILTER || indexType == IndexType.NGRAM_BF) { String indexColName = column.getName(); @@ -246,6 +274,10 @@ public void checkColumn(Column column, KeysType keysType, boolean enableUniqueKe + "invalid index: " + indexName); } + if (indexType == IndexType.ANN && !colType.isArrayType()) { + throw new AnalysisException("ANN index column must be array type"); + } + // In inverted index format v1, each subcolumn of a variant has its own index file, leading to high IOPS. // when the subcolumn type changes, it may result in missing files, causing link file failure. // There are two cases in which the inverted index format v1 is not supported: diff --git a/fe/fe-core/src/main/java/org/apache/doris/catalog/BuiltinScalarFunctions.java b/fe/fe-core/src/main/java/org/apache/doris/catalog/BuiltinScalarFunctions.java index 7b4290584c8e9e..a394bf7c1f97e3 100644 --- a/fe/fe-core/src/main/java/org/apache/doris/catalog/BuiltinScalarFunctions.java +++ b/fe/fe-core/src/main/java/org/apache/doris/catalog/BuiltinScalarFunctions.java @@ -236,6 +236,7 @@ import org.apache.doris.nereids.trees.expressions.functions.scalar.Ignore; import org.apache.doris.nereids.trees.expressions.functions.scalar.Initcap; import org.apache.doris.nereids.trees.expressions.functions.scalar.InnerProduct; +import org.apache.doris.nereids.trees.expressions.functions.scalar.InnerProductApproximate; import org.apache.doris.nereids.trees.expressions.functions.scalar.Instr; import org.apache.doris.nereids.trees.expressions.functions.scalar.InttoUuid; import org.apache.doris.nereids.trees.expressions.functions.scalar.Ipv4CIDRToRange; @@ -286,6 +287,7 @@ import org.apache.doris.nereids.trees.expressions.functions.scalar.JsonbValid; import org.apache.doris.nereids.trees.expressions.functions.scalar.L1Distance; import org.apache.doris.nereids.trees.expressions.functions.scalar.L2Distance; +import org.apache.doris.nereids.trees.expressions.functions.scalar.L2DistanceApproximate; import org.apache.doris.nereids.trees.expressions.functions.scalar.LastDay; import org.apache.doris.nereids.trees.expressions.functions.scalar.LastQueryId; import org.apache.doris.nereids.trees.expressions.functions.scalar.Lcm; @@ -724,6 +726,7 @@ public class BuiltinScalarFunctions implements FunctionHelper { scalar(If.class, "if"), scalar(Ignore.class, "ignore"), scalar(Initcap.class, "initcap"), + scalar(InnerProductApproximate.class, "inner_product_approximate"), scalar(InnerProduct.class, "inner_product"), scalar(Instr.class, "instr"), scalar(InttoUuid.class, "int_to_uuid"), @@ -778,6 +781,7 @@ public class BuiltinScalarFunctions implements FunctionHelper { scalar(JsonContains.class, "json_contains"), scalar(JsonKeys.class, "json_keys", "jsonb_keys"), scalar(L1Distance.class, "l1_distance"), + scalar(L2DistanceApproximate.class, "l2_distance_approximate"), scalar(L2Distance.class, "l2_distance"), scalar(LastDay.class, "last_day"), scalar(Lcm.class, "lcm"), diff --git a/fe/fe-core/src/main/java/org/apache/doris/catalog/Index.java b/fe/fe-core/src/main/java/org/apache/doris/catalog/Index.java index 5bc891a40e1e81..fbb39fb6ed8cb8 100644 --- a/fe/fe-core/src/main/java/org/apache/doris/catalog/Index.java +++ b/fe/fe-core/src/main/java/org/apache/doris/catalog/Index.java @@ -333,6 +333,10 @@ public OlapFile.TabletIndexPB toPb(Map columnMap, List builder.setIndexType(OlapFile.IndexType.BLOOMFILTER); break; + case ANN: + builder.setIndexType(OlapFile.IndexType.ANN); + break; + default: throw new RuntimeException("indexType " + indexType + " is not processed in toPb"); } diff --git a/fe/fe-core/src/main/java/org/apache/doris/nereids/glue/translator/PhysicalPlanTranslator.java b/fe/fe-core/src/main/java/org/apache/doris/nereids/glue/translator/PhysicalPlanTranslator.java index 14acf3f5f83a1c..e07cc4851dea62 100644 --- a/fe/fe-core/src/main/java/org/apache/doris/nereids/glue/translator/PhysicalPlanTranslator.java +++ b/fe/fe-core/src/main/java/org/apache/doris/nereids/glue/translator/PhysicalPlanTranslator.java @@ -910,6 +910,25 @@ private PlanFragment computePhysicalOlapScan(PhysicalOlapScan olapScan, PlanTran olapScanNode.setScoreSortLimit(olapScan.getScoreLimit().get()); } + // translate ann topn info + if (!olapScan.getAnnOrderKeys().isEmpty()) { + TupleDescriptor annSortTuple = olapScanNode.getTupleDesc(); + List orderingExprs = Lists.newArrayList(); + List ascOrders = Lists.newArrayList(); + List nullsFirstParams = Lists.newArrayList(); + List annOrderKeys = olapScan.getAnnOrderKeys(); + annOrderKeys.forEach(k -> { + orderingExprs.add(ExpressionTranslator.translate(k.getExpr(), context)); + ascOrders.add(k.isAsc()); + nullsFirstParams.add(k.isNullFirst()); + }); + SortInfo annSortInfo = new SortInfo(orderingExprs, ascOrders, nullsFirstParams, annSortTuple); + olapScanNode.setAnnSortInfo(annSortInfo); + } + if (olapScan.getAnnLimit().isPresent()) { + olapScanNode.setAnnSortLimit(olapScan.getAnnLimit().get()); + } + // TODO: move all node set cardinality into one place if (olapScan.getStats() != null) { // NOTICE: we should not set stats row count diff --git a/fe/fe-core/src/main/java/org/apache/doris/nereids/jobs/executor/Rewriter.java b/fe/fe-core/src/main/java/org/apache/doris/nereids/jobs/executor/Rewriter.java index 167ff51cf03d30..a3f2385c591f1a 100644 --- a/fe/fe-core/src/main/java/org/apache/doris/nereids/jobs/executor/Rewriter.java +++ b/fe/fe-core/src/main/java/org/apache/doris/nereids/jobs/executor/Rewriter.java @@ -138,6 +138,7 @@ import org.apache.doris.nereids.rules.rewrite.PushDownTopNThroughJoin; import org.apache.doris.nereids.rules.rewrite.PushDownTopNThroughUnion; import org.apache.doris.nereids.rules.rewrite.PushDownTopNThroughWindow; +import org.apache.doris.nereids.rules.rewrite.PushDownVectorTopNIntoOlapScan; import org.apache.doris.nereids.rules.rewrite.PushDownVirtualColumnsIntoOlapScan; import org.apache.doris.nereids.rules.rewrite.PushFilterInsideJoin; import org.apache.doris.nereids.rules.rewrite.PushProjectIntoUnion; @@ -536,6 +537,7 @@ public class Rewriter extends AbstractBatchJobExecutor { new MergeProjectable() )), custom(RuleType.ELIMINATE_UNNECESSARY_PROJECT, EliminateUnnecessaryProject::new), + topDown(new PushDownVectorTopNIntoOlapScan()), topDown(new PushDownVirtualColumnsIntoOlapScan()), topic("score optimize", topDown(new PushDownScoreTopNIntoOlapScan(), diff --git a/fe/fe-core/src/main/java/org/apache/doris/nereids/parser/LogicalPlanBuilder.java b/fe/fe-core/src/main/java/org/apache/doris/nereids/parser/LogicalPlanBuilder.java index e5eecb1ea4cf3e..ce0404c27c285c 100644 --- a/fe/fe-core/src/main/java/org/apache/doris/nereids/parser/LogicalPlanBuilder.java +++ b/fe/fe-core/src/main/java/org/apache/doris/nereids/parser/LogicalPlanBuilder.java @@ -5628,6 +5628,8 @@ public Command visitCreateIndex(CreateIndexContext ctx) { indexType = "NGRAM_BF"; } else if (ctx.INVERTED() != null) { indexType = "INVERTED"; + } else if (ctx.ANN() != null) { + indexType = "ANN"; } String comment = ctx.STRING_LITERAL() == null ? "" : stripQuotes(ctx.STRING_LITERAL().getText()); // change BITMAP index to INVERTED index diff --git a/fe/fe-core/src/main/java/org/apache/doris/nereids/rules/RuleType.java b/fe/fe-core/src/main/java/org/apache/doris/nereids/rules/RuleType.java index 9755e736313a0a..6fc962cd211d21 100644 --- a/fe/fe-core/src/main/java/org/apache/doris/nereids/rules/RuleType.java +++ b/fe/fe-core/src/main/java/org/apache/doris/nereids/rules/RuleType.java @@ -320,6 +320,7 @@ public enum RuleType { PUSH_CONJUNCTS_INTO_ES_SCAN(RuleTypeClass.REWRITE), PUSH_DOWN_VIRTUAL_COLUMNS_INTO_OLAP_SCAN(RuleTypeClass.REWRITE), PUSH_DOWN_SCORE_TOPN_INTO_OLAP_SCAN(RuleTypeClass.REWRITE), + PUSH_DOWN_VECTOR_TOPN_INTO_OLAP_SCAN(RuleTypeClass.REWRITE), CHECK_SCORE_USAGE(RuleTypeClass.REWRITE), OLAP_SCAN_TABLET_PRUNE(RuleTypeClass.REWRITE), PUSH_AGGREGATE_TO_OLAP_SCAN(RuleTypeClass.REWRITE), diff --git a/fe/fe-core/src/main/java/org/apache/doris/nereids/rules/implementation/LogicalOlapScanToPhysicalOlapScan.java b/fe/fe-core/src/main/java/org/apache/doris/nereids/rules/implementation/LogicalOlapScanToPhysicalOlapScan.java index 96a4d970600ae2..dd7f83ecd0796e 100644 --- a/fe/fe-core/src/main/java/org/apache/doris/nereids/rules/implementation/LogicalOlapScanToPhysicalOlapScan.java +++ b/fe/fe-core/src/main/java/org/apache/doris/nereids/rules/implementation/LogicalOlapScanToPhysicalOlapScan.java @@ -65,7 +65,9 @@ public Rule build() { olapScan.getOperativeSlots(), olapScan.getVirtualColumns(), olapScan.getScoreOrderKeys(), - olapScan.getScoreLimit()) + olapScan.getScoreLimit(), + olapScan.getAnnOrderKeys(), + olapScan.getAnnLimit()) ).toRule(RuleType.LOGICAL_OLAP_SCAN_TO_PHYSICAL_OLAP_SCAN_RULE); } diff --git a/fe/fe-core/src/main/java/org/apache/doris/nereids/rules/rewrite/PushDownVectorTopNIntoOlapScan.java b/fe/fe-core/src/main/java/org/apache/doris/nereids/rules/rewrite/PushDownVectorTopNIntoOlapScan.java new file mode 100644 index 00000000000000..271a043b48a83b --- /dev/null +++ b/fe/fe-core/src/main/java/org/apache/doris/nereids/rules/rewrite/PushDownVectorTopNIntoOlapScan.java @@ -0,0 +1,167 @@ +// 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. + +package org.apache.doris.nereids.rules.rewrite; + +import org.apache.doris.analysis.IndexDef.IndexType; +import org.apache.doris.catalog.Column; +import org.apache.doris.catalog.Index; +import org.apache.doris.catalog.TableIf; +import org.apache.doris.nereids.rules.Rule; +import org.apache.doris.nereids.rules.RuleType; +import org.apache.doris.nereids.trees.expressions.Alias; +import org.apache.doris.nereids.trees.expressions.Cast; +import org.apache.doris.nereids.trees.expressions.Expression; +import org.apache.doris.nereids.trees.expressions.NamedExpression; +import org.apache.doris.nereids.trees.expressions.SlotReference; +import org.apache.doris.nereids.trees.expressions.functions.scalar.InnerProductApproximate; +import org.apache.doris.nereids.trees.expressions.functions.scalar.L2DistanceApproximate; +import org.apache.doris.nereids.trees.plans.Plan; +import org.apache.doris.nereids.trees.plans.logical.LogicalFilter; +import org.apache.doris.nereids.trees.plans.logical.LogicalOlapScan; +import org.apache.doris.nereids.trees.plans.logical.LogicalProject; +import org.apache.doris.nereids.trees.plans.logical.LogicalTopN; +import org.apache.doris.nereids.util.ExpressionUtils; + +import com.google.common.collect.ImmutableList; +import com.google.common.collect.Maps; + +import java.util.List; +import java.util.Map; +import java.util.Optional; +import java.util.Set; + +/** + * extract virtual column from filter and push down them into olap scan. + */ +public class PushDownVectorTopNIntoOlapScan implements RewriteRuleFactory { + @Override + public List buildRules() { + return ImmutableList.of( + logicalTopN(logicalProject(logicalOlapScan())).when(t -> t.getOrderKeys().size() == 1).then(topN -> { + LogicalProject project = topN.child(); + LogicalOlapScan scan = project.child(); + return pushDown(topN, project, scan, Optional.empty()); + }).toRule(RuleType.PUSH_DOWN_VIRTUAL_COLUMNS_INTO_OLAP_SCAN), + logicalTopN(logicalProject(logicalFilter(logicalOlapScan()))) + .when(t -> t.getOrderKeys().size() == 1).then(topN -> { + LogicalProject> project = topN.child(); + LogicalFilter filter = project.child(); + LogicalOlapScan scan = filter.child(); + return pushDown(topN, project, scan, Optional.of(filter)); + }).toRule(RuleType.PUSH_DOWN_VIRTUAL_COLUMNS_INTO_OLAP_SCAN) + ); + } + + private Plan pushDown( + LogicalTopN topN, + LogicalProject project, + LogicalOlapScan scan, + Optional> optionalFilter) { + // Retrives the expression used for ordering in the TopN. + Expression orderKey = topN.getOrderKeys().get(0).getExpr(); + // The order key must be a SlotReference corresponding to an expr. + if (!(orderKey instanceof SlotReference)) { + return null; + } + SlotReference keySlot = (SlotReference) orderKey; + Expression orderKeyExpr = null; + Alias orderKeyAlias = null; + // Find the corresponding expression in the project that matches the keySlot. + for (NamedExpression projection : project.getProjects()) { + if (projection.toSlot().equals(keySlot) && projection instanceof Alias) { + orderKeyExpr = ((Alias) projection).child(); + orderKeyAlias = (Alias) projection; + break; + } + } + if (orderKeyExpr == null) { + return null; + } + + boolean l2Dist; + boolean innerProduct; + l2Dist = orderKeyExpr instanceof L2DistanceApproximate; + innerProduct = orderKeyExpr instanceof InnerProductApproximate; + if (!(l2Dist) && !(innerProduct)) { + return null; + } + + Expression left = null; + if (l2Dist) { + L2DistanceApproximate l2DistanceApproximate = (L2DistanceApproximate) orderKeyExpr; + left = l2DistanceApproximate.left(); + } else { + InnerProductApproximate innerProductApproximate = (InnerProductApproximate) orderKeyExpr; + left = innerProductApproximate.left(); + } + + while (left instanceof Cast) { + left = ((Cast) left).child(); + } + + if (l2Dist) { + if (!(left instanceof SlotReference && ((L2DistanceApproximate) orderKeyExpr).right().isConstant())) { + return null; + } + } else { + if (!(left instanceof SlotReference && ((InnerProductApproximate) orderKeyExpr).right().isConstant())) { + return null; + } + } + + SlotReference leftInput = (SlotReference) left; + if (!leftInput.getOriginalColumn().isPresent() || !leftInput.getOriginalTable().isPresent()) { + return null; + } + TableIf table = leftInput.getOriginalTable().get(); + Column column = leftInput.getOriginalColumn().get(); + boolean hasAnnIndexOnColumn = false; + for (Index index : table.getTableIndexes().getIndexes()) { + if (index.getIndexType() == IndexType.ANN) { + if (index.getColumns().size() != 1) { + continue; + } + if (index.getColumns().get(0).equalsIgnoreCase(column.getName())) { + hasAnnIndexOnColumn = true; + break; + } + } + } + if (!hasAnnIndexOnColumn) { + return null; + } + + Plan plan = scan.withVirtualColumnsAndTopN( + ImmutableList.of(orderKeyAlias), + topN.getOrderKeys(), Optional.of(topN.getLimit() + topN.getOffset()), + ImmutableList.of(), Optional.empty()); + + Map replaceMap = Maps.newHashMap(); + replaceMap.put(orderKeyAlias, orderKeyAlias.toSlot()); + replaceMap.put(orderKeyExpr, orderKeyAlias.toSlot()); + if (optionalFilter.isPresent()) { + LogicalFilter filter = optionalFilter.get(); + Set newConjuncts = ExpressionUtils.replace(filter.getConjuncts(), replaceMap); + plan = filter.withConjunctsAndChild(newConjuncts, plan); + } + List newProjections = ExpressionUtils + .replaceNamedExpressions(project.getProjects(), replaceMap); + LogicalProject newProject = project.withProjectsAndChild(newProjections, plan); + return topN.withChildren(newProject); + } +} diff --git a/fe/fe-core/src/main/java/org/apache/doris/nereids/trees/expressions/functions/scalar/InnerProductApproximate.java b/fe/fe-core/src/main/java/org/apache/doris/nereids/trees/expressions/functions/scalar/InnerProductApproximate.java new file mode 100644 index 00000000000000..be779f633b60a1 --- /dev/null +++ b/fe/fe-core/src/main/java/org/apache/doris/nereids/trees/expressions/functions/scalar/InnerProductApproximate.java @@ -0,0 +1,70 @@ +// 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. + +package org.apache.doris.nereids.trees.expressions.functions.scalar; + +import org.apache.doris.catalog.FunctionSignature; +import org.apache.doris.nereids.trees.expressions.Expression; +import org.apache.doris.nereids.trees.expressions.functions.AlwaysNullable; +import org.apache.doris.nereids.trees.expressions.functions.ExplicitlyCastableSignature; +import org.apache.doris.nereids.trees.expressions.shape.BinaryExpression; +import org.apache.doris.nereids.trees.expressions.visitor.ExpressionVisitor; +import org.apache.doris.nereids.types.ArrayType; +import org.apache.doris.nereids.types.DoubleType; + +import com.google.common.base.Preconditions; +import com.google.common.collect.ImmutableList; + +import java.util.List; + +/** + * inner_product function + */ +public class InnerProductApproximate extends ScalarFunction implements ExplicitlyCastableSignature, + BinaryExpression, AlwaysNullable { + + public static final List SIGNATURES = ImmutableList.of( + FunctionSignature.ret(DoubleType.INSTANCE) + .args(ArrayType.of(DoubleType.INSTANCE), ArrayType.of(DoubleType.INSTANCE)) + ); + + /** + * constructor with 1 argument. + */ + public InnerProductApproximate(Expression arg0, Expression arg1) { + super("inner_product_approximate", arg0, arg1); + } + + /** + * withChildren. + */ + @Override + public InnerProductApproximate withChildren(List children) { + Preconditions.checkArgument(children.size() == 2); + return new InnerProductApproximate(children.get(0), children.get(1)); + } + + @Override + public R accept(ExpressionVisitor visitor, C context) { + return visitor.visitInnerProductApproximate(this, context); + } + + @Override + public List getSignatures() { + return SIGNATURES; + } +} diff --git a/fe/fe-core/src/main/java/org/apache/doris/nereids/trees/expressions/functions/scalar/L2DistanceApproximate.java b/fe/fe-core/src/main/java/org/apache/doris/nereids/trees/expressions/functions/scalar/L2DistanceApproximate.java new file mode 100644 index 00000000000000..7b4058479ab75d --- /dev/null +++ b/fe/fe-core/src/main/java/org/apache/doris/nereids/trees/expressions/functions/scalar/L2DistanceApproximate.java @@ -0,0 +1,70 @@ +// 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. + +package org.apache.doris.nereids.trees.expressions.functions.scalar; + +import org.apache.doris.catalog.FunctionSignature; +import org.apache.doris.nereids.trees.expressions.Expression; +import org.apache.doris.nereids.trees.expressions.functions.AlwaysNullable; +import org.apache.doris.nereids.trees.expressions.functions.ExplicitlyCastableSignature; +import org.apache.doris.nereids.trees.expressions.shape.BinaryExpression; +import org.apache.doris.nereids.trees.expressions.visitor.ExpressionVisitor; +import org.apache.doris.nereids.types.ArrayType; +import org.apache.doris.nereids.types.DoubleType; + +import com.google.common.base.Preconditions; +import com.google.common.collect.ImmutableList; + +import java.util.List; + +/** + * l2_distance_approximate function + */ +public class L2DistanceApproximate extends ScalarFunction implements ExplicitlyCastableSignature, + BinaryExpression, AlwaysNullable { + + public static final List SIGNATURES = ImmutableList.of( + FunctionSignature.ret(DoubleType.INSTANCE) + .args(ArrayType.of(DoubleType.INSTANCE), ArrayType.of(DoubleType.INSTANCE)) + ); + + /** + * constructor with 1 argument. + */ + public L2DistanceApproximate(Expression arg0, Expression arg1) { + super("l2_distance_approximate", arg0, arg1); + } + + /** + * withChildren. + */ + @Override + public L2DistanceApproximate withChildren(List children) { + Preconditions.checkArgument(children.size() == 2); + return new L2DistanceApproximate(children.get(0), children.get(1)); + } + + @Override + public R accept(ExpressionVisitor visitor, C context) { + return visitor.visitL2DistanceApproximate(this, context); + } + + @Override + public List getSignatures() { + return SIGNATURES; + } +} diff --git a/fe/fe-core/src/main/java/org/apache/doris/nereids/trees/expressions/visitor/ScalarFunctionVisitor.java b/fe/fe-core/src/main/java/org/apache/doris/nereids/trees/expressions/visitor/ScalarFunctionVisitor.java index 33ded080fae842..0d2d9d3ec59e2f 100644 --- a/fe/fe-core/src/main/java/org/apache/doris/nereids/trees/expressions/visitor/ScalarFunctionVisitor.java +++ b/fe/fe-core/src/main/java/org/apache/doris/nereids/trees/expressions/visitor/ScalarFunctionVisitor.java @@ -240,6 +240,7 @@ import org.apache.doris.nereids.trees.expressions.functions.scalar.Ignore; import org.apache.doris.nereids.trees.expressions.functions.scalar.Initcap; import org.apache.doris.nereids.trees.expressions.functions.scalar.InnerProduct; +import org.apache.doris.nereids.trees.expressions.functions.scalar.InnerProductApproximate; import org.apache.doris.nereids.trees.expressions.functions.scalar.Instr; import org.apache.doris.nereids.trees.expressions.functions.scalar.InttoUuid; import org.apache.doris.nereids.trees.expressions.functions.scalar.Ipv4CIDRToRange; @@ -290,6 +291,7 @@ import org.apache.doris.nereids.trees.expressions.functions.scalar.JsonbValid; import org.apache.doris.nereids.trees.expressions.functions.scalar.L1Distance; import org.apache.doris.nereids.trees.expressions.functions.scalar.L2Distance; +import org.apache.doris.nereids.trees.expressions.functions.scalar.L2DistanceApproximate; import org.apache.doris.nereids.trees.expressions.functions.scalar.LastDay; import org.apache.doris.nereids.trees.expressions.functions.scalar.LastQueryId; import org.apache.doris.nereids.trees.expressions.functions.scalar.Lcm; @@ -1374,6 +1376,10 @@ default R visitInnerProduct(InnerProduct innerProduct, C context) { return visitScalarFunction(innerProduct, context); } + default R visitInnerProductApproximate(InnerProductApproximate innerProductApproximate, C context) { + return visitScalarFunction(innerProductApproximate, context); + } + default R visitInstr(Instr instr, C context) { return visitScalarFunction(instr, context); } @@ -1574,6 +1580,10 @@ default R visitL2Distance(L2Distance l2Distance, C context) { return visitScalarFunction(l2Distance, context); } + default R visitL2DistanceApproximate(L2DistanceApproximate l2DistanceApproximate, C context) { + return visitScalarFunction(l2DistanceApproximate, context); + } + default R visitLastDay(LastDay lastDay, C context) { return visitScalarFunction(lastDay, context); } diff --git a/fe/fe-core/src/main/java/org/apache/doris/nereids/trees/plans/commands/info/BuildIndexOp.java b/fe/fe-core/src/main/java/org/apache/doris/nereids/trees/plans/commands/info/BuildIndexOp.java index 90bc95d4f63eb8..b3183e39ce8b0e 100644 --- a/fe/fe-core/src/main/java/org/apache/doris/nereids/trees/plans/commands/info/BuildIndexOp.java +++ b/fe/fe-core/src/main/java/org/apache/doris/nereids/trees/plans/commands/info/BuildIndexOp.java @@ -125,6 +125,10 @@ public void validate(ConnectContext ctx) throws UserException { + " is not partitioned, cannot build index with partitions."); } } + if (indexDef.getIndexType() == IndexDef.IndexType.ANN) { + throw new AnalysisException( + "ANN index can only be created during table creation, not through BUILD INDEX."); + } indexDef.validate(); this.index = existedIdx.clone(); } diff --git a/fe/fe-core/src/main/java/org/apache/doris/nereids/trees/plans/commands/info/CreateIndexOp.java b/fe/fe-core/src/main/java/org/apache/doris/nereids/trees/plans/commands/info/CreateIndexOp.java index 1e00aa99fbc582..332859b536dcdc 100644 --- a/fe/fe-core/src/main/java/org/apache/doris/nereids/trees/plans/commands/info/CreateIndexOp.java +++ b/fe/fe-core/src/main/java/org/apache/doris/nereids/trees/plans/commands/info/CreateIndexOp.java @@ -20,6 +20,7 @@ import org.apache.doris.alter.AlterOpType; import org.apache.doris.analysis.AlterTableClause; import org.apache.doris.analysis.CreateIndexClause; +import org.apache.doris.analysis.IndexDef; import org.apache.doris.catalog.Index; import org.apache.doris.common.AnalysisException; import org.apache.doris.common.UserException; @@ -77,6 +78,12 @@ public void validate(ConnectContext ctx) throws UserException { if (tableName != null) { tableName.analyze(ctx); } + + if (indexDef.getIndexType() == IndexDef.IndexType.ANN) { + throw new AnalysisException( + "ANN index can only be created during table creation, not through CREATE INDEX."); + } + indexDef.validate(); index = indexDef.translateToCatalogStyle(); } diff --git a/fe/fe-core/src/main/java/org/apache/doris/nereids/trees/plans/commands/info/IndexDefinition.java b/fe/fe-core/src/main/java/org/apache/doris/nereids/trees/plans/commands/info/IndexDefinition.java index 4899fc56c00963..5aec7d5d0bfdb2 100644 --- a/fe/fe-core/src/main/java/org/apache/doris/nereids/trees/plans/commands/info/IndexDefinition.java +++ b/fe/fe-core/src/main/java/org/apache/doris/nereids/trees/plans/commands/info/IndexDefinition.java @@ -17,6 +17,7 @@ package org.apache.doris.nereids.trees.plans.commands.info; +import org.apache.doris.analysis.AnnIndexPropertiesChecker; import org.apache.doris.analysis.IndexDef; import org.apache.doris.analysis.IndexDef.IndexType; import org.apache.doris.analysis.InvertedIndexUtil; @@ -80,6 +81,10 @@ public IndexDefinition(String name, boolean ifNotExists, List cols, Stri this.indexType = IndexType.NGRAM_BF; break; } + case "ANN": { + this.indexType = IndexType.ANN; + break; + } default: throw new AnalysisException("unknown index type " + indexTypeName); } @@ -132,6 +137,26 @@ public static boolean isSupportIdxType(DataType columnType) { public void checkColumn(ColumnDefinition column, KeysType keysType, boolean enableUniqueKeyMergeOnWrite, TInvertedIndexFileStorageFormat invertedIndexFileStorageFormat) throws AnalysisException { + if (indexType == IndexType.ANN) { + if (column.isNullable()) { + throw new AnalysisException("ANN index must be built on a column that is not nullable"); + } + String indexColName = column.getName(); + caseSensitivityCols.add(indexColName); + DataType colType = column.getType(); + if (!colType.isArrayType()) { + throw new AnalysisException("ANN index column must be array type, invalid index: " + name); + } + DataType itemType = ((ArrayType) colType).getItemType(); + if (!itemType.isFloatType()) { + throw new AnalysisException("ANN index column item type must be float type, invalid index: " + name); + } + if (keysType != KeysType.DUP_KEYS) { + throw new AnalysisException("ANN index can only be used in DUP_KEYS table"); + } + return; + } + if (indexType == IndexType.BITMAP || indexType == IndexType.INVERTED || indexType == IndexType.BLOOMFILTER || indexType == IndexType.NGRAM_BF) { String indexColName = column.getName(); @@ -219,6 +244,10 @@ public void validate() { return; } + if (indexType == IndexDef.IndexType.ANN) { + AnnIndexPropertiesChecker.checkProperties(this.properties); + } + if (indexType == IndexDef.IndexType.BITMAP || indexType == IndexDef.IndexType.INVERTED) { if (cols == null || cols.size() != 1) { throw new AnalysisException( diff --git a/fe/fe-core/src/main/java/org/apache/doris/nereids/trees/plans/logical/LogicalOlapScan.java b/fe/fe-core/src/main/java/org/apache/doris/nereids/trees/plans/logical/LogicalOlapScan.java index bd9d0cb1e6e326..b6ec1b53a7e969 100644 --- a/fe/fe-core/src/main/java/org/apache/doris/nereids/trees/plans/logical/LogicalOlapScan.java +++ b/fe/fe-core/src/main/java/org/apache/doris/nereids/trees/plans/logical/LogicalOlapScan.java @@ -145,6 +145,9 @@ public class LogicalOlapScan extends LogicalCatalogRelation implements OlapScan private final List scoreOrderKeys; private final Optional scoreLimit; + // use for ann push down + private final List annOrderKeys; + private final Optional annLimit; public LogicalOlapScan(RelationId id, OlapTable table) { this(id, table, ImmutableList.of()); @@ -160,7 +163,7 @@ public LogicalOlapScan(RelationId id, OlapTable table, List qualifier) { -1, false, PreAggStatus.unset(), ImmutableList.of(), ImmutableList.of(), Maps.newHashMap(), Optional.empty(), false, ImmutableMap.of(), ImmutableList.of(), ImmutableList.of(), ImmutableList.of(), - ImmutableList.of(), Optional.empty()); + ImmutableList.of(), Optional.empty(), ImmutableList.of(), Optional.empty()); } public LogicalOlapScan(RelationId id, OlapTable table, List qualifier, List tabletIds, @@ -169,7 +172,7 @@ public LogicalOlapScan(RelationId id, OlapTable table, List qualifier, L table.getPartitionIds(), false, tabletIds, -1, false, PreAggStatus.unset(), ImmutableList.of(), hints, Maps.newHashMap(), tableSample, false, ImmutableMap.of(), ImmutableList.of(), operativeSlots, - ImmutableList.of(), ImmutableList.of(), Optional.empty()); + ImmutableList.of(), ImmutableList.of(), Optional.empty(), ImmutableList.of(), Optional.empty()); } /** @@ -182,9 +185,13 @@ public LogicalOlapScan(RelationId id, OlapTable table, List qualifier, L specifiedPartitions, false, tabletIds, -1, false, PreAggStatus.unset(), specifiedPartitions, hints, Maps.newHashMap(), tableSample, false, ImmutableMap.of(), ImmutableList.of(), operativeSlots, - ImmutableList.of(), ImmutableList.of(), Optional.empty()); + ImmutableList.of(), ImmutableList.of(), Optional.empty(), + ImmutableList.of(), Optional.empty()); } + /** + * constructor. + */ public LogicalOlapScan(RelationId id, OlapTable table, List qualifier, List tabletIds, List selectedPartitionIds, long selectedIndexId, PreAggStatus preAggStatus, List specifiedPartitions, List hints, Optional tableSample, @@ -193,7 +200,8 @@ public LogicalOlapScan(RelationId id, OlapTable table, List qualifier, L selectedPartitionIds, false, tabletIds, selectedIndexId, true, preAggStatus, specifiedPartitions, hints, Maps.newHashMap(), tableSample, true, ImmutableMap.of(), - ImmutableList.of(), operativeSlots, ImmutableList.of(), ImmutableList.of(), Optional.empty()); + ImmutableList.of(), operativeSlots, ImmutableList.of(), ImmutableList.of(), Optional.empty(), + ImmutableList.of(), Optional.empty()); } /** @@ -208,7 +216,8 @@ public LogicalOlapScan(RelationId id, Table table, List qualifier, Optional tableSample, boolean directMvScan, Map>> colToSubPathsMap, List specifiedTabletIds, Collection operativeSlots, List virtualColumns, - List scoreOrderKeys, Optional scoreLimit) { + List scoreOrderKeys, Optional scoreLimit, + List annOrderKeys, Optional annLimit) { super(id, PlanType.LOGICAL_OLAP_SCAN, table, qualifier, operativeSlots, virtualColumns, groupExpression, logicalProperties); Preconditions.checkArgument(selectedPartitionIds != null, @@ -242,6 +251,8 @@ public LogicalOlapScan(RelationId id, Table table, List qualifier, this.subPathToSlotMap = Maps.newHashMap(); this.scoreOrderKeys = Utils.fastToImmutableList(scoreOrderKeys); this.scoreLimit = scoreLimit; + this.annOrderKeys = Utils.fastToImmutableList(annOrderKeys); + this.annLimit = annLimit; } public List getSelectedPartitionIds() { @@ -310,7 +321,9 @@ public boolean equals(Object o) { && Objects.equals(hints, that.hints) && Objects.equals(tableSample, that.tableSample) && Objects.equals(scoreOrderKeys, that.scoreOrderKeys) - && Objects.equals(scoreLimit, that.scoreLimit); + && Objects.equals(scoreLimit, that.scoreLimit) + && Objects.equals(annOrderKeys, that.annOrderKeys) + && Objects.equals(annLimit, that.annLimit); } @Override @@ -325,7 +338,7 @@ public LogicalOlapScan withGroupExpression(Optional groupExpres selectedPartitionIds, partitionPruned, selectedTabletIds, selectedIndexId, indexSelected, preAggStatus, manuallySpecifiedPartitions, hints, cacheSlotWithSlotName, tableSample, directMvScan, colToSubPathsMap, manuallySpecifiedTabletIds, - operativeSlots, virtualColumns, scoreOrderKeys, scoreLimit); + operativeSlots, virtualColumns, scoreOrderKeys, scoreLimit, annOrderKeys, annLimit); } @Override @@ -335,7 +348,7 @@ public Plan withGroupExprLogicalPropChildren(Optional groupExpr selectedPartitionIds, partitionPruned, selectedTabletIds, selectedIndexId, indexSelected, preAggStatus, manuallySpecifiedPartitions, hints, cacheSlotWithSlotName, tableSample, directMvScan, colToSubPathsMap, manuallySpecifiedTabletIds, - operativeSlots, virtualColumns, scoreOrderKeys, scoreLimit); + operativeSlots, virtualColumns, scoreOrderKeys, scoreLimit, annOrderKeys, annLimit); } /** @@ -347,7 +360,7 @@ public LogicalOlapScan withSelectedPartitionIds(List selectedPartitionIds) selectedPartitionIds, true, selectedTabletIds, selectedIndexId, indexSelected, preAggStatus, manuallySpecifiedPartitions, hints, cacheSlotWithSlotName, tableSample, directMvScan, colToSubPathsMap, manuallySpecifiedTabletIds, - operativeSlots, virtualColumns, scoreOrderKeys, scoreLimit); + operativeSlots, virtualColumns, scoreOrderKeys, scoreLimit, annOrderKeys, annLimit); } /** @@ -361,7 +374,7 @@ public LogicalOlapScan withMaterializedIndexSelected(long indexId) { selectedPartitionIds, partitionPruned, selectedTabletIds, indexId, true, PreAggStatus.unset(), manuallySpecifiedPartitions, hints, cacheSlotWithSlotName, tableSample, directMvScan, colToSubPathsMap, manuallySpecifiedTabletIds, - operativeSlots, virtualColumns, scoreOrderKeys, scoreLimit); + operativeSlots, virtualColumns, scoreOrderKeys, scoreLimit, annOrderKeys, annLimit); } /** @@ -373,7 +386,7 @@ public LogicalOlapScan withSelectedTabletIds(List selectedTabletIds) { selectedPartitionIds, partitionPruned, selectedTabletIds, selectedIndexId, indexSelected, preAggStatus, manuallySpecifiedPartitions, hints, cacheSlotWithSlotName, tableSample, directMvScan, colToSubPathsMap, manuallySpecifiedTabletIds, - operativeSlots, virtualColumns, scoreOrderKeys, scoreLimit); + operativeSlots, virtualColumns, scoreOrderKeys, scoreLimit, annOrderKeys, annLimit); } /** @@ -385,7 +398,7 @@ public LogicalOlapScan withPreAggStatus(PreAggStatus preAggStatus) { selectedPartitionIds, partitionPruned, selectedTabletIds, selectedIndexId, indexSelected, preAggStatus, manuallySpecifiedPartitions, hints, cacheSlotWithSlotName, tableSample, directMvScan, colToSubPathsMap, manuallySpecifiedTabletIds, - operativeSlots, virtualColumns, scoreOrderKeys, scoreLimit); + operativeSlots, virtualColumns, scoreOrderKeys, scoreLimit, annOrderKeys, annLimit); } /** @@ -397,7 +410,7 @@ public LogicalOlapScan withColToSubPathsMap(Map>> colTo selectedPartitionIds, partitionPruned, selectedTabletIds, selectedIndexId, indexSelected, preAggStatus, manuallySpecifiedPartitions, hints, cacheSlotWithSlotName, tableSample, directMvScan, colToSubPathsMap, manuallySpecifiedTabletIds, - operativeSlots, virtualColumns, scoreOrderKeys, scoreLimit); + operativeSlots, virtualColumns, scoreOrderKeys, scoreLimit, annOrderKeys, annLimit); } /** @@ -409,7 +422,7 @@ public LogicalOlapScan withManuallySpecifiedTabletIds(List manuallySpecifi selectedPartitionIds, partitionPruned, selectedTabletIds, selectedIndexId, indexSelected, preAggStatus, manuallySpecifiedPartitions, hints, cacheSlotWithSlotName, tableSample, directMvScan, colToSubPathsMap, manuallySpecifiedTabletIds, - operativeSlots, virtualColumns, scoreOrderKeys, scoreLimit); + operativeSlots, virtualColumns, scoreOrderKeys, scoreLimit, annOrderKeys, annLimit); } @Override @@ -420,7 +433,7 @@ public LogicalOlapScan withRelationId(RelationId relationId) { selectedPartitionIds, false, selectedTabletIds, selectedIndexId, indexSelected, preAggStatus, manuallySpecifiedPartitions, hints, Maps.newHashMap(), tableSample, directMvScan, colToSubPathsMap, selectedTabletIds, - operativeSlots, virtualColumns, scoreOrderKeys, scoreLimit); + operativeSlots, virtualColumns, scoreOrderKeys, scoreLimit, annOrderKeys, annLimit); } /** @@ -439,7 +452,8 @@ public LogicalOlapScan withVirtualColumns(List virtualColumns) selectedPartitionIds, partitionPruned, selectedTabletIds, selectedIndexId, indexSelected, preAggStatus, manuallySpecifiedPartitions, hints, cacheSlotWithSlotName, tableSample, directMvScan, colToSubPathsMap, - manuallySpecifiedTabletIds, operativeSlots, virtualColumns, scoreOrderKeys, scoreLimit); + manuallySpecifiedTabletIds, operativeSlots, virtualColumns, scoreOrderKeys, scoreLimit, + annOrderKeys, annLimit); } /** @@ -462,7 +476,8 @@ public LogicalOlapScan withVirtualColumnsAndTopN( selectedPartitionIds, partitionPruned, selectedTabletIds, selectedIndexId, indexSelected, preAggStatus, manuallySpecifiedPartitions, hints, cacheSlotWithSlotName, tableSample, directMvScan, colToSubPathsMap, - manuallySpecifiedTabletIds, operativeSlots, virtualColumns, scoreOrderKeys, scoreLimit); + manuallySpecifiedTabletIds, operativeSlots, virtualColumns, scoreOrderKeys, scoreLimit, + annOrderKeys, annLimit); } @Override @@ -623,6 +638,14 @@ public Optional getScoreLimit() { return scoreLimit; } + public List getAnnOrderKeys() { + return annOrderKeys; + } + + public Optional getAnnLimit() { + return annLimit; + } + private List createSlotsVectorized(List columns) { List qualified = qualified(); SlotReference[] slots = new SlotReference[columns.size()]; @@ -784,7 +807,8 @@ public CatalogRelation withOperativeSlots(Collection operativeSlots) { selectedPartitionIds, partitionPruned, selectedTabletIds, selectedIndexId, indexSelected, preAggStatus, manuallySpecifiedPartitions, hints, cacheSlotWithSlotName, tableSample, directMvScan, colToSubPathsMap, - manuallySpecifiedTabletIds, operativeSlots, virtualColumns, scoreOrderKeys, scoreLimit); + manuallySpecifiedTabletIds, operativeSlots, virtualColumns, scoreOrderKeys, scoreLimit, + annOrderKeys, annLimit); } private Map constructReplaceMap(MTMV mtmv) { diff --git a/fe/fe-core/src/main/java/org/apache/doris/nereids/trees/plans/physical/PhysicalLazyMaterializeOlapScan.java b/fe/fe-core/src/main/java/org/apache/doris/nereids/trees/plans/physical/PhysicalLazyMaterializeOlapScan.java index 7ca6c8472abac2..aa1d67ce65b677 100644 --- a/fe/fe-core/src/main/java/org/apache/doris/nereids/trees/plans/physical/PhysicalLazyMaterializeOlapScan.java +++ b/fe/fe-core/src/main/java/org/apache/doris/nereids/trees/plans/physical/PhysicalLazyMaterializeOlapScan.java @@ -58,7 +58,10 @@ public PhysicalLazyMaterializeOlapScan(PhysicalOlapScan physicalOlapScan, physicalOlapScan.getOperativeSlots(), physicalOlapScan.getVirtualColumns(), physicalOlapScan.getScoreOrderKeys(), - physicalOlapScan.getScoreLimit()); + physicalOlapScan.getScoreLimit(), + physicalOlapScan.getAnnOrderKeys(), + physicalOlapScan.getAnnLimit() + ); this.scan = physicalOlapScan; this.rowId = rowId; this.lazySlots = ImmutableList.copyOf(lazySlots); diff --git a/fe/fe-core/src/main/java/org/apache/doris/nereids/trees/plans/physical/PhysicalOlapScan.java b/fe/fe-core/src/main/java/org/apache/doris/nereids/trees/plans/physical/PhysicalOlapScan.java index dbea157825f51e..f9c8ed54f2f733 100644 --- a/fe/fe-core/src/main/java/org/apache/doris/nereids/trees/plans/physical/PhysicalOlapScan.java +++ b/fe/fe-core/src/main/java/org/apache/doris/nereids/trees/plans/physical/PhysicalOlapScan.java @@ -70,6 +70,9 @@ public class PhysicalOlapScan extends PhysicalCatalogRelation implements OlapSca private final List scoreOrderKeys; private final Optional scoreLimit; + // use for ann push down + private final List annOrderKeys; + private final Optional annLimit; /** * Constructor for PhysicalOlapScan. @@ -79,12 +82,14 @@ public PhysicalOlapScan(RelationId id, OlapTable olapTable, List qualifi PreAggStatus preAggStatus, List baseOutputs, Optional groupExpression, LogicalProperties logicalProperties, Optional tableSample, List operativeSlots, List virtualColumns, - List scoreOrderKeys, Optional scoreLimit) { + List scoreOrderKeys, Optional scoreLimit, + List annOrderKeys, Optional annLimit) { this(id, olapTable, qualifier, selectedIndexId, selectedTabletIds, selectedPartitionIds, distributionSpec, preAggStatus, baseOutputs, groupExpression, logicalProperties, null, - null, tableSample, operativeSlots, virtualColumns, scoreOrderKeys, scoreLimit); + null, tableSample, operativeSlots, virtualColumns, scoreOrderKeys, scoreLimit, + annOrderKeys, annLimit); } /** @@ -97,7 +102,8 @@ public PhysicalOlapScan(RelationId id, OlapTable olapTable, List qualifi PhysicalProperties physicalProperties, Statistics statistics, Optional tableSample, Collection operativeSlots, List virtualColumns, - List scoreOrderKeys, Optional scoreLimit) { + List scoreOrderKeys, Optional scoreLimit, + List annOrderKeys, Optional annLimit) { super(id, PlanType.PHYSICAL_OLAP_SCAN, olapTable, qualifier, groupExpression, logicalProperties, physicalProperties, statistics, operativeSlots); this.selectedIndexId = selectedIndexId; @@ -111,6 +117,8 @@ public PhysicalOlapScan(RelationId id, OlapTable olapTable, List qualifi this.virtualColumns = ImmutableList.copyOf(virtualColumns); this.scoreOrderKeys = ImmutableList.copyOf(scoreOrderKeys); this.scoreLimit = scoreLimit; + this.annOrderKeys = ImmutableList.copyOf(annOrderKeys); + this.annLimit = annLimit; } @Override @@ -166,6 +174,14 @@ public Optional getScoreLimit() { return scoreLimit; } + public List getAnnOrderKeys() { + return annOrderKeys; + } + + public Optional getAnnLimit() { + return annLimit; + } + @Override public String getFingerprint() { String partitions = ""; @@ -241,7 +257,9 @@ public boolean equals(Object o) { && Objects.equals(operativeSlots, olapScan.operativeSlots) && Objects.equals(virtualColumns, olapScan.virtualColumns) && Objects.equals(scoreOrderKeys, olapScan.scoreOrderKeys) - && Objects.equals(scoreLimit, olapScan.scoreLimit); + && Objects.equals(scoreLimit, olapScan.scoreLimit) + && Objects.equals(annOrderKeys, olapScan.annOrderKeys) + && Objects.equals(annLimit, olapScan.annLimit); } @Override @@ -259,7 +277,7 @@ public PhysicalOlapScan withGroupExpression(Optional groupExpre return new PhysicalOlapScan(relationId, getTable(), qualifier, selectedIndexId, selectedTabletIds, selectedPartitionIds, distributionSpec, preAggStatus, baseOutputs, groupExpression, getLogicalProperties(), tableSample, operativeSlots, virtualColumns, - scoreOrderKeys, scoreLimit); + scoreOrderKeys, scoreLimit, annOrderKeys, annLimit); } @Override @@ -268,7 +286,7 @@ public Plan withGroupExprLogicalPropChildren(Optional groupExpr return new PhysicalOlapScan(relationId, getTable(), qualifier, selectedIndexId, selectedTabletIds, selectedPartitionIds, distributionSpec, preAggStatus, baseOutputs, groupExpression, logicalProperties.get(), tableSample, operativeSlots, virtualColumns, - scoreOrderKeys, scoreLimit); + scoreOrderKeys, scoreLimit, annOrderKeys, annLimit); } @Override @@ -277,7 +295,7 @@ public PhysicalOlapScan withPhysicalPropertiesAndStats( return new PhysicalOlapScan(relationId, getTable(), qualifier, selectedIndexId, selectedTabletIds, selectedPartitionIds, distributionSpec, preAggStatus, baseOutputs, groupExpression, getLogicalProperties(), physicalProperties, statistics, tableSample, operativeSlots, - virtualColumns, scoreOrderKeys, scoreLimit); + virtualColumns, scoreOrderKeys, scoreLimit, annOrderKeys, annLimit); } @Override @@ -303,7 +321,8 @@ public CatalogRelation withOperativeSlots(Collection operativeSlots) { return new PhysicalOlapScan(relationId, (OlapTable) table, qualifier, selectedIndexId, selectedTabletIds, selectedPartitionIds, distributionSpec, preAggStatus, baseOutputs, groupExpression, getLogicalProperties(), getPhysicalProperties(), statistics, - tableSample, operativeSlots, virtualColumns, scoreOrderKeys, scoreLimit); + tableSample, operativeSlots, virtualColumns, scoreOrderKeys, scoreLimit, + annOrderKeys, annLimit); } @Override diff --git a/fe/fe-core/src/main/java/org/apache/doris/planner/OlapScanNode.java b/fe/fe-core/src/main/java/org/apache/doris/planner/OlapScanNode.java index 5d9ba14682a54f..670af2e28a3ef3 100644 --- a/fe/fe-core/src/main/java/org/apache/doris/planner/OlapScanNode.java +++ b/fe/fe-core/src/main/java/org/apache/doris/planner/OlapScanNode.java @@ -196,6 +196,8 @@ public class OlapScanNode extends ScanNode { protected List rewrittenProjectList; private long maxVersion = -1L; + private SortInfo annSortInfo = null; + private long annSortLimit = -1; private SortInfo scoreSortInfo = null; private long scoreSortLimit = -1; @@ -258,6 +260,14 @@ public void setNereidsPrunedTabletIds(Set nereidsPrunedTabletIds) { this.nereidsPrunedTabletIds = nereidsPrunedTabletIds; } + public long getTotalTabletsNum() { + return totalTabletsNum; + } + + public boolean getForceOpenPreAgg() { + return forceOpenPreAgg; + } + public ArrayList getScanTabletIds() { return scanTabletIds; } @@ -282,6 +292,14 @@ public void setScoreSortLimit(long scoreSortLimit) { this.scoreSortLimit = scoreSortLimit; } + public void setAnnSortInfo(SortInfo annSortInfo) { + this.annSortInfo = annSortInfo; + } + + public void setAnnSortLimit(long annSortLimit) { + this.annSortLimit = annSortLimit; + } + public Collection getSelectedPartitionIds() { return selectedPartitionIds; } @@ -1019,6 +1037,17 @@ public String getNodeExplainString(String prefix, TExplainLevel detailLevel) { if (scoreSortLimit != -1) { output.append(prefix).append("SCORE SORT LIMIT: ").append(scoreSortLimit).append("\n"); } + + if (annSortInfo != null) { + output.append(prefix).append("ANN SORT INFO:\n"); + annSortInfo.getOrderingExprs().forEach(expr -> { + output.append(prefix).append(prefix).append(expr.toSql()).append("\n"); + }); + } + if (annSortLimit != -1) { + output.append(prefix).append("ANN SORT LIMIT: ").append(annSortLimit).append("\n"); + } + if (useTopnFilter()) { String topnFilterSources = String.join(",", topnFilterSortNodes.stream() @@ -1181,6 +1210,19 @@ protected void toThrift(TPlanNode msg) { if (scoreSortLimit != -1) { msg.olap_scan_node.setScoreSortLimit(scoreSortLimit); } + if (annSortInfo != null) { + TSortInfo tAnnSortInfo = new TSortInfo( + Expr.treesToThrift(annSortInfo.getOrderingExprs()), + annSortInfo.getIsAscOrder(), + annSortInfo.getNullsFirst()); + if (annSortInfo.getSortTupleSlotExprs() != null) { + tAnnSortInfo.setSortTupleSlotExprs(Expr.treesToThrift(annSortInfo.getSortTupleSlotExprs())); + } + msg.olap_scan_node.setAnnSortInfo(tAnnSortInfo); + } + if (annSortLimit != -1) { + msg.olap_scan_node.setAnnSortLimit(annSortLimit); + } msg.olap_scan_node.setKeyType(olapTable.getKeysType().toThrift()); String tableName = olapTable.getName(); if (selectedIndexId != -1) { diff --git a/fe/fe-core/src/main/java/org/apache/doris/qe/SessionVariable.java b/fe/fe-core/src/main/java/org/apache/doris/qe/SessionVariable.java index 9e402ef1d5ed54..9ed8a45b7215f8 100644 --- a/fe/fe-core/src/main/java/org/apache/doris/qe/SessionVariable.java +++ b/fe/fe-core/src/main/java/org/apache/doris/qe/SessionVariable.java @@ -786,6 +786,9 @@ public static double getHotValueThreshold() { public static final String ENABLE_STRICT_CAST = "enable_strict_cast"; public static final String DEFAULT_LLM_RESOURCE = "default_llm_resource"; + public static final String HNSW_EF_SEARCH = "hnsw_ef_search"; + public static final String HNSW_CHECK_RELATIVE_DISTANCE = "hnsw_check_relative_distance"; + public static final String HNSW_BOUNDED_QUEUE = "hnsw_bounded_queue"; public static final String DEFAULT_VARIANT_MAX_SUBCOLUMNS_COUNT = "default_variant_max_subcolumns_count"; @@ -2790,6 +2793,21 @@ public boolean isEnableESParallelScroll() { }) public boolean enableAddIndexForNewData = false; + @VariableMgr.VarAttr(name = HNSW_EF_SEARCH, needForward = true, + description = {"HNSW索引的EF搜索参数，控制搜索的精度和速度", + "HNSW index EF search parameter, controls the precision and speed of the search"}) + public int hnswEFSearch = 32; + + @VariableMgr.VarAttr(name = HNSW_CHECK_RELATIVE_DISTANCE, needForward = true, + description = {"是否启用相对距离检查机制，以提升HNSW搜索的准确性", + "Enable relative distance checking to improve HNSW search accuracy"}) + public boolean hnswCheckRelativeDistance = true; + + @VariableMgr.VarAttr(name = HNSW_BOUNDED_QUEUE, needForward = true, + description = {"是否使用有界优先队列来优化HNSW的搜索性能", + "Whether to use a bounded priority queue to optimize HNSW search performance"}) + public boolean hnswBoundedQueue = true; + @VariableMgr.VarAttr( name = DEFAULT_VARIANT_MAX_SUBCOLUMNS_COUNT, needForward = true, @@ -4433,6 +4451,11 @@ public TQueryOptions toThrift() { tResult.setExchangeMultiBlocksByteSize(exchangeMultiBlocksByteSize); tResult.setEnableStrictCast(enableStrictCast); tResult.setNewVersionUnixTimestamp(true); // once FE upgraded, always use new version + + tResult.setHnswEfSearch(hnswEFSearch); + tResult.setHnswCheckRelativeDistance(hnswCheckRelativeDistance); + tResult.setHnswBoundedQueue(hnswBoundedQueue); + return tResult; } diff --git a/fe/fe-core/src/test/java/org/apache/doris/nereids/glue/translator/PhysicalPlanTranslatorTest.java b/fe/fe-core/src/test/java/org/apache/doris/nereids/glue/translator/PhysicalPlanTranslatorTest.java index 4143b1bb2284ba..a31b1abdaa1ca0 100644 --- a/fe/fe-core/src/test/java/org/apache/doris/nereids/glue/translator/PhysicalPlanTranslatorTest.java +++ b/fe/fe-core/src/test/java/org/apache/doris/nereids/glue/translator/PhysicalPlanTranslatorTest.java @@ -68,7 +68,8 @@ public void testOlapPrune(@Injectable LogicalProperties placeHolder) throws Exce PhysicalOlapScan scan = new PhysicalOlapScan(StatementScopeIdGenerator.newRelationId(), t1, qualifier, t1.getBaseIndexId(), Collections.emptyList(), Collections.emptyList(), null, PreAggStatus.on(), ImmutableList.of(), Optional.empty(), t1Properties, Optional.empty(), - ImmutableList.of(), ImmutableList.of(), ImmutableList.of(), Optional.empty()); + ImmutableList.of(), ImmutableList.of(), ImmutableList.of(), Optional.empty(), + ImmutableList.of(), Optional.empty()); Literal t1FilterRight = new IntegerLiteral(1); Expression t1FilterExpr = new GreaterThan(col1, t1FilterRight); PhysicalFilter filter = diff --git a/fe/fe-core/src/test/java/org/apache/doris/nereids/postprocess/MergeProjectPostProcessTest.java b/fe/fe-core/src/test/java/org/apache/doris/nereids/postprocess/MergeProjectPostProcessTest.java index 2401ed14e02b01..3177cd97316ef0 100644 --- a/fe/fe-core/src/test/java/org/apache/doris/nereids/postprocess/MergeProjectPostProcessTest.java +++ b/fe/fe-core/src/test/java/org/apache/doris/nereids/postprocess/MergeProjectPostProcessTest.java @@ -79,7 +79,8 @@ public void testMergeProj(@Injectable LogicalProperties placeHolder, @Injectable PhysicalOlapScan scan = new PhysicalOlapScan(RelationId.createGenerator().getNextId(), t1, qualifier, 0L, Collections.emptyList(), Collections.emptyList(), null, PreAggStatus.on(), ImmutableList.of(), Optional.empty(), t1Properties, Optional.empty(), ImmutableList.of(), - ImmutableList.of(), ImmutableList.of(), Optional.empty()); + ImmutableList.of(), ImmutableList.of(), Optional.empty(), + ImmutableList.of(), Optional.empty()); Alias x = new Alias(a, "x"); List projList3 = Lists.newArrayList(x, b, c); PhysicalProject proj3 = new PhysicalProject(projList3, placeHolder, scan); diff --git a/fe/fe-core/src/test/java/org/apache/doris/nereids/postprocess/PushDownFilterThroughProjectTest.java b/fe/fe-core/src/test/java/org/apache/doris/nereids/postprocess/PushDownFilterThroughProjectTest.java index e7ae66335df676..76e9bb0722f3ad 100644 --- a/fe/fe-core/src/test/java/org/apache/doris/nereids/postprocess/PushDownFilterThroughProjectTest.java +++ b/fe/fe-core/src/test/java/org/apache/doris/nereids/postprocess/PushDownFilterThroughProjectTest.java @@ -93,7 +93,8 @@ public void testPushFilter(@Injectable LogicalProperties placeHolder, PhysicalOlapScan scan = new PhysicalOlapScan(RelationId.createGenerator().getNextId(), t1, qualifier, 0L, Collections.emptyList(), Collections.emptyList(), null, PreAggStatus.on(), ImmutableList.of(), Optional.empty(), t1Properties, - Optional.empty(), ImmutableList.of(), ImmutableList.of(), ImmutableList.of(), Optional.empty()); + Optional.empty(), ImmutableList.of(), ImmutableList.of(), ImmutableList.of(), Optional.empty(), + ImmutableList.of(), Optional.empty()); Alias x = new Alias(a, "x"); List projList3 = Lists.newArrayList(x, b, c); PhysicalProject proj3 = new PhysicalProject(projList3, placeHolder, scan); @@ -132,8 +133,8 @@ public void testNotPushFilterWithNonfoldable(@Injectable LogicalProperties place PhysicalOlapScan scan = new PhysicalOlapScan(RelationId.createGenerator().getNextId(), t1, qualifier, 0L, Collections.emptyList(), Collections.emptyList(), null, PreAggStatus.on(), ImmutableList.of(), Optional.empty(), t1Properties, - Optional.empty(), new ArrayList<>(), ImmutableList.of(), ImmutableList.of(), Optional.empty()); - + Optional.empty(), new ArrayList<>(), ImmutableList.of(), ImmutableList.of(), Optional.empty(), + ImmutableList.of(), Optional.empty()); Alias x = new Alias(a, "x"); List projList3 = Lists.newArrayList(x, b, c); PhysicalProject proj3 = new PhysicalProject(projList3, placeHolder, scan); diff --git a/fe/fe-core/src/test/java/org/apache/doris/nereids/trees/plans/PlanEqualsTest.java b/fe/fe-core/src/test/java/org/apache/doris/nereids/trees/plans/PlanEqualsTest.java index e309fc5ded525a..5e3152219c9a48 100644 --- a/fe/fe-core/src/test/java/org/apache/doris/nereids/trees/plans/PlanEqualsTest.java +++ b/fe/fe-core/src/test/java/org/apache/doris/nereids/trees/plans/PlanEqualsTest.java @@ -343,20 +343,23 @@ void testPhysicalOlapScan( 1L, selectedTabletId, olapTable.getPartitionIds(), distributionSpecHash, PreAggStatus.on(), ImmutableList.of(), Optional.empty(), logicalProperties, Optional.empty(), - ImmutableList.of(), ImmutableList.of(), ImmutableList.of(), Optional.empty()); + ImmutableList.of(), ImmutableList.of(), ImmutableList.of(), Optional.empty(), + ImmutableList.of(), Optional.empty()); PhysicalOlapScan expected = new PhysicalOlapScan(id, olapTable, Lists.newArrayList("a"), 1L, selectedTabletId, olapTable.getPartitionIds(), distributionSpecHash, PreAggStatus.on(), ImmutableList.of(), Optional.empty(), logicalProperties, Optional.empty(), - ImmutableList.of(), ImmutableList.of(), ImmutableList.of(), Optional.empty()); + ImmutableList.of(), ImmutableList.of(), ImmutableList.of(), Optional.empty(), + ImmutableList.of(), Optional.empty()); Assertions.assertEquals(expected, actual); PhysicalOlapScan unexpected = new PhysicalOlapScan(id, olapTable, Lists.newArrayList("b"), 12345L, selectedTabletId, olapTable.getPartitionIds(), distributionSpecHash, PreAggStatus.on(), ImmutableList.of(), Optional.empty(), logicalProperties, Optional.empty(), - ImmutableList.of(), ImmutableList.of(), ImmutableList.of(), Optional.empty()); + ImmutableList.of(), ImmutableList.of(), ImmutableList.of(), Optional.empty(), + ImmutableList.of(), Optional.empty()); Assertions.assertNotEquals(unexpected, actual); } diff --git a/gensrc/proto/olap_file.proto b/gensrc/proto/olap_file.proto index 165ab898a22bb9..6e7dea461504b5 100644 --- a/gensrc/proto/olap_file.proto +++ b/gensrc/proto/olap_file.proto @@ -371,6 +371,7 @@ enum IndexType { INVERTED = 1; BLOOMFILTER = 2; NGRAM_BF = 3; + ANN = 4; } enum InvertedIndexStorageFormatPB { diff --git a/gensrc/thrift/Descriptors.thrift b/gensrc/thrift/Descriptors.thrift index 84a2a70a666e14..f7577749152648 100644 --- a/gensrc/thrift/Descriptors.thrift +++ b/gensrc/thrift/Descriptors.thrift @@ -170,7 +170,8 @@ enum TIndexType { BITMAP = 0, INVERTED = 1, BLOOMFILTER = 2, - NGRAM_BF = 3 + NGRAM_BF = 3, + ANN = 4 } enum TPartialUpdateNewRowPolicy { diff --git a/gensrc/thrift/PaloInternalService.thrift b/gensrc/thrift/PaloInternalService.thrift index 438b54fc306021..54cf0dc75e2ea2 100644 --- a/gensrc/thrift/PaloInternalService.thrift +++ b/gensrc/thrift/PaloInternalService.thrift @@ -398,6 +398,10 @@ struct TQueryOptions { 166: optional bool enable_strict_cast = false 167: optional bool new_version_unix_timestamp = false + 168: optional i32 hnsw_ef_search = 32; + 169: optional bool hnsw_check_relative_distance = true; + 170: optional bool hnsw_bounded_queue = true; + // For cloud, to control if the content would be written into file cache // In write path, to control if the content would be written into file cache. // In read path, read from file cache or remote storage when execute query. diff --git a/gensrc/thrift/PlanNodes.thrift b/gensrc/thrift/PlanNodes.thrift index cf4e88a0906a37..998f75cf40a1fc 100644 --- a/gensrc/thrift/PlanNodes.thrift +++ b/gensrc/thrift/PlanNodes.thrift @@ -825,6 +825,8 @@ struct TOlapScanNode { 18: optional list topn_filter_source_node_ids //deprecated, move to TPlanNode.106 19: optional TSortInfo score_sort_info 20: optional i64 score_sort_limit + 21: optional TSortInfo ann_sort_info + 22: optional i64 ann_sort_limit } struct TEqJoinCondition { diff --git a/regression-test/data/ann_index_p0/ann_index_basic.out b/regression-test/data/ann_index_p0/ann_index_basic.out new file mode 100644 index 00000000000000..b0787f6b52c2e1 --- /dev/null +++ b/regression-test/data/ann_index_p0/ann_index_basic.out @@ -0,0 +1,46 @@ +-- This file is automatically generated. You should know what you did if you want to edit this +-- !sql_l2_insert -- +3 + +-- !sql_l2_query -- +1 0.0 +2 0.5196152329444885 +3 13.928388595581055 + +-- !sql_ip_insert -- +3 + +-- !sql_ip_query -- +3 +2 +1 + +-- !sql_l2_threshold -- +1 +2 + +-- !sql_l2_desc -- +3 +2 + +-- !sql_ip_asc -- +1 +2 + +-- !sql_l2_small_pred -- +1 +2 +3 +4 + +-- !sql_l2_large_pred -- +1 +2 +3 +4 +5 + +-- !sql_compound -- +1 +2 + diff --git a/regression-test/data/ann_index_p0/ann_with_fulltext.out b/regression-test/data/ann_index_p0/ann_with_fulltext.out new file mode 100644 index 00000000000000..6701752255c15f --- /dev/null +++ b/regression-test/data/ann_index_p0/ann_with_fulltext.out @@ -0,0 +1,20 @@ +-- This file is automatically generated. You should know what you did if you want to edit this +-- !q0 -- +0 [39.906116, 10.495334, 54.08394, 88.67262, 55.243687, 10.162686, 36.335983, 38.684258] This example illustrates how subtle differences can influence perception. It's more about interpretation than right or wrong. 100 +1 [62.759315, 97.15586, 25.832521, 39.604908, 88.76715, 72.64085, 9.688437, 17.721428] Thanks for all the comments, good and bad. They help us refine our test. Keep in mind that we're attempting to figure you out in 40 pairs of pictures. We did this so that lots of people could take it, just to introduce the idea.

A real test would have more like 200 pairs, which is what the YC founders took when we assessed their attributes in the first place. 101 +2 [15.447449, 59.7771, 65.54516, 12.973712, 99.685135, 72.080734, 85.71118, 99.35976] At a glance, these might seem obvious, but there’s nuance in every choice. Don’t rush. 102 +3 [72.26747, 46.42257, 32.368374, 80.50209, 5.777631, 98.803314, 7.0915947, 68.62693] We're testing how consistent your judgments are over a range of visual impressions. There's no single 'correct' answer. 103 +4 [22.098177, 74.10027, 63.634556, 4.710955, 12.405106, 79.39356, 63.014366, 68.67834] Some pairs are meant to be tricky. Your intuition is part of what we're analyzing. 104 +5 [27.53003, 72.1106, 50.891026, 38.459953, 68.30715, 20.610682, 94.806274, 45.181377] This data will help us identify patterns in how people perceive attributes such as trustworthiness or confidence. 105 +6 [77.73215, 64.42907, 71.50025, 43.85641, 94.42648, 50.04773, 65.12575, 68.58207] Sometimes people see entirely different things in the same image. That's part of the exploration. 106 +7 [2.1537063, 82.667885, 16.171143, 71.126656, 5.335274, 40.286068, 11.943586, 3.69409] Don't worry if you’re unsure. The ambiguity is intentional — that’s what makes this interesting. 107 +8 [54.435013, 56.800594, 59.335514, 55.829235, 85.46627, 33.388138, 11.076194, 20.480877] Your reactions help us understand which features people subconsciously favor or avoid. 108 +9 [76.197945, 60.623528, 84.229805, 31.652937, 71.82595, 48.04684, 71.29212, 30.282396] This task isn’t about right answers, but about consistency in your judgments over time. 109 + +-- !q1 -- +0 [39.906116, 10.495334, 54.08394, 88.67262, 55.243687, 10.162686, 36.335983, 38.684258] This example illustrates how subtle differences can influence perception. It's more about interpretation than right or wrong. 100 + +-- !q2 -- +0 [39.906116, 10.495334, 54.08394, 88.67262, 55.243687, 10.162686, 36.335983, 38.684258] This example illustrates how subtle differences can influence perception. It's more about interpretation than right or wrong. 100 81.69190979003906 +9 [76.197945, 60.623528, 84.229805, 31.652937, 71.82595, 48.04684, 71.29212, 30.282396] This task isn’t about right answers, but about consistency in your judgments over time. 109 122.17068481445312 + diff --git a/regression-test/data/ann_index_p0/delete_where.out b/regression-test/data/ann_index_p0/delete_where.out new file mode 100644 index 00000000000000..9f88b492f3fc62 --- /dev/null +++ b/regression-test/data/ann_index_p0/delete_where.out @@ -0,0 +1,14 @@ +-- This file is automatically generated. You should know what you did if you want to edit this +-- !sql_1 -- +1 [1, 2, 3] 11 +2 [4, 5, 6] 22 +3 [7, 8, 9] 33 + +-- !sql_2 -- +2 [4, 5, 6] 22 +3 [7, 8, 9] 33 + +-- !sql_3 -- +2 5.196152422706632 +3 10.392304845413264 + diff --git a/regression-test/data/ann_index_p0/insert_with_invalid_array.out b/regression-test/data/ann_index_p0/insert_with_invalid_array.out new file mode 100644 index 00000000000000..194b6b4b55354a --- /dev/null +++ b/regression-test/data/ann_index_p0/insert_with_invalid_array.out @@ -0,0 +1,4 @@ +-- This file is automatically generated. You should know what you did if you want to edit this +-- !sql -- +1 [1, 2, 3] + diff --git a/regression-test/data/ann_index_p0/memtbl_on_sink.out b/regression-test/data/ann_index_p0/memtbl_on_sink.out new file mode 100644 index 00000000000000..c4d26f6e5a9ef2 --- /dev/null +++ b/regression-test/data/ann_index_p0/memtbl_on_sink.out @@ -0,0 +1,10 @@ +-- This file is automatically generated. You should know what you did if you want to edit this +-- !sql_0 -- +1 [1, 2, 3] + +-- !sql_1 -- +1 [1, 2, 3] + +-- !sql_2 -- +1 0.0 + diff --git a/regression-test/data/ann_index_p0/mow_with_ann.out b/regression-test/data/ann_index_p0/mow_with_ann.out new file mode 100644 index 00000000000000..c3cc035245f3e6 --- /dev/null +++ b/regression-test/data/ann_index_p0/mow_with_ann.out @@ -0,0 +1,27 @@ +-- This file is automatically generated. You should know what you did if you want to edit this +-- !ann_only_topn -- +0 81.69190833232089 +5 90.8576121918621 +6 111.23396339841393 + +-- !ann_only_range -- +2 + +-- !inv_only_score -- +9 2.2117803 +1 1.9442263 +0 1.7344174 + +-- !ann_with_filter -- +0 81.69190833232089 +9 122.1706825093796 +1 163.05757214930287 + +-- !hybrid_search -- +0 \N 81.69190833232089 ann +5 \N 90.8576121918621 ann +6 \N 111.23396339841393 ann +9 2.211780309677124 \N score +1 1.9442262649536133 \N score +0 1.73441743850708 \N score + diff --git a/regression-test/suites/ann_index_p0/ann_index_basic.groovy b/regression-test/suites/ann_index_p0/ann_index_basic.groovy new file mode 100644 index 00000000000000..b64e795c993aea --- /dev/null +++ b/regression-test/suites/ann_index_p0/ann_index_basic.groovy @@ -0,0 +1,142 @@ +// 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. + +// Doris ANN Index Basic Test Suite +// Includes range search, topn search, and compound search + +suite ("ann_index_basic") { + sql "set enable_common_expr_pushdown=true;" + + // 1) Basic L2 ANN table: dim=3 + sql "drop table if exists tbl_ann_l2" + sql """ + CREATE TABLE tbl_ann_l2 ( + id INT NOT NULL, + embedding ARRAY NOT NULL, + INDEX idx_emb (`embedding`) USING ANN PROPERTIES( + "index_type"="hnsw", + "metric_type"="l2_distance", + "dim"="3" + ) + ) ENGINE=OLAP + DUPLICATE KEY(id) + DISTRIBUTED BY HASH(id) BUCKETS 2 + PROPERTIES ("replication_num" = "1"); + """ + + qt_sql_l2_insert """ + INSERT INTO tbl_ann_l2 VALUES + (1, [1.0, 2.0, 3.0]), + (2, [0.5, 2.1, 2.9]), + (3, [10.0, 10.0, 10.0]); + """ + + // Query: l2 distance ascending (closest first) + qt_sql_l2_query "select id, l2_distance_approximate(embedding, [1.0,2.0,3.0]) as dist from tbl_ann_l2 order by dist limit 3;" + + // 2) Basic inner_product ANN table: dim=4 + sql "drop table if exists tbl_ann_ip" + sql """ + CREATE TABLE tbl_ann_ip ( + id INT NOT NULL, + embedding ARRAY NOT NULL, + INDEX idx_emb (`embedding`) USING ANN PROPERTIES( + "index_type"="hnsw", + "metric_type"="inner_product", + "dim"="4" + ) + ) ENGINE=OLAP + DUPLICATE KEY(id) + DISTRIBUTED BY HASH(id) BUCKETS 2 + PROPERTIES ("replication_num" = "1"); + """ + + qt_sql_ip_insert """ + INSERT INTO tbl_ann_ip VALUES + (1, [0.1, 0.2, 0.3, 0.4]), + (2, [0.5, 0.6, 0.7, 0.8]), + (3, [1.0, 1.0, 1.0, 1.0]); + """ + + // Query: inner product descending (higher score first) + qt_sql_ip_query "select id from tbl_ann_ip order by inner_product_approximate(embedding, [0.1,0.2,0.3,0.4]) desc limit 3;" + + // 3) Simple threshold filter using l2_distance_approximate + qt_sql_l2_threshold "select id from tbl_ann_l2 where l2_distance_approximate(embedding, [1.0,2.0,3.0]) < 5.0 order by id;" + + // 4) Descending l2 order (should exercise path where Desc topn for l2/cosine cannot be evaluated by ann index) + qt_sql_l2_desc "select id from tbl_ann_l2 order by l2_distance_approximate(embedding, [1.0,2.0,3.0]) desc limit 2;" + + // 5) Ascending inner_product order (should exercise path where Asc topn for inner product cannot be evaluated by ann index) + qt_sql_ip_asc "select id from tbl_ann_ip order by inner_product_approximate(embedding, [0.1,0.2,0.3,0.4]) asc limit 2;" + + // 6) Large table to exercise predicate-input-ratio check (create many rows and run topn with small-range predicate) + sql "drop table if exists tbl_ann_l2_large" + sql """ + CREATE TABLE tbl_ann_l2_large ( + id INT NOT NULL, + embedding ARRAY NOT NULL, + INDEX idx_emb (`embedding`) USING ANN PROPERTIES( + "index_type"="hnsw", + "metric_type"="l2_distance", + "dim"="3" + ) + ) ENGINE=OLAP + DUPLICATE KEY(id) + DISTRIBUTED BY HASH(id) BUCKETS 2 + PROPERTIES ("replication_num" = "1"); + """ + + // insert 50 rows with simple embeddings + sql "truncate table tbl_ann_l2_large" + def values = [] + for (i in 1..50) { + def a = i * 1.0 + def b = (i + 1) * 1.0 + def c = (i + 2) * 1.0 + values.add("(${i}, [${a}, ${b}, ${c}])") + } + sql "INSERT INTO tbl_ann_l2_large VALUES ${values.join(',')};" + + // topn with small predicate (id < 5) -> selects 4/50 = 8% (<30%), should exercise "will not use ann index" path + qt_sql_l2_small_pred "select id from tbl_ann_l2_large where id < 5 order by l2_distance_approximate(embedding, [1.0,2.0,3.0]) limit 5;" + + // topn with large predicate (id < 40) -> selects 39/50 = 78% (>30%), more likely to use ann index + qt_sql_l2_large_pred "select id from tbl_ann_l2_large where id < 40 order by l2_distance_approximate(embedding, [1.0,2.0,3.0]) limit 5;" + + // 7) Compound search: inverted index + ann index + sql "drop table if exists ann_compound" + sql """ + CREATE TABLE ann_compound ( + id INT NOT NULL, + embedding ARRAY NOT NULL, + txt STRING NOT NULL, + INDEX idx_txt(`txt`) USING INVERTED PROPERTIES("parser"="english"), + INDEX idx_ann(`embedding`) USING ANN PROPERTIES( + "index_type"="hnsw", + "metric_type"="l2_distance", + "dim"="3" + ) + ) DUPLICATE KEY(id) + DISTRIBUTED BY HASH(id) BUCKETS 1 + PROPERTIES ("replication_num" = "1"); + """ + + sql "INSERT INTO ann_compound VALUES (1, [1.0,2.0,3.0], 'quick brown fox'), (2, [2.0,3.0,4.0], 'lazy dog fox'), (3, [10.0,10.0,10.0], 'unrelated text');" + + qt_sql_compound "select id from ann_compound where txt match_any 'fox' order by l2_distance_approximate(embedding, [1.0,2.0,3.0]) limit 3;" +} diff --git a/regression-test/suites/ann_index_p0/ann_with_fulltext.groovy b/regression-test/suites/ann_index_p0/ann_with_fulltext.groovy new file mode 100644 index 00000000000000..06bc158c5f9b03 --- /dev/null +++ b/regression-test/suites/ann_index_p0/ann_with_fulltext.groovy @@ -0,0 +1,65 @@ +// 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. + +suite("ann_with_fulltext") { + sql "drop table if exists ann_with_fulltext" + sql "set profile_level=2;" + sql "set enable_common_expr_pushdown=true;" + + sql """ + create table ann_with_fulltext ( + id int not null, + embedding array not null, + comment String not null, + value int null, + INDEX idx_comment(`comment`) USING INVERTED PROPERTIES("parser" = "english") COMMENT 'inverted index for comment', + INDEX ann_embedding(`embedding`) USING ANN PROPERTIES("index_type"="hnsw","metric_type"="l2_distance","dim"="8") + ) duplicate key (`id`) + distributed by hash(`id`) buckets 1 + properties("replication_num"="1"); + """ + + sql """ + INSERT INTO ann_with_fulltext (id, embedding, comment, value) VALUES + (0, [39.906116, 10.495334, 54.08394, 88.67262, 55.243687, 10.162686, 36.335983, 38.684258], "This example illustrates how subtle differences can influence perception. It's more about interpretation than right or wrong.", 100), + (1, [62.759315, 97.15586, 25.832521, 39.604908, 88.76715, 72.64085, 9.688437, 17.721428], "Thanks for all the comments, good and bad. They help us refine our test. Keep in mind that we're attempting to figure you out in 40 pairs of pictures. We did this so that lots of people could take it, just to introduce the idea.

A real test would have more like 200 pairs, which is what the YC founders took when we assessed their attributes in the first place.", 101), + (2, [15.447449, 59.7771, 65.54516, 12.973712, 99.685135, 72.080734, 85.71118, 99.35976], "At a glance, these might seem obvious, but there’s nuance in every choice. Don’t rush.", 102), + (3, [72.26747, 46.42257, 32.368374, 80.50209, 5.777631, 98.803314, 7.0915947, 68.62693], "We're testing how consistent your judgments are over a range of visual impressions. There's no single 'correct' answer.", 103), + (4, [22.098177, 74.10027, 63.634556, 4.710955, 12.405106, 79.39356, 63.014366, 68.67834], "Some pairs are meant to be tricky. Your intuition is part of what we're analyzing.", 104), + (5, [27.53003, 72.1106, 50.891026, 38.459953, 68.30715, 20.610682, 94.806274, 45.181377], "This data will help us identify patterns in how people perceive attributes such as trustworthiness or confidence.", 105), + (6, [77.73215, 64.42907, 71.50025, 43.85641, 94.42648, 50.04773, 65.12575, 68.58207], "Sometimes people see entirely different things in the same image. That's part of the exploration.", 106), + (7, [2.1537063, 82.667885, 16.171143, 71.126656, 5.335274, 40.286068, 11.943586, 3.69409], "Don't worry if you’re unsure. The ambiguity is intentional — that’s what makes this interesting.", 107), + (8, [54.435013, 56.800594, 59.335514, 55.829235, 85.46627, 33.388138, 11.076194, 20.480877], "Your reactions help us understand which features people subconsciously favor or avoid.", 108), + (9, [76.197945, 60.623528, 84.229805, 31.652937, 71.82595, 48.04684, 71.29212, 30.282396], "This task isn’t about right answers, but about consistency in your judgments over time.", 109); + """ + qt_q0 """ + select * from ann_with_fulltext ORDER BY id; + """ + + qt_q1 """ + select * from ann_with_fulltext where comment match_any "illustrates comments answers" and l2_distance_approximate(embedding, [26.360261917114258,7.05784273147583,32.361351013183594,86.39714050292969,58.79527282714844,27.189321517944336,99.38946533203125,80.19270324707031]) < 105.66439056396484 + """ + + qt_q2 """ + select *, l2_distance_approximate(embedding, [26.360261917114258,7.05784273147583,32.361351013183594,86.39714050292969,58.79527282714844,27.189321517944336,99.38946533203125,80.19270324707031]) as dist from ann_with_fulltext where comment match_any "illustrates comments answers" order by dist limit 2 + """ + // This case should error because score() + another ordering expression is not allowed for TopN push down + test { + sql "select score() as score, l2_distance_approximate(embedding, [26.360261917114258,7.05784273147583,32.361351013183594,86.39714050292969,58.79527282714844,27.189321517944336,99.38946533203125,80.19270324707031]) as dist from ann_with_fulltext where comment match_any \"illustrates comments answers\" order by score, dist limit 2" + exception "TopN must have exactly one ordering expression for score() push down optimization" + } +} \ No newline at end of file diff --git a/regression-test/suites/ann_index_p0/create_ann_index_test.groovy b/regression-test/suites/ann_index_p0/create_ann_index_test.groovy new file mode 100644 index 00000000000000..c36592302bbb37 --- /dev/null +++ b/regression-test/suites/ann_index_p0/create_ann_index_test.groovy @@ -0,0 +1,325 @@ +// 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. + +suite("create_ann_index_test") { + sql "set enable_common_expr_pushdown=true;" + // Test that CREATE INDEX for ANN is not supported + sql "drop table if exists tbl_not_null" + sql """ + CREATE TABLE `tbl_not_null` ( + `id` int NOT NULL COMMENT "", + `embedding` array NOT NULL COMMENT "" + ) ENGINE=OLAP + DUPLICATE KEY(`id`) COMMENT "OLAP" + DISTRIBUTED BY HASH(`id`) BUCKETS 2 + PROPERTIES ( + "replication_num" = "1" + ); + """ + + test { + sql """ + CREATE INDEX idx_test_ann ON tbl_not_null(`embedding`) USING ANN PROPERTIES( + "index_type"="hnsw", + "metric_type"="l2_distance", + "dim"="1" + ); + """ + exception "ANN index can only be created during table creation, not through CREATE INDEX" + } + + // Test cases for creating tables with ANN indexes + + // 1. Case for nullable column + sql "drop table if exists tbl_nullable_ann" + test { + sql """ + CREATE TABLE tbl_nullable_ann ( + id INT NOT NULL COMMENT "", + embedding ARRAY NULL COMMENT "", + INDEX idx_nullable_ann (`embedding`) USING ANN PROPERTIES( + "index_type"="hnsw", + "metric_type"="l2_distance", + "dim"="1" + ) + ) ENGINE=OLAP + DUPLICATE KEY(id) COMMENT "OLAP" + DISTRIBUTED BY HASH(id) BUCKETS 2 + PROPERTIES ( + "replication_num" = "1" + ); + """ + exception "ANN index must be built on a column that is not nullable" + } + + // 2. Invalid properties cases + // dim is not a positive integer + sql "drop table if exists tbl_ann_invalid_dim" + test { + sql """ + CREATE TABLE tbl_ann_invalid_dim ( + id INT NOT NULL COMMENT "", + embedding ARRAY NOT NULL COMMENT "", + INDEX idx_test_ann (`embedding`) USING ANN PROPERTIES( + "index_type"="hnsw", + "metric_type"="l2_distance", + "dim"="-1" + ) + ) ENGINE=OLAP + DUPLICATE KEY(id) COMMENT "OLAP" + DISTRIBUTED BY HASH(id) BUCKETS 2 + PROPERTIES ( + "replication_num" = "1" + ); + """ + exception "dim of ann index must be a positive integer" + } + + // dim is not a number + sql "drop table if exists tbl_ann_invalid_dim_str" + test { + sql """ + CREATE TABLE tbl_ann_invalid_dim_str ( + id INT NOT NULL COMMENT "", + embedding ARRAY NOT NULL COMMENT "", + INDEX idx_test_ann (`embedding`) USING ANN PROPERTIES( + "index_type"="hnsw", + "metric_type"="l2_distance", + "dim"="abc" + ) + ) ENGINE=OLAP + DUPLICATE KEY(id) COMMENT "OLAP" + DISTRIBUTED BY HASH(id) BUCKETS 2 + PROPERTIES ( + "replication_num" = "1" + ); + """ + exception "dim of ann index must be a positive integer" + } + + // dim is missing + sql "drop table if exists tbl_ann_missing_dim" + test { + sql """ + CREATE TABLE tbl_ann_missing_dim ( + id INT NOT NULL COMMENT "", + embedding ARRAY NOT NULL COMMENT "", + INDEX idx_test_ann (`embedding`) USING ANN PROPERTIES( + "index_type"="hnsw", + "metric_type"="l2_distance" + ) + ) ENGINE=OLAP + DUPLICATE KEY(id) COMMENT "OLAP" + DISTRIBUTED BY HASH(id) BUCKETS 2 + PROPERTIES ( + "replication_num" = "1" + ); + """ + exception "dim of ann index must be specified" + } + + // index_type is missing + sql "drop table if exists tbl_ann_missing_index_type" + test { + sql """ + CREATE TABLE tbl_ann_missing_index_type ( + id INT NOT NULL COMMENT "", + embedding ARRAY NOT NULL COMMENT "", + INDEX idx_test_ann (`embedding`) USING ANN PROPERTIES( + "metric_type"="l2_distance", + "dim"="1" + ) + ) ENGINE=OLAP + DUPLICATE KEY(id) COMMENT "OLAP" + DISTRIBUTED BY HASH(id) BUCKETS 2 + PROPERTIES ( + "replication_num" = "1" + ); + """ + exception "index_type of ann index be specified." + } + + // metric_type is missing + sql "drop table if exists tbl_ann_missing_metric_type" + test { + sql """ + CREATE TABLE tbl_ann_missing_metric_type ( + id INT NOT NULL COMMENT "", + embedding ARRAY NOT NULL COMMENT "", + INDEX idx_test_ann (`embedding`) USING ANN PROPERTIES( + "index_type"="hnsw", + "dim"="1" + ) + ) ENGINE=OLAP + DUPLICATE KEY(id) COMMENT "OLAP" + DISTRIBUTED BY HASH(id) BUCKETS 2 + PROPERTIES ( + "replication_num" = "1" + ); + """ + exception "metric_type of ann index must be specified." + } + + // index_type is incorrect + sql "drop table if exists tbl_ann_invalid_index_type" + test { + sql """ + CREATE TABLE tbl_ann_invalid_index_type ( + id INT NOT NULL COMMENT "", + embedding ARRAY NOT NULL COMMENT "", + INDEX idx_test_ann (`embedding`) USING ANN PROPERTIES( + "index_type"="ivf", + "metric_type"="l2_distance", + "dim"="1" + ) + ) ENGINE=OLAP + DUPLICATE KEY(id) COMMENT "OLAP" + DISTRIBUTED BY HASH(id) BUCKETS 2 + PROPERTIES ( + "replication_num" = "1" + ); + """ + exception "only support ann index with type hnsw" + } + + // metric_type is incorrect + sql "drop table if exists tbl_ann_invalid_metric_type" + test { + sql """ + CREATE TABLE tbl_ann_invalid_metric_type ( + id INT NOT NULL COMMENT "", + embedding ARRAY NOT NULL COMMENT "", + INDEX idx_test_ann (`embedding`) USING ANN PROPERTIES( + "index_type"="hnsw", + "metric_type"="cosine", + "dim"="1" + ) + ) ENGINE=OLAP + DUPLICATE KEY(id) COMMENT "OLAP" + DISTRIBUTED BY HASH(id) BUCKETS 2 + PROPERTIES ( + "replication_num" = "1" + ); + """ + exception "only support ann index with metric l2_distance or inner_product" + } + + // 不支持的属性 quantization (已移除) + sql "drop table if exists tbl_ann_invalid_quantization" + test { + sql """ + CREATE TABLE tbl_ann_invalid_quantization ( + id INT NOT NULL COMMENT "", + embedding ARRAY NOT NULL COMMENT "", + INDEX idx_test_ann (`embedding`) USING ANN PROPERTIES( + "index_type"="hnsw", + "metric_type"="l2_distance", + "dim"="1", + "quantization"="flat" + ) + ) ENGINE=OLAP + DUPLICATE KEY(id) COMMENT "OLAP" + DISTRIBUTED BY HASH(id) BUCKETS 2 + PROPERTIES ( + "replication_num" = "1" + ); + """ + exception "unknown ann index property" + } + + // Unknown property + sql "drop table if exists tbl_ann_unknown_property" + test { + sql """ + CREATE TABLE tbl_ann_unknown_property ( + id INT NOT NULL COMMENT "", + embedding ARRAY NOT NULL COMMENT "", + INDEX idx_test_ann (`embedding`) USING ANN PROPERTIES( + "index_type"="hnsw", + "metric_type"="l2_distance", + "dim"="1", + "unknown"="xxx" + ) + ) ENGINE=OLAP + DUPLICATE KEY(id) COMMENT "OLAP" + DISTRIBUTED BY HASH(id) BUCKETS 2 + PROPERTIES ( + "replication_num" = "1" + ); + """ + exception "unknown ann index property: unknown" + } + + // 3. Valid CREATE TABLE with ANN index (l2_distance) + sql "drop table if exists tbl_ann_valid_l2" + sql """ + CREATE TABLE tbl_ann_valid_l2 ( + id INT NOT NULL COMMENT "", + embedding ARRAY NOT NULL COMMENT "", + INDEX idx_test_ann (`embedding`) USING ANN PROPERTIES( + "index_type"="hnsw", + "metric_type"="l2_distance", + "dim"="128" + ) + ) ENGINE=OLAP + DUPLICATE KEY(id) COMMENT "OLAP" + DISTRIBUTED BY HASH(id) BUCKETS 2 + PROPERTIES ( + "replication_num" = "1" + ); + """ + + // 4. Valid CREATE TABLE with ANN index (inner_product) + sql "drop table if exists tbl_ann_valid_inner_product" + sql """ + CREATE TABLE tbl_ann_valid_inner_product ( + id INT NOT NULL COMMENT "", + embedding ARRAY NOT NULL COMMENT "", + INDEX idx_test_ann (`embedding`) USING ANN PROPERTIES( + "index_type"="hnsw", + "metric_type"="inner_product", + "dim"="128" + ) + ) ENGINE=OLAP + DUPLICATE KEY(id) COMMENT "OLAP" + DISTRIBUTED BY HASH(id) BUCKETS 2 + PROPERTIES ( + "replication_num" = "1" + ); + """ + + sql "drop table if exists tbl_ann_unique_key" + test { + sql """ + CREATE TABLE tbl_ann_unique_key ( + id INT NOT NULL COMMENT "", + embedding ARRAY NOT NULL COMMENT "", + INDEX idx_test_ann (`embedding`) USING ANN PROPERTIES( + "index_type"="hnsw", + "metric_type"="inner_product", + "dim"="128" + ) + ) ENGINE=OLAP + UNIQUE KEY(id) COMMENT "OLAP" + DISTRIBUTED BY HASH(id) BUCKETS 2 + PROPERTIES ( + "replication_num" = "1" + ); + """ + exception "ANN index can only be used in DUP_KEYS table" + } +} \ No newline at end of file diff --git a/regression-test/suites/ann_index_p0/create_tbl_with_ann_index_test.groovy b/regression-test/suites/ann_index_p0/create_tbl_with_ann_index_test.groovy new file mode 100644 index 00000000000000..99925c354d8113 --- /dev/null +++ b/regression-test/suites/ann_index_p0/create_tbl_with_ann_index_test.groovy @@ -0,0 +1,297 @@ +// 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. + +suite("create_tbl_with_ann_index_test") { + sql "set enable_common_expr_pushdown=true;" + sql "drop table if exists ann_tbl1" + test { + sql """ + CREATE TABLE ann_tbl1 ( + id INT NOT NULL COMMENT "", + vec ARRAY NOT NULL COMMENT "", + INDEX ann_idx1 (vec) USING ANN PROPERTIES( + "index_type" = "hnsw", + "metric_type" = "l2_distance", + "dim" = "128" + ) + ) ENGINE=OLAP + DUPLICATE KEY(id) COMMENT "OLAP" + DISTRIBUTED BY HASH(id) BUCKETS 2 + PROPERTIES ( + "replication_num" = "1" + ); + """ + } + + sql "drop table if exists ann_tbl2" + test { + sql """ + CREATE TABLE ann_tbl2 ( + id INT NOT NULL COMMENT "", + vec ARRAY NOT NULL COMMENT "", + INDEX ann_idx2 (vec) USING ANN PROPERTIES( + "index_type" = "ivf", + "metric_type" = "l2_distance", + "dim" = "128" + ) + ) ENGINE=OLAP + DUPLICATE KEY(id) COMMENT "OLAP" + DISTRIBUTED BY HASH(id) BUCKETS 2 + PROPERTIES ( + "replication_num" = "1" + ); + """ + exception "only support ann index with type hnsw" + } + + // metric_type 错误 + sql "drop table if exists ann_tbl3" + test { + sql """ + CREATE TABLE ann_tbl3 ( + id INT NOT NULL COMMENT "", + vec ARRAY NOT NULL COMMENT "", + INDEX ann_idx3 (vec) USING ANN PROPERTIES( + "index_type" = "hnsw", + "metric_type" = "cosine", + "dim" = "128" + ) + ) ENGINE=OLAP + DUPLICATE KEY(id) COMMENT "OLAP" + DISTRIBUTED BY HASH(id) BUCKETS 2 + PROPERTIES ( + "replication_num" = "1" + ); + """ + exception "only support ann index with metric l2_distance or inner_product" + } + + // dim 非正整数 + sql "drop table if exists ann_tbl4" + test { + sql """ + CREATE TABLE ann_tbl4 ( + id INT NOT NULL COMMENT "", + vec ARRAY NOT NULL COMMENT "", + INDEX ann_idx4 (vec) USING ANN PROPERTIES( + "index_type" = "hnsw", + "metric_type" = "l2_distance", + "dim" = "-1" + ) + ) ENGINE=OLAP + DUPLICATE KEY(id) COMMENT "OLAP" + DISTRIBUTED BY HASH(id) BUCKETS 2 + PROPERTIES ( + "replication_num" = "1" + ); + """ + exception "dim of ann index must be a positive integer" + } + + // dim 非数字 + sql "drop table if exists ann_tbl5" + test { + sql """ + CREATE TABLE ann_tbl5 ( + id INT NOT NULL COMMENT "", + vec ARRAY NOT NULL COMMENT "", + INDEX ann_idx5 (vec) USING ANN PROPERTIES( + "index_type" = "hnsw", + "metric_type" = "l2_distance", + "dim" = "abc" + ) + ) ENGINE=OLAP + DUPLICATE KEY(id) COMMENT "OLAP" + DISTRIBUTED BY HASH(id) BUCKETS 2 + PROPERTIES ( + "replication_num" = "1" + ); + """ + exception "dim of ann index must be a positive integer" + } + + // 不支持的属性 quantization (已移除) + sql "drop table if exists ann_tbl6" + test { + sql """ + CREATE TABLE ann_tbl6 ( + id INT NOT NULL COMMENT "", + vec ARRAY NOT NULL COMMENT "", + INDEX ann_idx6 (vec) USING ANN PROPERTIES( + "index_type" = "hnsw", + "metric_type" = "l2_distance", + "dim" = "128", + "quantization" = "flat" + ) + ) ENGINE=OLAP + DUPLICATE KEY(id) COMMENT "OLAP" + DISTRIBUTED BY HASH(id) BUCKETS 2 + PROPERTIES ( + "replication_num" = "1" + ); + """ + exception "unknown ann index property" + } + + // 缺少 index_type + sql "drop table if exists ann_tbl7" + test { + sql """ + CREATE TABLE ann_tbl7 ( + id INT NOT NULL COMMENT "", + vec ARRAY NOT NULL COMMENT "", + INDEX ann_idx7 (vec) USING ANN PROPERTIES( + "metric_type" = "l2_distance", + "dim" = "128" + ) + ) ENGINE=OLAP + DUPLICATE KEY(id) COMMENT "OLAP" + DISTRIBUTED BY HASH(id) BUCKETS 2 + PROPERTIES ( + "replication_num" = "1" + ); + """ + exception "index_type of ann index be specified." + } + + // 缺少 metric_type + sql "drop table if exists ann_tbl8" + test { + sql """ + CREATE TABLE ann_tbl8 ( + id INT NOT NULL COMMENT "", + vec ARRAY NOT NULL COMMENT "", + INDEX ann_idx8 (vec) USING ANN PROPERTIES( + "index_type" = "hnsw", + "dim" = "128" + ) + ) ENGINE=OLAP + DUPLICATE KEY(id) COMMENT "OLAP" + DISTRIBUTED BY HASH(id) BUCKETS 2 + PROPERTIES ( + "replication_num" = "1" + ); + """ + exception "metric_type of ann index must be specified." + } + + // 缺少 dim + sql "drop table if exists ann_tbl9" + test { + sql """ + CREATE TABLE ann_tbl9 ( + id INT NOT NULL COMMENT "", + vec ARRAY NOT NULL COMMENT "", + INDEX ann_idx9 (vec) USING ANN PROPERTIES( + "index_type" = "hnsw", + "metric_type" = "l2_distance" + ) + ) ENGINE=OLAP + DUPLICATE KEY(id) COMMENT "OLAP" + DISTRIBUTED BY HASH(id) BUCKETS 2 + PROPERTIES ( + "replication_num" = "1" + ); + """ + exception "dim of ann index must be specified" + } + + // 未知属性 + sql "drop table if exists ann_tbl10" + test { + sql """ + CREATE TABLE ann_tbl10 ( + id INT NOT NULL COMMENT "", + vec ARRAY NOT NULL COMMENT "", + INDEX ann_idx10 (vec) USING ANN PROPERTIES( + "index_type" = "hnsw", + "metric_type" = "l2_distance", + "dim" = "128", + "unknown" = "xxx" + ) + ) ENGINE=OLAP + DUPLICATE KEY(id) COMMENT "OLAP" + DISTRIBUTED BY HASH(id) BUCKETS 2 + PROPERTIES ( + "replication_num" = "1" + ); + """ + exception "unknown ann index property" + } + + // 不支持的属性 quantization (已移除) + sql "drop table if exists ann_tbl12" + test { + sql """ + CREATE TABLE ann_tbl12 ( + id INT NOT NULL COMMENT "", + vec ARRAY NOT NULL COMMENT "", + INDEX ann_idx12 (vec) USING ANN PROPERTIES( + "index_type" = "hnsw", + "metric_type" = "inner_product", + "dim" = "128", + "quantization" = "pq" + ) + ) ENGINE=OLAP + DUPLICATE KEY(id) COMMENT "OLAP" + DISTRIBUTED BY HASH(id) BUCKETS 2 + PROPERTIES ( + "replication_num" = "1" + ); + """ + exception "unknown ann index property" + } + + // 成功创建 ANN 索引 - 基本配置 + sql "drop table if exists ann_tbl13" + sql """ + CREATE TABLE ann_tbl13 ( + id INT NOT NULL COMMENT "", + vec ARRAY NOT NULL COMMENT "", + INDEX ann_idx13 (vec) USING ANN PROPERTIES( + "index_type" = "hnsw", + "metric_type" = "l2_distance", + "dim" = "128" + ) + ) ENGINE=OLAP + DUPLICATE KEY(id) COMMENT "OLAP" + DISTRIBUTED BY HASH(id) BUCKETS 2 + PROPERTIES ( + "replication_num" = "1" + ); + """ + + // 成功创建 ANN 索引 - inner_product + sql "drop table if exists ann_tbl14" + sql """ + CREATE TABLE ann_tbl14 ( + id INT NOT NULL COMMENT "", + vec ARRAY NOT NULL COMMENT "", + INDEX ann_idx14 (vec) USING ANN PROPERTIES( + "index_type" = "hnsw", + "metric_type" = "inner_product", + "dim" = "256" + ) + ) ENGINE=OLAP + DUPLICATE KEY(id) COMMENT "OLAP" + DISTRIBUTED BY HASH(id) BUCKETS 2 + PROPERTIES ( + "replication_num" = "1" + ); + """ + +} \ No newline at end of file diff --git a/regression-test/suites/ann_index_p0/delete_where.groovy b/regression-test/suites/ann_index_p0/delete_where.groovy new file mode 100644 index 00000000000000..677f713e6f9105 --- /dev/null +++ b/regression-test/suites/ann_index_p0/delete_where.groovy @@ -0,0 +1,50 @@ +// 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. + +suite("delete_where_with_ann") { + sql "set enable_common_expr_pushdown=true;" + sql "drop table if exists delete_where_with_ann" + test { + sql """ + CREATE TABLE delete_where_with_ann ( + id INT NOT NULL COMMENT "", + vec ARRAY NOT NULL COMMENT "", + value INT NULL COMMENT "", + INDEX ann_idx (vec) USING ANN PROPERTIES( + "index_type" = "hnsw", + "metric_type" = "l2_distance", + "dim" = "3" + ) + ) ENGINE=OLAP + DUPLICATE KEY(id) COMMENT "OLAP" + DISTRIBUTED BY HASH(id) BUCKETS AUTO + PROPERTIES ( + "replication_num" = "1" + ); + """ + } + + sql "insert into delete_where_with_ann values (1, [1.0, 2.0, 3.0], 11),(2, [4.0, 5.0, 6.0], 22),(3, [7.0, 8.0, 9.0], 33)" + + qt_sql_1 "select * from delete_where_with_ann order by id" + + sql "delete from delete_where_with_ann where id = 1" + + qt_sql_2 "select * from delete_where_with_ann order by id" + + qt_sql_3 "select id, l2_distance_approximate(vec, [1.0, 2.0, 3.0]) as dist from delete_where_with_ann order by dist limit 2;" +} \ No newline at end of file diff --git a/regression-test/suites/ann_index_p0/insert_with_invalid_array.groovy b/regression-test/suites/ann_index_p0/insert_with_invalid_array.groovy new file mode 100644 index 00000000000000..193392d92236ba --- /dev/null +++ b/regression-test/suites/ann_index_p0/insert_with_invalid_array.groovy @@ -0,0 +1,51 @@ +// 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. + +suite("insert_with_invalid_array") { + sql "set enable_common_expr_pushdown=true;" + sql "drop table if exists insert_with_invalid_array" + test { + sql """ + CREATE TABLE insert_with_invalid_array ( + id INT NOT NULL COMMENT "", + vec ARRAY NOT NULL COMMENT "", + INDEX ann_idx (vec) USING ANN PROPERTIES( + "index_type" = "hnsw", + "metric_type" = "l2_distance", + "dim" = "3" + ) + ) ENGINE=OLAP + DUPLICATE KEY(id) COMMENT "OLAP" + DISTRIBUTED BY HASH(id) BUCKETS AUTO + PROPERTIES ( + "replication_num" = "1" + ); + """ + } + + sql "insert into insert_with_invalid_array values (1, [1.0, 2.0, 3.0])" + + qt_sql "select * from insert_with_invalid_array order by id" + + // Insert with invalid array + test { + sql """ + INSERT INTO insert_with_invalid_array VALUES (1, [1.0]) + """ + exception "[INVALID_ARGUMENT]" + } +} \ No newline at end of file diff --git a/regression-test/suites/ann_index_p0/memtbl_on_sink.groovy b/regression-test/suites/ann_index_p0/memtbl_on_sink.groovy new file mode 100644 index 00000000000000..efa88a3a5dc279 --- /dev/null +++ b/regression-test/suites/ann_index_p0/memtbl_on_sink.groovy @@ -0,0 +1,69 @@ +// 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. + +suite("memtbl_on_sink") { + sql "set enable_common_expr_pushdown=true;" + sql "drop table if exists memtbl_on_sink" + test { + sql """ + CREATE TABLE memtbl_on_sink ( + id INT NOT NULL COMMENT "", + vec ARRAY NOT NULL COMMENT "", + INDEX ann_idx (vec) USING ANN PROPERTIES( + "index_type" = "hnsw", + "metric_type" = "l2_distance", + "dim" = "3" + ) + ) ENGINE=OLAP + DUPLICATE KEY(id) COMMENT "OLAP" + DISTRIBUTED BY HASH(id) BUCKETS AUTO + PROPERTIES ( + "replication_num" = "1" + ); + """ + } + + sql "insert into memtbl_on_sink values (1, [1.0, 2.0, 3.0])" + + qt_sql_0 "select * from memtbl_on_sink order by id" + + // Insert with invalid array + test { + sql """ + INSERT INTO memtbl_on_sink VALUES (1, [1.0]) + """ + exception "[INVALID_ARGUMENT]" + } + + sql "truncate table memtbl_on_sink;" + sql "set enable_memtable_on_sink_node=false;" + + sql "insert into memtbl_on_sink values (1, [1.0, 2.0, 3.0])" + + qt_sql_1 "select * from memtbl_on_sink order by id" + + // Insert with invalid array + test { + sql """ + INSERT INTO memtbl_on_sink VALUES (1, [1.0]) + """ + exception "[INVALID_ARGUMENT]" + } + + qt_sql_2 "select id, l2_distance_approximate(vec, [1.0, 2.0, 3.0]) as dist from memtbl_on_sink order by dist limit 2;" + +} \ No newline at end of file