feat: spawn io with spawn service

Signed-off-by: Ion Koutsouris <15728914+ion-elgreco@users.noreply.github.com>

Author: ion-elgreco

Cargo.toml

@@ -43,7 +43,7 @@ arrow-ord = { version = "55" }
 arrow-row = { version = "55" }
 arrow-schema = { version = "55" }
 arrow-select = { version = "55" }
-object_store = { version = "0.12.0" }
+object_store = { version = "0.12.1" }
 parquet = { version = "55" }
 
 # datafusion

crates/aws/src/storage.rs

@@ -18,7 +18,9 @@ use deltalake_core::{DeltaResult, DeltaTableError, ObjectStoreError, Path};
 use futures::stream::BoxStream;
 use futures::Future;
 use object_store::aws::AmazonS3;
+use object_store::client::SpawnedReqwestConnector;
 use object_store::RetryConfig;
+use tokio::runtime::Handle;
 use tracing::log::*;
 use url::Url;
 
@@ -39,13 +41,19 @@ impl ObjectStoreFactory for S3ObjectStoreFactory {
         url: &Url,
         storage_options: &HashMap<String, String>,
         retry: &RetryConfig,
+        handle: Option<Handle>,
     ) -> DeltaResult<(ObjectStoreRef, Path)> {
         let options = self.with_env_s3(storage_options);
 
         // All S3-likes should start their builder the same way
         let mut builder = AmazonS3Builder::new()
             .with_url(url.to_string())
             .with_retry(retry.clone());
+
+        if let Some(handle) = handle {
+            builder = builder.with_http_connector(SpawnedReqwestConnector::new(handle));
+        }
+
         for (key, value) in options.iter() {
             if let Ok(key) = AmazonS3ConfigKey::from_str(&key.to_ascii_lowercase()) {
                 builder = builder.with_config(key, value.clone());

crates/aws/tests/integration_s3_dynamodb.rs

@@ -107,7 +107,8 @@ async fn test_create_s3_table() -> TestResult<()> {
         deltalake_aws::constants::AWS_FORCE_CREDENTIAL_LOAD.into() => "true".into(),
         deltalake_aws::constants::AWS_ENDPOINT_URL.into()  => "http://localhost:4566".into(),
     };
-    let log_store = logstore_for(Url::parse(&table_uri)?, storage_options, None)?;
+    let storage_config = StorageConfig::parse_options(storage_options)?;
+    let log_store = logstore_for(Url::parse(&table_uri)?, storage_config)?;
 
     let payload = PutPayload::from_static(b"test-drivin");
     let _put = log_store

crates/azure/src/lib.rs

@@ -8,7 +8,9 @@ use deltalake_core::logstore::{
 };
 use deltalake_core::{DeltaResult, DeltaTableError, Path};
 use object_store::azure::{AzureConfigKey, MicrosoftAzureBuilder};
+use object_store::client::SpawnedReqwestConnector;
 use object_store::{ObjectStoreScheme, RetryConfig};
+use tokio::runtime::Handle;
 use url::Url;
 
 mod config;
@@ -40,12 +42,18 @@ impl ObjectStoreFactory for AzureFactory {
         url: &Url,
         options: &HashMap<String, String>,
         retry: &RetryConfig,
+        handle: Option<Handle>,
     ) -> DeltaResult<(ObjectStoreRef, Path)> {
         let config = config::AzureConfigHelper::try_new(options.as_azure_options())?.build()?;
 
         let mut builder = MicrosoftAzureBuilder::new()
             .with_url(url.to_string())
             .with_retry(retry.clone());
+
+        if let Some(handle) = handle {
+            builder = builder.with_http_connector(SpawnedReqwestConnector::new(handle));
+        }
+
         for (key, value) in config.iter() {
             builder = builder.with_config(*key, value.clone());
         }

crates/catalog-unity/src/lib.rs

@@ -17,6 +17,7 @@ use std::collections::HashMap;
 use std::future::Future;
 use std::str::FromStr;
 use std::sync::Arc;
+use tokio::runtime::Handle;
 
 use crate::credential::{
     AzureCliCredential, ClientSecretOAuthProvider, CredentialProvider, WorkspaceOAuthProvider,
@@ -840,6 +841,7 @@ impl ObjectStoreFactory for UnityCatalogFactory {
         table_uri: &Url,
         options: &HashMap<String, String>,
         _retry: &RetryConfig,
+        handle: Option<Handle>,
     ) -> DeltaResult<(ObjectStoreRef, Path)> {
         let (table_path, temp_creds) = UnityCatalogBuilder::execute_uc_future(
             UnityCatalogBuilder::get_uc_location_and_token(table_uri.as_str()),
@@ -850,8 +852,11 @@ impl ObjectStoreFactory for UnityCatalogFactory {
 
         // TODO(roeap): we should not have to go through the table here.
         // ideally we just create the right storage ...
-        let mut builder =
-            DeltaTableBuilder::from_uri(&table_path).with_io_runtime(IORuntime::default());
+        let mut builder = DeltaTableBuilder::from_uri(&table_path);
+
+        if let Some(handle) = handle {
+            builder = builder.with_io_runtime(IORuntime::RT(handle));
+        }
         if !storage_options.is_empty() {
             builder = builder.with_storage_options(storage_options.clone());
         }

crates/core/src/logstore/config.rs

@@ -10,10 +10,9 @@ use std::collections::HashMap;
 
 use ::object_store::RetryConfig;
 use object_store::{path::Path, prefix::PrefixStore, ObjectStore, ObjectStoreScheme};
-use tokio::runtime::Handle;
 
-use super::storage::runtime::RuntimeConfig;
 use super::storage::LimitConfig;
+use super::{storage::runtime::RuntimeConfig, IORuntime};
 use crate::{DeltaResult, DeltaTableError};
 
 pub trait TryUpdateKey: Default {
@@ -91,8 +90,9 @@ where
 pub struct StorageConfig {
     /// Runtime configuration.
     ///
-    /// Configuration to set up a dedicated IO runtime to execute IO related operations.
-    pub runtime: Option<RuntimeConfig>,
+    /// Configuration to set up a dedicated IO runtime to execute IO related operations or
+    /// dedicated handle.
+    pub runtime: Option<IORuntime>,
 
     pub retry: ::object_store::RetryConfig,
 
@@ -119,21 +119,12 @@ impl StorageConfig {
     /// Depending on the configuration, the following layers may be added:
     /// - Retry layer: Adds retry logic to the object store.
     /// - Limit layer: Limits the number of concurrent requests to the object store.
-    /// - Runtime layer: Executes IO related operations on a dedicated runtime.
     pub fn decorate_store<T: ObjectStore + Clone>(
         &self,
         store: T,
         table_root: &url::Url,
-        handle: Option<Handle>,
     ) -> DeltaResult<Box<dyn ObjectStore>> {
-        let inner = if let Some(runtime) = &self.runtime {
-            Box::new(runtime.decorate(store, handle)) as Box<dyn ObjectStore>
-        } else {
-            Box::new(store) as Box<dyn ObjectStore>
-        };
-
-        let inner = Self::decorate_prefix(inner, table_root)?;
-
+        let inner = Self::decorate_prefix(store, table_root)?;
         Ok(inner)
     }
 
@@ -169,7 +160,9 @@ where
         };
 
         let result = ParseResult::<RuntimeConfig>::from_iter(&config.raw);
-        config.runtime = (!result.is_default).then_some(result.config);
+        if let Some(runtime_config) = (!result.is_default).then_some(result.config) {
+            config.runtime = Some(IORuntime::Config(runtime_config));
+        };
 
         let result = ParseResult::<LimitConfig>::from_iter(result.unparsed);
         config.limit = (!result.is_default).then_some(result.config);
@@ -218,7 +211,7 @@ impl StorageConfig {
         let (runtime, remainder): (RuntimeConfig, _) = try_parse_impl(&props.raw)?;
         // NOTE: we only want to assign an actual runtime config we consumed an option
         if props.raw.len() > remainder.len() {
-            props.runtime = Some(runtime);
+            props.runtime = Some(IORuntime::Config(runtime));
         }
 
         let result = ParseResult::<LimitConfig>::from_iter(remainder);
@@ -235,6 +228,12 @@ impl StorageConfig {
         props.unknown_properties = remainder;
         Ok(props)
     }
+
+    // Provide an IO Runtime directly
+    pub fn with_io_runtime(mut self, rt: IORuntime) -> Self {
+        self.runtime = Some(rt);
+        self
+    }
 }
 
 pub(super) fn try_parse_impl<T, K, V, I>(options: I) -> DeltaResult<(T, HashMap<String, String>)>

crates/core/src/logstore/factories.rs

@@ -4,11 +4,12 @@ use std::{
 };
 
 use dashmap::DashMap;
-use object_store::path::Path;
 use object_store::RetryConfig;
+use object_store::{path::Path, DynObjectStore};
+use tokio::runtime::Handle;
 use url::Url;
 
-use super::{default_logstore, LogStore, ObjectStoreRef, StorageConfig};
+use super::{default_logstore, DeltaIOStorageBackend, LogStore, ObjectStoreRef, StorageConfig};
 use crate::{DeltaResult, DeltaTableError};
 
 /// Factory registry to manage [`ObjectStoreFactory`] instances
@@ -29,6 +30,7 @@ pub trait ObjectStoreFactory: Send + Sync {
         url: &Url,
         options: &HashMap<String, String>,
         retry: &RetryConfig,
+        handle: Option<Handle>,
     ) -> DeltaResult<(ObjectStoreRef, Path)>;
 }
 
@@ -41,8 +43,14 @@ impl ObjectStoreFactory for DefaultObjectStoreFactory {
         url: &Url,
         options: &HashMap<String, String>,
         _retry: &RetryConfig,
+        handle: Option<Handle>,
     ) -> DeltaResult<(ObjectStoreRef, Path)> {
-        default_parse_url_opts(url, options)
+        let (mut store, path) = default_parse_url_opts(url, options)?;
+
+        if let Some(handle) = handle {
+            store = Arc::new(DeltaIOStorageBackend::new(store, handle)) as Arc<DynObjectStore>;
+        }
+        Ok((store, path))
     }
 }
 
@@ -85,8 +93,8 @@ where
     let storage_config = StorageConfig::parse_options(options)?;
     if let Some(factory) = object_store_factories().get(&scheme) {
         let (store, _prefix) =
-            factory.parse_url_opts(url, &storage_config.raw, &storage_config.retry)?;
-        let store = storage_config.decorate_store(store, url, None)?;
+            factory.parse_url_opts(url, &storage_config.raw, &storage_config.retry, None)?;
+        let store = storage_config.decorate_store(store, url)?;
         Ok(Arc::new(store))
     } else {
         Err(DeltaTableError::InvalidTableLocation(url.clone().into()))