feat(python): expose index cache size (#655)

This is to enable https://github.com/lancedb/lancedb/issues/641.
Should be merged after https://github.com/lancedb/lance/pull/1587 is
released.

python/lancedb/remote/table.py

@@ -71,6 +71,7 @@ class RemoteTable(Table):
         vector_column_name: str = VECTOR_COLUMN_NAME,
         replace: bool = True,
         accelerator: Optional[str] = None,
+        index_cache_size: Optional[int] = None,
     ):
         """Create an index on the table.
         Currently, the only parameters that matter are
@@ -88,6 +89,11 @@ class RemoteTable(Table):
             The name of the vector column. Default is "vector".
         replace : bool
             Whether to replace the existing index. Default is True.
+        accelerator : str, optional
+            If set, use the given accelerator to create the index.
+            Default is None. Currently not supported.
+        index_cache_size : int, optional
+            The size of the index cache in number of entries. Default value is 256.
 
         Examples
         --------
@@ -115,6 +121,7 @@ class RemoteTable(Table):
             "column": vector_column_name,
             "index_type": index_type,
             "metric_type": metric,
+            "index_cache_size": index_cache_size,
         }
         resp = self._conn._loop.run_until_complete(
             self._conn._client.post(f"/v1/table/{self._name}/create_index/", data=data)

python/lancedb/table.py

@@ -188,6 +188,7 @@ class Table(ABC):
         vector_column_name: str = VECTOR_COLUMN_NAME,
         replace: bool = True,
         accelerator: Optional[str] = None,
+        index_cache_size: Optional[int] = None,
     ):
         """Create an index on the table.
 
@@ -212,6 +213,8 @@ class Table(ABC):
         accelerator: str, default None
             If set, use the given accelerator to create the index.
             Only support "cuda" for now.
+        index_cache_size : int, optional
+            The size of the index cache in number of entries. Default value is 256.
         """
         raise NotImplementedError
 
@@ -556,6 +559,7 @@ class LanceTable(Table):
         vector_column_name=VECTOR_COLUMN_NAME,
         replace: bool = True,
         accelerator: Optional[str] = None,
+        index_cache_size: Optional[int] = None,
     ):
         """Create an index on the table."""
         self._dataset.create_index(
@@ -566,6 +570,7 @@ class LanceTable(Table):
             num_sub_vectors=num_sub_vectors,
             replace=replace,
             accelerator=accelerator,
+            index_cache_size=index_cache_size,
         )
         self._reset_dataset()
         register_event("create_index")

python/tests/test_db.py

@@ -301,6 +301,7 @@ def test_replace_index(tmp_path):
         num_partitions=2,
         num_sub_vectors=4,
         replace=True,
+        index_cache_size=10,
     )
 
 

python/tests/test_table.py

@@ -213,6 +213,7 @@ def test_create_index_method():
                 num_sub_vectors=96,
                 vector_column_name="vector",
                 replace=True,
+                index_cache_size=256,
             )
 
             # Check that the _dataset.create_index method was called
@@ -225,6 +226,7 @@ def test_create_index_method():
                 num_sub_vectors=96,
                 replace=True,
                 accelerator=None,
+                index_cache_size=256,
             )