Enhance LanceDB index creation and search functionality with improved error handling and configuration options

Author: tae898

bindings/python/examples/11_vector_index_build.py

@@ -546,16 +546,54 @@ def create_index_lancedb(
     table,
     max_connections: int,
     beam_width: int,
-) -> None:
-    partitions = max(1, min(256, int(max_connections) * 8))
-    table.create_index(
-        metric="cosine",
-        vector_column_name="vector",
-        index_type="IVF_HNSW_SQ",
-        num_partitions=partitions,
-        m=int(max_connections),
-        ef_construction=int(beam_width),
-    )
+) -> dict:
+    common_kwargs = {
+        "metric": "cosine",
+        "vector_column_name": "vector",
+        "m": int(max_connections),
+        "ef_construction": int(beam_width),
+    }
+    attempts = [
+        ("HNSW", {}),
+        ("IVF_HNSW_SQ", {"num_partitions": 1}),
+    ]
+    last_error: Exception | None = None
+
+    for index_type, extra_kwargs in attempts:
+        try:
+            table.create_index(
+                index_type=index_type,
+                **common_kwargs,
+                **extra_kwargs,
+            )
+            config = {
+                "metric": "cosine",
+                "index_type": index_type,
+                "hnsw_m": int(max_connections),
+                "hnsw_ef_construct": int(beam_width),
+            }
+            if "num_partitions" in extra_kwargs:
+                config["num_partitions"] = int(extra_kwargs["num_partitions"])
+                config["quantization"] = "SQ"
+            else:
+                config["quantization"] = "NONE"
+
+            print(
+                "Using LanceDB index mode: "
+                f"{index_type}"
+                + (
+                    f" (num_partitions={extra_kwargs['num_partitions']})"
+                    if "num_partitions" in extra_kwargs
+                    else ""
+                )
+            )
+            return config
+        except (AttributeError, RuntimeError, TypeError, ValueError) as exc:
+            last_error = exc
+
+    raise RuntimeError(
+        "Unable to create LanceDB index in HNSW-like mode"
+    ) from last_error
 
 
 def vector_to_pg_literal(vec: np.ndarray) -> str:
@@ -1825,6 +1863,7 @@ def record(name: str, result, dur: float, rss_start: float, rss_end: float):
             phases[name].update(result)
 
     runtime_versions: dict[str, str | None] = {}
+    lancedb_index_config: dict[str, object] | None = None
 
     if args.backend == "arcadedb_sql":
         stop_cpu = start_cpu_logger(2)
@@ -1969,15 +2008,15 @@ def record(name: str, result, dur: float, rss_start: float, rss_end: float):
             )
             record("ingest", {"ingested": int(ingested)}, dur, r0, r1)
 
-            (_, dur, r0, r1) = timed_section(
+            (lancedb_index_config, dur, r0, r1) = timed_section(
                 "create_index",
                 lambda: create_index_lancedb(
                     table,
                     max_connections=args.max_connections,
                     beam_width=args.beam_width,
                 ),
             )
-            record("create_index", {}, dur, r0, r1)
+            record("create_index", lancedb_index_config, dur, r0, r1)
 
             close_db_fn = getattr(db, "close", None)
             (_, dur, r0, r1) = timed_section(
@@ -2379,8 +2418,18 @@ def server_pid_provider() -> int | None:
             "lancedb": {
                 "data_dir": str(db_path / "lancedb-data"),
                 "table": "vectordata",
-                "metric": "cosine",
-                "index_type": "IVF_HNSW_SQ",
+                "metric": (
+                    str((lancedb_index_config or {}).get("metric"))
+                    if lancedb_index_config
+                    else "cosine"
+                ),
+                "index_type": (
+                    str((lancedb_index_config or {}).get("index_type"))
+                    if lancedb_index_config
+                    else "IVF_HNSW_SQ"
+                ),
+                "num_partitions": ((lancedb_index_config or {}).get("num_partitions")),
+                "quantization": ((lancedb_index_config or {}).get("quantization")),
                 "hnsw_m": args.max_connections,
                 "hnsw_ef_construct": args.beam_width,
             },

bindings/python/examples/12_vector_search.py

@@ -18,7 +18,8 @@
 - ArcadeDB/jvector uses `overquery_factor` directly.
 - HNSW backends expose `ef_search`.
 - These knobs are not semantically identical, so this benchmark uses a fixed
-    normalization: `ef_search = 0.5 * k * overquery_factor` for pgvector/qdrant/milvus.
+    normalization: `ef_search = 0.5 * k * overquery_factor` for
+    faiss/pgvector/qdrant/milvus and LanceDB HNSW-like search tuning.
 """
 
 from __future__ import annotations
@@ -608,18 +609,48 @@ def search_lancedb(
     gt_full: dict[int, List[int]],
     k: int,
     overquery_factor: float,
+    build_config: dict,
 ) -> dict:
     latencies_ms: List[float] = []
     recalls: List[float] = []
 
-    nprobes = max(1, int(round(overquery_factor * 4)))
+    lancedb_cfg = build_config.get("lancedb") if isinstance(build_config, dict) else {}
+    if not isinstance(lancedb_cfg, dict):
+        lancedb_cfg = {}
+
+    index_type = str(lancedb_cfg.get("index_type") or "IVF_HNSW_SQ").upper()
+    ef_search = overquery_to_ef_search(k, overquery_factor)
+    nprobes = None
+    if index_type.startswith("IVF_"):
+        nprobes = (
+            1
+            if int(lancedb_cfg.get("num_partitions") or 1) <= 1
+            else max(1, int(round(overquery_factor * 4)))
+        )
+
+    applied_ef_search = None
+    applied_nprobes = None
 
     for q_idx, qid in enumerate(qids):
         start = time.perf_counter()
         search = table.search(queries[q_idx].tolist()).metric("cosine").limit(int(k))
-        if hasattr(search, "nprobes"):
+        if hasattr(search, "ef"):
+            try:
+                search = search.ef(int(ef_search))
+                applied_ef_search = int(ef_search)
+            except Exception:
+                pass
+        elif hasattr(search, "ef_search"):
+            try:
+                search = search.ef_search(int(ef_search))
+                applied_ef_search = int(ef_search)
+            except Exception:
+                pass
+
+        if nprobes is not None and hasattr(search, "nprobes"):
             try:
                 search = search.nprobes(int(nprobes))
+                applied_nprobes = int(nprobes)
             except Exception:
                 pass
         rows = search.to_list()
@@ -646,6 +677,8 @@ def search_lancedb(
         "latency_ms_mean": lat_mean,
         "latency_ms_p95": lat_p95,
         "recall_count": len(recalls),
+        "effective_ef_search": applied_ef_search,
+        "effective_nprobes": applied_nprobes,
     }
 
 
@@ -1957,7 +1990,8 @@ def main() -> None:
         default="1,2,3,4,6,8",
         help=(
             "Sweep values. For arcadedb: overquery factors. "
-            "For pgvector/qdrant/milvus: ef_search = 0.5 * k * factor. "
+            "For faiss/pgvector/qdrant/milvus and LanceDB HNSW-like tuning: "
+            "ef_search = 0.5 * k * factor. "
             "Default: 1,2,3,4,6,8"
         ),
     )
@@ -2267,6 +2301,7 @@ def record_phase(
                             gt_full,
                             k=args.k,
                             overquery_factor=overquery,
+                            build_config=build_config,
                         ),
                         queries=queries,
                         qids=qids,
@@ -2287,7 +2322,8 @@ def record_phase(
                 sweeps.append(
                     {
                         "overquery_factor": overquery,
-                        "effective_nprobes": max(1, int(round(overquery * 4))),
+                        "effective_ef_search": stats.get("effective_ef_search"),
+                        "effective_nprobes": stats.get("effective_nprobes"),
                         "phases": phases,
                         "recall_mean": stats.get("recall_mean"),
                         "recall_count": stats.get("recall_count"),
@@ -2765,6 +2801,27 @@ def server_pid_provider() -> int | None:
                 }
                 for factor in overqueries
             ],
+            "lancedb_search_mapping": (
+                [
+                    {
+                        "factor": factor,
+                        "effective_ef_search": overquery_to_ef_search(args.k, factor),
+                        "effective_nprobes": (
+                            1
+                            if str(
+                                ((build_config.get("lancedb") or {}).get("index_type"))
+                                or ""
+                            )
+                            .upper()
+                            .startswith("IVF_")
+                            else None
+                        ),
+                    }
+                    for factor in overqueries
+                ]
+                if args.backend == "lancedb"
+                else None
+            ),
             "sweeps": sweeps,
         },
         "peak_rss_mb": peak_rss_mb,
@@ -2877,10 +2934,13 @@ def server_pid_provider() -> int | None:
         recall_text = f"{recall:.4f}" if recall is not None else "n/a"
         lat_text = f"{lat:.2f}" if lat is not None else "n/a"
         p95_text = f"{p95:.2f}" if p95 is not None else "n/a"
-        if args.backend in {"pgvector", "qdrant", "milvus", "faiss"}:
+        if args.backend in {"pgvector", "qdrant", "milvus", "faiss", "lancedb"}:
             ef_text = str(sweep.get("effective_ef_search"))
+            extra = ""
+            if args.backend == "lancedb" and sweep.get("effective_nprobes") is not None:
+                extra = f" | nprobes={sweep.get('effective_nprobes')}"
             print(
-                f"factor={oq:>4} | ef_search={ef_text} | "
+                f"factor={oq:>4} | ef_search={ef_text}{extra} | "
                 f"recall@{args.k}={recall_text} | latency_mean_ms={lat_text} | "
                 f"latency_p95_ms={p95_text}"
             )

bindings/python/examples/scripts/run_11_vector_index_build_matrix.sh

@@ -15,7 +15,7 @@ source "$HELPERS_SH"
 # Large         10,000  32GB    16
 # X-Large       25,000  64GB    32
 
-DATASET="stackoverflow-large"
+DATASET="stackoverflow-xlarge"
 BATCH_SIZE=10000
 MEM_LIMIT="32g"
 THREADS=4
@@ -49,8 +49,8 @@ MILVUS_PORT=19530
 MILVUS_COMPOSE_VERSION="v2.6.10"
 MILVUS_COLLECTION="vectordata"
 
-# BACKENDS_RAW="arcadedb_sql,lancedb,faiss,pgvector,qdrant,milvus"
-BACKENDS_RAW="lancedb"
+BACKENDS_RAW="lancedb,pgvector,faiss,qdrant,milvus"
+# BACKENDS_RAW="arcadedb_sql"
 LABEL_PREFIX="sweep11"
 
 if [[ $# -gt 0 ]]; then