fix: wrap future streams also in io_runtime

Cargo.toml

@@ -76,5 +76,6 @@ uuid = { version = "1" }
 async-trait = { version = "0.1" }
 futures = { version = "0.3" }
 tokio = { version = "1" }
+tokio-stream = { version = "0.1" }
 num_cpus = { version = "1" }
 

crates/benchmarks/src/bin/merge.rs

@@ -72,7 +72,7 @@ pub async fn convert_tpcds_web_returns(input_path: String, table_path: String) -
         .await
         .unwrap();
 
-    DeltaOps::try_from_uri(table_path)
+    DeltaOps::try_from_uri(table_path, None)
         .await
         .unwrap()
         .write(table.collect().await.unwrap())
@@ -564,7 +564,7 @@ async fn main() {
             )
             .unwrap();
 
-            DeltaOps::try_from_uri(output)
+            DeltaOps::try_from_uri(output, None)
                 .await
                 .unwrap()
                 .write(vec![batch])

crates/core/Cargo.toml

@@ -75,6 +75,7 @@ tokio = { workspace = true, features = [
     "fs",
     "parking_lot",
 ] }
+tokio-stream = { workspace = true}
 
 # other deps (these should be organized and pulled into workspace.dependencies as necessary)
 cfg-if = "1"

crates/core/src/operations/create.rs

@@ -440,7 +440,7 @@ mod tests {
             "./{}",
             tmp_dir.path().file_name().unwrap().to_str().unwrap()
         );
-        let table = DeltaOps::try_from_uri(relative_path)
+        let table = DeltaOps::try_from_uri(relative_path, None)
             .await
             .unwrap()
             .create()

crates/core/src/operations/load_cdf.rs

@@ -493,7 +493,7 @@ pub(crate) mod tests {
     #[tokio::test]
     async fn test_load_local() -> TestResult {
         let ctx = SessionContext::new();
-        let table = DeltaOps::try_from_uri("../test/tests/data/cdf-table")
+        let table = DeltaOps::try_from_uri("../test/tests/data/cdf-table", None)
             .await?
             .load_cdf()
             .with_session_ctx(ctx.clone())
@@ -543,7 +543,7 @@ pub(crate) mod tests {
     async fn test_load_local_datetime() -> TestResult {
         let ctx = SessionContext::new();
         let starting_timestamp = NaiveDateTime::from_str("2023-12-22T17:10:21.675").unwrap();
-        let table = DeltaOps::try_from_uri("../test/tests/data/cdf-table")
+        let table = DeltaOps::try_from_uri("../test/tests/data/cdf-table", None)
             .await?
             .load_cdf()
             .with_session_ctx(ctx.clone())
@@ -589,7 +589,7 @@ pub(crate) mod tests {
     #[tokio::test]
     async fn test_load_local_non_partitioned() -> TestResult {
         let ctx = SessionContext::new();
-        let table = DeltaOps::try_from_uri("../test/tests/data/cdf-table-non-partitioned")
+        let table = DeltaOps::try_from_uri("../test/tests/data/cdf-table-non-partitioned", None)
             .await?
             .load_cdf()
             .with_session_ctx(ctx.clone())
@@ -641,7 +641,7 @@ pub(crate) mod tests {
 
     #[tokio::test]
     async fn test_load_bad_version_range() -> TestResult {
-        let table = DeltaOps::try_from_uri("../test/tests/data/cdf-table-non-partitioned")
+        let table = DeltaOps::try_from_uri("../test/tests/data/cdf-table-non-partitioned", None)
             .await?
             .load_cdf()
             .with_starting_version(4)
@@ -660,7 +660,7 @@ pub(crate) mod tests {
 
     #[tokio::test]
     async fn test_load_version_out_of_range() -> TestResult {
-        let table = DeltaOps::try_from_uri("../test/tests/data/cdf-table-non-partitioned")
+        let table = DeltaOps::try_from_uri("../test/tests/data/cdf-table-non-partitioned", None)
             .await?
             .load_cdf()
             .with_starting_version(5)
@@ -678,7 +678,7 @@ pub(crate) mod tests {
 
     #[tokio::test]
     async fn test_load_version_out_of_range_with_flag() -> TestResult {
-        let table = DeltaOps::try_from_uri("../test/tests/data/cdf-table-non-partitioned")
+        let table = DeltaOps::try_from_uri("../test/tests/data/cdf-table-non-partitioned", None)
             .await?
             .load_cdf()
             .with_starting_version(5)
@@ -702,7 +702,7 @@ pub(crate) mod tests {
     #[tokio::test]
     async fn test_load_timestamp_out_of_range() -> TestResult {
         let ending_timestamp = NaiveDateTime::from_str("2033-12-22T17:10:21.675").unwrap();
-        let table = DeltaOps::try_from_uri("../test/tests/data/cdf-table-non-partitioned")
+        let table = DeltaOps::try_from_uri("../test/tests/data/cdf-table-non-partitioned", None)
             .await?
             .load_cdf()
             .with_starting_timestamp(ending_timestamp.and_utc())
@@ -721,7 +721,7 @@ pub(crate) mod tests {
     #[tokio::test]
     async fn test_load_timestamp_out_of_range_with_flag() -> TestResult {
         let ending_timestamp = NaiveDateTime::from_str("2033-12-22T17:10:21.675").unwrap();
-        let table = DeltaOps::try_from_uri("../test/tests/data/cdf-table-non-partitioned")
+        let table = DeltaOps::try_from_uri("../test/tests/data/cdf-table-non-partitioned", None)
             .await?
             .load_cdf()
             .with_starting_timestamp(ending_timestamp.and_utc())
@@ -744,7 +744,7 @@ pub(crate) mod tests {
 
     #[tokio::test]
     async fn test_load_non_cdf() -> TestResult {
-        let table = DeltaOps::try_from_uri("../test/tests/data/simple_table")
+        let table = DeltaOps::try_from_uri("../test/tests/data/simple_table", None)
             .await?
             .load_cdf()
             .with_starting_version(0)
@@ -764,7 +764,7 @@ pub(crate) mod tests {
     async fn test_load_vacuumed_table() -> TestResult {
         let ending_timestamp = NaiveDateTime::from_str("2024-01-06T15:44:59.570")?;
         let ctx = SessionContext::new();
-        let table = DeltaOps::try_from_uri("../test/tests/data/checkpoint-cdf-table")
+        let table = DeltaOps::try_from_uri("../test/tests/data/checkpoint-cdf-table", None)
             .await?
             .load_cdf()
             .with_session_ctx(ctx.clone())

crates/core/src/operations/mod.rs

@@ -33,6 +33,7 @@ use self::{
 };
 use crate::errors::{DeltaResult, DeltaTableError};
 use crate::logstore::LogStoreRef;
+use crate::storage::IORuntime;
 use crate::table::builder::DeltaTableBuilder;
 use crate::DeltaTable;
 
@@ -130,11 +131,19 @@ impl DeltaOps {
     /// use deltalake_core::DeltaOps;
     ///
     /// async {
-    ///     let ops = DeltaOps::try_from_uri("memory://").await.unwrap();
+    ///     let ops = DeltaOps::try_from_uri("memory://", None).await.unwrap();
     /// };
     /// ```
-    pub async fn try_from_uri(uri: impl AsRef<str>) -> DeltaResult<Self> {
-        let mut table = DeltaTableBuilder::from_uri(uri).build()?;
+    pub async fn try_from_uri(
+        uri: impl AsRef<str>,
+        io_runtime: Option<IORuntime>,
+    ) -> DeltaResult<Self> {
+        let mut builder = DeltaTableBuilder::from_uri(uri);
+
+        if let Some(io_runtime) = io_runtime {
+            builder = builder.with_io_runtime(io_runtime)
+        };
+        let mut table = builder.build()?;
         // We allow for uninitialized locations, since we may want to create the table
         match table.load().await {
             Ok(_) => Ok(table.into()),
@@ -147,10 +156,15 @@ impl DeltaOps {
     pub async fn try_from_uri_with_storage_options(
         uri: impl AsRef<str>,
         storage_options: HashMap<String, String>,
+        io_runtime: Option<IORuntime>,
     ) -> DeltaResult<Self> {
-        let mut table = DeltaTableBuilder::from_uri(uri)
-            .with_storage_options(storage_options)
-            .build()?;
+        let mut builder = DeltaTableBuilder::from_uri(uri).with_storage_options(storage_options);
+
+        if let Some(io_runtime) = io_runtime {
+            builder = builder.with_io_runtime(io_runtime)
+        };
+        let mut table = builder.build()?;
+
         // We allow for uninitialized locations, since we may want to create the table
         match table.load().await {
             Ok(_) => Ok(table.into()),
@@ -183,7 +197,7 @@ impl DeltaOps {
     /// use deltalake_core::DeltaOps;
     ///
     /// async {
-    ///     let ops = DeltaOps::try_from_uri("memory://").await.unwrap();
+    ///     let ops = DeltaOps::try_from_uri("memory://", None).await.unwrap();
     ///     let table = ops.create().with_table_name("my_table").await.unwrap();
     ///     assert_eq!(table.version(), 0);
     /// };

crates/core/src/operations/optimize.rs

@@ -1728,7 +1728,7 @@ pub(super) mod zorder {
 
             // Run optimize
             let (_, metrics) =
-                DeltaOps::try_from_uri(tmp_dir.path().join(table_name).to_str().unwrap())
+                DeltaOps::try_from_uri(tmp_dir.path().join(table_name).to_str().unwrap(), None)
                     .await
                     .unwrap()
                     .optimize()

crates/core/src/operations/transaction/application.rs

@@ -17,7 +17,7 @@ mod tests {
         let tmp_path = std::fs::canonicalize(tmp_dir.path()).unwrap();
 
         let batch = get_record_batch(None, false);
-        let table = DeltaOps::try_from_uri(tmp_path.to_str().unwrap())
+        let table = DeltaOps::try_from_uri(tmp_path.to_str().unwrap(), None)
             .await
             .unwrap()
             .write(vec![batch.clone()])

crates/core/src/operations/write.rs

@@ -20,7 +20,7 @@
 //! let schema = Arc::new(arrow::datatypes::Schema::new(vec![id_field]));
 //! let ids = arrow::array::Int32Array::from(vec![1, 2, 3, 4, 5]);
 //! let batch = RecordBatch::try_new(schema, vec![Arc::new(ids)])?;
-//! let ops = DeltaOps::try_from_uri("../path/to/empty/dir").await?;
+//! let ops = DeltaOps::try_from_uri("../path/to/empty/dir", Some(IORuntime::default())).await?;
 //! let table = ops.write(vec![batch]).await?;
 //! ````
 
@@ -1904,7 +1904,7 @@ mod tests {
         let batch = RecordBatch::try_new(schema, vec![Arc::new(str_values), Arc::new(data_values)])
             .unwrap();
 
-        let ops = DeltaOps::try_from_uri(tmp_path.as_os_str().to_str().unwrap())
+        let ops = DeltaOps::try_from_uri(tmp_path.as_os_str().to_str().unwrap(), None)
             .await
             .unwrap();
 

crates/core/src/protocol/checkpoints.rs

@@ -1135,7 +1135,7 @@ mod tests {
         let table_schema = get_delta_schema();
         let temp_dir = tempfile::tempdir()?;
         let table_path = temp_dir.path().to_str().unwrap();
-        let mut table = DeltaOps::try_from_uri(&table_path)
+        let mut table = DeltaOps::try_from_uri(&table_path, None)
             .await?
             .create()
             .with_columns(table_schema.fields().cloned())
@@ -1210,7 +1210,7 @@ mod tests {
         )
         .unwrap();
 
-        let mut table = DeltaOps::try_from_uri(tmp_path.as_os_str().to_str().unwrap())
+        let mut table = DeltaOps::try_from_uri(tmp_path.as_os_str().to_str().unwrap(), None)
             .await?
             .write(vec![batch])
             .await?;
@@ -1229,7 +1229,7 @@ mod tests {
         )
         .unwrap();
 
-        let table = DeltaOps::try_from_uri(tmp_path.as_os_str().to_str().unwrap())
+        let table = DeltaOps::try_from_uri(tmp_path.as_os_str().to_str().unwrap(), None)
             .await?
             .write(vec![batch])
             .with_save_mode(SaveMode::Overwrite)