Upgrade DataFusion to arrow-rs/parquet 58.0.0 / object_store 0.13.0

Cargo.toml

@@ -91,19 +91,19 @@ ahash = { version = "0.8", default-features = false, features = [
     "runtime-rng",
 ] }
 apache-avro = { version = "0.21", default-features = false }
-arrow = { version = "57.3.0", features = [
+arrow = { version = "58.0.0", features = [
     "prettyprint",
     "chrono-tz",
 ] }
-arrow-buffer = { version = "57.2.0", default-features = false }
-arrow-flight = { version = "57.3.0", features = [
+arrow-buffer = { version = "58.0.0", default-features = false }
+arrow-flight = { version = "58.0.0", features = [
     "flight-sql-experimental",
 ] }
-arrow-ipc = { version = "57.2.0", default-features = false, features = [
+arrow-ipc = { version = "58.0.0", default-features = false, features = [
     "lz4",
 ] }
-arrow-ord = { version = "57.2.0", default-features = false }
-arrow-schema = { version = "57.2.0", default-features = false }
+arrow-ord = { version = "58.0.0", default-features = false }
+arrow-schema = { version = "58.0.0", default-features = false }
 async-trait = "0.1.89"
 bigdecimal = "0.4.8"
 bytes = "1.11"
@@ -165,9 +165,9 @@ liblzma = { version = "0.4.6", features = ["static"] }
 log = "^0.4"
 memchr = "2.8.0"
 num-traits = { version = "0.2" }
-object_store = { version = "0.12.5", default-features = false }
+object_store = { version = "0.13.1", default-features = false }
 parking_lot = "0.12"
-parquet = { version = "57.3.0", default-features = false, features = [
+parquet = { version = "58.0.0", default-features = false, features = [
     "arrow",
     "async",
     "object_store",

datafusion-cli/src/exec.rs

@@ -521,6 +521,7 @@ mod tests {
     use datafusion::common::plan_err;
 
     use datafusion::prelude::SessionContext;
+    use datafusion_common::assert_contains;
     use url::Url;
 
     async fn create_external_table_test(location: &str, sql: &str) -> Result<()> {
@@ -714,25 +715,24 @@ mod tests {
         let err = create_external_table_test(location, &sql)
             .await
             .unwrap_err();
-        assert!(err.to_string().contains("os error 2"));
+        assert_contains!(err.to_string(), "os error 2");
 
         // for service_account_key
         let sql = format!(
             "CREATE EXTERNAL TABLE test STORED AS PARQUET OPTIONS('gcp.service_account_key' '{service_account_key}') LOCATION '{location}'"
         );
         let err = create_external_table_test(location, &sql)
             .await
-            .unwrap_err()
-            .to_string();
-        assert!(err.contains("No RSA key found in pem file"), "{err}");
+            .unwrap_err();
+        assert_contains!(err.to_string(), "Error reading pem file: no items found");
 
         // for application_credentials_path
         let sql = format!("CREATE EXTERNAL TABLE test STORED AS PARQUET
             OPTIONS('gcp.application_credentials_path' '{application_credentials_path}') LOCATION '{location}'");
         let err = create_external_table_test(location, &sql)
             .await
             .unwrap_err();
-        assert!(err.to_string().contains("os error 2"));
+        assert_contains!(err.to_string(), "os error 2");
 
         Ok(())
     }

datafusion-cli/src/main.rs

@@ -617,8 +617,8 @@ mod tests {
         | filename                          | file_size_bytes | metadata_size_bytes | hits | extra            |
         +-----------------------------------+-----------------+---------------------+------+------------------+
         | alltypes_plain.parquet            | 1851            | 8882                | 2    | page_index=false |
-        | alltypes_tiny_pages.parquet       | 454233          | 269266              | 2    | page_index=true  |
-        | lz4_raw_compressed_larger.parquet | 380836          | 1347                | 2    | page_index=false |
+        | alltypes_tiny_pages.parquet       | 454233          | 269074              | 2    | page_index=true  |
+        | lz4_raw_compressed_larger.parquet | 380836          | 1339                | 2    | page_index=false |
         +-----------------------------------+-----------------+---------------------+------+------------------+
         ");
 
@@ -648,8 +648,8 @@ mod tests {
         | filename                          | file_size_bytes | metadata_size_bytes | hits | extra            |
         +-----------------------------------+-----------------+---------------------+------+------------------+
         | alltypes_plain.parquet            | 1851            | 8882                | 5    | page_index=false |
-        | alltypes_tiny_pages.parquet       | 454233          | 269266              | 2    | page_index=true  |
-        | lz4_raw_compressed_larger.parquet | 380836          | 1347                | 3    | page_index=false |
+        | alltypes_tiny_pages.parquet       | 454233          | 269074              | 2    | page_index=true  |
+        | lz4_raw_compressed_larger.parquet | 380836          | 1339                | 3    | page_index=false |
         +-----------------------------------+-----------------+---------------------+------+------------------+
         ");
 
@@ -841,8 +841,8 @@ mod tests {
         +---------------------+-----------+-----------------+------+
         | metadata_size_bytes | filename  | file_size_bytes | etag |
         +---------------------+-----------+-----------------+------+
-        | 212                 | 0.parquet | 3645            | 0    |
-        | 212                 | 1.parquet | 3645            | 1    |
+        | 212                 | 0.parquet | 3642            | 0    |
+        | 212                 | 1.parquet | 3642            | 1    |
         +---------------------+-----------+-----------------+------+
         ");
 

datafusion-cli/src/object_storage/instrumented.rs

@@ -36,10 +36,11 @@ use datafusion::{
     execution::object_store::{DefaultObjectStoreRegistry, ObjectStoreRegistry},
 };
 use futures::stream::{BoxStream, Stream};
+use futures::{StreamExt, TryStreamExt};
 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, ObjectStoreExt, PutMultipartOptions, PutOptions, PutPayload,
+    PutResult, Result, path::Path,
 };
 use parking_lot::{Mutex, RwLock};
 use url::Url;
@@ -230,40 +231,57 @@ impl InstrumentedObjectStore {
         let timestamp = Utc::now();
         let range = options.range.clone();
 
+        let head = options.head;
         let start = Instant::now();
         let ret = self.inner.get_opts(location, options).await?;
         let elapsed = start.elapsed();
 
+        let (op, size) = if head {
+            (Operation::Head, None)
+        } else {
+            (
+                Operation::Get,
+                Some((ret.range.end - ret.range.start) as usize),
+            )
+        };
+
         self.requests.lock().push(RequestDetails {
-            op: Operation::Get,
+            op,
             path: location.clone(),
             timestamp,
             duration: Some(elapsed),
-            size: Some((ret.range.end - ret.range.start) as usize),
+            size,
             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 requests_captured = Arc::clone(&self.requests);
+
         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(())
+        self.inner
+            .delete_stream(locations)
+            .and_then(move |location| {
+                let elapsed = start.elapsed();
+                requests_captured.lock().push(RequestDetails {
+                    op: Operation::Delete,
+                    path: location.clone(),
+                    timestamp,
+                    duration: Some(elapsed),
+                    size: None,
+                    range: None,
+                    extra_display: None,
+                });
+                futures::future::ok(location)
+            })
+            .boxed()
     }
 
     fn instrumented_list(
@@ -361,25 +379,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,12 +428,15 @@ impl ObjectStore for InstrumentedObjectStore {
         self.inner.get_opts(location, options).await
     }
 
-    async fn delete(&self, location: &Path) -> Result<()> {
+    fn delete_stream(
+        &self,
+        locations: BoxStream<'static, Result<Path>>,
+    ) -> BoxStream<'static, Result<Path>> {
         if self.enabled() {
-            return self.instrumented_delete(location).await;
+            return self.instrumented_delete_stream(locations);
         }
 
-        self.inner.delete(location).await
+        self.inner.delete_stream(locations)
     }
 
     fn list(&self, prefix: Option<&Path>) -> BoxStream<'static, Result<ObjectMeta>> {
@@ -453,28 +455,24 @@ 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<()> {
-        if self.enabled() {
-            return self.instrumented_copy_if_not_exists(from, to).await;
-        }
-
-        self.inner.copy_if_not_exists(from, to).await
-    }
-
-    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 match options.mode {
+                object_store::CopyMode::Create => {
+                    self.instrumented_copy_if_not_exists(from, to).await
+                }
+                object_store::CopyMode::Overwrite => {
+                    self.instrumented_copy(from, to).await
+                }
+            };
         }
 
-        self.inner.head(location).await
+        self.inner.copy_opts(from, to, options).await
     }
 }
 

datafusion-examples/examples/custom_data_source/adapter_serialization.rs

@@ -69,7 +69,7 @@ use datafusion_proto::protobuf::{
 };
 use object_store::memory::InMemory;
 use object_store::path::Path;
-use object_store::{ObjectStore, PutPayload};
+use object_store::{ObjectStore, ObjectStoreExt, PutPayload};
 use serde::{Deserialize, Serialize};
 
 /// Example showing how to preserve custom adapter information during plan serialization.

datafusion-examples/examples/custom_data_source/csv_json_opener.rs

@@ -36,7 +36,7 @@ use datafusion::{
 use datafusion::datasource::physical_plan::FileScanConfigBuilder;
 use datafusion_examples::utils::datasets::ExampleDataset;
 use futures::StreamExt;
-use object_store::{ObjectStore, local::LocalFileSystem, memory::InMemory};
+use object_store::{ObjectStoreExt, 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

@@ -40,7 +40,7 @@ use datafusion_physical_expr_adapter::{
 };
 use object_store::memory::InMemory;
 use object_store::path::Path;
-use object_store::{ObjectStore, PutPayload};
+use object_store::{ObjectStore, ObjectStoreExt, PutPayload};
 
 // Example showing how to implement custom casting rules to adapt file schemas.
 // This example enforces that casts must be strictly widening: if the file type is Int64 and the table type is Int32, it will error

datafusion-examples/examples/custom_data_source/default_column_values.rs

@@ -48,7 +48,7 @@ use datafusion_physical_expr_adapter::{
 use futures::StreamExt;
 use object_store::memory::InMemory;
 use object_store::path::Path;
-use object_store::{ObjectStore, PutPayload};
+use object_store::{ObjectStore, ObjectStoreExt, PutPayload};
 
 // Metadata key for storing default values in field metadata
 const DEFAULT_VALUE_METADATA_KEY: &str = "example.default_value";
@@ -79,7 +79,7 @@ pub async fn default_column_values() -> Result<()> {
         let mut buf = vec![];
 
         let props = WriterProperties::builder()
-            .set_max_row_group_size(2)
+            .set_max_row_group_row_count(Some(2))
             .build();
 
         let mut writer =

datafusion-examples/examples/data_io/json_shredding.rs

@@ -47,7 +47,7 @@ use datafusion_physical_expr_adapter::{
 };
 use object_store::memory::InMemory;
 use object_store::path::Path;
-use object_store::{ObjectStore, PutPayload};
+use object_store::{ObjectStoreExt, PutPayload};
 
 // Example showing how to implement custom filter rewriting for JSON shredding.
 //
@@ -76,7 +76,7 @@ pub async fn json_shredding() -> Result<()> {
         let mut buf = vec![];
 
         let props = WriterProperties::builder()
-            .set_max_row_group_size(2)
+            .set_max_row_group_row_count(Some(2))
             .build();
 
         let mut writer = ArrowWriter::try_new(&mut buf, batch.schema(), Some(props))

datafusion-examples/examples/data_io/parquet_advanced_index.rs

@@ -43,7 +43,7 @@ use datafusion::parquet::arrow::arrow_reader::{
     ArrowReaderOptions, ParquetRecordBatchReaderBuilder, RowSelection, RowSelector,
 };
 use datafusion::parquet::arrow::async_reader::{AsyncFileReader, ParquetObjectReader};
-use datafusion::parquet::file::metadata::ParquetMetaData;
+use datafusion::parquet::file::metadata::{PageIndexPolicy, ParquetMetaData};
 use datafusion::parquet::file::properties::{EnabledStatistics, WriterProperties};
 use datafusion::parquet::schema::types::ColumnPath;
 use datafusion::physical_expr::PhysicalExpr;
@@ -410,7 +410,7 @@ impl IndexedFile {
         let options = ArrowReaderOptions::new()
             // Load the page index when reading metadata to cache
             // so it is available to interpret row selections
-            .with_page_index(true);
+            .with_page_index_policy(PageIndexPolicy::Required);
         let reader =
             ParquetRecordBatchReaderBuilder::try_new_with_options(file, options)?;
         let metadata = reader.metadata().clone();
@@ -567,7 +567,7 @@ impl ParquetFileReaderFactory for CachedParquetFileReaderFactory {
             .object_meta
             .location
             .parts()
-            .last()
+            .next_back()
             .expect("No path in location")
             .as_ref()
             .to_string();
@@ -659,7 +659,7 @@ fn make_demo_file(path: impl AsRef<Path>, value_range: Range<i32>) -> Result<()>
     // enable page statistics for the tag column,
     // for everything else.
     let props = WriterProperties::builder()
-        .set_max_row_group_size(100)
+        .set_max_row_group_row_count(Some(100))
         // compute column chunk (per row group) statistics by default
         .set_statistics_enabled(EnabledStatistics::Chunk)
         // compute column page statistics for the tag column

datafusion/common/src/file_options/mod.rs

@@ -84,7 +84,7 @@ mod tests {
         .build();
 
         // Verify the expected options propagated down to parquet crate WriterProperties struct
-        assert_eq!(properties.max_row_group_size(), 123);
+        assert_eq!(properties.max_row_group_row_count(), Some(123));
         assert_eq!(properties.data_page_size_limit(), 123);
         assert_eq!(properties.write_batch_size(), 123);
         assert_eq!(properties.writer_version(), WriterVersion::PARQUET_2_0);