From 51ab7bcb571570ffad19235e116890a410de8bef Mon Sep 17 00:00:00 2001 From: jukejian Date: Sun, 23 Aug 2026 17:42:56 +0800 Subject: [PATCH] feat: expose shared session APIs --- README.md | 28 ++++++++ include/lance/lance.h | 65 +++++++++++++++++++ include/lance/lance.hpp | 59 ++++++++++++++--- src/dataset.rs | 39 ++++++++++- src/lib.rs | 2 + src/session.rs | 128 +++++++++++++++++++++++++++++++++++++ tests/c_api_test.rs | 80 +++++++++++++++++++++++ tests/cpp/test_c_api.c | 25 ++++++++ tests/cpp/test_cpp_api.cpp | 16 +++++ 9 files changed, 433 insertions(+), 9 deletions(-) create mode 100644 src/session.rs diff --git a/README.md b/README.md index 436c2f1..c5f298a 100644 --- a/README.md +++ b/README.md @@ -169,6 +169,34 @@ ds.scan() // consume stream... ``` +### Share metadata and index caches + +Use a session to share Lance's metadata and index caches across dataset +handles. Cache limits are bytes; zero requests zero capacity for the +corresponding cache. + +```c +LanceSession* session = lance_session_new( + 6ULL * 1024 * 1024 * 1024, + 1ULL * 1024 * 1024 * 1024); +LanceDataset* ds = + lance_dataset_open_with_session("data.lance", NULL, 0, session); + +LanceSessionCacheStats stats; +lance_session_get_cache_stats(session, &stats); + +lance_session_close(session); /* ds remains valid */ +lance_dataset_close(ds); +``` + +```cpp +lance::Session session( + 6ULL * 1024 * 1024 * 1024, + 1ULL * 1024 * 1024 * 1024); +auto ds = lance::Dataset::open_with_session(session, "data.lance"); +auto stats = session.cache_stats(); +``` + ### Open at a specific version `lance_dataset_open` takes a `version` argument — `0` means the latest, any diff --git a/include/lance/lance.h b/include/lance/lance.h index db22c62..298cf76 100644 --- a/include/lance/lance.h +++ b/include/lance/lance.h @@ -181,11 +181,57 @@ void lance_free_string(const char* s); typedef struct LanceDataset LanceDataset; typedef struct LanceScanner LanceScanner; typedef struct LanceBatch LanceBatch; +typedef struct LanceSession LanceSession; typedef struct LanceVersions LanceVersions; typedef struct LanceDataStatistics LanceDataStatistics; typedef struct LanceIndexSegmentBuilder LanceIndexSegmentBuilder; typedef struct LanceIndexSegmentMetadata LanceIndexSegmentMetadata; +/* ─── Shared session ─── */ + +/** + * Snapshot of a shared session's cache statistics. + * + * Cache sizes are the bytes currently retained, not their configured limits. + */ +typedef struct LanceSessionCacheStats { + uint64_t index_cache_hits; + uint64_t index_cache_misses; + uint64_t index_cache_entries; + uint64_t index_cache_size_bytes; + uint64_t metadata_cache_hits; + uint64_t metadata_cache_misses; + uint64_t metadata_cache_entries; + uint64_t metadata_cache_size_bytes; +} LanceSessionCacheStats; + +/** + * Create a session that can share metadata and index caches across datasets. + * + * Cache limits are specified in bytes. Pass 0 to request zero capacity. + * @return Session handle, or NULL on error + */ +LanceSession* lance_session_new( + uint64_t index_cache_size_bytes, + uint64_t metadata_cache_size_bytes +); + +/** + * Close a session handle. Safe to call with NULL. Datasets previously opened + * with the session remain valid and retain the shared cache state. + */ +void lance_session_close(LanceSession* session); + +/** + * Copy current cache statistics to `out_stats`. + * + * @return 0 on success, -1 on error + */ +int32_t lance_session_get_cache_stats( + const LanceSession* session, + LanceSessionCacheStats* out_stats +); + /* ─── Dataset lifecycle ─── */ /** @@ -207,6 +253,25 @@ LanceDataset* lance_dataset_open( uint64_t version ); +/** + * Open a Lance dataset using a shared session. + * + * The dataset retains the shared session state and remains valid if the caller + * subsequently closes `session`. + * + * @param uri Dataset path (file://, s3://, memory://, etc.) + * @param storage_opts NULL-terminated key-value pairs ["k1","v1",NULL], or NULL + * @param version Version to open (0 = latest) + * @param session Shared session; must not be NULL + * @return Dataset handle, or NULL on error + */ +LanceDataset* lance_dataset_open_with_session( + const char* uri, + const char* const* storage_opts, + uint64_t version, + const LanceSession* session +); + /** Close and free a dataset handle. Safe to call with NULL. */ void lance_dataset_close(LanceDataset* dataset); diff --git a/include/lance/lance.hpp b/include/lance/lance.hpp index 6d25953..26cb953 100644 --- a/include/lance/lance.hpp +++ b/include/lance/lance.hpp @@ -163,11 +163,43 @@ struct SqlColumn { std::string expression; }; +// ─── Shared Session ────────────────────────────────────────────────────────── + +class Session { + Handle handle_; + +public: + Session(uint64_t index_cache_size_bytes, uint64_t metadata_cache_size_bytes) + : handle_(lance_session_new(index_cache_size_bytes, metadata_cache_size_bytes)) { + if (!handle_) check_error(); + } + + LanceSessionCacheStats cache_stats() const { + LanceSessionCacheStats stats{}; + if (lance_session_get_cache_stats(handle_.get(), &stats) != 0) + check_error(); + return stats; + } + + const LanceSession* c_handle() const { return handle_.get(); } +}; + // ─── Dataset ───────────────────────────────────────────────────────────────── class Dataset { Handle handle_; + static std::vector to_c_storage_options( + const std::vector>& storage_opts) { + std::vector kv; + for (auto& [k, v] : storage_opts) { + kv.push_back(k.c_str()); + kv.push_back(v.c_str()); + } + kv.push_back(nullptr); + return kv; + } + public: /// Open a dataset at the given URI. Pass `version` = 0 (the default) for /// the latest, or a specific version id from `versions()` to check out @@ -177,14 +209,7 @@ class Dataset { const std::vector>& storage_opts = {}, uint64_t version = 0) { - // Build NULL-terminated key-value array for storage options. - std::vector kv; - for (auto& [k, v] : storage_opts) { - kv.push_back(k.c_str()); - kv.push_back(v.c_str()); - } - kv.push_back(nullptr); - + auto kv = to_c_storage_options(storage_opts); const char* const* opts_ptr = storage_opts.empty() ? nullptr : kv.data(); @@ -193,6 +218,24 @@ class Dataset { return Dataset(ds); } + /// Open a dataset with caches owned by `session`. The returned dataset + /// remains valid if the Session wrapper is destroyed first. + static Dataset open_with_session( + const Session& session, + const std::string& uri, + const std::vector>& storage_opts = {}, + uint64_t version = 0) { + + auto kv = to_c_storage_options(storage_opts); + const char* const* opts_ptr = + storage_opts.empty() ? nullptr : kv.data(); + + auto* ds = lance_dataset_open_with_session( + uri.c_str(), opts_ptr, version, session.c_handle()); + if (!ds) check_error(); + return Dataset(ds); + } + /// Write an Arrow record batch stream to a Lance dataset and return the /// open dataset at the committed version. /// diff --git a/src/dataset.rs b/src/dataset.rs index 1397b74..cc1f87c 100644 --- a/src/dataset.rs +++ b/src/dataset.rs @@ -17,6 +17,7 @@ use lance_core::Result; use crate::error::{ffi_try, swallow_unwind}; use crate::helpers; use crate::runtime::block_on; +use crate::session::LanceSession; use crate::stream_guard::guarded_ffi_stream_from_reader; /// Opaque handle representing an opened Lance dataset. @@ -121,15 +122,48 @@ pub unsafe extern "C" fn lance_dataset_open( version: u64, ) -> *mut LanceDataset { ffi_try!( - unsafe { open_dataset_inner(uri, storage_options, version) }, + unsafe { open_dataset_inner(uri, storage_options, version, None) }, null ) } +/// Open a Lance dataset using a shared session. +/// +/// The dataset retains shared ownership of the session state, so the caller +/// may close the session handle after this function returns successfully. +#[unsafe(no_mangle)] +pub unsafe extern "C" fn lance_dataset_open_with_session( + uri: *const c_char, + storage_options: *const *const c_char, + version: u64, + session: *const LanceSession, +) -> *mut LanceDataset { + ffi_try!( + unsafe { open_dataset_with_session_inner(uri, storage_options, version, session) }, + null + ) +} + +unsafe fn open_dataset_with_session_inner( + uri: *const c_char, + storage_options: *const *const c_char, + version: u64, + session: *const LanceSession, +) -> Result<*mut LanceDataset> { + if session.is_null() { + return Err(lance_core::Error::invalid_input_source( + "session must not be NULL".into(), + )); + } + let session = unsafe { &*session }; + unsafe { open_dataset_inner(uri, storage_options, version, Some(session)) } +} + unsafe fn open_dataset_inner( uri: *const c_char, storage_options: *const *const c_char, version: u64, + session: Option<&LanceSession>, ) -> Result<*mut LanceDataset> { let uri_str = unsafe { helpers::parse_c_string(uri)? } .ok_or_else(|| lance_core::Error::invalid_input_source("uri must not be NULL".into()))?; @@ -143,6 +177,9 @@ unsafe fn open_dataset_inner( if version != 0 { builder = builder.with_version(version); } + if let Some(session) = session { + builder = builder.with_session(session.inner.clone()); + } let dataset = block_on(builder.load())?; let handle = LanceDataset { diff --git a/src/lib.rs b/src/lib.rs index d58e875..33a7253 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -39,6 +39,7 @@ mod merge_insert; mod restore; pub mod runtime; mod scanner; +mod session; pub mod stream_guard; mod update; mod versions; @@ -63,6 +64,7 @@ pub use index_segment::*; pub use merge_insert::*; pub use restore::*; pub use scanner::*; +pub use session::*; pub use update::*; pub use versions::*; pub use writer::*; diff --git a/src/session.rs b/src/session.rs new file mode 100644 index 0000000..60a1623 --- /dev/null +++ b/src/session.rs @@ -0,0 +1,128 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright The Lance Authors + +//! Shared Lance session C API. + +use std::sync::Arc; + +use lance::session::Session; +use lance_core::Result; + +use crate::error::{ffi_try, swallow_unwind}; +use crate::runtime::block_on; + +/// Opaque handle for sharing Lance metadata and index caches across datasets. +pub struct LanceSession { + pub(crate) inner: Arc, +} + +/// Snapshot of a session's metadata and index cache statistics. +#[repr(C)] +#[derive(Clone, Copy, Debug, Default, PartialEq, Eq)] +pub struct LanceSessionCacheStats { + pub index_cache_hits: u64, + pub index_cache_misses: u64, + pub index_cache_entries: u64, + pub index_cache_size_bytes: u64, + pub metadata_cache_hits: u64, + pub metadata_cache_misses: u64, + pub metadata_cache_entries: u64, + pub metadata_cache_size_bytes: u64, +} + +/// Create a shared Lance session with byte-based cache limits. +/// +/// A zero limit requests zero capacity for the corresponding cache. +#[unsafe(no_mangle)] +pub extern "C" fn lance_session_new( + index_cache_size_bytes: u64, + metadata_cache_size_bytes: u64, +) -> *mut LanceSession { + ffi_try!( + session_new_inner(index_cache_size_bytes, metadata_cache_size_bytes), + null + ) +} + +fn session_new_inner( + index_cache_size_bytes: u64, + metadata_cache_size_bytes: u64, +) -> Result<*mut LanceSession> { + let index_cache_size_bytes = u64_to_usize(index_cache_size_bytes, "index_cache_size_bytes")?; + let metadata_cache_size_bytes = + u64_to_usize(metadata_cache_size_bytes, "metadata_cache_size_bytes")?; + let session = Session::new( + index_cache_size_bytes, + metadata_cache_size_bytes, + Default::default(), + ); + Ok(Box::into_raw(Box::new(LanceSession { + inner: Arc::new(session), + }))) +} + +/// Close a session handle. +/// +/// Datasets opened with this session retain shared ownership of its runtime +/// state and remain valid after this handle is closed. +#[unsafe(no_mangle)] +pub unsafe extern "C" fn lance_session_close(session: *mut LanceSession) { + if !session.is_null() { + swallow_unwind("lance_session_close", || unsafe { + let _ = Box::from_raw(session); + }); + } +} + +/// Copy the current cache statistics into `out_stats`. +#[unsafe(no_mangle)] +pub unsafe extern "C" fn lance_session_get_cache_stats( + session: *const LanceSession, + out_stats: *mut LanceSessionCacheStats, +) -> i32 { + ffi_try!( + unsafe { session_get_cache_stats_inner(session, out_stats) }, + neg + ) +} + +unsafe fn session_get_cache_stats_inner( + session: *const LanceSession, + out_stats: *mut LanceSessionCacheStats, +) -> Result { + if session.is_null() || out_stats.is_null() { + return Err(lance_core::Error::invalid_input_source( + "session and out_stats must not be NULL".into(), + )); + } + let session = unsafe { &*session }; + let (index, metadata) = block_on(async { + let index = session.inner.index_cache_stats().await; + let metadata = session.inner.metadata_cache_stats().await; + (index, metadata) + }); + unsafe { + std::ptr::write_unaligned( + out_stats, + LanceSessionCacheStats { + index_cache_hits: index.hits, + index_cache_misses: index.misses, + index_cache_entries: index.num_entries as u64, + index_cache_size_bytes: index.size_bytes as u64, + metadata_cache_hits: metadata.hits, + metadata_cache_misses: metadata.misses, + metadata_cache_entries: metadata.num_entries as u64, + metadata_cache_size_bytes: metadata.size_bytes as u64, + }, + ); + } + Ok(0) +} + +fn u64_to_usize(value: u64, field: &'static str) -> Result { + usize::try_from(value).map_err(|_| { + lance_core::Error::invalid_input_source( + format!("{field}={value} exceeds usize::MAX on this target").into(), + ) + }) +} diff --git a/tests/c_api_test.rs b/tests/c_api_test.rs index 1a1c78c..8b6b194 100644 --- a/tests/c_api_test.rs +++ b/tests/c_api_test.rs @@ -233,6 +233,86 @@ fn test_open_close() { unsafe { lance_dataset_close(ptr::null_mut()) }; } +#[test] +fn test_shared_session_open_and_cache_stats() { + let (_tmp, uri) = create_test_dataset(); + let c_uri = c_str(&uri); + let session = lance_session_new(0, 16 * 1024 * 1024); + assert!(!session.is_null(), "session creation should succeed"); + + let mut initial = LanceSessionCacheStats::default(); + assert_eq!( + unsafe { lance_session_get_cache_stats(session, &mut initial) }, + 0 + ); + assert_eq!(initial.metadata_cache_hits, 0); + assert_eq!(initial.metadata_cache_misses, 0); + + let first = unsafe { lance_dataset_open_with_session(c_uri.as_ptr(), ptr::null(), 0, session) }; + assert!(!first.is_null(), "first shared-session open should succeed"); + assert_eq!( + scan_all_rows(first) + .iter() + .map(|batch| batch.num_rows()) + .sum::(), + 5 + ); + unsafe { lance_dataset_close(first) }; + + let second = + unsafe { lance_dataset_open_with_session(c_uri.as_ptr(), ptr::null(), 0, session) }; + assert!( + !second.is_null(), + "second shared-session open should succeed" + ); + assert_eq!( + scan_all_rows(second) + .iter() + .map(|batch| batch.num_rows()) + .sum::(), + 5 + ); + + // Cache activity depends on whether the backend supplies manifest sizes. + let mut stats = LanceSessionCacheStats::default(); + assert_eq!( + unsafe { lance_session_get_cache_stats(session, &mut stats) }, + 0 + ); + + unsafe { lance_session_close(session) }; + assert_eq!(unsafe { lance_dataset_count_rows(second) }, 5); + unsafe { lance_dataset_close(second) }; +} + +#[test] +fn test_shared_session_rejects_null_inputs() { + let (_tmp, uri) = create_test_dataset(); + let c_uri = c_str(&uri); + let ds = + unsafe { lance_dataset_open_with_session(c_uri.as_ptr(), ptr::null(), 0, ptr::null()) }; + assert!(ds.is_null()); + assert_eq!(lance_last_error_code(), LanceErrorCode::InvalidArgument); + + let session = lance_session_new(0, 0); + assert!(!session.is_null()); + let mut stats = LanceSessionCacheStats::default(); + assert_eq!( + unsafe { lance_session_get_cache_stats(ptr::null(), &mut stats) }, + -1 + ); + assert_eq!(lance_last_error_code(), LanceErrorCode::InvalidArgument); + assert_eq!( + unsafe { lance_session_get_cache_stats(session, ptr::null_mut()) }, + -1 + ); + assert_eq!(lance_last_error_code(), LanceErrorCode::InvalidArgument); + unsafe { + lance_session_close(session); + lance_session_close(ptr::null_mut()); + } +} + #[test] fn test_open_nonexistent() { let c_uri = c_str("memory://nonexistent_dataset_xyz"); diff --git a/tests/cpp/test_c_api.c b/tests/cpp/test_c_api.c index 4499cb9..c49ecfa 100644 --- a/tests/cpp/test_c_api.c +++ b/tests/cpp/test_c_api.c @@ -102,6 +102,30 @@ static void test_open_and_metadata(const char *uri) { printf("OK\n"); } +static void test_shared_session(const char *uri) { + printf(" test_shared_session... "); + + LanceSession *session = lance_session_new(0, 16 * 1024 * 1024); + ASSERT(session != NULL, "session creation failed"); + + LanceDataset *ds = lance_dataset_open_with_session(uri, NULL, 0, session); + ASSERT(ds != NULL, "shared-session dataset open failed"); + + LanceSessionCacheStats stats; + memset(&stats, 0, sizeof(stats)); + int32_t rc = lance_session_get_cache_stats(session, &stats); + ASSERT(rc == 0, "session cache stats failed"); + + /* The dataset retains the shared state after the caller drops its handle. */ + lance_session_close(session); + ASSERT(lance_dataset_count_rows(ds) > 0, + "dataset should remain valid after session close"); + + lance_dataset_close(ds); + printf("metadata_entries=%llu... OK\n", + (unsigned long long)stats.metadata_cache_entries); +} + static void test_scan(const char *uri) { printf(" test_scan... "); @@ -948,6 +972,7 @@ int main(int argc, char **argv) { printf("Running C API tests with dataset: %s\n", uri); test_open_and_metadata(uri); + test_shared_session(uri); test_scan(uri); test_scan_with_limit(uri); test_versions(uri); diff --git a/tests/cpp/test_cpp_api.cpp b/tests/cpp/test_cpp_api.cpp index b5d090a..17b1ab6 100644 --- a/tests/cpp/test_cpp_api.cpp +++ b/tests/cpp/test_cpp_api.cpp @@ -84,6 +84,21 @@ static void test_dataset_open(const std::string& uri) { PASS(); } +static void test_shared_session(const std::string& uri) { + TEST(test_shared_session); + + auto session = std::make_unique(0, 16 * 1024 * 1024); + auto ds = lance::Dataset::open_with_session(*session, uri); + auto stats = session->cache_stats(); + + session.reset(); + assert(ds.count_rows() > 0); + + printf("metadata_entries=%llu... ", + (unsigned long long)stats.metadata_cache_entries); + PASS(); +} + static void test_dataset_schema(const std::string& uri) { TEST(test_dataset_schema); @@ -906,6 +921,7 @@ int main(int argc, char** argv) { printf("Running C++ API tests with dataset: %s\n", uri.c_str()); test_dataset_open(uri); + test_shared_session(uri); test_dataset_schema(uri); test_scanner_fluent(uri); test_scanner_async_stream_ownership(uri);