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>

python/CLAUDE.md

@@ -14,3 +14,6 @@ Common commands:
 Before committing changes, run lints and then formatting.
 
 When you change the Rust code, you will need to recompile the Python bindings: `make develop`.
+
+When you export new types from Rust to Python, you must manually update `python/lancedb/_lancedb.pyi`
+with the corresponding type hints. You can run `pyright` to check for type errors in the Python code.

python/python/lancedb/__init__.py

@@ -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__",
 ]
 

python/python/lancedb/_lancedb.pyi

@@ -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:

python/python/lancedb/db.py

@@ -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:

python/python/tests/test_session.py

@@ -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

python/src/connection.rs

@@ -179,7 +179,7 @@ impl Connection {
 }
 
 #[pyfunction]
-#[pyo3(signature = (uri, api_key=None, region=None, host_override=None, read_consistency_interval=None, client_config=None, storage_options=None))]
+#[pyo3(signature = (uri, api_key=None, region=None, host_override=None, read_consistency_interval=None, client_config=None, storage_options=None, session=None))]
 #[allow(clippy::too_many_arguments)]
 pub fn connect(
     py: Python,
@@ -190,6 +190,7 @@ pub fn connect(
     read_consistency_interval: Option<f64>,
     client_config: Option<PyClientConfig>,
     storage_options: Option<HashMap<String, String>>,
+    session: Option<crate::session::Session>,
 ) -> PyResult<Bound<'_, PyAny>> {
     future_into_py(py, async move {
         let mut builder = lancedb::connect(&uri);
@@ -213,6 +214,9 @@ pub fn connect(
         if let Some(client_config) = client_config {
             builder = builder.client_config(client_config.into());
         }
+        if let Some(session) = session {
+            builder = builder.session(session.inner.clone());
+        }
         Ok(Connection::new(builder.execute().await.infer_error()?))
     })
 }

python/src/lib.rs

@@ -11,6 +11,7 @@ use pyo3::{
     wrap_pyfunction, Bound, PyResult, Python,
 };
 use query::{FTSQuery, HybridQuery, Query, VectorQuery};
+use session::Session;
 use table::{
     AddColumnsResult, AddResult, AlterColumnsResult, DeleteResult, DropColumnsResult, MergeResult,
     Table, UpdateResult,
@@ -21,6 +22,7 @@ pub mod connection;
 pub mod error;
 pub mod index;
 pub mod query;
+pub mod session;
 pub mod table;
 pub mod util;
 
@@ -31,6 +33,7 @@ pub fn _lancedb(_py: Python, m: &Bound<'_, PyModule>) -> PyResult<()> {
         .write_style("LANCEDB_LOG_STYLE");
     env_logger::init_from_env(env);
     m.add_class::<Connection>()?;
+    m.add_class::<Session>()?;
     m.add_class::<Table>()?;
     m.add_class::<IndexConfig>()?;
     m.add_class::<Query>()?;

python/src/session.rs

@@ -0,0 +1,107 @@
+// SPDX-License-Identifier: Apache-2.0
+// SPDX-FileCopyrightText: Copyright The LanceDB Authors
+
+use std::sync::Arc;
+
+use lancedb::{ObjectStoreRegistry, Session as LanceSession};
+use pyo3::{pyclass, pymethods, PyResult};
+
+/// A session for managing caches and object stores across LanceDB operations.
+///
+/// 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.
+#[pyclass]
+#[derive(Clone)]
+pub struct Session {
+    pub(crate) inner: Arc<LanceSession>,
+}
+
+impl Default for Session {
+    fn default() -> Self {
+        Self {
+            inner: Arc::new(LanceSession::default()),
+        }
+    }
+}
+
+#[pymethods]
+impl Session {
+    /// Create a new session with custom cache sizes.
+    ///
+    /// Parameters
+    /// ----------
+    /// index_cache_size_bytes : int, optional
+    ///     The size of the index cache in bytes.
+    ///     Index data is stored in memory in this cache to speed up queries.
+    ///     Default: 6GB (6 * 1024 * 1024 * 1024 bytes)
+    /// metadata_cache_size_bytes : int, optional
+    ///     The size of the metadata cache in bytes.
+    ///     The metadata cache stores file metadata and schema information in memory.
+    ///     This cache improves scan and write performance.
+    ///     Default: 1GB (1024 * 1024 * 1024 bytes)
+    #[new]
+    #[pyo3(signature = (index_cache_size_bytes=None, metadata_cache_size_bytes=None))]
+    pub fn new(
+        index_cache_size_bytes: Option<usize>,
+        metadata_cache_size_bytes: Option<usize>,
+    ) -> PyResult<Self> {
+        let index_cache_size = index_cache_size_bytes.unwrap_or(6 * 1024 * 1024 * 1024); // 6GB default
+        let metadata_cache_size = metadata_cache_size_bytes.unwrap_or(1024 * 1024 * 1024); // 1GB default
+
+        let session = LanceSession::new(
+            index_cache_size,
+            metadata_cache_size,
+            Arc::new(ObjectStoreRegistry::default()),
+        );
+
+        Ok(Self {
+            inner: Arc::new(session),
+        })
+    }
+
+    /// Create a session with default cache sizes.
+    ///
+    /// This is equivalent to creating a session with 6GB index cache
+    /// and 1GB metadata cache.
+    ///
+    /// Returns
+    /// -------
+    /// Session
+    ///     A new Session with default cache sizes
+    #[staticmethod]
+    #[allow(clippy::should_implement_trait)]
+    pub fn default() -> Self {
+        Default::default()
+    }
+
+    /// Get the current size of the session caches in bytes.
+    ///
+    /// Returns
+    /// -------
+    /// int
+    ///     The total size of all caches in the session
+    #[getter]
+    pub fn size_bytes(&self) -> u64 {
+        self.inner.size_bytes()
+    }
+
+    /// Get the approximate number of items cached in the session.
+    ///
+    /// Returns
+    /// -------
+    /// int
+    ///     The number of cached items across all caches
+    #[getter]
+    pub fn approx_num_items(&self) -> usize {
+        self.inner.approx_num_items()
+    }
+
+    fn __repr__(&self) -> String {
+        format!(
+            "Session(size_bytes={}, approx_num_items={})",
+            self.size_bytes(),
+            self.approx_num_items()
+        )
+    }
+}