feat: Add Store.replace_prefix

Cargo.toml

@@ -26,7 +26,7 @@ chrono = "0.4.44"
 futures = "0.3.31"
 http = "1.4"
 indexmap = "2"
-object_store = { version = "0.14.0", features = [
+object_store = { git = "https://github.com/apache/arrow-rs-object-store", rev = "978fb48b9ea6ebe3f8f4833f656e793cde79ab15", features = [
     "aws",
     "azure",
     "gcp",

obstore/python/obstore/_store/__init__.pyi

@@ -185,6 +185,33 @@ class LocalStore:
 
     def __eq__(self, value: object) -> bool: ...
     def __getnewargs_ex__(self): ...
+    def replace_prefix(self, prefix: str | Path | None) -> Self:
+        """Construct a new store with a different prefix.
+
+        All other configuration is inherited from this store. Note that unlike the
+        remote stores, there's no underlying connection pool to reuse here, so this is
+        equivalent to constructing a new [`LocalStore`][obstore.store.LocalStore]
+        directly.
+
+        The new prefix fully replaces the existing one; it is not appended to it.
+
+        **Example:**
+
+        ```py
+        store = LocalStore(prefix="/data/2024")
+        store_2025 = store.replace_prefix("/data/2025")
+        ```
+
+        Args:
+            prefix: The prefix to apply to all operations in the new store. Pass `None`
+                to construct a store with no prefix. If `mkdir` is `True` on this store,
+                the directory at `prefix` will attempt to be created.
+
+        Returns:
+            LocalStore
+
+        """
+
     @property
     def prefix(self) -> Path | None:
         """Get the prefix applied to all operations in this store, if any."""

obstore/python/obstore/_store/_aws.pyi

@@ -620,6 +620,35 @@ class S3Store:
 
     def __eq__(self, value: object) -> bool: ...
     def __getnewargs_ex__(self): ...
+    def replace_prefix(self, prefix: str | None) -> Self:
+        """Construct a new store with a different prefix.
+
+        The new store **shares the underlying HTTP client and connection pool** with
+        this store, so no new connections need to be established, and any already-open
+        connections stay warm. This makes it cheap to create many stores pointing at
+        different prefixes of the same bucket.
+
+        All other configuration is inherited from this store.
+
+        The new prefix fully replaces the existing one; it is not appended to it. It is
+        always interpreted relative to the root of the bucket.
+
+        **Example:**
+
+        ```py
+        store = S3Store("bucket", prefix="data/2024")
+        store_2025 = store.replace_prefix("data/2025")
+        ```
+
+        Args:
+            prefix: The prefix to apply to all operations in the new store. Pass `None`
+                to construct a store with no prefix.
+
+        Returns:
+            S3Store
+
+        """
+
     @property
     def prefix(self) -> str | None:
         """Get the prefix applied to all operations in this store, if any."""

obstore/python/obstore/_store/_azure.pyi

@@ -448,6 +448,35 @@ class AzureStore:
 
     def __eq__(self, value: object) -> bool: ...
     def __getnewargs_ex__(self): ...
+    def replace_prefix(self, prefix: str | None) -> Self:
+        """Construct a new store with a different prefix.
+
+        The new store **shares the underlying HTTP client and connection pool** with
+        this store, so no new connections need to be established, and any already-open
+        connections stay warm. This makes it cheap to create many stores pointing at
+        different prefixes of the same container.
+
+        All other configuration is inherited from this store.
+
+        The new prefix fully replaces the existing one; it is not appended to it. It is
+        always interpreted relative to the root of the container.
+
+        **Example:**
+
+        ```py
+        store = AzureStore("container", prefix="data/2024")
+        store_2025 = store.replace_prefix("data/2025")
+        ```
+
+        Args:
+            prefix: The prefix to apply to all operations in the new store. Pass `None`
+                to construct a store with no prefix.
+
+        Returns:
+            AzureStore
+
+        """
+
     @property
     def prefix(self) -> str | None:
         """Get the prefix applied to all operations in this store, if any."""

obstore/python/obstore/_store/_gcs.pyi

@@ -226,6 +226,35 @@ class GCSStore:
 
     def __eq__(self, value: object) -> bool: ...
     def __getnewargs_ex__(self): ...
+    def replace_prefix(self, prefix: str | None) -> Self:
+        """Construct a new store with a different prefix.
+
+        The new store **shares the underlying HTTP client and connection pool** with
+        this store, so no new connections need to be established, and any already-open
+        connections stay warm. This makes it cheap to create many stores pointing at
+        different prefixes of the same bucket.
+
+        All other configuration is inherited from this store.
+
+        The new prefix fully replaces the existing one; it is not appended to it. It is
+        always interpreted relative to the root of the bucket.
+
+        **Example:**
+
+        ```py
+        store = GCSStore("bucket", prefix="data/2024")
+        store_2025 = store.replace_prefix("data/2025")
+        ```
+
+        Args:
+            prefix: The prefix to apply to all operations in the new store. Pass `None`
+                to construct a store with no prefix.
+
+        Returns:
+            GCSStore
+
+        """
+
     @property
     def prefix(self) -> str | None:
         """Get the prefix applied to all operations in this store, if any."""

pyo3-object_store/Cargo.toml

@@ -28,7 +28,7 @@ humantime = "2.1"
 http = "1"
 # This is already an object_store dependency
 itertools = "0.15.0"
-object_store = { version = "0.14.0", features = [
+object_store = { git = "https://github.com/apache/arrow-rs-object-store", rev = "978fb48b9ea6ebe3f8f4833f656e793cde79ab15", features = [
     "aws",
     "azure",
     "gcp",

pyo3-object_store/src/aws/store.rs

@@ -57,6 +57,16 @@ impl S3Config {
 
         PyTuple::new(py, [args, kwargs.into_bound_py_any(py)?])
     }
+
+    fn replace_prefix(&self, new_prefix: Option<PyPath>) -> Self {
+        Self {
+            prefix: new_prefix,
+            config: self.config.clone(),
+            client_options: self.client_options.clone(),
+            retry_config: self.retry_config.clone(),
+            credential_provider: self.credential_provider.clone(),
+        }
+    }
 }
 
 /// A Python-facing wrapper around an [`AmazonS3`].
@@ -216,6 +226,13 @@ impl PyS3Store {
         self.config.credential_provider.as_ref()
     }
 
+    fn replace_prefix(&self, prefix: Option<PyPath>) -> PyObjectStoreResult<Self> {
+        Ok(Self {
+            store: Arc::new(self.store.replace_prefix(prefix.clone())),
+            config: self.config.replace_prefix(prefix),
+        })
+    }
+
     #[getter]
     fn retry_config(&self) -> Option<&PyRetryConfig> {
         self.config.retry_config.as_ref()

pyo3-object_store/src/azure/store.rs

@@ -63,6 +63,16 @@ impl AzureConfig {
 
         PyTuple::new(py, [args, kwargs.into_bound_py_any(py)?])
     }
+
+    fn replace_prefix(&self, new_prefix: Option<PyPath>) -> Self {
+        Self {
+            prefix: new_prefix,
+            config: self.config.clone(),
+            client_options: self.client_options.clone(),
+            retry_config: self.retry_config.clone(),
+            credential_provider: self.credential_provider.clone(),
+        }
+    }
 }
 
 /// A Python-facing wrapper around a [`MicrosoftAzure`].
@@ -238,6 +248,13 @@ impl PyAzureStore {
         self.config.credential_provider.as_ref()
     }
 
+    fn replace_prefix(&self, prefix: Option<PyPath>) -> PyObjectStoreResult<Self> {
+        Ok(Self {
+            store: Arc::new(self.store.replace_prefix(prefix.clone())),
+            config: self.config.replace_prefix(prefix),
+        })
+    }
+
     #[getter]
     fn retry_config(&self) -> Option<&PyRetryConfig> {
         self.config.retry_config.as_ref()

pyo3-object_store/src/gcp/store.rs

@@ -55,6 +55,16 @@ impl GCSConfig {
 
         PyTuple::new(py, [args, kwargs.into_bound_py_any(py)?])
     }
+
+    fn replace_prefix(&self, new_prefix: Option<PyPath>) -> Self {
+        Self {
+            prefix: new_prefix,
+            config: self.config.clone(),
+            client_options: self.client_options.clone(),
+            retry_config: self.retry_config.clone(),
+            credential_provider: self.credential_provider.clone(),
+        }
+    }
 }
 
 /// A Python-facing wrapper around a [`GoogleCloudStorage`].
@@ -202,6 +212,13 @@ impl PyGCSStore {
         self.config.credential_provider.as_ref()
     }
 
+    fn replace_prefix(&self, prefix: Option<PyPath>) -> PyObjectStoreResult<Self> {
+        Ok(Self {
+            store: Arc::new(self.store.replace_prefix(prefix.clone())),
+            config: self.config.replace_prefix(prefix),
+        })
+    }
+
     #[getter]
     fn retry_config(&self) -> Option<&PyRetryConfig> {
         self.config.retry_config.as_ref()

pyo3-object_store/src/local.rs

@@ -140,4 +140,11 @@ impl PyLocalStore {
             py.None().into_bound_py_any(py)
         }
     }
+
+    fn replace_prefix(&self, prefix: Option<std::path::PathBuf>) -> PyObjectStoreResult<Self> {
+        // Here we use Self::new instead of `replace_prefix` as on the other stores because 1) this
+        // doesn't use a MaybePrefixedStore wrapper and 2) there's no underlying connection pool we
+        // need to reuse.
+        Self::new(prefix, self.config.automatic_cleanup, self.config.mkdir)
+    }
 }

pyo3-object_store/src/prefix.rs

@@ -88,6 +88,16 @@ impl<T: ObjectStore> MaybePrefixedStore<T> {
     }
 }
 
+impl<T: ObjectStore + Clone> MaybePrefixedStore<T> {
+    /// Create a new instance of [`MaybePrefixedStore`] with a new prefix
+    pub fn replace_prefix(&self, new_prefix: Option<impl Into<Path>>) -> Self {
+        Self {
+            prefix: new_prefix.map(|x| x.into()),
+            inner: self.inner.clone(),
+        }
+    }
+}
+
 // Note: This is a relative hack to move these two functions to pure functions so they don't rely
 // on the `self` lifetime. Expected to be cleaned up before merge.
 //