Add hnsw.ef_search variable to allow changing ef at runtime (#199)

lanterndata · Oct 25, 2023 · be86766 · be86766
1 parent e68e849
commit be86766
Show file tree

Hide file tree

Showing 7 changed files with 289 additions and 7 deletions.
diff --git a/src/hnsw/options.c b/src/hnsw/options.c
@@ -28,6 +28,7 @@
 static relopt_kind ldb_hnsw_index_withopts;
 
 int ldb_hnsw_init_k;
+int ldb_hnsw_ef_search;
 
 // this variable is only set during testing and controls whether
 // certain elog() calls are made
@@ -188,7 +189,7 @@ void _PG_init(void)
 
     add_int_reloption(ldb_hnsw_index_withopts,
                       "ef",
-                      "HNSW ef-construction hyperparameter",
+                      "HNSW ef-search hyperparameter",
                       HNSW_DEFAULT_EF,
                       1,
                       HNSW_MAX_EF
@@ -220,6 +221,19 @@ void _PG_init(void)
                             NULL,
                             NULL);
 
+    DefineCustomIntVariable("lantern_hnsw.ef_search",
+                            "Expansion factor to use during vector search in a scan",
+                            "Valid values are in range [1, 400]",
+                            &ldb_hnsw_ef_search,
+                            USEARCH_SEARCH_EF_INVALID_VALUE,
+                            1,
+                            HNSW_MAX_EF,
+                            PGC_USERSET,
+                            0,
+                            NULL,
+                            NULL,
+                            NULL);
+
     DefineCustomBoolVariable("_lantern_internal.is_test",
                              "Whether or not the DB is in a regression test",
                              "set this to 1 to enable extra logging for use in lanterndb regression tests",

diff --git a/src/hnsw/options.h b/src/hnsw/options.h
@@ -52,6 +52,7 @@ usearch_metric_kind_t ldb_HnswGetMetricKind(Relation index);
 bytea* ldb_amoptions(Datum reloptions, bool validate);
 
 extern int  ldb_hnsw_init_k;
+extern int  ldb_hnsw_ef_search;
 extern bool ldb_is_test;
 
 #endif  // LDB_HNSW_OPTIONS_H
diff --git a/src/hnsw/scan.c b/src/hnsw/scan.c
@@ -160,6 +160,8 @@ bool ldb_amgettuple(IndexScanDesc scan, ScanDirection dir)
     // about the furtheest neighbors
     Assert(ScanDirectionIsForward(dir));
 
+    int ef = ldb_hnsw_ef_search;  // 0 if not set, but we pass it into usearch_custom_ef anyway since 0 is also a
+                                  // sentinel value there
     if(scanstate->first) {
         int             num_returned;
         Datum           value;
@@ -199,8 +201,14 @@ bool ldb_amgettuple(IndexScanDesc scan, ScanDirection dir)
                  k,
                  "index size exceeded work_mem during scan, consider increasing work_mem");
         ldb_dlog("LANTERN querying index for %d elements", k);
-        num_returned = usearch_search(
-            scanstate->usearch_index, vec, usearch_scalar_f32_k, k, scanstate->labels, scanstate->distances, &error);
+        num_returned = usearch_search(scanstate->usearch_index,
+                                      vec,
+                                      usearch_scalar_f32_k,
+                                      k,
+                                      ef,
+                                      scanstate->labels,
+                                      scanstate->distances,
+                                      &error);
         ldb_wal_retriever_area_reset(scanstate->retriever_ctx, NULL);
 
         scanstate->count = num_returned;
@@ -240,8 +248,14 @@ bool ldb_amgettuple(IndexScanDesc scan, ScanDirection dir)
                  "index size exceeded work_mem during scan, consider increasing work_mem");
 
         ldb_dlog("LANTERN - querying index for %d elements", k);
-        num_returned = usearch_search(
-            scanstate->usearch_index, vec, usearch_scalar_f32_k, k, scanstate->labels, scanstate->distances, &error);
+        num_returned = usearch_search(scanstate->usearch_index,
+                                      vec,
+                                      usearch_scalar_f32_k,
+                                      k,
+                                      ef,
+                                      scanstate->labels,
+                                      scanstate->distances,
+                                      &error);
         ldb_wal_retriever_area_reset(scanstate->retriever_ctx, NULL);
 
         scanstate->count = num_returned;

diff --git a/test/expected/hnsw_ef_search.out b/test/expected/hnsw_ef_search.out
@@ -0,0 +1,197 @@
+------------------------------------------------------------------------------
+-- Test changing lantern_hnsw.ef_search variable at runtime 
+------------------------------------------------------------------------------
+\ir utils/sift1k_array.sql
+CREATE TABLE IF NOT EXISTS sift_base1k (
+    id SERIAL,
+    v REAL[]
+);
+COPY sift_base1k (v) FROM '/tmp/lantern/vector_datasets/sift_base1k_arrays.csv' WITH csv;
+CREATE INDEX hnsw_l2_index ON sift_base1k USING hnsw (v) WITH (_experimental_index_path='/tmp/lantern/files/index-sift1k-l2.usearch');
+INFO:  done init usearch index
+INFO:  done loading usearch index
+INFO:  done saving 1000 vectors
+SELECT * FROM ldb_get_indexes('sift_base1k');
+   indexname   |  size  |                                                                   indexdef                                                                   | total_index_size 
+---------------+--------+----------------------------------------------------------------------------------------------------------------------------------------------+------------------
+ hnsw_l2_index | 720 kB | CREATE INDEX hnsw_l2_index ON public.sift_base1k USING hnsw (v) WITH (_experimental_index_path='/tmp/lantern/files/index-sift1k-l2.usearch') | 720 kB
+(1 row)
+
+INSERT INTO sift_base1k (id, v) VALUES 
+(1001, array_fill(1, ARRAY[128])),
+(1002, array_fill(2, ARRAY[128]));
+-- Validate error on invalid ef_search values
+\set ON_ERROR_STOP off
+SET lantern_hnsw.ef_search = -1;
+ERROR:  -1 is outside the valid range for parameter "lantern_hnsw.ef_search" (1 .. 400)
+SET lantern_hnsw.ef_search = 0;
+ERROR:  0 is outside the valid range for parameter "lantern_hnsw.ef_search" (1 .. 400)
+SET lantern_hnsw.ef_search = 401;
+ERROR:  401 is outside the valid range for parameter "lantern_hnsw.ef_search" (1 .. 400)
+\set ON_ERROR_STOP on
+-- Repeat the same query while varying ef parameter
+-- NOTE: it is not entirely known if the results of these are deterministic
+SET enable_seqscan = false;
+SELECT v AS v1001 FROM sift_base1k WHERE id = 1001 \gset
+-- Queries below have the same result
+SET lantern_hnsw.ef_search = 1;
+SELECT ROUND(l2sq_dist(v, :'v1001')::numeric, 2) FROM sift_base1k order by v <-> :'v1001' LIMIT 10;
+   round   
+-----------
+      0.00
+    128.00
+ 249249.00
+ 249285.00
+ 249418.00
+ 249515.00
+ 249589.00
+ 249647.00
+ 249652.00
+ 249675.00
+(10 rows)
+
+SET lantern_hnsw.ef_search = 2;
+SELECT ROUND(l2sq_dist(v, :'v1001')::numeric, 2) FROM sift_base1k order by v <-> :'v1001' LIMIT 10;
+   round   
+-----------
+      0.00
+    128.00
+ 249249.00
+ 249285.00
+ 249418.00
+ 249515.00
+ 249589.00
+ 249647.00
+ 249652.00
+ 249675.00
+(10 rows)
+
+SET lantern_hnsw.ef_search = 4;
+SELECT ROUND(l2sq_dist(v, :'v1001')::numeric, 2) FROM sift_base1k order by v <-> :'v1001' LIMIT 10;
+   round   
+-----------
+      0.00
+    128.00
+ 249249.00
+ 249285.00
+ 249418.00
+ 249515.00
+ 249589.00
+ 249647.00
+ 249652.00
+ 249675.00
+(10 rows)
+
+SET lantern_hnsw.ef_search = 8;
+SELECT ROUND(l2sq_dist(v, :'v1001')::numeric, 2) FROM sift_base1k order by v <-> :'v1001' LIMIT 10;
+   round   
+-----------
+      0.00
+    128.00
+ 249249.00
+ 249285.00
+ 249418.00
+ 249515.00
+ 249589.00
+ 249647.00
+ 249652.00
+ 249675.00
+(10 rows)
+
+SET lantern_hnsw.ef_search = 16;
+SELECT ROUND(l2sq_dist(v, :'v1001')::numeric, 2) FROM sift_base1k order by v <-> :'v1001' LIMIT 10;
+   round   
+-----------
+      0.00
+    128.00
+ 249249.00
+ 249285.00
+ 249418.00
+ 249515.00
+ 249589.00
+ 249647.00
+ 249652.00
+ 249675.00
+(10 rows)
+
+-- Queries below have the same result, which is different from above
+SET lantern_hnsw.ef_search = 32;
+SELECT ROUND(l2sq_dist(v, :'v1001')::numeric, 2) FROM sift_base1k order by v <-> :'v1001' LIMIT 10;
+   round   
+-----------
+      0.00
+    128.00
+ 249249.00
+ 249285.00
+ 249418.00
+ 249457.00
+ 249515.00
+ 249589.00
+ 249647.00
+ 249652.00
+(10 rows)
+
+SET lantern_hnsw.ef_search = 64;
+SELECT ROUND(l2sq_dist(v, :'v1001')::numeric, 2) FROM sift_base1k order by v <-> :'v1001' LIMIT 10;
+   round   
+-----------
+      0.00
+    128.00
+ 249249.00
+ 249285.00
+ 249418.00
+ 249457.00
+ 249515.00
+ 249589.00
+ 249647.00
+ 249652.00
+(10 rows)
+
+SET lantern_hnsw.ef_search = 128;
+SELECT ROUND(l2sq_dist(v, :'v1001')::numeric, 2) FROM sift_base1k order by v <-> :'v1001' LIMIT 10;
+   round   
+-----------
+      0.00
+    128.00
+ 249249.00
+ 249285.00
+ 249418.00
+ 249457.00
+ 249515.00
+ 249589.00
+ 249647.00
+ 249652.00
+(10 rows)
+
+SET lantern_hnsw.ef_search = 256;
+SELECT ROUND(l2sq_dist(v, :'v1001')::numeric, 2) FROM sift_base1k order by v <-> :'v1001' LIMIT 10;
+   round   
+-----------
+      0.00
+    128.00
+ 249249.00
+ 249285.00
+ 249418.00
+ 249457.00
+ 249515.00
+ 249589.00
+ 249647.00
+ 249652.00
+(10 rows)
+
+SET lantern_hnsw.ef_search = 400;
+SELECT ROUND(l2sq_dist(v, :'v1001')::numeric, 2) FROM sift_base1k order by v <-> :'v1001' LIMIT 10;
+   round   
+-----------
+      0.00
+    128.00
+ 249249.00
+ 249285.00
+ 249418.00
+ 249457.00
+ 249515.00
+ 249589.00
+ 249647.00
+ 249652.00
+(10 rows)
+
diff --git a/test/schedule.txt b/test/schedule.txt
@@ -4,4 +4,4 @@
 # - 'test' lines may have multiple space-separated tests. All tests in a single 'test' line will be run in parallel
 
 test_pgvector: hnsw_vector
-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
+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_ef_search
diff --git a/test/sql/hnsw_ef_search.sql b/test/sql/hnsw_ef_search.sql
@@ -0,0 +1,56 @@
+------------------------------------------------------------------------------
+-- Test changing lantern_hnsw.ef_search variable at runtime 
+------------------------------------------------------------------------------
+
+\ir utils/sift1k_array.sql
+
+CREATE INDEX hnsw_l2_index ON sift_base1k USING hnsw (v) WITH (_experimental_index_path='/tmp/lantern/files/index-sift1k-l2.usearch');
+SELECT * FROM ldb_get_indexes('sift_base1k');
+
+INSERT INTO sift_base1k (id, v) VALUES 
+(1001, array_fill(1, ARRAY[128])),
+(1002, array_fill(2, ARRAY[128]));
+
+-- Validate error on invalid ef_search values
+\set ON_ERROR_STOP off
+SET lantern_hnsw.ef_search = -1;
+SET lantern_hnsw.ef_search = 0;
+SET lantern_hnsw.ef_search = 401;
+\set ON_ERROR_STOP on
+
+-- Repeat the same query while varying ef parameter
+-- NOTE: it is not entirely known if the results of these are deterministic
+SET enable_seqscan = false;
+SELECT v AS v1001 FROM sift_base1k WHERE id = 1001 \gset
+
+-- Queries below have the same result
+SET lantern_hnsw.ef_search = 1;
+SELECT ROUND(l2sq_dist(v, :'v1001')::numeric, 2) FROM sift_base1k order by v <-> :'v1001' LIMIT 10;
+
+SET lantern_hnsw.ef_search = 2;
+SELECT ROUND(l2sq_dist(v, :'v1001')::numeric, 2) FROM sift_base1k order by v <-> :'v1001' LIMIT 10;
+
+SET lantern_hnsw.ef_search = 4;
+SELECT ROUND(l2sq_dist(v, :'v1001')::numeric, 2) FROM sift_base1k order by v <-> :'v1001' LIMIT 10;
+
+SET lantern_hnsw.ef_search = 8;
+SELECT ROUND(l2sq_dist(v, :'v1001')::numeric, 2) FROM sift_base1k order by v <-> :'v1001' LIMIT 10;
+
+SET lantern_hnsw.ef_search = 16;
+SELECT ROUND(l2sq_dist(v, :'v1001')::numeric, 2) FROM sift_base1k order by v <-> :'v1001' LIMIT 10;
+
+-- Queries below have the same result, which is different from above
+SET lantern_hnsw.ef_search = 32;
+SELECT ROUND(l2sq_dist(v, :'v1001')::numeric, 2) FROM sift_base1k order by v <-> :'v1001' LIMIT 10;
+
+SET lantern_hnsw.ef_search = 64;
+SELECT ROUND(l2sq_dist(v, :'v1001')::numeric, 2) FROM sift_base1k order by v <-> :'v1001' LIMIT 10;
+
+SET lantern_hnsw.ef_search = 128;
+SELECT ROUND(l2sq_dist(v, :'v1001')::numeric, 2) FROM sift_base1k order by v <-> :'v1001' LIMIT 10;
+
+SET lantern_hnsw.ef_search = 256;
+SELECT ROUND(l2sq_dist(v, :'v1001')::numeric, 2) FROM sift_base1k order by v <-> :'v1001' LIMIT 10;
+
+SET lantern_hnsw.ef_search = 400;
+SELECT ROUND(l2sq_dist(v, :'v1001')::numeric, 2) FROM sift_base1k order by v <-> :'v1001' LIMIT 10;
diff --git a/third_party/usearch b/third_party/usearch
+32 −9		c/lib.cpp
+2 −1		c/usearch.h
+17 −0		include/usearch/index_punned_dense.hpp