Add weighted vector search

Author: Ngalstyan4

sql/lantern.sql

@@ -110,7 +110,6 @@ $$ LANGUAGE plpgsql VOLATILE;
 -- Create access method
 DO $BODY$
 DECLARE
-	hnsw_am_exists boolean;
 	pgvector_exists boolean;
 BEGIN
 	-- Check if the vector type from pgvector exists
@@ -593,4 +592,90 @@ $async_tasks_related$ LANGUAGE plpgsql;
 SELECT _lantern_internal.maybe_setup_lantern_tasks();
 DROP FUNCTION _lantern_internal.maybe_setup_lantern_tasks();
 
+-- Weighted vector search
 -- Asynchronous task scheduling BEGIN
+CREATE OR REPLACE FUNCTION _lantern_internal.maybe_setup_weighted_vector_search() RETURNS VOID AS
+$weighted_vector_search$
+DECLARE
+	pgvector_exists boolean;
+BEGIN
+	-- Check if the vector type from pgvector exists
+	SELECT EXISTS (
+		SELECT 1
+		FROM pg_type
+		WHERE typname = 'vector'
+	) INTO pgvector_exists;
+
+  IF NOT pgvector_exists THEN
+    RAISE NOTICE 'pgvector extension not found. Skipping lantern weighted vector search setup';
+    RETURN;
+  END IF;
+
+  CREATE OR REPLACE FUNCTION lantern.weighted_vector_search(
+    relation_type anyelement,
+    max_dist numeric,
+    w1 numeric,
+    col1 text,
+    vec1 vector,
+    w2 numeric= 0,
+    col2 text = NULL,
+    vec2 vector = NULL,
+    w3 numeric = 0,
+    col3 text = NULL,
+    vec3 vector = NULL)
+    RETURNS TABLE ("row" anyelement) AS
+  $$
+  DECLARE
+    joint_condition text;
+    query_base text;
+    query1 text;
+    query2 text;
+    query3 text;
+    old_hnsw_ef_search numeric;
+    qlimit integer := 1000;
+  BEGIN
+    qlimit := 500;
+    -- Joint similarity metric condition
+    joint_condition := format('(%s * (%I <=> %L))', w1, col1, vec1);
+    IF w2 > 0 AND col2 IS NOT NULL AND vec2 IS NOT NULL THEN
+      joint_condition := joint_condition || format(' + (%s * (%I <=> %L))', w2, col2, vec2);
+    END IF;
+    IF w3 > 0 AND col3 IS NOT NULL AND vec3 IS NOT NULL THEN
+      joint_condition := joint_condition || format(' + (%s * (%I <=> %L))', w3, col3, vec3);
+    END IF;
+
+    -- Base query with joint similarity metric
+    query_base := format('SELECT *
+    FROM %s WHERE %s < %L', pg_typeof(relation_type), joint_condition, max_dist);
+
+    -- extend hnsw depth to the maximum
+   -- old_hnsw_ef_search := hnsw.ef_search;
+    EXECUTE format('SET hnsw.ef_search TO %L', qlimit);
+
+
+    -- Query 1: Order by first condition's weighted similarity
+    query1 := format('%s ORDER BY %I <=> %L LIMIT %L', query_base, col1, vec1, qlimit);
+    RETURN QUERY EXECUTE query1;
+
+
+    -- Query 2: Order by other conditions' weighted similarity, if applicable
+
+    IF w2 > 0 AND col2 IS NOT NULL AND vec2 IS NOT NULL THEN
+      query2 := format('%s ORDER BY %I <=> %L LIMIT %L', query_base, col2, vec2, qlimit);
+      RETURN QUERY EXECUTE query2;
+    END IF;
+
+    IF w3 > 0 AND col3 IS NOT NULL AND vec3 IS NOT NULL THEN
+      query3 := format('%s ORDER BY %I <=> %L LIMIT %L', query_base, col3, vec3, qlimit);
+      RETURN QUERY EXECUTE query3;
+    END IF;
+
+    -- SET hnsw.ef_search = old_hnsw_ef_search;
+  END
+  $$ LANGUAGE plpgsql;
+END
+$weighted_vector_search$ LANGUAGE plpgsql;
+
+SELECT _lantern_internal.maybe_setup_weighted_vector_search();
+DROP FUNCTION _lantern_internal.maybe_setup_weighted_vector_search;
+

test/expected/hnsw_vector.out

@@ -5,9 +5,6 @@
 -- pgvector if it is present on initialization
 DROP EXTENSION IF EXISTS lantern;
 CREATE EXTENSION IF NOT EXISTS vector;
--- Setting min messages to ERROR so the WARNING about existing hnsw access method is NOT printed
--- in tests. This makes sure that regression tests pass on pgvector <=0.4.4 as well as >=0.5.0
-SET client_min_messages=ERROR;
 CREATE EXTENSION lantern;
 RESET client_min_messages;
 SET lantern.pgvector_compat=FALSE;

test/expected/weighted_search.out

@@ -0,0 +1,61 @@
+DROP EXTENSION IF EXISTS lantern;
+CREATE EXTENSION IF NOT EXISTS vector;
+\ir utils/sift1k_vector.sql
+CREATE TABLE IF NOT EXISTS sift_base1k (
+    id SERIAL,
+    v VECTOR(128)
+);
+COPY sift_base1k (v) FROM '/tmp/lantern/vector_datasets/sift_base1k.csv' WITH CSV;
+CREATE INDEX ON sift_base1k USING hnsw (v vector_l2_ops) WITH (M=5, ef_construction=20);
+-- Note: We drop the Lantern extension and re-create it because Lantern only supports
+-- pgvector if it is present on initialization
+CREATE EXTENSION lantern;
+-- create variables with 4th and 444th vector and find closest 10 IDs to each
+SELECT v as v4 FROM sift_Base1k WHERE id = 4 \gset
+SELECT v as v444 FROM sift_Base1k WHERE id = 444 \gset
+SELECT id, ROUND((v <-> :'v4')::numeric, 2) as dist FROM sift_Base1k ORDER BY v <-> :'v4' LIMIT 10;
+ id  |  dist  
+-----+--------
+   4 |   0.00
+   2 | 122.45
+  15 | 141.39
+   8 | 226.66
+ 163 | 244.31
+   6 | 249.80
+  63 | 257.77
+ 183 | 259.18
+ 254 | 263.45
+ 116 | 264.64
+(10 rows)
+
+SELECT id, ROUND((v <-> :'v444')::numeric, 2) as dist FROM sift_Base1k ORDER BY v <-> :'v444' LIMIT 10;
+ id  |  dist  
+-----+--------
+ 444 |   0.00
+ 830 | 214.16
+ 854 | 298.42
+ 557 | 302.76
+  62 | 305.77
+  58 | 306.00
+  30 | 307.27
+ 358 | 308.79
+  60 | 309.94
+ 591 | 310.15
+(10 rows)
+
+-- SELECT id, v <-> :'v4' as v4_dist, v <-> :'v444' v444_dist, weighted_dist FROM lantern.weighted_vector_search(CAST(NULL as "sift_base1k"), max_dist => 20., w1=> 1., col1=>'v'::text, vec1=>:'v4'::vector, w2=> 1., col2=>'v'::text, vec2=>:'v444'::vector) as v4_weighted_search LIMIT 10
+SELECT id, ROUND((v <-> :'v4')::numeric, 2) as v4_dist, ROUND((v <-> :'v444')::numeric, 2) v444_dist FROM lantern.weighted_vector_search(CAST(NULL as "sift_base1k"), max_dist => 20., w1=> 1., col1=>'v'::text, vec1=>:'v4'::vector, w2=> 1., col2=>'v'::text, vec2=>:'v444'::vector) as v4_weighted_search LIMIT 10
+ id  | v4_dist | v444_dist 
+-----+---------+-----------
+   4 |    0.00 |    390.54
+   2 |  122.45 |    380.85
+  15 |  141.39 |    404.62
+   8 |  226.66 |    345.44
+ 163 |  244.31 |    401.58
+   6 |  249.80 |    417.95
+  63 |  257.77 |    338.20
+ 183 |  259.18 |    383.49
+ 254 |  263.45 |    392.81
+ 116 |  264.64 |    398.24
+(10 rows)
+

test/schedule.txt

@@ -6,5 +6,5 @@
 test:   hnsw_config hnsw_correct hnsw_create hnsw_create_expr hnsw_dist_func hnsw_insert hnsw_select hnsw_todo hnsw_index_from_file hnsw_cost_estimate ext_relocation                hnsw_failure_point hnsw_operators hnsw_blockmap_create hnsw_create_unlogged hnsw_insert_unlogged hnsw_logged_unlogged missing_outer_snapshot_portal hnsw_pq
 test_pg_cron: async_tasks
 test: hnsw_pq_index
-test_pgvector: hnsw_vector
+test_pgvector: hnsw_vector weighted_search
 test_extras: hnsw_extras

test/sql/hnsw_vector.sql

@@ -6,9 +6,6 @@
 -- pgvector if it is present on initialization
 DROP EXTENSION IF EXISTS lantern;
 CREATE EXTENSION IF NOT EXISTS vector;
--- Setting min messages to ERROR so the WARNING about existing hnsw access method is NOT printed
--- in tests. This makes sure that regression tests pass on pgvector <=0.4.4 as well as >=0.5.0
-SET client_min_messages=ERROR;
 CREATE EXTENSION lantern;
 RESET client_min_messages;
 SET lantern.pgvector_compat=FALSE;

test/sql/weighted_search.sql

@@ -0,0 +1,19 @@
+DROP EXTENSION IF EXISTS lantern;
+CREATE EXTENSION IF NOT EXISTS vector;
+
+\ir utils/sift1k_vector.sql
+CREATE INDEX ON sift_base1k USING hnsw (v vector_l2_ops) WITH (M=5, ef_construction=20);
+
+-- Note: We drop the Lantern extension and re-create it because Lantern only supports
+-- pgvector if it is present on initialization
+CREATE EXTENSION lantern;
+
+-- create variables with 4th and 444th vector and find closest 10 IDs to each
+SELECT v as v4 FROM sift_Base1k WHERE id = 4 \gset
+SELECT v as v444 FROM sift_Base1k WHERE id = 444 \gset
+
+SELECT id, ROUND((v <-> :'v4')::numeric, 2) as dist FROM sift_Base1k ORDER BY v <-> :'v4' LIMIT 10;
+
+SELECT id, ROUND((v <-> :'v444')::numeric, 2) as dist FROM sift_Base1k ORDER BY v <-> :'v444' LIMIT 10;
+-- SELECT id, v <-> :'v4' as v4_dist, v <-> :'v444' v444_dist, weighted_dist FROM lantern.weighted_vector_search(CAST(NULL as "sift_base1k"), max_dist => 20., w1=> 1., col1=>'v'::text, vec1=>:'v4'::vector, w2=> 1., col2=>'v'::text, vec2=>:'v444'::vector) as v4_weighted_search LIMIT 10
+SELECT id, ROUND((v <-> :'v4')::numeric, 2) as v4_dist, ROUND((v <-> :'v444')::numeric, 2) v444_dist FROM lantern.weighted_vector_search(CAST(NULL as "sift_base1k"), max_dist => 20., w1=> 1., col1=>'v'::text, vec1=>:'v4'::vector, w2=> 1., col2=>'v'::text, vec2=>:'v444'::vector) as v4_weighted_search LIMIT 10