docs: add doc and test for 4bit PQ (#3212)

lancedb · Dec 10, 2024 · ef9d0c2 · ef9d0c2
1 parent 10c31b3
commit ef9d0c2
Show file tree

Hide file tree

Showing 2 changed files with 18 additions and 0 deletions.
diff --git a/python/python/lance/dataset.py b/python/python/lance/dataset.py
@@ -1668,6 +1668,9 @@ def create_index(
         Optional parameters for "IVF_PQ":
             ivf_centroids :
                 K-mean centroids for IVF clustering.
+            num_bits : int, optional
+                The number of bits for PQ (Product Quantization). Default is 8.
+                Only 4, 8 are supported.
 
         Optional parameters for "IVF_HNSW_*":
             max_level : int

diff --git a/python/python/tests/test_vector_index.py b/python/python/tests/test_vector_index.py
@@ -386,6 +386,21 @@ def test_create_dot_index(dataset, tmp_path):
     assert ann_ds.has_index
 
 
+def test_create_4bit_ivf_pq_index(dataset, tmp_path):
+    assert not dataset.has_index
+    ann_ds = lance.write_dataset(dataset.to_table(), tmp_path / "indexed.lance")
+    ann_ds = ann_ds.create_index(
+        "vector",
+        index_type="IVF_PQ",
+        num_partitions=1,
+        num_sub_vectors=16,
+        num_bits=4,
+        metric="l2",
+    )
+    index = ann_ds.stats.index_stats("vector_idx")
+    assert index["indices"][0]["sub_index"]["nbits"] == 4
+
+
 def test_create_ivf_hnsw_pq_index(dataset, tmp_path):
     assert not dataset.has_index
     ann_ds = lance.write_dataset(dataset.to_table(), tmp_path / "indexed.lance")