feat!: move object store registry to the session, re-use stores (#3689)

BREAKING CHANGE: removes `object_store_registry` from `WriteParams` and
`ReadParams`. The registry is now taken from the `Session`, which is
already on those parameters. Also, most `ObjectStore` constructors now
return `Arc<ObjectStore>`.

Closes #3684

* Add a cache of in-use datasets within the registry.
* Move `ObjectStoreRegistry` onto the `Session` object. Combined with
the cache, this lets datasets using the same session share object
stores, as long as they use the same parameters.

Author: wjones127

.github/workflows/rust.yml

@@ -173,7 +173,7 @@ jobs:
           - nightly
     defaults:
       run:
-        working-directory: ./rust/lance
+        working-directory: ./rust
     steps:
       - uses: actions/checkout@v4
       - uses: Swatinem/rust-cache@v2
@@ -200,7 +200,7 @@ jobs:
     runs-on: windows-latest
     defaults:
       run:
-        working-directory: rust/lance
+        working-directory: rust
     steps:
       - uses: actions/checkout@v4
       - uses: Swatinem/rust-cache@v2

java/core/lance-jni/src/blocking_dataset.rs

@@ -116,7 +116,6 @@ impl BlockingDataset {
         read_version: Option<u64>,
         storage_options: HashMap<String, String>,
     ) -> Result<Self> {
-        let object_store_registry = Arc::new(ObjectStoreRegistry::default());
         let inner = RT.block_on(Dataset::commit(
             uri,
             operation,
@@ -126,7 +125,7 @@ impl BlockingDataset {
                 ..Default::default()
             }),
             None,
-            object_store_registry,
+            Default::default(),
             false, // TODO: support enable_v2_manifest_paths
         ))?;
         Ok(Self { inner })

java/core/lance-jni/src/file_reader.rs

@@ -89,7 +89,6 @@ fn inner_open<'local>(env: &mut JNIEnv<'local>, file_uri: JString) -> Result<JOb
 
     let reader = RT.block_on(async move {
         let (obj_store, path) = ObjectStore::from_uri(&file_uri_str).await?;
-        let obj_store = Arc::new(obj_store);
         let config = SchedulerConfig::max_bandwidth(&obj_store);
         let scan_scheduler = ScanScheduler::new(obj_store, config);
 

python/src/dataset.rs

@@ -2196,11 +2196,7 @@ impl PyFullTextQuery {
 
     #[staticmethod]
     #[pyo3(signature = (positive, negative,negative_boost=None))]
-    fn boost_query(
-        positive: PyFullTextQuery,
-        negative: PyFullTextQuery,
-        negative_boost: Option<f32>,
-    ) -> PyResult<Self> {
+    fn boost_query(positive: Self, negative: Self, negative_boost: Option<f32>) -> PyResult<Self> {
         Ok(Self {
             inner: BoostQuery::new(positive.inner, negative.inner, negative_boost).into(),
         })

python/src/file.rs

@@ -333,7 +333,7 @@ fn path_to_parent(path: &Path) -> PyResult<(Path, String)> {
 
 pub async fn object_store_from_uri_or_path_no_options(
     uri_or_path: impl AsRef<str>,
-) -> PyResult<(ObjectStore, Path)> {
+) -> PyResult<(Arc<ObjectStore>, Path)> {
     object_store_from_uri_or_path(uri_or_path, None).await
 }
 
@@ -344,7 +344,7 @@ pub async fn object_store_from_uri_or_path_no_options(
 pub async fn object_store_from_uri_or_path(
     uri_or_path: impl AsRef<str>,
     storage_options: Option<HashMap<String, String>>,
-) -> PyResult<(ObjectStore, Path)> {
+) -> PyResult<(Arc<ObjectStore>, Path)> {
     if let Ok(mut url) = Url::parse(uri_or_path.as_ref()) {
         if url.scheme().len() > 1 {
             let path = object_store::path::Path::parse(url.path()).map_err(|e| {
@@ -376,7 +376,7 @@ pub async fn object_store_from_uri_or_path(
     let path = Path::parse(uri_or_path.as_ref()).map_err(|e| {
         PyIOError::new_err(format!("Invalid path `{}`: {}", uri_or_path.as_ref(), e))
     })?;
-    let object_store = ObjectStore::local();
+    let object_store = Arc::new(ObjectStore::local());
     Ok((object_store, path))
 }
 
@@ -393,7 +393,7 @@ impl LanceFileReader {
         let (object_store, path) =
             object_store_from_uri_or_path(uri_or_path, storage_options).await?;
         let scheduler = ScanScheduler::new(
-            Arc::new(object_store),
+            object_store,
             SchedulerConfig {
                 io_buffer_size_bytes: 2 * 1024 * 1024 * 1024,
             },

python/src/indices.rs

@@ -242,6 +242,7 @@ pub fn transform_vectors(
     )?
 }
 
+#[allow(deprecated)]
 async fn do_shuffle_transformed_vectors(
     unsorted_filenames: Vec<String>,
     dir_path: &str,

rust/lance-file/benches/reader.rs

@@ -33,8 +33,12 @@ fn bench_reader(c: &mut Criterion) {
 
         let tempdir = tempfile::tempdir().unwrap();
         let test_path = tempdir.path();
-        let (object_store, base_path) =
-            ObjectStore::from_path(test_path.as_os_str().to_str().unwrap()).unwrap();
+        let (object_store, base_path) = rt
+            .block_on(ObjectStore::from_uri(
+                test_path.as_os_str().to_str().unwrap(),
+            ))
+            .unwrap();
+
         let file_path = base_path.child("foo.lance");
         let object_writer = rt.block_on(object_store.create(&file_path)).unwrap();
 
@@ -59,7 +63,7 @@ fn bench_reader(c: &mut Criterion) {
                 let data = &data;
                 rt.block_on(async move {
                     let store_scheduler = ScanScheduler::new(
-                        Arc::new(object_store.clone()),
+                        object_store.clone(),
                         SchedulerConfig::default_for_testing(),
                     );
                     let scheduler = store_scheduler.open_file(file_path).await.unwrap();
@@ -125,8 +129,11 @@ fn bench_random_access(c: &mut Criterion) {
 
         let tempdir = tempfile::tempdir().unwrap();
         let test_path = tempdir.path();
-        let (object_store, base_path) =
-            ObjectStore::from_path(test_path.as_os_str().to_str().unwrap()).unwrap();
+        let (object_store, base_path) = rt
+            .block_on(ObjectStore::from_uri(
+                test_path.as_os_str().to_str().unwrap(),
+            ))
+            .unwrap();
         let file_path = base_path.child("foo.lance");
         let object_writer = rt.block_on(object_store.create(&file_path)).unwrap();
 
@@ -150,10 +157,8 @@ fn bench_random_access(c: &mut Criterion) {
         let object_store = &object_store;
         let file_path = &file_path;
         let reader = rt.block_on(async move {
-            let store_scheduler = ScanScheduler::new(
-                Arc::new(object_store.clone()),
-                SchedulerConfig::default_for_testing(),
-            );
+            let store_scheduler =
+                ScanScheduler::new(object_store.clone(), SchedulerConfig::default_for_testing());
             let scheduler = store_scheduler.open_file(file_path).await.unwrap();
             Arc::new(
                 FileReader::try_open(

rust/lance-file/src/reader.rs

@@ -766,6 +766,7 @@ mod tests {
     };
     use arrow_array::{BooleanArray, Int32Array};
     use arrow_schema::{Field as ArrowField, Fields as ArrowFields, Schema as ArrowSchema};
+    use lance_io::object_store::ObjectStoreParams;
 
     #[tokio::test]
     async fn test_take() {
@@ -1364,8 +1365,17 @@ mod tests {
 
     #[tokio::test]
     async fn test_take_boolean_beyond_chunk() {
-        let mut store = ObjectStore::memory();
-        store.set_block_size(256);
+        let store = ObjectStore::from_uri_and_params(
+            Arc::new(Default::default()),
+            "memory://",
+            &ObjectStoreParams {
+                block_size: Some(256),
+                ..Default::default()
+            },
+        )
+        .await
+        .unwrap()
+        .0;
         let path = Path::from("/take_bools");
 
         let arrow_schema = Arc::new(ArrowSchema::new(vec![ArrowField::new(

rust/lance-index/benches/inverted.rs

@@ -34,7 +34,7 @@ fn bench_inverted(c: &mut Criterion) {
     let index_dir = Path::from_filesystem_path(tempdir.path()).unwrap();
     let store = rt.block_on(async {
         Arc::new(LanceIndexStore::new(
-            ObjectStore::local(),
+            Arc::new(ObjectStore::local()),
             index_dir,
             FileMetadataCache::no_cache(),
         ))

rust/lance-index/benches/ngram.rs

@@ -31,7 +31,7 @@ fn bench_ngram(c: &mut Criterion) {
     let index_dir = Path::from_filesystem_path(tempdir.path()).unwrap();
     let store = rt.block_on(async {
         Arc::new(LanceIndexStore::new(
-            ObjectStore::local(),
+            Arc::new(ObjectStore::local()),
             index_dir,
             FileMetadataCache::no_cache(),
         ))