Complete upgrade to object_store 0.13.0

Author: alamb

datafusion-cli/src/object_storage/instrumented.rs

@@ -35,11 +35,14 @@ use datafusion::{
     error::DataFusionError,
     execution::object_store::{DefaultObjectStoreRegistry, ObjectStoreRegistry},
 };
-use futures::stream::{BoxStream, Stream};
+use futures::{
+    StreamExt,
+    stream::{BoxStream, Stream},
+};
 use object_store::{
-    GetOptions, GetRange, GetResult, ListResult, MultipartUpload, ObjectMeta,
-    ObjectStore, PutMultipartOptions, PutOptions, PutPayload, PutResult, Result,
-    path::Path,
+    CopyOptions, GetOptions, GetRange, GetResult, ListResult, MultipartUpload,
+    ObjectMeta, ObjectStore, PutMultipartOptions, PutOptions, PutPayload, PutResult,
+    Result, path::Path,
 };
 use parking_lot::{Mutex, RwLock};
 use url::Url;
@@ -228,42 +231,56 @@ impl InstrumentedObjectStore {
         options: GetOptions,
     ) -> Result<GetResult> {
         let timestamp = Utc::now();
+        let operation = if options.head {
+            Operation::Head
+        } else {
+            Operation::Get
+        };
         let range = options.range.clone();
 
         let start = Instant::now();
         let ret = self.inner.get_opts(location, options).await?;
         let elapsed = start.elapsed();
 
         self.requests.lock().push(RequestDetails {
-            op: Operation::Get,
+            op: operation,
             path: location.clone(),
             timestamp,
             duration: Some(elapsed),
-            size: Some((ret.range.end - ret.range.start) as usize),
+            size: (operation == Operation::Get)
+                .then_some((ret.range.end - ret.range.start) as usize),
             range,
             extra_display: None,
         });
 
         Ok(ret)
     }
 
-    async fn instrumented_delete(&self, location: &Path) -> Result<()> {
+    fn instrumented_delete_stream(
+        &self,
+        locations: BoxStream<'static, Result<Path>>,
+    ) -> BoxStream<'static, Result<Path>> {
         let timestamp = Utc::now();
         let start = Instant::now();
-        self.inner.delete(location).await?;
-        let elapsed = start.elapsed();
-
-        self.requests.lock().push(RequestDetails {
-            op: Operation::Delete,
-            path: location.clone(),
-            timestamp,
-            duration: Some(elapsed),
-            size: None,
-            range: None,
-            extra_display: None,
-        });
-
-        Ok(())
+        let requests = Arc::clone(&self.requests);
+
+        self.inner
+            .delete_stream(locations)
+            .inspect(move |result| {
+                if let Ok(path) = result {
+                    let elapsed = start.elapsed();
+                    requests.lock().push(RequestDetails {
+                        op: Operation::Delete,
+                        path: path.clone(),
+                        timestamp,
+                        duration: Some(elapsed),
+                        size: None,
+                        range: None,
+                        extra_display: None,
+                    });
+                }
+            })
+            .boxed()
     }
 
     fn instrumented_list(
@@ -320,33 +337,15 @@ impl InstrumentedObjectStore {
         Ok(ret)
     }
 
-    async fn instrumented_copy(&self, from: &Path, to: &Path) -> Result<()> {
-        let timestamp = Utc::now();
-        let start = Instant::now();
-        self.inner.copy(from, to).await?;
-        let elapsed = start.elapsed();
-
-        self.requests.lock().push(RequestDetails {
-            op: Operation::Copy,
-            path: from.clone(),
-            timestamp,
-            duration: Some(elapsed),
-            size: None,
-            range: None,
-            extra_display: Some(format!("copy_to: {to}")),
-        });
-
-        Ok(())
-    }
-
-    async fn instrumented_copy_if_not_exists(
+    async fn instrumented_copy_opts(
         &self,
         from: &Path,
         to: &Path,
+        options: CopyOptions,
     ) -> Result<()> {
         let timestamp = Utc::now();
         let start = Instant::now();
-        self.inner.copy_if_not_exists(from, to).await?;
+        self.inner.copy_opts(from, to, options).await?;
         let elapsed = start.elapsed();
 
         self.requests.lock().push(RequestDetails {
@@ -361,25 +360,6 @@ impl InstrumentedObjectStore {
 
         Ok(())
     }
-
-    async fn instrumented_head(&self, location: &Path) -> Result<ObjectMeta> {
-        let timestamp = Utc::now();
-        let start = Instant::now();
-        let ret = self.inner.head(location).await?;
-        let elapsed = start.elapsed();
-
-        self.requests.lock().push(RequestDetails {
-            op: Operation::Head,
-            path: location.clone(),
-            timestamp,
-            duration: Some(elapsed),
-            size: None,
-            range: None,
-            extra_display: None,
-        });
-
-        Ok(ret)
-    }
 }
 
 impl fmt::Display for InstrumentedObjectStore {
@@ -429,14 +409,6 @@ impl ObjectStore for InstrumentedObjectStore {
         self.inner.get_opts(location, options).await
     }
 
-    async fn delete(&self, location: &Path) -> Result<()> {
-        if self.enabled() {
-            return self.instrumented_delete(location).await;
-        }
-
-        self.inner.delete(location).await
-    }
-
     fn list(&self, prefix: Option<&Path>) -> BoxStream<'static, Result<ObjectMeta>> {
         if self.enabled() {
             return self.instrumented_list(prefix);
@@ -453,28 +425,28 @@ impl ObjectStore for InstrumentedObjectStore {
         self.inner.list_with_delimiter(prefix).await
     }
 
-    async fn copy(&self, from: &Path, to: &Path) -> Result<()> {
-        if self.enabled() {
-            return self.instrumented_copy(from, to).await;
-        }
-
-        self.inner.copy(from, to).await
-    }
-
-    async fn copy_if_not_exists(&self, from: &Path, to: &Path) -> Result<()> {
+    fn delete_stream(
+        &self,
+        locations: BoxStream<'static, Result<Path>>,
+    ) -> BoxStream<'static, Result<Path>> {
         if self.enabled() {
-            return self.instrumented_copy_if_not_exists(from, to).await;
+            return self.instrumented_delete_stream(locations);
         }
 
-        self.inner.copy_if_not_exists(from, to).await
+        self.inner.delete_stream(locations)
     }
 
-    async fn head(&self, location: &Path) -> Result<ObjectMeta> {
+    async fn copy_opts(
+        &self,
+        from: &Path,
+        to: &Path,
+        options: CopyOptions,
+    ) -> Result<()> {
         if self.enabled() {
-            return self.instrumented_head(location).await;
+            return self.instrumented_copy_opts(from, to, options).await;
         }
 
-        self.inner.head(location).await
+        self.inner.copy_opts(from, to, options).await
     }
 }
 
@@ -824,6 +796,7 @@ impl ObjectStoreRegistry for InstrumentedObjectStoreRegistry {
 #[cfg(test)]
 mod tests {
     use futures::StreamExt;
+    use object_store::ObjectStoreExt;
     use object_store::WriteMultipart;
 
     use super::*;

datafusion-examples/examples/custom_data_source/csv_json_opener.rs

@@ -36,7 +36,8 @@ use datafusion::{
 
 use datafusion::datasource::physical_plan::FileScanConfigBuilder;
 use futures::StreamExt;
-use object_store::{ObjectStore, local::LocalFileSystem, memory::InMemory};
+use object_store::ObjectStoreExt;
+use object_store::{local::LocalFileSystem, memory::InMemory};
 
 /// This example demonstrates using the low level [`FileStream`] / [`FileOpener`] APIs to directly
 /// read data from (CSV/JSON) into Arrow RecordBatches.

datafusion-examples/examples/custom_data_source/custom_file_casts.rs

@@ -38,6 +38,7 @@ use datafusion::prelude::SessionConfig;
 use datafusion_physical_expr_adapter::{
     DefaultPhysicalExprAdapterFactory, PhysicalExprAdapter, PhysicalExprAdapterFactory,
 };
+use object_store::ObjectStoreExt;
 use object_store::memory::InMemory;
 use object_store::path::Path;
 use object_store::{ObjectStore, PutPayload};

datafusion-examples/examples/custom_data_source/default_column_values.rs

@@ -46,6 +46,7 @@ use datafusion_physical_expr_adapter::{
     replace_columns_with_literals,
 };
 use futures::StreamExt;
+use object_store::ObjectStoreExt;
 use object_store::memory::InMemory;
 use object_store::path::Path;
 use object_store::{ObjectStore, PutPayload};

datafusion-examples/examples/data_io/json_shredding.rs

@@ -45,9 +45,10 @@ use datafusion::scalar::ScalarValue;
 use datafusion_physical_expr_adapter::{
     DefaultPhysicalExprAdapterFactory, PhysicalExprAdapter, PhysicalExprAdapterFactory,
 };
+use object_store::ObjectStoreExt;
+use object_store::PutPayload;
 use object_store::memory::InMemory;
 use object_store::path::Path;
-use object_store::{ObjectStore, PutPayload};
 
 // Example showing how to implement custom filter rewriting for JSON shredding.
 //

datafusion-examples/examples/data_io/parquet_advanced_index.rs

@@ -567,7 +567,7 @@ impl ParquetFileReaderFactory for CachedParquetFileReaderFactory {
             .object_meta
             .location
             .parts()
-            .last()
+            .next_back()
             .expect("No path in location")
             .as_ref()
             .to_string();

datafusion/core/benches/push_down_filter.rs

@@ -26,6 +26,7 @@ use datafusion_physical_optimizer::PhysicalOptimizerRule;
 use datafusion_physical_optimizer::filter_pushdown::FilterPushdown;
 use datafusion_physical_plan::ExecutionPlan;
 use object_store::ObjectStore;
+use object_store::ObjectStoreExt;
 use object_store::memory::InMemory;
 use object_store::path::Path;
 use parquet::arrow::ArrowWriter;

datafusion/core/benches/sql_query_with_io.rs

@@ -30,6 +30,7 @@ use datafusion::{
 };
 use datafusion_execution::runtime_env::RuntimeEnv;
 use itertools::Itertools;
+use object_store::ObjectStoreExt;
 use object_store::{
     ObjectStore,
     memory::InMemory,

datafusion/core/src/datasource/file_format/csv.rs

@@ -60,12 +60,14 @@ mod tests {
     use futures::StreamExt;
     use futures::stream::BoxStream;
     use insta::assert_snapshot;
+    use object_store::ObjectStoreExt;
     use object_store::chunked::ChunkedStore;
     use object_store::local::LocalFileSystem;
     use object_store::path::Path;
     use object_store::{
-        Attributes, GetOptions, GetResult, GetResultPayload, ListResult, MultipartUpload,
-        ObjectMeta, ObjectStore, PutMultipartOptions, PutOptions, PutPayload, PutResult,
+        Attributes, CopyOptions, GetOptions, GetResult, GetResultPayload, ListResult,
+        MultipartUpload, ObjectMeta, ObjectStore, PutMultipartOptions, PutOptions,
+        PutPayload, PutResult,
     };
     use regex::Regex;
     use rstest::*;
@@ -159,10 +161,19 @@ mod tests {
 
         fn delete_stream(
             &self,
-            locations: BoxStream<'static, object_store::Result<Path>>,
+            _locations: BoxStream<'static, object_store::Result<Path>>,
         ) -> BoxStream<'static, object_store::Result<Path>> {
             unimplemented!()
         }
+
+        async fn copy_opts(
+            &self,
+            _from: &Path,
+            _to: &Path,
+            _options: CopyOptions,
+        ) -> object_store::Result<()> {
+            unimplemented!()
+        }
     }
 
     impl VariableStream {