mirror of
https://github.com/lancedb/lancedb.git
synced 2025-12-27 15:12:53 +00:00
feat: allow Python and Typescript users to create Sessions (#2530)
## Summary - Exposes `Session` in Python and Typescript so users can set the `index_cache_size_bytes` and `metadata_cache_size_bytes` * The `Session` is attached to the `Connection`, and thus shared across all tables in that connection. - Adds deprecation warnings for table-level cache configuration 🤖 Generated with [Claude Code](https://claude.ai/code) --------- Co-authored-by: Claude <noreply@anthropic.com>
This commit is contained in:
@@ -18,6 +18,7 @@ from .remote import ClientConfig
|
||||
from .remote.db import RemoteDBConnection
|
||||
from .schema import vector
|
||||
from .table import AsyncTable
|
||||
from ._lancedb import Session
|
||||
|
||||
|
||||
def connect(
|
||||
@@ -30,6 +31,7 @@ def connect(
|
||||
request_thread_pool: Optional[Union[int, ThreadPoolExecutor]] = None,
|
||||
client_config: Union[ClientConfig, Dict[str, Any], None] = None,
|
||||
storage_options: Optional[Dict[str, str]] = None,
|
||||
session: Optional[Session] = None,
|
||||
**kwargs: Any,
|
||||
) -> DBConnection:
|
||||
"""Connect to a LanceDB database.
|
||||
@@ -64,6 +66,12 @@ def connect(
|
||||
storage_options: dict, optional
|
||||
Additional options for the storage backend. See available options at
|
||||
<https://lancedb.github.io/lancedb/guides/storage/>
|
||||
session: Session, optional
|
||||
(For LanceDB OSS only)
|
||||
A session to use for this connection. Sessions allow you to configure
|
||||
cache sizes for index and metadata caches, which can significantly
|
||||
impact memory use and performance. They can also be re-used across
|
||||
multiple connections to share the same cache state.
|
||||
|
||||
Examples
|
||||
--------
|
||||
@@ -113,6 +121,7 @@ def connect(
|
||||
uri,
|
||||
read_consistency_interval=read_consistency_interval,
|
||||
storage_options=storage_options,
|
||||
session=session,
|
||||
)
|
||||
|
||||
|
||||
@@ -125,6 +134,7 @@ async def connect_async(
|
||||
read_consistency_interval: Optional[timedelta] = None,
|
||||
client_config: Optional[Union[ClientConfig, Dict[str, Any]]] = None,
|
||||
storage_options: Optional[Dict[str, str]] = None,
|
||||
session: Optional[Session] = None,
|
||||
) -> AsyncConnection:
|
||||
"""Connect to a LanceDB database.
|
||||
|
||||
@@ -158,6 +168,12 @@ async def connect_async(
|
||||
storage_options: dict, optional
|
||||
Additional options for the storage backend. See available options at
|
||||
<https://lancedb.github.io/lancedb/guides/storage/>
|
||||
session: Session, optional
|
||||
(For LanceDB OSS only)
|
||||
A session to use for this connection. Sessions allow you to configure
|
||||
cache sizes for index and metadata caches, which can significantly
|
||||
impact memory use and performance. They can also be re-used across
|
||||
multiple connections to share the same cache state.
|
||||
|
||||
Examples
|
||||
--------
|
||||
@@ -197,6 +213,7 @@ async def connect_async(
|
||||
read_consistency_interval_secs,
|
||||
client_config,
|
||||
storage_options,
|
||||
session,
|
||||
)
|
||||
)
|
||||
|
||||
@@ -212,6 +229,7 @@ __all__ = [
|
||||
"DBConnection",
|
||||
"LanceDBConnection",
|
||||
"RemoteDBConnection",
|
||||
"Session",
|
||||
"__version__",
|
||||
]
|
||||
|
||||
|
||||
@@ -6,6 +6,19 @@ import pyarrow as pa
|
||||
from .index import BTree, IvfFlat, IvfPq, Bitmap, LabelList, HnswPq, HnswSq, FTS
|
||||
from .remote import ClientConfig
|
||||
|
||||
class Session:
|
||||
def __init__(
|
||||
self,
|
||||
index_cache_size_bytes: Optional[int] = None,
|
||||
metadata_cache_size_bytes: Optional[int] = None,
|
||||
): ...
|
||||
@staticmethod
|
||||
def default() -> "Session": ...
|
||||
@property
|
||||
def size_bytes(self) -> int: ...
|
||||
@property
|
||||
def approx_num_items(self) -> int: ...
|
||||
|
||||
class Connection(object):
|
||||
uri: str
|
||||
async def table_names(
|
||||
@@ -89,6 +102,7 @@ async def connect(
|
||||
read_consistency_interval: Optional[float],
|
||||
client_config: Optional[Union[ClientConfig, Dict[str, Any]]],
|
||||
storage_options: Optional[Dict[str, str]],
|
||||
session: Optional[Session],
|
||||
) -> Connection: ...
|
||||
|
||||
class RecordBatchStream:
|
||||
|
||||
@@ -37,6 +37,7 @@ if TYPE_CHECKING:
|
||||
from ._lancedb import Connection as LanceDbConnection
|
||||
from .common import DATA, URI
|
||||
from .embeddings import EmbeddingFunctionConfig
|
||||
from ._lancedb import Session
|
||||
|
||||
|
||||
class DBConnection(EnforceOverrides):
|
||||
@@ -247,6 +248,9 @@ class DBConnection(EnforceOverrides):
|
||||
name: str
|
||||
The name of the table.
|
||||
index_cache_size: int, default 256
|
||||
**Deprecated**: Use session-level cache configuration instead.
|
||||
Create a Session with custom cache sizes and pass it to lancedb.connect().
|
||||
|
||||
Set the size of the index cache, specified as a number of entries
|
||||
|
||||
The exact meaning of an "entry" will depend on the type of index:
|
||||
@@ -354,6 +358,7 @@ class LanceDBConnection(DBConnection):
|
||||
*,
|
||||
read_consistency_interval: Optional[timedelta] = None,
|
||||
storage_options: Optional[Dict[str, str]] = None,
|
||||
session: Optional[Session] = None,
|
||||
):
|
||||
if not isinstance(uri, Path):
|
||||
scheme = get_uri_scheme(uri)
|
||||
@@ -367,6 +372,7 @@ class LanceDBConnection(DBConnection):
|
||||
self._entered = False
|
||||
self.read_consistency_interval = read_consistency_interval
|
||||
self.storage_options = storage_options
|
||||
self.session = session
|
||||
|
||||
if read_consistency_interval is not None:
|
||||
read_consistency_interval_secs = read_consistency_interval.total_seconds()
|
||||
@@ -382,6 +388,7 @@ class LanceDBConnection(DBConnection):
|
||||
read_consistency_interval_secs,
|
||||
None,
|
||||
storage_options,
|
||||
session,
|
||||
)
|
||||
|
||||
self._conn = AsyncConnection(LOOP.run(do_connect()))
|
||||
@@ -475,6 +482,17 @@ class LanceDBConnection(DBConnection):
|
||||
-------
|
||||
A LanceTable object representing the table.
|
||||
"""
|
||||
if index_cache_size is not None:
|
||||
import warnings
|
||||
|
||||
warnings.warn(
|
||||
"index_cache_size is deprecated. Use session-level cache "
|
||||
"configuration instead. Create a Session with custom cache sizes "
|
||||
"and pass it to lancedb.connect().",
|
||||
DeprecationWarning,
|
||||
stacklevel=2,
|
||||
)
|
||||
|
||||
return LanceTable.open(
|
||||
self,
|
||||
name,
|
||||
@@ -820,6 +838,9 @@ class AsyncConnection(object):
|
||||
See available options at
|
||||
<https://lancedb.github.io/lancedb/guides/storage/>
|
||||
index_cache_size: int, default 256
|
||||
**Deprecated**: Use session-level cache configuration instead.
|
||||
Create a Session with custom cache sizes and pass it to lancedb.connect().
|
||||
|
||||
Set the size of the index cache, specified as a number of entries
|
||||
|
||||
The exact meaning of an "entry" will depend on the type of index:
|
||||
|
||||
38
python/python/tests/test_session.py
Normal file
38
python/python/tests/test_session.py
Normal file
@@ -0,0 +1,38 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright The LanceDB Authors
|
||||
|
||||
import lancedb
|
||||
|
||||
|
||||
def test_session_cache_configuration(tmp_path):
|
||||
"""Test Session cache configuration and basic functionality."""
|
||||
# Create session with small cache limits for testing
|
||||
index_cache_size = 1024 * 1024 # 1MB
|
||||
metadata_cache_size = 512 * 1024 # 512KB
|
||||
|
||||
session = lancedb.Session(
|
||||
index_cache_size_bytes=index_cache_size,
|
||||
metadata_cache_size_bytes=metadata_cache_size,
|
||||
)
|
||||
|
||||
# Record initial cache state
|
||||
initial_cache_size = session.size_bytes
|
||||
initial_cache_items = session.approx_num_items
|
||||
|
||||
# Test session works with database connection
|
||||
db = lancedb.connect(tmp_path, session=session)
|
||||
|
||||
# Create and use a table to exercise the session
|
||||
data = [{"id": i, "text": f"item {i}"} for i in range(100)]
|
||||
table = db.create_table("test", data)
|
||||
results = list(table.to_arrow().to_pylist())
|
||||
|
||||
assert len(results) == 100
|
||||
|
||||
# Verify cache usage increased after operations
|
||||
final_cache_size = session.size_bytes
|
||||
final_cache_items = session.approx_num_items
|
||||
|
||||
assert final_cache_size > initial_cache_size # Cache should have grown
|
||||
assert final_cache_items >= initial_cache_items # Items should not decrease
|
||||
assert initial_cache_size < index_cache_size + metadata_cache_size
|
||||
Reference in New Issue
Block a user