Optimize lateral VSS joins using HNSW index

.github/workflows/MainDistributionPipeline.yml

@@ -28,6 +28,7 @@ jobs:
     with:
       duckdb_version: main
       extension_name: vss
+      ci_tools_version: main
 
   duckdb-stable-deploy:
     name: Deploy extension binaries

src/hnsw/CMakeLists.txt

@@ -9,8 +9,9 @@ set(EXTENSION_SOURCES
         ${CMAKE_CURRENT_SOURCE_DIR}/hnsw_index_scan.cpp
         ${CMAKE_CURRENT_SOURCE_DIR}/hnsw_plan_index_create.cpp
         ${CMAKE_CURRENT_SOURCE_DIR}/hnsw_topk_operator.cpp
-        ${CMAKE_CURRENT_SOURCE_DIR}/hnsw_optimize_topk.cpp
         ${CMAKE_CURRENT_SOURCE_DIR}/hnsw_optimize_expr.cpp
+        ${CMAKE_CURRENT_SOURCE_DIR}/hnsw_optimize_join.cpp
+        ${CMAKE_CURRENT_SOURCE_DIR}/hnsw_optimize_topk.cpp
         ${CMAKE_CURRENT_SOURCE_DIR}/hnsw_optimize_scan.cpp
         PARENT_SCOPE
 )

src/hnsw/hnsw_index.cpp

@@ -3,12 +3,11 @@
 #include "duckdb/common/serializer/binary_deserializer.hpp"
 #include "duckdb/common/serializer/binary_serializer.hpp"
 #include "duckdb/execution/index/fixed_size_allocator.hpp"
-#include "duckdb/storage/table/scan_state.hpp"
 #include "duckdb/planner/operator/logical_get.hpp"
+#include "duckdb/storage/table/scan_state.hpp"
 #include "hnsw/hnsw.hpp"
 
 namespace duckdb {
-
 //------------------------------------------------------------------------------
 // Linked Blocks
 //------------------------------------------------------------------------------
@@ -209,6 +208,8 @@ HNSWIndex::HNSWIndex(const string &name, IndexConstraintType index_constraint_ty
 		index.reserve(MinValue(static_cast<idx_t>(32), estimated_cardinality));
 	}
 	index_size = index.size();
+
+	function_matcher = MakeFunctionMatcher();
 }
 
 idx_t HNSWIndex::GetVectorSize() const {
@@ -228,24 +229,6 @@ string HNSWIndex::GetMetric() const {
 	}
 }
 
-bool HNSWIndex::MatchesDistanceFunction(const string &name) const {
-	if ((name == "array_distance" || name == "<->") &&
-	    index.metric().metric_kind() == unum::usearch::metric_kind_t::l2sq_k) {
-		// Note: usearch uses l2sq, for their metric, but its functionally equivalent to sqrt(l2sq)
-		return true;
-	}
-	if ((name == "array_cosine_distance" || name == "<=>") &&
-	    index.metric().metric_kind() == unum::usearch::metric_kind_t::cos_k) {
-		return true;
-	}
-	if ((name == "array_negative_inner_product" || name == "<#>") &&
-	    index.metric().metric_kind() == unum::usearch::metric_kind_t::ip_k) {
-		// Note: usearch uses (1.0 - ip) for their metric, but its functionally equivalent to (-ip)
-		return true;
-	}
-	return false;
-}
-
 const case_insensitive_map_t<unum::usearch::metric_kind_t> HNSWIndex::METRIC_KIND_MAP = {
     {"l2sq", unum::usearch::metric_kind_t::l2sq_k},
     {"cosine", unum::usearch::metric_kind_t::cos_k},
@@ -324,11 +307,11 @@ unique_ptr<IndexScanState> HNSWIndex::InitializeScan(float *query_vector, idx_t
 	return std::move(state);
 }
 
-idx_t HNSWIndex::Scan(IndexScanState &state, Vector &result) {
+idx_t HNSWIndex::Scan(IndexScanState &state, Vector &result, idx_t result_offset) {
 	auto &scan_state = state.Cast<HNSWIndexScanState>();
 
 	idx_t count = 0;
-	auto row_ids = FlatVector::GetData<row_t>(result);
+	auto row_ids = FlatVector::GetData<row_t>(result) + result_offset;
 
 	// Push the row ids into the result vector, up to STANDARD_VECTOR_SIZE or the
 	// end of the result set
@@ -339,6 +322,58 @@ idx_t HNSWIndex::Scan(IndexScanState &state, Vector &result) {
 	return count;
 }
 
+struct MultiScanState final : IndexScanState {
+	Vector vec;
+	vector<row_t> row_ids;
+	size_t ef_search;
+	MultiScanState(size_t ef_search_p) : vec(LogicalType::ROW_TYPE, nullptr), ef_search(ef_search_p) {
+	}
+};
+
+unique_ptr<IndexScanState> HNSWIndex::InitializeMultiScan(ClientContext &context) {
+	// Try to get the ef_search parameter from the database or use the default value
+	auto ef_search = index.expansion_search();
+
+	Value hnsw_ef_search_opt;
+	if (context.TryGetCurrentSetting("hnsw_ef_search", hnsw_ef_search_opt)) {
+		if (!hnsw_ef_search_opt.IsNull() && hnsw_ef_search_opt.type() == LogicalType::BIGINT) {
+			const auto val = hnsw_ef_search_opt.GetValue<int64_t>();
+			if (val > 0) {
+				ef_search = static_cast<idx_t>(val);
+			}
+		}
+	}
+	// Return the new state
+	return make_uniq<MultiScanState>(ef_search);
+}
+
+idx_t HNSWIndex::ExecuteMultiScan(IndexScanState &state_p, float *query_vector, idx_t limit) {
+	auto &state = state_p.Cast<MultiScanState>();
+
+	USearchIndexType::search_result_t search_result;
+	{
+		auto lock = rwlock.GetSharedLock();
+		search_result = index.ef_search(query_vector, limit, state.ef_search);
+	}
+
+	const auto offset = state.row_ids.size();
+	state.row_ids.resize(state.row_ids.size() + search_result.size());
+	search_result.dump_to(state.row_ids.data() + offset);
+
+	return search_result.size();
+}
+
+const Vector &HNSWIndex::GetMultiScanResult(IndexScanState &state) {
+	auto &scan_state = state.Cast<MultiScanState>();
+	FlatVector::SetData(scan_state.vec, (data_ptr_t)scan_state.row_ids.data());
+	return scan_state.vec;
+}
+
+void HNSWIndex::ResetMultiScan(IndexScanState &state) {
+	auto &scan_state = state.Cast<MultiScanState>();
+	scan_state.row_ids.clear();
+}
+
 void HNSWIndex::CommitDrop(IndexLock &index_lock) {
 	// Acquire an exclusive lock to drop the index
 	auto lock = rwlock.GetExclusiveLock();
@@ -367,13 +402,15 @@ void HNSWIndex::Construct(DataChunk &input, Vector &row_ids, idx_t thread_idx) {
 	auto vec_child_data = FlatVector::GetData<float>(vec_child_vec);
 	auto rowid_data = FlatVector::GetData<row_t>(row_ids);
 
+	auto to_add_count = FlatVector::Validity(vec_vec).CountValid(count);
+
 	// Check if we need to resize the index
 	// We keep the size of the index in a separate atomic to avoid
 	// locking exclusively when checking
 	bool needs_resize = false;
 	{
 		auto lock = rwlock.GetSharedLock();
-		if (index_size.fetch_add(count) + count > index.capacity()) {
+		if (index_size.fetch_add(to_add_count) + to_add_count > index.capacity()) {
 			needs_resize = true;
 		}
 	}
@@ -394,6 +431,11 @@ void HNSWIndex::Construct(DataChunk &input, Vector &row_ids, idx_t thread_idx) {
 		// Now we can be sure that we have enough space in the index
 		auto lock = rwlock.GetSharedLock();
 		for (idx_t out_idx = 0; out_idx < count; out_idx++) {
+			if(FlatVector::IsNull(vec_vec, out_idx)) {
+				// Dont add nulls
+				continue;
+			}
+
 			auto rowid = rowid_data[out_idx];
 			auto result = index.add(rowid, vec_child_data + (out_idx * array_size), thread_idx);
 			if (!result) {
@@ -472,7 +514,6 @@ void HNSWIndex::PersistToDisk() {
 	}
 
 	// Write
-
 	if (root_block_ptr.Get() == 0) {
 		root_block_ptr = linked_block_allocator->New();
 	}
@@ -536,35 +577,85 @@ void HNSWIndex::VerifyAllocations(IndexLock &state) {
 //------------------------------------------------------------------------------
 // Can rewrite index expression?
 //------------------------------------------------------------------------------
-static void RewriteIndexExpression(const Index &index, LogicalGet &get, Expression &expr, bool &rewrite_possible,
-                                   bool &any_column_ref) {
+static void TryBindIndexExpressionInternal(Expression &expr, idx_t table_idx, const vector<column_t> &index_columns,
+                                           const vector<column_t> &table_columns, bool &success, bool &found) {
+
 	if (expr.type == ExpressionType::BOUND_COLUMN_REF) {
-		any_column_ref = true;
-		auto &bound_colref = expr.Cast<BoundColumnRefExpression>();
-		// bound column ref: rewrite to fit in the current set of bound column ids
-		bound_colref.binding.table_index = get.table_index;
-		auto &column_ids = index.GetColumnIds();
-		auto &get_column_ids = get.GetColumnIds();
-		column_t referenced_column = column_ids[bound_colref.binding.column_index];
-		// search for the referenced column in the set of column_ids
-		for (idx_t i = 0; i < get_column_ids.size(); i++) {
-			if (get_column_ids[i] == referenced_column) {
-				bound_colref.binding.column_index = i;
+		found = true;
+		auto &ref = expr.Cast<BoundColumnRefExpression>();
+
+		// Rewrite the column reference to fit in the current set of bound column ids
+		ref.binding.table_index = table_idx;
+
+		const auto referenced_column = index_columns[ref.binding.column_index];
+		for (idx_t i = 0; i < table_columns.size(); i++) {
+			if (table_columns[i] == referenced_column) {
+				ref.binding.column_index = i;
 				return;
 			}
 		}
-		// column id not found in bound columns in the LogicalGet: rewrite not possible
-		rewrite_possible = false;
+		success = false;
 	}
-	ExpressionIterator::EnumerateChildren(
-	    expr, [&](Expression &child) { RewriteIndexExpression(index, get, child, rewrite_possible, any_column_ref); });
+
+	ExpressionIterator::EnumerateChildren(expr, [&](Expression &child) {
+		TryBindIndexExpressionInternal(child, table_idx, index_columns, table_columns, success, found);
+	});
 }
 
-bool HNSWIndex::CanRewriteIndexExpression(LogicalGet &get, Expression &column_ref) const {
-	bool rewrite_possible = true;
-	bool any_column_ref = false;
-	RewriteIndexExpression(*this, get, column_ref, rewrite_possible, any_column_ref);
-	return any_column_ref && rewrite_possible;
+bool HNSWIndex::TryBindIndexExpression(LogicalGet &get, unique_ptr<Expression> &result) const {
+	auto expr_ptr = unbound_expressions.back()->Copy();
+
+	auto &expr = *expr_ptr;
+	auto &index_columns = GetColumnIds();
+	auto &table_columns = get.GetColumnIds();
+
+	auto success = true;
+	auto found = false;
+
+	TryBindIndexExpressionInternal(expr, get.table_index, index_columns, table_columns, success, found);
+
+	if (success && found) {
+		result = std::move(expr_ptr);
+		return true;
+	}
+	return false;
+}
+
+bool HNSWIndex::TryMatchDistanceFunction(const unique_ptr<Expression> &expr,
+                                         vector<reference<Expression>> &bindings) const {
+	return function_matcher->Match(*expr, bindings);
+}
+
+unique_ptr<ExpressionMatcher> HNSWIndex::MakeFunctionMatcher() const {
+	unordered_set<string> distance_functions;
+	switch (index.metric().metric_kind()) {
+	case unum::usearch::metric_kind_t::l2sq_k:
+		distance_functions = {"array_distance", "<->"};
+		break;
+	case unum::usearch::metric_kind_t::cos_k:
+		distance_functions = {"array_cosine_distance", "<=>"};
+		break;
+	case unum::usearch::metric_kind_t::ip_k:
+		distance_functions = {"array_negative_inner_product", "<#>"};
+		break;
+	default:
+		throw NotImplementedException("Unknown metric kind");
+	}
+
+	auto matcher = make_uniq<FunctionExpressionMatcher>();
+	matcher->function = make_uniq<ManyFunctionMatcher>(distance_functions);
+	matcher->expr_type = make_uniq<SpecificExpressionTypeMatcher>(ExpressionType::BOUND_FUNCTION);
+	matcher->policy = SetMatcher::Policy::UNORDERED;
+
+	auto lhs_matcher = make_uniq<ExpressionMatcher>();
+	lhs_matcher->type = make_uniq<SpecificTypeMatcher>(LogicalType::ARRAY(LogicalType::FLOAT, GetVectorSize()));
+	matcher->matchers.push_back(std::move(lhs_matcher));
+
+	auto rhs_matcher = make_uniq<ExpressionMatcher>();
+	rhs_matcher->type = make_uniq<SpecificTypeMatcher>(LogicalType::ARRAY(LogicalType::FLOAT, GetVectorSize()));
+	matcher->matchers.push_back(std::move(rhs_matcher));
+
+	return std::move(matcher);
 }
 
 //------------------------------------------------------------------------------