diff --git a/docs/index.rst b/docs/index.rst index e761263..fa2c261 100644 --- a/docs/index.rst +++ b/docs/index.rst @@ -38,6 +38,23 @@ This example shows how to create a LanceDB table with vector data using Apache A LanceDBConnectBuilder* builder = lancedb_connect("./my_database"); LanceDBConnection* db = lancedb_connect_builder_execute(builder); +.. note:: + + You can optionally configure session cache sizes and attach the session to + the connection builder: + + .. code-block:: c + + LanceDBSessionOptions options = {0}; + options.index_cache_bytes = 512 * 1024 * 1024; + options.metadata_cache_bytes = 256 * 1024 * 1024; + + LanceDBSession* session = lancedb_session_new(&options); + + LanceDBConnectBuilder* builder = lancedb_connect("./my_database"); + builder = lancedb_connect_builder_session(builder, session); + LanceDBConnection* db = lancedb_connect_builder_execute(builder); + **2. Create Arrow schema (using Arrow C++ API):** .. code-block:: cpp diff --git a/include/lancedb.h b/include/lancedb.h index 4aa47a9..6dea360 100644 --- a/include/lancedb.h +++ b/include/lancedb.h @@ -7,6 +7,7 @@ #define LANCEDB_H #include +#include #ifdef __cplusplus extern "C" { @@ -47,6 +48,11 @@ typedef struct LanceDBVectorQuery LanceDBVectorQuery; */ typedef struct LanceDBQueryResult LanceDBQueryResult; +/** + * Opaque handle to a LanceDB Session + */ +typedef struct LanceDBSession LanceDBSession; + /** * Opaque handle to Arrow RecordBatchReader */ @@ -212,6 +218,24 @@ typedef struct { int when_not_matched_insert_all; // Insert all new records (1 = true, 0 = false) } LanceDBMergeInsertConfig; +/** + * Session creation options + */ +typedef struct { + size_t index_cache_bytes; // Index cache size in bytes (0 = default) + size_t metadata_cache_bytes; // Metadata cache size in bytes (0 = default) +} LanceDBSessionOptions; + +/** + * Session cache statistics + */ +typedef struct { + uint64_t hits; // Number of cache hits + uint64_t misses; // Number of cache misses + size_t num_entries; // Number of entries in cache + size_t size_bytes; // Cache size in bytes +} LanceDBSessionCacheStats; + /** * Create a ConnectBuilder for the given URI * @@ -253,6 +277,18 @@ LanceDBConnection* lancedb_connect_builder_execute(LanceDBConnectBuilder* builde */ LanceDBConnectBuilder* lancedb_connect_builder_storage_option(LanceDBConnectBuilder* builder, const char* key, const char* value); +/** + * Set session for the connection builder + * + * @param builder - pointer to LanceDBConnectBuilder returned from lancedb_connect() + * @param session - pointer to LanceDBSession, or NULL to keep current/default session behavior + * @return Non-null pointer to LanceDBConnectBuilder on success, NULL on failure + * + * The builder is consumed by this function and must not be used after calling. + * Passing NULL session is a no-op. + */ +LanceDBConnectBuilder* lancedb_connect_builder_session(LanceDBConnectBuilder* builder, const LanceDBSession* session); + /** * Free a ConnectBuilder * @@ -523,6 +559,59 @@ void lancedb_free_namespace_list(char** namespaces, size_t count); */ void lancedb_connection_free(LanceDBConnection* connection); +/** + * Create a new session + * + * @param options - pointer to LanceDBSessionOptions, or NULL for defaults + * @return Non-null pointer to LanceDBSession on success, NULL on failure + * + * The returned session must be freed with lancedb_session_free(). + */ +LanceDBSession* lancedb_session_new(const LanceDBSessionOptions* options); + +/** + * Get index cache stats for a session + * + * @param session - pointer to LanceDBSession + * @param out_stats - pointer to receive session cache statistics + * @param error_message - optional pointer to receive detailed error message (NULL to ignore) + * @return Error code indicating success or failure + * + * If error_message is provided and an error occurs, the caller must free + * the error message with lancedb_free_string(). + */ +LanceDBError lancedb_session_index_cache_stats( + const LanceDBSession* session, + LanceDBSessionCacheStats* out_stats, + char** error_message +); + +/** + * Get metadata cache stats for a session + * + * @param session - pointer to LanceDBSession + * @param out_stats - pointer to receive session cache statistics + * @param error_message - optional pointer to receive detailed error message (NULL to ignore) + * @return Error code indicating success or failure + * + * If error_message is provided and an error occurs, the caller must free + * the error message with lancedb_free_string(). + */ +LanceDBError lancedb_session_metadata_cache_stats( + const LanceDBSession* session, + LanceDBSessionCacheStats* out_stats, + char** error_message +); + +/** + * Free a Session + * + * @param session - pointer to LanceDBSession + * + * After calling this function, the session pointer must not be used. + */ +void lancedb_session_free(LanceDBSession* session); + /** * Free a Table * diff --git a/src/connection.rs b/src/connection.rs index 17a6b77..a91aee1 100644 --- a/src/connection.rs +++ b/src/connection.rs @@ -46,6 +46,31 @@ pub struct LanceDBTableNamesBuilder { inner: Box, } +/// Opaque handle to a Session +#[repr(C)] +pub struct LanceDBSession { + inner: Arc, +} + +/// Session creation options +#[repr(C)] +pub struct LanceDBSessionOptions { + pub index_cache_bytes: usize, + pub metadata_cache_bytes: usize, +} + +/// Cache statistics for a session cache +#[repr(C)] +pub struct LanceDBSessionCacheStats { + pub hits: u64, + pub misses: u64, + pub num_entries: usize, + pub size_bytes: usize, +} + +const DEFAULT_INDEX_CACHE_SIZE_BYTES: usize = 6 * 1024 * 1024 * 1024; +const DEFAULT_METADATA_CACHE_SIZE_BYTES: usize = 1024 * 1024 * 1024; + /// Runtime to handle async operations static RUNTIME: OnceLock = OnceLock::new(); @@ -157,6 +182,42 @@ pub unsafe extern "C" fn lancedb_connect_builder_storage_option( Box::into_raw(boxed_builder) } +/// Set a session for the connection builder +/// +/// # Safety +/// - `builder` must be a valid pointer returned from `lancedb_connect` +/// - `builder` will be consumed and must not be used after calling this function +/// - `session` can be NULL (no-op) or a valid pointer returned from `lancedb_session_new` +/// +/// # Returns +/// - A new pointer to LanceDBConnectBuilder on success +/// - Null pointer on failure +#[no_mangle] +pub unsafe extern "C" fn lancedb_connect_builder_session( + builder: *mut LanceDBConnectBuilder, + session: *const LanceDBSession, +) -> *mut LanceDBConnectBuilder { + if builder.is_null() { + return ptr::null_mut(); + } + + let builder_box = Box::from_raw(builder); + let connect_builder = *builder_box.inner; + + let updated_builder = if session.is_null() { + connect_builder + } else { + let session_ref = &*session; + connect_builder.session(session_ref.inner.clone()) + }; + + let boxed_builder = Box::new(LanceDBConnectBuilder { + inner: Box::new(updated_builder), + }); + + Box::into_raw(boxed_builder) +} + /// Free a ConnectBuilder /// /// # Safety @@ -199,6 +260,116 @@ pub unsafe extern "C" fn lancedb_connection_uri( cached_uri.as_ptr() } +/// Create a new session +/// +/// # Safety +/// - `options` can be NULL to use default session configuration +/// +/// # Returns +/// - Non-null pointer to LanceDBSession on success +/// - Null pointer on failure +#[no_mangle] +pub unsafe extern "C" fn lancedb_session_new( + options: *const LanceDBSessionOptions, +) -> *mut LanceDBSession { + let session = if options.is_null() { + Arc::new(lancedb::Session::default()) + } else { + let session_options = &*options; + let index_cache_bytes = if session_options.index_cache_bytes == 0 { + DEFAULT_INDEX_CACHE_SIZE_BYTES + } else { + session_options.index_cache_bytes + }; + let metadata_cache_bytes = if session_options.metadata_cache_bytes == 0 { + DEFAULT_METADATA_CACHE_SIZE_BYTES + } else { + session_options.metadata_cache_bytes + }; + Arc::new(lancedb::Session::new( + index_cache_bytes, + metadata_cache_bytes, + Arc::new(lancedb::ObjectStoreRegistry::default()), + )) + }; + + Box::into_raw(Box::new(LanceDBSession { inner: session })) +} + +/// Get index cache stats for a session +/// +/// # Safety +/// - `session` must be a valid pointer returned from `lancedb_session_new` +/// - `out_stats` must be a valid pointer to receive cache stats +/// - `error_message` can be NULL to ignore detailed error messages +/// +/// # Returns +/// - Error code indicating success or failure +#[no_mangle] +pub unsafe extern "C" fn lancedb_session_index_cache_stats( + session: *const LanceDBSession, + out_stats: *mut LanceDBSessionCacheStats, + error_message: *mut *mut c_char, +) -> LanceDBError { + if session.is_null() || out_stats.is_null() { + set_invalid_argument_message(error_message); + return LanceDBError::InvalidArgument; + } + + let session_ref = &*session; + let stats = get_runtime().block_on(session_ref.inner.index_cache_stats()); + *out_stats = LanceDBSessionCacheStats { + hits: stats.hits, + misses: stats.misses, + num_entries: stats.num_entries, + size_bytes: stats.size_bytes, + }; + LanceDBError::Success +} + +/// Get metadata cache stats for a session +/// +/// # Safety +/// - `session` must be a valid pointer returned from `lancedb_session_new` +/// - `out_stats` must be a valid pointer to receive cache stats +/// - `error_message` can be NULL to ignore detailed error messages +/// +/// # Returns +/// - Error code indicating success or failure +#[no_mangle] +pub unsafe extern "C" fn lancedb_session_metadata_cache_stats( + session: *const LanceDBSession, + out_stats: *mut LanceDBSessionCacheStats, + error_message: *mut *mut c_char, +) -> LanceDBError { + if session.is_null() || out_stats.is_null() { + set_invalid_argument_message(error_message); + return LanceDBError::InvalidArgument; + } + + let session_ref = &*session; + let stats = get_runtime().block_on(session_ref.inner.metadata_cache_stats()); + *out_stats = LanceDBSessionCacheStats { + hits: stats.hits, + misses: stats.misses, + num_entries: stats.num_entries, + size_bytes: stats.size_bytes, + }; + LanceDBError::Success +} + +/// Free a Session +/// +/// # Safety +/// - `session` must be a valid pointer returned from `lancedb_session_new` +/// - `session` must not be used after calling this function +#[no_mangle] +pub unsafe extern "C" fn lancedb_session_free(session: *mut LanceDBSession) { + if !session.is_null() { + let _ = Box::from_raw(session); + } +} + /// Create a new table with Arrow schema and data /// /// # Safety diff --git a/tests/test_connection.cpp b/tests/test_connection.cpp index 605524a..c6844d5 100644 --- a/tests/test_connection.cpp +++ b/tests/test_connection.cpp @@ -23,6 +23,8 @@ TEST_CASE_METHOD(BaseFixture, "LanceDB Connection Builder", "[connection]") { REQUIRE(builder != nullptr); builder = lancedb_connect_builder_storage_option(builder, "hello", "world"); REQUIRE(builder != nullptr); + builder = lancedb_connect_builder_session(builder, nullptr); + REQUIRE(builder != nullptr); LanceDBConnection* db = lancedb_connect_builder_execute(builder); REQUIRE(db != nullptr); lancedb_connection_free(db); @@ -60,6 +62,81 @@ TEST_CASE_METHOD(BaseFixture, "LanceDB Connection Builder", "[connection]") { builder = lancedb_connect_builder_storage_option(builder, "hello", NON_UTF8); REQUIRE(builder == nullptr); } + SECTION("Attach session to connect builder") { + LanceDBSession* session = lancedb_session_new(nullptr); + REQUIRE(session != nullptr); + LanceDBConnectBuilder* builder = lancedb_connect(uri.c_str()); + REQUIRE(builder != nullptr); + builder = lancedb_connect_builder_session(builder, session); + REQUIRE(builder != nullptr); + LanceDBConnection* db = lancedb_connect_builder_execute(builder); + REQUIRE(db != nullptr); + lancedb_connection_free(db); + lancedb_session_free(session); + } + SECTION("NULL session on builder should be allowed and use default session") { + LanceDBConnectBuilder* builder = lancedb_connect(uri.c_str()); + REQUIRE(builder != nullptr); + builder = lancedb_connect_builder_session(builder, nullptr); + REQUIRE(builder != nullptr); + LanceDBConnection* db = lancedb_connect_builder_execute(builder); + REQUIRE(db != nullptr); + lancedb_connection_free(db); + } +} + +TEST_CASE_METHOD(BaseFixture, "LanceDB Session", "[connection]") { + SECTION("Create and free default session") { + LanceDBSession* session = lancedb_session_new(nullptr); + REQUIRE(session != nullptr); + lancedb_session_free(session); + } + SECTION("Create session with options") { + LanceDBSessionOptions options{}; + options.index_cache_bytes = 1024 * 1024; + options.metadata_cache_bytes = 2 * 1024 * 1024; + LanceDBSession* session = lancedb_session_new(&options); + REQUIRE(session != nullptr); + lancedb_session_free(session); + } + SECTION("Zero options fallback to defaults") { + LanceDBSessionOptions options{}; + options.index_cache_bytes = 0; + options.metadata_cache_bytes = 0; + LanceDBSession* session = lancedb_session_new(&options); + REQUIRE(session != nullptr); + lancedb_session_free(session); + } + SECTION("Get session cache stats") { + LanceDBSession* session = lancedb_session_new(nullptr); + REQUIRE(session != nullptr); + LanceDBSessionCacheStats index_stats{}; + LanceDBSessionCacheStats metadata_stats{}; + char* error_message = nullptr; + auto index_result = lancedb_session_index_cache_stats(session, &index_stats, &error_message); + REQUIRE(index_result == LANCEDB_SUCCESS); + REQUIRE(error_message == nullptr); + auto metadata_result = lancedb_session_metadata_cache_stats(session, &metadata_stats, &error_message); + REQUIRE(metadata_result == LANCEDB_SUCCESS); + REQUIRE(error_message == nullptr); + lancedb_session_free(session); + } + SECTION("Get session cache stats with invalid args should fail") { + LanceDBSession* session = lancedb_session_new(nullptr); + REQUIRE(session != nullptr); + LanceDBSessionCacheStats stats{}; + char* error_message = nullptr; + auto result = lancedb_session_index_cache_stats(nullptr, &stats, &error_message); + REQUIRE(result == LANCEDB_INVALID_ARGUMENT); + REQUIRE(error_message != nullptr); + lancedb_free_string(error_message); + error_message = nullptr; + result = lancedb_session_metadata_cache_stats(session, nullptr, &error_message); + REQUIRE(result == LANCEDB_INVALID_ARGUMENT); + REQUIRE(error_message != nullptr); + lancedb_free_string(error_message); + lancedb_session_free(session); + } } TEST_CASE_METHOD(LanceDBFixture, "LanceDB Tables", "[connection]") { @@ -401,4 +478,3 @@ TEST_CASE_METHOD(LanceDBFixture, "LanceDB Namespaces", "[connection]") { lancedb_free_string(error_message); } } -