feat(commit): let registered providers declare their commit handler

rust/lance-io/src/object_store.rs

@@ -141,7 +141,7 @@ fn stream_copy_error(
     }))
 }
 
-pub use providers::{ObjectStoreProvider, ObjectStoreRegistry};
+pub use providers::{CommitHandlerType, ObjectStoreProvider, ObjectStoreRegistry};
 pub use read_dir::ReadDirOptions;
 pub use storage_options::{
     BASE_SCOPED_OPTION_PREFIX, BaseScopedStorageOptionsProvider, EXPIRES_AT_MILLIS_KEY,
@@ -224,6 +224,11 @@ pub struct ObjectStore {
     /// The backend's paginated listing API, when it has one. `None` means
     /// [`Self::read_dir_page`] has to list a directory in full to page through it.
     pub(crate) paginated_lister: Option<Arc<dyn PaginatedListStore>>,
+    /// The commit-handler capability declared by the provider that built this
+    /// store. It is captured at construction time so the default commit handler
+    /// is bound to the exact store, not re-derived from a registry that may
+    /// have changed.
+    commit_handler_type: CommitHandlerType,
 }
 
 // Hand-written because `PaginatedListStore` is not `Debug`.
@@ -244,6 +249,7 @@ impl std::fmt::Debug for ObjectStore {
             .field("io_tracker", &self.io_tracker)
             .field("store_prefix", &self.store_prefix)
             .field("paginated_lister", &self.paginated_lister.is_some())
+            .field("commit_handler_type", &self.commit_handler_type)
             .finish()
     }
 }
@@ -677,6 +683,9 @@ impl ObjectStore {
                 store_prefix,
                 // Type-erased on the way in, so there is no telling if it can paginate.
                 paginated_lister: None,
+                // A caller-supplied store has no provider to declare a
+                // capability, so fall back to the conflict-safe default.
+                commit_handler_type: CommitHandlerType::ConditionalPut,
             };
             let path = Path::parse(path.path())?;
             return Ok((Arc::new(store), path));
@@ -779,6 +788,13 @@ impl ObjectStore {
         &self.scheme
     }
 
+    /// The commit-handler capability declared by the provider that built this
+    /// store. Commit-handler resolution uses this to pick the default handler,
+    /// so the choice stays bound to the store the provider actually produced.
+    pub fn commit_handler_type(&self) -> CommitHandlerType {
+        self.commit_handler_type
+    }
+
     pub fn block_size(&self) -> usize {
         self.block_size
     }
@@ -1797,6 +1813,13 @@ impl ObjectStore {
             store_prefix,
             // Type-erased on the way in, so there is no telling if it can paginate.
             paginated_lister: None,
+            // Derive the capability from the registered provider for this
+            // scheme when one exists; otherwise fall back to the conflict-safe
+            // default.
+            commit_handler_type: DEFAULT_OBJECT_STORE_REGISTRY
+                .get_provider(scheme)
+                .map(|provider| provider.commit_handler())
+                .unwrap_or_default(),
         }
     }
 }

rust/lance-io/src/object_store/providers.rs

@@ -38,10 +38,40 @@ pub mod tencent;
 #[cfg(feature = "tos")]
 pub mod tos;
 
+/// Which built-in commit handler a provider's stores should use.
+///
+/// Commit-handler selection is otherwise inferred from the URL scheme. A
+/// registered out-of-tree provider serves a scheme that the built-in resolver
+/// does not recognize, so this lets the provider declare the guarantee its
+/// store actually offers instead of falling back to a scheme-string guess.
+#[derive(Debug, Clone, Copy, PartialEq, Eq, Default)]
+pub enum CommitHandlerType {
+    /// Atomic put-if-not-exists. Safe for concurrent writers: a second writer
+    /// racing on the same manifest version gets a conflict instead of silently
+    /// overwriting. This is the default; it fits any store whose backend
+    /// supports create-if-not-exists (most object stores).
+    #[default]
+    ConditionalPut,
+    /// Blind overwrite with no concurrency protection. A store that cannot
+    /// offer atomic create-if-not-exists should return this so callers get the
+    /// documented "concurrent writes may lose data" behavior rather than an
+    /// error at commit time.
+    Unsafe,
+}
+
 #[async_trait::async_trait]
 pub trait ObjectStoreProvider: std::fmt::Debug + Sync + Send {
     async fn new_store(&self, base_path: Url, params: &ObjectStoreParams) -> Result<ObjectStore>;
 
+    /// The commit handler this provider's stores should use.
+    ///
+    /// Defaults to the conflict-safe [`CommitHandlerType::ConditionalPut`].
+    /// Override it to return [`CommitHandlerType::Unsafe`] for a store that
+    /// cannot support atomic create-if-not-exists.
+    fn commit_handler(&self) -> CommitHandlerType {
+        CommitHandlerType::ConditionalPut
+    }
+
     /// Extract the path relative to the base of the store.
     ///
     /// For example, in S3 the path is relative to the bucket. So a URL of
@@ -214,6 +244,11 @@ impl ObjectStoreRegistry {
     ) -> Result<Arc<ObjectStore>> {
         let mut store = provider.new_store(base_path, params).await?;
 
+        // Capture the provider's declared commit-handler capability on the
+        // exact store it built, so commit-handler resolution is bound to this
+        // store rather than a later registry lookup.
+        store.commit_handler_type = provider.commit_handler();
+
         store.inner = store.inner.traced();
 
         // Label metrics by the store's unique prefix (e.g. `s3$bucket`,
@@ -694,4 +729,91 @@ mod tests {
         let stats = registry.stats();
         assert_eq!((stats.hits, stats.misses, stats.active_stores), (0, 0, 0));
     }
+
+    /// A provider that builds a working in-memory store yet declares it can
+    /// only offer unsafe commits, used to prove the declared capability is
+    /// stamped onto the store the registry builds even for a known scheme.
+    #[derive(Debug)]
+    struct UnsafeMemoryProvider;
+
+    #[async_trait::async_trait]
+    impl ObjectStoreProvider for UnsafeMemoryProvider {
+        async fn new_store(
+            &self,
+            base_path: Url,
+            params: &ObjectStoreParams,
+        ) -> Result<ObjectStore> {
+            MemoryStoreProvider.new_store(base_path, params).await
+        }
+
+        fn commit_handler(&self) -> CommitHandlerType {
+            CommitHandlerType::Unsafe
+        }
+
+        fn calculate_object_store_prefix(
+            &self,
+            _url: &Url,
+            _storage_options: Option<&HashMap<String, String>>,
+        ) -> Result<String> {
+            Ok("memory".to_string())
+        }
+    }
+
+    #[tokio::test]
+    async fn test_provider_commit_handler_travels_with_store() {
+        let url = Url::parse("memory://bucket/ds").unwrap();
+
+        // A built-in provider declares the conflict-safe default, and the
+        // registry stamps that capability onto the store it builds.
+        let registry = ObjectStoreRegistry::default();
+        let store = registry
+            .new_store(url.clone(), &ObjectStoreParams::default())
+            .await
+            .unwrap();
+        assert_eq!(
+            store.commit_handler_type(),
+            CommitHandlerType::ConditionalPut
+        );
+
+        // Overriding a known scheme with a provider that declares Unsafe makes
+        // the store it builds carry Unsafe: the capability is bound to the
+        // exact store, not inferred from the scheme name.
+        let registry = ObjectStoreRegistry::default();
+        registry.insert("memory", Arc::new(UnsafeMemoryProvider));
+        let store = registry
+            .new_store(url.clone(), &ObjectStoreParams::default())
+            .await
+            .unwrap();
+        assert_eq!(store.commit_handler_type(), CommitHandlerType::Unsafe);
+
+        // A cached store keeps the capability it was built with.
+        let cached = registry
+            .get_store(url, &ObjectStoreParams::default())
+            .await
+            .unwrap();
+        assert_eq!(cached.commit_handler_type(), CommitHandlerType::Unsafe);
+    }
+
+    #[test]
+    fn test_builtin_providers_keep_conditional_put_capability() {
+        // Regression guard for existing providers: the declared capability now
+        // drives commit-handler selection, so every provider the default
+        // registry ships must keep declaring the conflict-safe put-if-not-exists
+        // handler it resolved to before this change. A built-in that silently
+        // switched to Unsafe would change the commit handler of every store on
+        // that scheme.
+        let registry = ObjectStoreRegistry::default();
+        let providers = registry
+            .providers
+            .read()
+            .expect("ObjectStoreRegistry lock poisoned");
+        assert!(!providers.is_empty());
+        for (scheme, provider) in providers.iter() {
+            assert_eq!(
+                provider.commit_handler(),
+                CommitHandlerType::ConditionalPut,
+                "built-in provider for `{scheme}` changed its declared commit handler",
+            );
+        }
+    }
 }

rust/lance-io/src/object_store/providers/aws.rs

@@ -266,6 +266,7 @@ impl ObjectStoreProvider for AwsStoreProvider {
             store_prefix: self
                 .calculate_object_store_prefix(&base_path, params.storage_options())?,
             paginated_lister,
+            commit_handler_type: crate::object_store::CommitHandlerType::ConditionalPut,
         })
     }
 }

rust/lance-io/src/object_store/providers/azure.rs

@@ -306,6 +306,7 @@ impl ObjectStoreProvider for AzureBlobStoreProvider {
             store_prefix: self
                 .calculate_object_store_prefix(&base_path, params.storage_options())?,
             paginated_lister,
+            commit_handler_type: crate::object_store::CommitHandlerType::ConditionalPut,
         })
     }
 

rust/lance-io/src/object_store/providers/gcp.rs

@@ -350,6 +350,7 @@ impl ObjectStoreProvider for GcsStoreProvider {
             store_prefix: self
                 .calculate_object_store_prefix(&base_path, params.storage_options())?,
             paginated_lister,
+            commit_handler_type: crate::object_store::CommitHandlerType::ConditionalPut,
         })
     }
 }

rust/lance-io/src/object_store/providers/goosefs.rs

@@ -193,6 +193,7 @@ impl ObjectStoreProvider for GooseFsStoreProvider {
                 .calculate_object_store_prefix(&base_path, params.storage_options())?,
             // Listed in full: no paginated lister covers OpenDAL yet.
             paginated_lister: None,
+            commit_handler_type: crate::object_store::CommitHandlerType::ConditionalPut,
         })
     }
 

rust/lance-io/src/object_store/providers/huggingface.rs

@@ -219,6 +219,7 @@ impl ObjectStoreProvider for HuggingfaceStoreProvider {
                 .calculate_object_store_prefix(&base_path, params.storage_options())?,
             // Listed in full: no paginated lister covers OpenDAL yet.
             paginated_lister: None,
+            commit_handler_type: crate::object_store::CommitHandlerType::ConditionalPut,
         })
     }
 

rust/lance-io/src/object_store/providers/local.rs

@@ -141,6 +141,7 @@ impl ObjectStoreProvider for FileStoreProvider {
             // Listed in full: reading a directory is one local walk whatever the page size,
             // so there is no request for a page to be pushed into.
             paginated_lister: None,
+            commit_handler_type: crate::object_store::CommitHandlerType::ConditionalPut,
         })
     }
 
@@ -204,6 +205,7 @@ mod tests {
             io_tracker: Default::default(),
             store_prefix: "file$rooted-test".to_owned(),
             paginated_lister: None,
+            commit_handler_type: crate::object_store::CommitHandlerType::ConditionalPut,
         }
     }
 

rust/lance-io/src/object_store/providers/memory.rs

@@ -37,6 +37,7 @@ impl ObjectStoreProvider for MemoryStoreProvider {
             // Listed in full: the store is already in memory, so a page costs no less than
             // the directory does.
             paginated_lister: None,
+            commit_handler_type: crate::object_store::CommitHandlerType::ConditionalPut,
         })
     }
 

rust/lance-io/src/object_store/providers/oss.rs

@@ -147,6 +147,7 @@ impl ObjectStoreProvider for OssStoreProvider {
             store_prefix: self.calculate_object_store_prefix(&url, params.storage_options())?,
             // Listed in full: no paginated lister covers OpenDAL yet.
             paginated_lister: None,
+            commit_handler_type: crate::object_store::CommitHandlerType::ConditionalPut,
         })
     }
 }

rust/lance-io/src/object_store/providers/tencent.rs

@@ -103,6 +103,7 @@ impl ObjectStoreProvider for TencentStoreProvider {
             store_prefix: self.calculate_object_store_prefix(&url, params.storage_options())?,
             // Listed in full: no paginated lister covers OpenDAL yet.
             paginated_lister: None,
+            commit_handler_type: crate::object_store::CommitHandlerType::ConditionalPut,
         })
     }
 }

rust/lance-io/src/object_store/providers/tos.rs

@@ -153,6 +153,7 @@ impl ObjectStoreProvider for TosStoreProvider {
             store_prefix: self.calculate_object_store_prefix(&url, params.storage_options())?,
             // Listed in full: no paginated lister covers OpenDAL yet.
             paginated_lister: None,
+            commit_handler_type: crate::object_store::CommitHandlerType::ConditionalPut,
         })
     }
 }

rust/lance-namespace-impls/src/dir/manifest.rs

@@ -1819,8 +1819,13 @@ impl ManifestNamespace {
     }
 
     /// Resolve the commit handler for the `__manifest` dataset's storage backend.
-    async fn manifest_commit_handler(&self) -> Result<Arc<dyn CommitHandler>> {
-        commit_handler_from_url(&self.root, &None)
+    /// The capability is read from the exact store the commit is written through,
+    /// which for stores like `memory://` can differ from `self.object_store`.
+    async fn manifest_commit_handler(
+        &self,
+        object_store: &ObjectStore,
+    ) -> Result<Arc<dyn CommitHandler>> {
+        commit_handler_from_url(&self.root, &None, object_store.commit_handler_type())
             .await
             .map_err(|e| {
                 lance_core::Error::from(NamespaceError::Internal {
@@ -1950,7 +1955,6 @@ impl ManifestNamespace {
         let max_retries = self.manifest_rewrite_commit_retries();
         let mut retries = 0;
         let build_indices = self.inline_optimization_enabled;
-        let commit_handler = self.manifest_commit_handler().await?;
 
         loop {
             let dataset_guard = self.manifest_dataset.get_refreshed().await?;
@@ -1963,6 +1967,9 @@ impl ManifestNamespace {
             // Staged files, indices, the commit, and cleanup must all use the dataset's
             // own object store (see `commit_manifest_overwrite`).
             let object_store = dataset.object_store(None).await?;
+            // Resolve the handler from that same committing store so an override
+            // registered for this scheme is honored (see manifest_commit_handler).
+            let commit_handler = self.manifest_commit_handler(&object_store).await?;
 
             let source = Self::manifest_projected_stream(&dataset).await?;
             let resolution = make_mutation().conflict_resolution();