feat: change dataset uri to return full qualified url instead of object store path (lancedb#2416)

Author: eddyxu

Cargo.toml

@@ -20,7 +20,7 @@ exclude = ["python"]
 resolver = "2"
 
 [workspace.package]
-version = "0.11.2"
+version = "0.12.0"
 edition = "2021"
 authors = ["Lance Devs <dev@lancedb.com>"]
 license = "Apache-2.0"
@@ -43,20 +43,20 @@ categories = [
 rust-version = "1.75"
 
 [workspace.dependencies]
-lance = { version = "=0.11.2", path = "./rust/lance" }
-lance-arrow = { version = "=0.11.2", path = "./rust/lance-arrow" }
-lance-core = { version = "=0.11.2", path = "./rust/lance-core" }
-lance-datafusion = { version = "=0.11.2", path = "./rust/lance-datafusion" }
-lance-datagen = { version = "=0.11.2", path = "./rust/lance-datagen" }
-lance-encoding = { version = "=0.11.2", path = "./rust/lance-encoding" }
-lance-encoding-datafusion = { version = "=0.11.2", path = "./rust/lance-encoding-datafusion" }
-lance-file = { version = "=0.11.2", path = "./rust/lance-file" }
-lance-index = { version = "=0.11.2", path = "./rust/lance-index" }
-lance-io = { version = "=0.11.2", path = "./rust/lance-io" }
-lance-linalg = { version = "=0.11.2", path = "./rust/lance-linalg" }
-lance-table = { version = "=0.11.2", path = "./rust/lance-table" }
-lance-test-macros = { version = "=0.11.2", path = "./rust/lance-test-macros" }
-lance-testing = { version = "=0.11.2", path = "./rust/lance-testing" }
+lance = { version = "=0.12.0", path = "./rust/lance" }
+lance-arrow = { version = "=0.12.0", path = "./rust/lance-arrow" }
+lance-core = { version = "=0.12.0", path = "./rust/lance-core" }
+lance-datafusion = { version = "=0.12.0", path = "./rust/lance-datafusion" }
+lance-datagen = { version = "=0.12.0", path = "./rust/lance-datagen" }
+lance-encoding = { version = "=0.12.0", path = "./rust/lance-encoding" }
+lance-encoding-datafusion = { version = "=0.12.0", path = "./rust/lance-encoding-datafusion" }
+lance-file = { version = "=0.12.0", path = "./rust/lance-file" }
+lance-index = { version = "=0.12.0", path = "./rust/lance-index" }
+lance-io = { version = "=0.12.0", path = "./rust/lance-io" }
+lance-linalg = { version = "=0.12.0", path = "./rust/lance-linalg" }
+lance-table = { version = "=0.12.0", path = "./rust/lance-table" }
+lance-test-macros = { version = "=0.12.0", path = "./rust/lance-test-macros" }
+lance-testing = { version = "=0.12.0", path = "./rust/lance-testing" }
 approx = "0.5.1"
 # Note that this one does not include pyarrow
 arrow = { version = "51.0.0", optional = false, features = ["prettyprint"] }
@@ -130,7 +130,7 @@ rustc_version = "0.4"
 serde = { version = "^1" }
 serde_json = { version = "1" }
 shellexpand = "3.0"
-snafu = "0.7.4"
+snafu = "0.7.5"
 tempfile = "3"
 test-log = { version = "0.2.15" }
 tokio = { version = "1.23", features = [

python/Cargo.toml

@@ -1,6 +1,6 @@
 [package]
 name = "pylance"
-version = "0.11.2"
+version = "0.12.0"
 edition = "2021"
 authors = ["Lance Devs <dev@lancedb.com>"]
 rust-version = "1.65"

python/src/fragment.rs

@@ -432,7 +432,7 @@ impl FragmentMetadata {
 
 #[pyfunction(name = "_cleanup_partial_writes")]
 pub fn cleanup_partial_writes(base_uri: &str, files: Vec<(String, String)>) -> PyResult<()> {
-    let (store, _) = RT
+    let (store, base_path) = RT
         .runtime
         .block_on(ObjectStore::from_uri(base_uri))
         .map_err(|err| PyIOError::new_err(format!("Failed to create object store: {}", err)))?;
@@ -443,15 +443,19 @@ pub fn cleanup_partial_writes(base_uri: &str, files: Vec<(String, String)>) -> P
         .collect();
 
     #[allow(clippy::map_identity)]
-    async fn inner(store: ObjectStore, files: Vec<(Path, String)>) -> Result<(), ::lance::Error> {
+    async fn inner(
+        store: ObjectStore,
+        base_path: Path,
+        files: Vec<(Path, String)>,
+    ) -> Result<(), ::lance::Error> {
         let files_iter = files
             .iter()
             .map(|(path, multipart_id)| (path, multipart_id));
-        lance::dataset::cleanup::cleanup_partial_writes(&store, files_iter).await
+        lance::dataset::cleanup::cleanup_partial_writes(&store, &base_path, files_iter).await
     }
 
     RT.runtime
-        .block_on(inner(store, files))
+        .block_on(inner(store, base_path, files))
         .map_err(|err| PyIOError::new_err(format!("Failed to cleanup files: {}", err)))
 }
 

rust/lance-index/src/vector/hnsw/builder.rs

@@ -766,6 +766,7 @@ mod tests {
     use lance_linalg::distance::DistanceType;
     use lance_table::format::SelfDescribingFileReader;
     use lance_testing::datagen::generate_random_array;
+    use object_store::path::Path;
 
     use crate::vector::{
         flat::storage::FlatStorage,
@@ -792,7 +793,7 @@ mod tests {
         .unwrap();
 
         let object_store = ObjectStore::memory();
-        let path = object_store.base_path().child("test_builder_write_load");
+        let path = Path::from("test_builder_write_load");
         let writer = object_store.create(&path).await.unwrap();
         let schema = Schema::new(vec![
             VECTOR_ID_FIELD.clone(),

rust/lance-io/src/object_store.rs

@@ -84,7 +84,6 @@ pub struct ObjectStore {
     // Inner object store
     pub inner: Arc<dyn OSObjectStore>,
     scheme: String,
-    base_path: Path,
     block_size: usize,
 }
 
@@ -94,9 +93,7 @@ impl DeepSizeOf for ObjectStore {
         // shouldn't be too big.  The only exception might be the write cache but, if
         // the writer cache has data, it means we're using it somewhere else that isn't
         // a cache and so that doesn't really count.
-        self.scheme.deep_size_of_children(context)
-            + self.base_path.as_ref().deep_size_of_children(context)
-            + self.block_size.deep_size_of_children(context)
+        self.scheme.deep_size_of_children(context) + self.block_size.deep_size_of_children(context)
     }
 }
 
@@ -356,15 +353,14 @@ impl ObjectStore {
         uri: &str,
         params: &ObjectStoreParams,
     ) -> Result<(Self, Path)> {
-        let (object_store, base_path) = match Url::parse(uri) {
+        let (object_store, path) = match Url::parse(uri) {
             Ok(url) if url.scheme().len() == 1 && cfg!(windows) => {
                 // On Windows, the drive is parsed as a scheme
                 Self::from_path(uri)
             }
             Ok(url) => {
                 let store = Self::new_from_url(url.clone(), params.clone()).await?;
-                let path = Path::from(url.path());
-                Ok((store, path))
+                Ok((store, Path::from(url.path())))
             }
             Err(_) => Self::from_path(uri),
         }?;
@@ -378,7 +374,7 @@ impl ObjectStore {
                     .unwrap_or(object_store.inner),
                 ..object_store
             },
-            base_path,
+            path,
         ))
     }
 
@@ -399,7 +395,6 @@ impl ObjectStore {
             Self {
                 inner: Arc::new(LocalFileSystem::new()).traced(),
                 scheme: String::from(scheme),
-                base_path: Path::from_absolute_path(expanded_path.as_path())?,
                 block_size: 4 * 1024, // 4KB block size
             },
             Path::from_absolute_path(expanded_path.as_path())?,
@@ -419,7 +414,6 @@ impl ObjectStore {
         Self {
             inner: Arc::new(LocalFileSystem::new()).traced(),
             scheme: String::from("file"),
-            base_path: Path::from("/"),
             block_size: 4 * 1024, // 4KB block size
         }
     }
@@ -429,7 +423,6 @@ impl ObjectStore {
         Self {
             inner: Arc::new(InMemory::new()).traced(),
             scheme: String::from("memory"),
-            base_path: Path::from("/"),
             block_size: 64 * 1024,
         }
     }
@@ -447,10 +440,6 @@ impl ObjectStore {
         self.block_size = new_size;
     }
 
-    pub fn base_path(&self) -> &Path {
-        &self.base_path
-    }
-
     /// Open a file for path.
     ///
     /// Parameters
@@ -747,11 +736,9 @@ async fn configure_store(url: &str, options: ObjectStoreParams) -> Result<Object
             Ok(ObjectStore {
                 inner: Arc::new(store),
                 scheme: String::from(url.scheme()),
-                base_path: Path::from(url.path()),
                 block_size: 64 * 1024,
             })
         }
-
         "gs" => {
             storage_options.with_env_gcs();
             let mut builder = GoogleCloudStorageBuilder::new().with_url(url.as_ref());
@@ -763,23 +750,21 @@ async fn configure_store(url: &str, options: ObjectStoreParams) -> Result<Object
             // object_store 0.10.0 is available.
             let store = PatchedGoogleCloudStorage(Arc::new(store));
             let store = Arc::new(store);
+
             Ok(ObjectStore {
                 inner: store,
                 scheme: String::from("gs"),
-                base_path: Path::from(url.path()),
                 block_size: 64 * 1024,
             })
         }
         "az" => {
             storage_options.with_env_azure();
-
             let (store, _) = parse_url_opts(&url, storage_options.as_azure_options())?;
             let store = Arc::new(store);
 
             Ok(ObjectStore {
                 inner: store,
                 scheme: String::from("az"),
-                base_path: Path::from(url.path()),
                 block_size: 64 * 1024,
             })
         }
@@ -794,7 +779,6 @@ async fn configure_store(url: &str, options: ObjectStoreParams) -> Result<Object
         "memory" => Ok(ObjectStore {
             inner: Arc::new(InMemory::new()).traced(),
             scheme: String::from("memory"),
-            base_path: Path::from(url.path()),
             block_size: 64 * 1024,
         }),
         unknow_scheme => {
@@ -824,7 +808,6 @@ impl ObjectStore {
         Self {
             inner: store,
             scheme: scheme.into(),
-            base_path: location.path().into(),
             block_size,
         }
     }
@@ -969,6 +952,26 @@ mod tests {
         }
     }
 
+    #[tokio::test]
+    async fn test_cloud_paths() {
+        let uri = "s3://bucket/foo.lance";
+        let (store, path) = ObjectStore::from_uri(uri).await.unwrap();
+        assert_eq!(store.scheme, "s3");
+        assert_eq!(path.to_string(), "foo.lance");
+
+        let (store, path) = ObjectStore::from_uri("s3+ddb://bucket/foo.lance")
+            .await
+            .unwrap();
+        assert_eq!(store.scheme, "s3");
+        assert_eq!(path.to_string(), "foo.lance");
+
+        let (store, path) = ObjectStore::from_uri("gs://bucket/foo.lance")
+            .await
+            .unwrap();
+        assert_eq!(store.scheme, "gs");
+        assert_eq!(path.to_string(), "foo.lance");
+    }
+
     #[tokio::test]
     async fn test_relative_paths() {
         let tmp_dir = tempfile::tempdir().unwrap();

rust/lance/src/dataset.rs

@@ -84,6 +84,11 @@ pub(crate) const DEFAULT_METADATA_CACHE_SIZE: usize = 256;
 pub struct Dataset {
     pub(crate) object_store: Arc<ObjectStore>,
     pub(crate) commit_handler: Arc<dyn CommitHandler>,
+    /// Uri of the dataset.
+    ///
+    /// On cloud storage, we can not use [Dataset::base] to build the full uri because the
+    /// `bucket` is swlloed in the inner [ObjectStore].
+    uri: String,
     pub(crate) base: Path,
     pub(crate) manifest: Arc<Manifest>,
     pub(crate) session: Arc<Session>,
@@ -282,6 +287,7 @@ impl Dataset {
         Self::checkout_manifest(
             self.object_store.clone(),
             base_path,
+            self.uri.to_string(),
             &manifest_location,
             self.session.clone(),
             self.commit_handler.clone(),
@@ -292,6 +298,7 @@ impl Dataset {
     async fn checkout_manifest(
         object_store: Arc<ObjectStore>,
         base_path: Path,
+        uri: String,
         manifest_location: &ManifestLocation,
         session: Arc<Session>,
         commit_handler: Arc<dyn CommitHandler>,
@@ -354,6 +361,7 @@ impl Dataset {
         Ok(Self {
             object_store,
             base: base_path,
+            uri,
             manifest: Arc::new(manifest),
             commit_handler,
             session,
@@ -495,6 +503,7 @@ impl Dataset {
         Ok(Self {
             object_store,
             base,
+            uri: uri.to_string(),
             manifest: Arc::new(manifest.clone()),
             session: Arc::new(Session::default()),
             commit_handler,
@@ -590,9 +599,9 @@ impl Dataset {
         self.append_impl(batches, params).await
     }
 
-    /// Get the base URI of the dataset.
-    pub fn uri(&self) -> &Path {
-        &self.base
+    /// Get the fully qualified URI of this dataset.
+    pub fn uri(&self) -> &str {
+        &self.uri
     }
 
     /// Get the full manifest of the dataset version.
@@ -789,6 +798,7 @@ impl Dataset {
         Ok(Self {
             object_store: Arc::new(object_store),
             base,
+            uri: base_uri.to_string(),
             manifest: Arc::new(manifest.clone()),
             session: Arc::new(Session::default()),
             commit_handler,
@@ -3378,4 +3388,38 @@ mod tests {
 
         assert!(!dataset_dir.exists());
     }
+
+    #[tokio::test]
+    async fn test_dataset_uri_roundtrips() {
+        let schema = Arc::new(ArrowSchema::new(vec![ArrowField::new(
+            "a",
+            DataType::Int32,
+            false,
+        )]));
+
+        let test_dir = tempdir().unwrap();
+        let test_uri = test_dir.path().to_str().unwrap();
+        let vectors = Arc::new(Int32Array::from_iter_values(vec![]));
+
+        let data = RecordBatch::try_new(schema.clone(), vec![vectors]);
+        let reader = RecordBatchIterator::new(vec![data.unwrap()].into_iter().map(Ok), schema);
+        let dataset = Dataset::write(
+            reader,
+            test_uri,
+            Some(WriteParams {
+                ..Default::default()
+            }),
+        )
+        .await
+        .unwrap();
+
+        let uri = dataset.uri();
+        assert_eq!(uri, test_uri);
+
+        let ds2 = Dataset::open(uri).await.unwrap();
+        assert_eq!(
+            ds2.latest_version_id().await.unwrap(),
+            dataset.latest_version_id().await.unwrap()
+        );
+    }
 }