feat(index): expose layered RaBitQ defaults across SDKs

python/python/lancedb/index.py

@@ -882,9 +882,13 @@ class IvfRq:
     num_partitions: int, default sqrt(num_rows)
         Number of IVF partitions to create.
 
-    num_bits: int, default 1
+    num_bits: int, default 5
         Number of bits to encode each dimension in the RabitQ codebook.
 
+    layered: Optional[bool], default None
+        Use prefix planes. None uses the product default: False in OSS and True
+        in Enterprise. Ignored for 1-bit and 2-bit codes.
+
     max_iterations: int, default 50
         Max iterations to train kmeans when computing IVF partitions.
 
@@ -897,13 +901,14 @@ class IvfRq:
 
     distance_type: Literal["l2", "cosine", "dot"] = "l2"
     num_partitions: Optional[int] = None
-    num_bits: int = 1
+    num_bits: int = 5
     max_iterations: int = 50
     sample_rate: int = 256
     target_partition_size: Optional[int] = None
     # Name of the accelerator (e.g. "cuda") to use for IVF training. When set,
     # create_index() dispatches to pylance to build the index on the accelerator.
     accelerator: Optional[str] = None
+    layered: Optional[bool] = None
 
 
 # The API reference renders this module with a single mkdocstrings directive,

python/python/lancedb/table.py

@@ -3491,6 +3491,7 @@ class LanceTable(Table):
                         replace=replace,
                         accelerator=acc,
                         num_bits=getattr(config, "num_bits", 8),
+                        layered=getattr(config, "layered", None),
                         m=getattr(config, "m", 20),
                         ef_construction=getattr(config, "ef_construction", 300),
                         target_partition_size=getattr(

python/python/tests/test_index.py

@@ -421,13 +421,29 @@ async def test_create_4bit_ivfpq_index(some_table: AsyncTable):
 
 
 @pytest.mark.asyncio
-async def test_create_ivfrq_index(some_table: AsyncTable):
-    await some_table.create_index("vector", config=IvfRq(num_bits=1))
+@pytest.mark.parametrize("num_bits", [1, 2, 5])
+@pytest.mark.parametrize("layered", [None, False, True])
+async def test_create_ivfrq_index(some_table: AsyncTable, num_bits, layered):
+    await some_table.create_index(
+        "vector", config=IvfRq(num_bits=num_bits, layered=layered)
+    )
     indices = await some_table.list_indices()
     assert len(indices) == 1
     assert indices[0].index_type == "IvfRq"
     assert indices[0].columns == ["vector"]
     assert indices[0].name == "vector_idx"
+    query = list(map(float, range(DIM)))
+    expected = (
+        await some_table.vector_search(query).bypass_vector_index().limit(10).to_arrow()
+    )
+    actual = (
+        await some_table.vector_search(query)
+        .nprobes(NROWS)
+        .refine_factor(26)
+        .limit(10)
+        .to_arrow()
+    )
+    assert len(set(actual["id"].to_pylist()) & set(expected["id"].to_pylist())) >= 5
 
 
 @pytest.mark.asyncio

python/python/tests/test_remote_db.py

@@ -1476,7 +1476,8 @@ def test_remote_json_fts_config(server_version):
     assert request["block_size"] == 256
 
 
-def test_remote_create_index_new_api(caplog):
+@pytest.mark.parametrize("layered", [None, False, True])
+def test_remote_create_index_new_api(caplog, layered):
     received_requests = []
 
     def handler(request):
@@ -1549,7 +1550,9 @@ def test_remote_create_index_new_api(caplog):
             table.create_index("category", config=BTree())
             table.create_index("text", config=FTS(block_size=256))
             # IvfRq via new API
-            table.create_index("vector", config=IvfRq(distance_type="l2"))
+            table.create_index(
+                "vector", config=IvfRq(distance_type="l2", layered=layered)
+            )
             table.create_index(
                 "vector", config=IvfPq(distance_type="l2"), replace=False
             )
@@ -1573,6 +1576,11 @@ def test_remote_create_index_new_api(caplog):
             "vector",
             "vector",
         ]
+        assert received_requests[3]["num_bits"] == 5
+        if layered is None:
+            assert "layered" not in received_requests[3]
+        else:
+            assert received_requests[3]["layered"] is layered
         assert received_requests[2]["block_size"] == 256
         assert received_requests[4]["replace"] is False
         assert received_requests[5]["replace"] is False

python/src/index.rs

@@ -138,6 +138,9 @@ pub fn extract_index_params(source: &Option<Bound<'_, PyAny>>) -> PyResult<Lance
                     .max_iterations(params.max_iterations)
                     .sample_rate(params.sample_rate)
                     .num_bits(params.num_bits);
+                if let Some(layered) = params.layered {
+                    ivf_rq_builder = ivf_rq_builder.layered(layered);
+                }
                 if let Some(num_partitions) = params.num_partitions {
                     ivf_rq_builder = ivf_rq_builder.num_partitions(num_partitions);
                 }
@@ -266,6 +269,7 @@ struct IvfRqParams {
     distance_type: String,
     num_partitions: Option<u32>,
     num_bits: u32,
+    layered: Option<bool>,
     max_iterations: u32,
     sample_rate: u32,
     target_partition_size: Option<u32>,