From 3d1f102087fcaf0deb15b6ef8ba5ba4ff57dd4d0 Mon Sep 17 00:00:00 2001 From: Will Jones Date: Thu, 24 Jul 2025 12:06:29 -0700 Subject: [PATCH] feat: allow Python and Typescript users to create `Session`s (#2530) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit ## 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 --- Cargo.lock | 4 +- docs/src/js/classes/Session.md | 84 +++++++++++++++ docs/src/js/functions/connect.md | 19 +++- docs/src/js/globals.md | 1 + docs/src/js/interfaces/ConnectionOptions.md | 11 ++ docs/src/js/interfaces/OpenTableOptions.md | 7 +- nodejs/__test__/session.test.ts | 46 +++++++++ nodejs/lancedb/connection.ts | 3 + nodejs/lancedb/index.ts | 28 +++-- nodejs/src/connection.rs | 4 + nodejs/src/lib.rs | 4 + nodejs/src/session.rs | 102 +++++++++++++++++++ python/CLAUDE.md | 3 + python/python/lancedb/__init__.py | 18 ++++ python/python/lancedb/_lancedb.pyi | 14 +++ python/python/lancedb/db.py | 21 ++++ python/python/tests/test_session.py | 38 +++++++ python/src/connection.rs | 6 +- python/src/lib.rs | 3 + python/src/session.rs | 107 ++++++++++++++++++++ rust/lancedb/src/lib.rs | 4 + 21 files changed, 514 insertions(+), 13 deletions(-) create mode 100644 docs/src/js/classes/Session.md create mode 100644 nodejs/__test__/session.test.ts create mode 100644 nodejs/src/session.rs create mode 100644 python/python/tests/test_session.py create mode 100644 python/src/session.rs diff --git a/Cargo.lock b/Cargo.lock index b88a5c67..f2bbcd0b 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -3819,9 +3819,9 @@ dependencies = [ [[package]] name = "io-uring" -version = "0.7.8" +version = "0.7.9" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "b86e202f00093dcba4275d4636b93ef9dd75d025ae560d2521b45ea28ab49013" +checksum = "d93587f37623a1a17d94ef2bc9ada592f5465fe7732084ab7beefabe5c77c0c4" dependencies = [ "bitflags 2.9.1", "cfg-if", diff --git a/docs/src/js/classes/Session.md b/docs/src/js/classes/Session.md new file mode 100644 index 00000000..110fbf68 --- /dev/null +++ b/docs/src/js/classes/Session.md @@ -0,0 +1,84 @@ +[**@lancedb/lancedb**](../README.md) • **Docs** + +*** + +[@lancedb/lancedb](../globals.md) / Session + +# Class: Session + +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 performance for large datasets. + +## Constructors + +### new Session() + +```ts +new Session(indexCacheSizeBytes?, metadataCacheSizeBytes?): Session +``` + +Create a new session with custom cache sizes. + +# Parameters + +- `index_cache_size_bytes`: The size of the index cache in bytes. + Defaults to 6GB if not specified. +- `metadata_cache_size_bytes`: The size of the metadata cache in bytes. + Defaults to 1GB if not specified. + +#### Parameters + +* **indexCacheSizeBytes?**: `null` \| `bigint` + +* **metadataCacheSizeBytes?**: `null` \| `bigint` + +#### Returns + +[`Session`](Session.md) + +## Methods + +### approxNumItems() + +```ts +approxNumItems(): number +``` + +Get the approximate number of items cached in the session. + +#### Returns + +`number` + +*** + +### sizeBytes() + +```ts +sizeBytes(): bigint +``` + +Get the current size of the session caches in bytes. + +#### Returns + +`bigint` + +*** + +### default() + +```ts +static default(): 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`](Session.md) diff --git a/docs/src/js/functions/connect.md b/docs/src/js/functions/connect.md index 054928c0..ce35bb80 100644 --- a/docs/src/js/functions/connect.md +++ b/docs/src/js/functions/connect.md @@ -6,10 +6,13 @@ # Function: connect() -## connect(uri, options) +## connect(uri, options, session) ```ts -function connect(uri, options?): Promise +function connect( + uri, + options?, + session?): Promise ``` Connect to a LanceDB instance at the given URI. @@ -29,6 +32,8 @@ Accepted formats: * **options?**: `Partial`<[`ConnectionOptions`](../interfaces/ConnectionOptions.md)> The options to use when connecting to the database +* **session?**: [`Session`](../classes/Session.md) + ### Returns `Promise`<[`Connection`](../classes/Connection.md)> @@ -77,7 +82,7 @@ Accepted formats: [ConnectionOptions](../interfaces/ConnectionOptions.md) for more details on the URI format. -### Example +### Examples ```ts const conn = await connect({ @@ -85,3 +90,11 @@ const conn = await connect({ storageOptions: {timeout: "60s"} }); ``` + +```ts +const session = Session.default(); +const conn = await connect({ + uri: "/path/to/database", + session: session +}); +``` diff --git a/docs/src/js/globals.md b/docs/src/js/globals.md index e7e6beac..11f3d899 100644 --- a/docs/src/js/globals.md +++ b/docs/src/js/globals.md @@ -29,6 +29,7 @@ - [Query](classes/Query.md) - [QueryBase](classes/QueryBase.md) - [RecordBatchIterator](classes/RecordBatchIterator.md) +- [Session](classes/Session.md) - [Table](classes/Table.md) - [TagContents](classes/TagContents.md) - [Tags](classes/Tags.md) diff --git a/docs/src/js/interfaces/ConnectionOptions.md b/docs/src/js/interfaces/ConnectionOptions.md index 4837c4c1..14b05b24 100644 --- a/docs/src/js/interfaces/ConnectionOptions.md +++ b/docs/src/js/interfaces/ConnectionOptions.md @@ -70,6 +70,17 @@ Defaults to 'us-east-1'. *** +### session? + +```ts +optional session: Session; +``` + +(For LanceDB OSS only): the session to use for this connection. Holds +shared caches and other session-specific state. + +*** + ### storageOptions? ```ts diff --git a/docs/src/js/interfaces/OpenTableOptions.md b/docs/src/js/interfaces/OpenTableOptions.md index dab6294d..009549ec 100644 --- a/docs/src/js/interfaces/OpenTableOptions.md +++ b/docs/src/js/interfaces/OpenTableOptions.md @@ -8,7 +8,7 @@ ## Properties -### indexCacheSize? +### ~~indexCacheSize?~~ ```ts optional indexCacheSize: number; @@ -16,6 +16,11 @@ optional indexCacheSize: number; Set the size of the index cache, specified as a number of entries +#### Deprecated + +Use session-level cache configuration instead. +Create a Session with custom cache sizes and pass it to the connect() function. + The exact meaning of an "entry" will depend on the type of index: - IVF: there is one entry for each IVF partition - BTREE: there is one entry for the entire index diff --git a/nodejs/__test__/session.test.ts b/nodejs/__test__/session.test.ts new file mode 100644 index 00000000..4b650d25 --- /dev/null +++ b/nodejs/__test__/session.test.ts @@ -0,0 +1,46 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright The LanceDB Authors + +import * as tmp from "tmp"; +import { Session, connect } from "../lancedb"; + +describe("Session", () => { + let tmpDir: tmp.DirResult; + beforeEach(() => { + tmpDir = tmp.dirSync({ unsafeCleanup: true }); + }); + afterEach(() => tmpDir.removeCallback()); + + it("should configure cache sizes and work with database operations", async () => { + // Create session with small cache limits for testing + const indexCacheSize = BigInt(1024 * 1024); // 1MB + const metadataCacheSize = BigInt(512 * 1024); // 512KB + + const session = new Session(indexCacheSize, metadataCacheSize); + + // Record initial cache state + const initialCacheSize = session.sizeBytes(); + const initialCacheItems = session.approxNumItems(); + + // Test session works with database connection + const db = await connect({ uri: tmpDir.name, session: session }); + + // Create and use a table to exercise the session + const data = Array.from({ length: 100 }, (_, i) => ({ + id: i, + text: `item ${i}`, + })); + const table = await db.createTable("test", data); + const results = await table.query().limit(5).toArray(); + + expect(results).toHaveLength(5); + + // Verify cache usage increased after operations + const finalCacheSize = session.sizeBytes(); + const finalCacheItems = session.approxNumItems(); + + expect(finalCacheSize).toBeGreaterThan(initialCacheSize); // Cache should have grown + expect(finalCacheItems).toBeGreaterThanOrEqual(initialCacheItems); // Items should not decrease + expect(initialCacheSize).toBeLessThan(indexCacheSize + metadataCacheSize); // Within limits + }); +}); diff --git a/nodejs/lancedb/connection.ts b/nodejs/lancedb/connection.ts index 53ff6aa5..373fb670 100644 --- a/nodejs/lancedb/connection.ts +++ b/nodejs/lancedb/connection.ts @@ -85,6 +85,9 @@ export interface OpenTableOptions { /** * Set the size of the index cache, specified as a number of entries * + * @deprecated Use session-level cache configuration instead. + * Create a Session with custom cache sizes and pass it to the connect() function. + * * The exact meaning of an "entry" will depend on the type of index: * - IVF: there is one entry for each IVF partition * - BTREE: there is one entry for the entire index diff --git a/nodejs/lancedb/index.ts b/nodejs/lancedb/index.ts index 0bf4f9f5..ce0f4546 100644 --- a/nodejs/lancedb/index.ts +++ b/nodejs/lancedb/index.ts @@ -10,6 +10,7 @@ import { import { ConnectionOptions, Connection as LanceDbConnection, + Session, } from "./native.js"; export { @@ -51,6 +52,8 @@ export { OpenTableOptions, } from "./connection"; +export { Session } from "./native.js"; + export { ExecutableQuery, Query, @@ -131,6 +134,7 @@ export { IntoSql, packBits } from "./util"; export async function connect( uri: string, options?: Partial, + session?: Session, ): Promise; /** * Connect to a LanceDB instance at the given URI. @@ -149,31 +153,43 @@ export async function connect( * storageOptions: {timeout: "60s"} * }); * ``` + * + * @example + * ```ts + * const session = Session.default(); + * const conn = await connect({ + * uri: "/path/to/database", + * session: session + * }); + * ``` */ export async function connect( options: Partial & { uri: string }, ): Promise; export async function connect( uriOrOptions: string | (Partial & { uri: string }), - options: Partial = {}, + options?: Partial, ): Promise { let uri: string | undefined; + let finalOptions: Partial = {}; + if (typeof uriOrOptions !== "string") { const { uri: uri_, ...opts } = uriOrOptions; uri = uri_; - options = opts; + finalOptions = opts; } else { uri = uriOrOptions; + finalOptions = options || {}; } if (!uri) { throw new Error("uri is required"); } - options = (options as ConnectionOptions) ?? {}; - (options).storageOptions = cleanseStorageOptions( - (options).storageOptions, + finalOptions = (finalOptions as ConnectionOptions) ?? {}; + (finalOptions).storageOptions = cleanseStorageOptions( + (finalOptions).storageOptions, ); - const nativeConn = await LanceDbConnection.new(uri, options); + const nativeConn = await LanceDbConnection.new(uri, finalOptions); return new LocalConnection(nativeConn); } diff --git a/nodejs/src/connection.rs b/nodejs/src/connection.rs index 49bbf004..d6313e4a 100644 --- a/nodejs/src/connection.rs +++ b/nodejs/src/connection.rs @@ -74,6 +74,10 @@ impl Connection { builder = builder.host_override(&host_override); } + if let Some(session) = options.session { + builder = builder.session(session.inner.clone()); + } + Ok(Self::inner_new(builder.execute().await.default_error()?)) } diff --git a/nodejs/src/lib.rs b/nodejs/src/lib.rs index 7674e370..1d5a4333 100644 --- a/nodejs/src/lib.rs +++ b/nodejs/src/lib.rs @@ -14,6 +14,7 @@ pub mod merge; mod query; pub mod remote; mod rerankers; +mod session; mod table; mod util; @@ -34,6 +35,9 @@ pub struct ConnectionOptions { /// /// The available options are described at https://lancedb.github.io/lancedb/guides/storage/ pub storage_options: Option>, + /// (For LanceDB OSS only): the session to use for this connection. Holds + /// shared caches and other session-specific state. + pub session: Option, /// (For LanceDB cloud only): configuration for the remote HTTP client. pub client_config: Option, diff --git a/nodejs/src/session.rs b/nodejs/src/session.rs new file mode 100644 index 00000000..c33cfb17 --- /dev/null +++ b/nodejs/src/session.rs @@ -0,0 +1,102 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright The LanceDB Authors + +use std::sync::Arc; + +use lancedb::{ObjectStoreRegistry, Session as LanceSession}; +use napi::bindgen_prelude::*; +use napi_derive::*; + +/// 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. +#[napi] +#[derive(Clone)] +pub struct Session { + pub(crate) inner: Arc, +} + +impl std::fmt::Debug for Session { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + f.debug_struct("Session") + .field("size_bytes", &self.inner.size_bytes()) + .field("approx_num_items", &self.inner.approx_num_items()) + .finish() + } +} + +#[napi] +impl Session { + /// Create a new session with custom cache sizes. + /// + /// # Parameters + /// + /// - `index_cache_size_bytes`: The size of the index cache in bytes. + /// Index data is stored in memory in this cache to speed up queries. + /// Defaults to 6GB if not specified. + /// - `metadata_cache_size_bytes`: 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. + /// Defaults to 1GB if not specified. + #[napi(constructor)] + pub fn new( + index_cache_size_bytes: Option, + metadata_cache_size_bytes: Option, + ) -> napi::Result { + let index_cache_size = index_cache_size_bytes + .map(|size| size.get_u64().1 as usize) + .unwrap_or(6 * 1024 * 1024 * 1024); // 6GB default + + let metadata_cache_size = metadata_cache_size_bytes + .map(|size| size.get_u64().1 as usize) + .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. + #[napi(factory)] + pub fn default() -> Self { + Self { + inner: Arc::new(LanceSession::default()), + } + } + + /// Get the current size of the session caches in bytes. + #[napi] + pub fn size_bytes(&self) -> BigInt { + BigInt::from(self.inner.size_bytes()) + } + + /// Get the approximate number of items cached in the session. + #[napi] + pub fn approx_num_items(&self) -> u32 { + self.inner.approx_num_items() as u32 + } +} + +// Implement FromNapiValue for Session to work with napi(object) +impl napi::bindgen_prelude::FromNapiValue for Session { + unsafe fn from_napi_value( + env: napi::sys::napi_env, + napi_val: napi::sys::napi_value, + ) -> napi::Result { + let object: napi::bindgen_prelude::ClassInstance = + napi::bindgen_prelude::ClassInstance::from_napi_value(env, napi_val)?; + let copy = object.clone(); + Ok(copy) + } +} diff --git a/python/CLAUDE.md b/python/CLAUDE.md index 07d78211..993e83cf 100644 --- a/python/CLAUDE.md +++ b/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. diff --git a/python/python/lancedb/__init__.py b/python/python/lancedb/__init__.py index f43a2b27..3227a9af 100644 --- a/python/python/lancedb/__init__.py +++ b/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 + 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 + 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__", ] diff --git a/python/python/lancedb/_lancedb.pyi b/python/python/lancedb/_lancedb.pyi index 8f4c5a06..b0f4ad7d 100644 --- a/python/python/lancedb/_lancedb.pyi +++ b/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: diff --git a/python/python/lancedb/db.py b/python/python/lancedb/db.py index d4f30b83..fedde25e 100644 --- a/python/python/lancedb/db.py +++ b/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 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: diff --git a/python/python/tests/test_session.py b/python/python/tests/test_session.py new file mode 100644 index 00000000..8923477c --- /dev/null +++ b/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 diff --git a/python/src/connection.rs b/python/src/connection.rs index 975a756c..2e2f64d3 100644 --- a/python/src/connection.rs +++ b/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, client_config: Option, storage_options: Option>, + session: Option, ) -> PyResult> { 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()?)) }) } diff --git a/python/src/lib.rs b/python/src/lib.rs index bd1b8d02..e83c7913 100644 --- a/python/src/lib.rs +++ b/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::()?; + m.add_class::()?; m.add_class::()?; m.add_class::()?; m.add_class::()?; diff --git a/python/src/session.rs b/python/src/session.rs new file mode 100644 index 00000000..abc499e3 --- /dev/null +++ b/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, +} + +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, + metadata_cache_size_bytes: Option, + ) -> PyResult { + 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() + ) + } +} diff --git a/rust/lancedb/src/lib.rs b/rust/lancedb/src/lib.rs index 9f80884f..e2dc1794 100644 --- a/rust/lancedb/src/lib.rs +++ b/rust/lancedb/src/lib.rs @@ -290,3 +290,7 @@ impl Display for DistanceType { /// Connect to a database pub use connection::connect; + +/// Re-export Lance Session and ObjectStoreRegistry for custom session creation +pub use lance::session::Session; +pub use lance_io::object_store::ObjectStoreRegistry;