feat: add 'target_partition_size' param (#2642)

this exposes the param `target_partition_size` from lance

---------

Signed-off-by: BubbleCal <bubble-cal@outlook.com>

python/python/lancedb/index.py

@@ -251,6 +251,13 @@ class HnswPq:
         results. In most cases, there is no benefit to setting this higher than 500.
         This value should be set to a value that is not less than `ef` in the
         search phase.
+
+    target_partition_size, default is 1,048,576
+
+        The target size of each partition.
+
+        This value controls the tradeoff between search performance and accuracy.
+        faster search but less accurate results as higher value.
     """
 
     distance_type: Literal["l2", "cosine", "dot"] = "l2"
@@ -261,6 +268,7 @@ class HnswPq:
     sample_rate: int = 256
     m: int = 20
     ef_construction: int = 300
+    target_partition_size: Optional[int] = None
 
 
 @dataclass
@@ -351,6 +359,12 @@ class HnswSq:
         This value should be set to a value that is not less than `ef` in the search
         phase.
 
+    target_partition_size, default is 1,048,576
+
+        The target size of each partition.
+
+        This value controls the tradeoff between search performance and accuracy.
+        faster search but less accurate results as higher value.
     """
 
     distance_type: Literal["l2", "cosine", "dot"] = "l2"
@@ -359,6 +373,7 @@ class HnswSq:
     sample_rate: int = 256
     m: int = 20
     ef_construction: int = 300
+    target_partition_size: Optional[int] = None
 
 
 @dataclass
@@ -444,12 +459,20 @@ class IvfFlat:
         cases the default should be sufficient.
 
         The default value is 256.
+
+    target_partition_size, default is 8192
+
+        The target size of each partition.
+
+        This value controls the tradeoff between search performance and accuracy.
+        faster search but less accurate results as higher value.
     """
 
     distance_type: Literal["l2", "cosine", "dot", "hamming"] = "l2"
     num_partitions: Optional[int] = None
     max_iterations: int = 50
     sample_rate: int = 256
+    target_partition_size: Optional[int] = None
 
 
 @dataclass
@@ -564,6 +587,13 @@ class IvfPq:
         cases the default should be sufficient.
 
         The default value is 256.
+
+    target_partition_size, default is 8192
+
+        The target size of each partition.
+
+        This value controls the tradeoff between search performance and accuracy.
+        faster search but less accurate results as higher value.
     """
 
     distance_type: Literal["l2", "cosine", "dot"] = "l2"
@@ -572,6 +602,7 @@ class IvfPq:
     num_bits: int = 8
     max_iterations: int = 50
     sample_rate: int = 256
+    target_partition_size: Optional[int] = None
 
 
 __all__ = [

python/python/lancedb/table.py

@@ -691,6 +691,7 @@ class Table(ABC):
         ef_construction: int = 300,
         name: Optional[str] = None,
         train: bool = True,
+        target_partition_size: Optional[int] = None,
     ):
         """Create an index on the table.
 
@@ -2002,6 +2003,7 @@ class LanceTable(Table):
         *,
         name: Optional[str] = None,
         train: bool = True,
+        target_partition_size: Optional[int] = None,
     ):
         """Create an index on the table."""
         if accelerator is not None:
@@ -2018,6 +2020,7 @@ class LanceTable(Table):
                 num_bits=num_bits,
                 m=m,
                 ef_construction=ef_construction,
+                target_partition_size=target_partition_size,
             )
             self.checkout_latest()
             return
@@ -2027,6 +2030,7 @@ class LanceTable(Table):
                 num_partitions=num_partitions,
                 max_iterations=max_iterations,
                 sample_rate=sample_rate,
+                target_partition_size=target_partition_size,
             )
         elif index_type == "IVF_PQ":
             config = IvfPq(
@@ -2036,6 +2040,7 @@ class LanceTable(Table):
                 num_bits=num_bits,
                 max_iterations=max_iterations,
                 sample_rate=sample_rate,
+                target_partition_size=target_partition_size,
             )
         elif index_type == "IVF_HNSW_PQ":
             config = HnswPq(
@@ -2047,6 +2052,7 @@ class LanceTable(Table):
                 sample_rate=sample_rate,
                 m=m,
                 ef_construction=ef_construction,
+                target_partition_size=target_partition_size,
             )
         elif index_type == "IVF_HNSW_SQ":
             config = HnswSq(
@@ -2056,6 +2062,7 @@ class LanceTable(Table):
                 sample_rate=sample_rate,
                 m=m,
                 ef_construction=ef_construction,
+                target_partition_size=target_partition_size,
             )
         else:
             raise ValueError(f"Unknown index type {index_type}")

python/python/tests/test_table.py

@@ -674,6 +674,45 @@ def test_create_index_method(mock_create_index, mem_db: DBConnection):
         "vector", replace=True, config=expected_config, name=None, train=True
     )
 
+    # Test with target_partition_size
+    table.create_index(
+        metric="l2",
+        num_sub_vectors=96,
+        vector_column_name="vector",
+        replace=True,
+        index_cache_size=256,
+        num_bits=4,
+        target_partition_size=8192,
+    )
+    expected_config = IvfPq(
+        distance_type="l2",
+        num_sub_vectors=96,
+        num_bits=4,
+        target_partition_size=8192,
+    )
+    mock_create_index.assert_called_with(
+        "vector", replace=True, config=expected_config, name=None, train=True
+    )
+
+    # target_partition_size has a default value,
+    # so `num_partitions` and `target_partition_size` are not required
+    table.create_index(
+        metric="l2",
+        num_sub_vectors=96,
+        vector_column_name="vector",
+        replace=True,
+        index_cache_size=256,
+        num_bits=4,
+    )
+    expected_config = IvfPq(
+        distance_type="l2",
+        num_sub_vectors=96,
+        num_bits=4,
+    )
+    mock_create_index.assert_called_with(
+        "vector", replace=True, config=expected_config, name=None, train=True
+    )
+
     table.create_index(
         vector_column_name="my_vector",
         metric="dot",