Support setting MariaDB HNSW index parameters

Setting the parameters through session variables
HugoWenTD · Apr 26, 2024 · e09784e · e09784e
1 parent 5a7f75f
commit e09784e
Show file tree

Hide file tree

Showing 2 changed files with 9 additions and 8 deletions.
diff --git a/ann_benchmarks/algorithms/mariadb/config.yml b/ann_benchmarks/algorithms/mariadb/config.yml
@@ -7,11 +7,11 @@ float:
     module: ann_benchmarks.algorithms.mariadb
     name: mariadb
     run_groups:
-      M-16:
-        arg_groups: [{M: 16, efConstruction: 200}]
+      M-50:
+        arg_groups: [{M: 50, efConstruction: 10}]
         args: {}
-        query_args: [[10, 20, 40, 80, 120, 200, 400, 800]]
-      M-24:
-        arg_groups: [{M: 24, efConstruction: 200}]
+        query_args: [[10]]
+      PGVector-default:
+        arg_groups: [{M: 16, efConstruction: 64}]
         args: {}
-        query_args: [[10, 20, 40, 80, 120, 200, 400, 800]]
+        query_args: [[40]]
diff --git a/ann_benchmarks/algorithms/mariadb/module.py b/ann_benchmarks/algorithms/mariadb/module.py
@@ -248,6 +248,8 @@ def fit(self, X):
         self._cur.execute("DROP DATABASE IF EXISTS ann")
         self._cur.execute("CREATE DATABASE ann")
         self._cur.execute("USE ann")
+        self._cur.execute("SET hnsw_max_connection_per_layer = %d" % self._m)
+        self._cur.execute("SET hnsw_ef_constructor = %d" % self._ef_construction)
         # Innodb create table with index is not supported with the latest commit of the develop branch.
         # Once all supported we could use:
         #self._cur.execute("CREATE TABLE t1 (id INT PRIMARY KEY, v BLOB NOT NULL, vector INDEX (v)) ENGINE=InnoDB;")
@@ -284,8 +286,7 @@ def fit(self, X):
     def set_query_arguments(self, ef_search):
         # Set ef_search
         self._ef_search = ef_search
-        # Not supported by MariaDB at the moment
-        #self._cur.execute("SET hnsw.ef_search = %d" % ef_search)
+        self._cur.execute("SET hnsw_ef_search = %d" % ef_search)
 
     def query(self, v, n):
         self._cur.execute("SELECT id FROM t1 ORDER by vec_distance(v, %s) LIMIT %d", (bytes(vector_to_hex(v)), n))