diff --git a/python/lancedb/remote/table.py b/python/lancedb/remote/table.py index 0c214d28..cb1a3229 100644 --- a/python/lancedb/remote/table.py +++ b/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) diff --git a/python/lancedb/table.py b/python/lancedb/table.py index ea2b5274..e21b3330 100644 --- a/python/lancedb/table.py +++ b/python/lancedb/table.py @@ -187,6 +187,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. @@ -211,6 +212,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 @@ -555,6 +558,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( @@ -565,6 +569,7 @@ class LanceTable(Table): num_sub_vectors=num_sub_vectors, replace=replace, accelerator=accelerator, + index_cache_size=index_cache_size, ) self._reset_dataset() diff --git a/python/tests/test_db.py b/python/tests/test_db.py index 90719d85..a7afa56e 100644 --- a/python/tests/test_db.py +++ b/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, ) diff --git a/python/tests/test_table.py b/python/tests/test_table.py index 155b13a0..258a929e 100644 --- a/python/tests/test_table.py +++ b/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, )