diff --git a/Cargo.lock b/Cargo.lock index 6d4c38d8d9d04..f6a2942d45658 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -1704,7 +1704,6 @@ version = "55.0.0" dependencies = [ "arrow", "arrow-schema", - "async-trait", "bytes", "bzip2", "chrono", @@ -1778,7 +1777,6 @@ name = "datafusion-benchmarks" version = "55.0.0" dependencies = [ "arrow", - "async-trait", "bytes", "clap", "criterion", @@ -1809,7 +1807,6 @@ name = "datafusion-catalog" version = "55.0.0" dependencies = [ "arrow", - "async-trait", "dashmap", "datafusion-common", "datafusion-common-runtime", @@ -1832,7 +1829,6 @@ name = "datafusion-catalog-listing" version = "55.0.0" dependencies = [ "arrow", - "async-trait", "chrono", "datafusion-catalog", "datafusion-common", @@ -1856,7 +1852,6 @@ name = "datafusion-cli" version = "55.0.0" dependencies = [ "arrow", - "async-trait", "aws-config", "aws-credential-types", "chrono", @@ -1928,7 +1923,6 @@ version = "55.0.0" dependencies = [ "arrow", "async-compression", - "async-trait", "bytes", "bzip2", "chrono", @@ -1967,7 +1961,6 @@ version = "55.0.0" dependencies = [ "arrow", "arrow-ipc", - "async-trait", "bytes", "chrono", "datafusion-common", @@ -1991,7 +1984,6 @@ version = "55.0.0" dependencies = [ "arrow", "arrow-avro", - "async-trait", "bytes", "datafusion-common", "datafusion-common-runtime", @@ -2013,7 +2005,6 @@ name = "datafusion-datasource-csv" version = "55.0.0" dependencies = [ "arrow", - "async-trait", "bytes", "datafusion-common", "datafusion-common-runtime", @@ -2035,7 +2026,6 @@ name = "datafusion-datasource-json" version = "55.0.0" dependencies = [ "arrow", - "async-trait", "bytes", "datafusion-common", "datafusion-common-runtime", @@ -2059,7 +2049,6 @@ version = "55.0.0" dependencies = [ "arrow", "arrow-schema", - "async-trait", "bytes", "chrono", "criterion", @@ -2099,7 +2088,6 @@ dependencies = [ "arrow", "arrow-flight", "arrow-schema", - "async-trait", "base64 0.23.1", "bytes", "dashmap", @@ -2139,7 +2127,6 @@ version = "55.0.0" dependencies = [ "arrow", "arrow-buffer", - "async-trait", "bytes", "chrono", "dashmap", @@ -2166,7 +2153,6 @@ version = "55.0.0" dependencies = [ "arrow", "arrow-schema", - "async-trait", "chrono", "ctor", "datafusion-common", @@ -2178,6 +2164,7 @@ dependencies = [ "datafusion-proto-common", "datafusion-proto-models", "env_logger", + "futures", "indexmap 2.14.2", "insta", "itertools 0.15.0", @@ -2205,7 +2192,6 @@ dependencies = [ "arrow", "arrow-schema", "async-ffi", - "async-trait", "chrono", "datafusion", "datafusion-catalog", @@ -2333,12 +2319,12 @@ name = "datafusion-functions-table" version = "55.0.0" dependencies = [ "arrow", - "async-trait", "datafusion-catalog", "datafusion-common", "datafusion-expr", "datafusion-physical-expr", "datafusion-physical-plan", + "futures", "parking_lot", ] @@ -2380,7 +2366,6 @@ name = "datafusion-optimizer" version = "55.0.0" dependencies = [ "arrow", - "async-trait", "chrono", "criterion", "ctor", @@ -2492,7 +2477,6 @@ dependencies = [ "arrow-ipc", "arrow-ord", "arrow-schema", - "async-trait", "bytes", "criterion", "datafusion-common", @@ -2531,7 +2515,6 @@ name = "datafusion-proto" version = "55.0.0" dependencies = [ "arrow", - "async-trait", "datafusion", "datafusion-catalog", "datafusion-catalog-listing", @@ -2555,6 +2538,7 @@ dependencies = [ "datafusion-proto-models", "doc-comment", "flate2", + "futures", "object_store", "pretty_assertions", "prost", @@ -2612,7 +2596,6 @@ name = "datafusion-session" version = "55.0.0" dependencies = [ "arrow-schema", - "async-trait", "datafusion-common", "datafusion-execution", "datafusion-expr", @@ -2682,7 +2665,6 @@ name = "datafusion-sqllogictest" version = "55.0.0" dependencies = [ "arrow", - "async-trait", "bigdecimal", "bytes", "chrono", @@ -2715,10 +2697,10 @@ name = "datafusion-substrait" version = "55.0.0" dependencies = [ "async-recursion", - "async-trait", "chrono", "datafusion", "datafusion-functions-aggregate", + "futures", "half", "insta", "itertools 0.15.0", diff --git a/Cargo.toml b/Cargo.toml index efeb580074138..3821041b4c2ad 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -114,7 +114,6 @@ arrow-ipc = { version = "59.2.0", default-features = false, features = [ ] } arrow-ord = { version = "59.2.0", default-features = false } arrow-schema = { version = "59.2.0", default-features = false } -async-trait = "0.1.89" bigdecimal = "0.4.8" bytes = "1.11" bzip2 = "0.6.1" diff --git a/benchmarks/Cargo.toml b/benchmarks/Cargo.toml index cd222c3a5e5b0..e85f70372f903 100644 --- a/benchmarks/Cargo.toml +++ b/benchmarks/Cargo.toml @@ -41,7 +41,6 @@ mimalloc_extended = ["libmimalloc-sys/extended"] [dependencies] arrow = { workspace = true } -async-trait = "0.1" bytes = { workspace = true } clap = { version = "4.6.0", features = ["derive", "env", "string"] } criterion = { workspace = true, features = ["async_tokio", "html_reports"] } diff --git a/benchmarks/src/util/latency_object_store.rs b/benchmarks/src/util/latency_object_store.rs index 9ef8d1b78b751..d6b86dfa151bd 100644 --- a/benchmarks/src/util/latency_object_store.rs +++ b/benchmarks/src/util/latency_object_store.rs @@ -23,10 +23,11 @@ //! - P99: ~150-200ms use std::fmt; +use std::future::Future; +use std::pin::Pin; use std::sync::atomic::{AtomicUsize, Ordering}; use std::time::Duration; -use async_trait::async_trait; use futures::StreamExt; use futures::stream::BoxStream; use object_store::path::Path; @@ -89,37 +90,67 @@ impl fmt::Display for LatencyObjectStore { } } -#[async_trait] impl ObjectStore for LatencyObjectStore { - async fn put_opts( - &self, - location: &Path, + fn put_opts<'life0, 'life1, 'async_trait>( + &'life0 self, + location: &'life1 Path, payload: PutPayload, opts: PutOptions, - ) -> Result { - self.inner.put_opts(location, payload, opts).await + ) -> Pin> + Send + 'async_trait>> + where + 'life0: 'async_trait, + 'life1: 'async_trait, + Self: 'async_trait, + { + Box::pin(async move { self.inner.put_opts(location, payload, opts).await }) } - async fn put_multipart_opts( - &self, - location: &Path, + fn put_multipart_opts<'life0, 'life1, 'async_trait>( + &'life0 self, + location: &'life1 Path, opts: PutMultipartOptions, - ) -> Result> { - self.inner.put_multipart_opts(location, opts).await + ) -> Pin< + Box>> + Send + 'async_trait>, + > + where + 'life0: 'async_trait, + 'life1: 'async_trait, + Self: 'async_trait, + { + Box::pin(async move { self.inner.put_multipart_opts(location, opts).await }) } - async fn get_opts(&self, location: &Path, options: GetOptions) -> Result { - tokio::time::sleep(self.next_get_latency()).await; - self.inner.get_opts(location, options).await + fn get_opts<'life0, 'life1, 'async_trait>( + &'life0 self, + location: &'life1 Path, + options: GetOptions, + ) -> Pin> + Send + 'async_trait>> + where + 'life0: 'async_trait, + 'life1: 'async_trait, + Self: 'async_trait, + { + Box::pin(async move { + tokio::time::sleep(self.next_get_latency()).await; + self.inner.get_opts(location, options).await + }) } - async fn get_ranges( - &self, - location: &Path, - ranges: &[std::ops::Range], - ) -> Result> { - tokio::time::sleep(self.next_get_latency()).await; - self.inner.get_ranges(location, ranges).await + fn get_ranges<'life0, 'life1, 'life2, 'async_trait>( + &'life0 self, + location: &'life1 Path, + ranges: &'life2 [std::ops::Range], + ) -> Pin>> + Send + 'async_trait>> + where + 'life0: 'async_trait, + 'life1: 'async_trait, + 'life2: 'async_trait, + Self: 'async_trait, + { + Box::pin(async move { + tokio::time::sleep(self.next_get_latency()).await; + self.inner.get_ranges(location, ranges).await + }) } fn delete_stream( @@ -141,17 +172,33 @@ impl ObjectStore for LatencyObjectStore { .boxed() } - async fn list_with_delimiter(&self, prefix: Option<&Path>) -> Result { - tokio::time::sleep(self.next_list_latency()).await; - self.inner.list_with_delimiter(prefix).await + fn list_with_delimiter<'life0, 'life1, 'async_trait>( + &'life0 self, + prefix: Option<&'life1 Path>, + ) -> Pin> + Send + 'async_trait>> + where + 'life0: 'async_trait, + 'life1: 'async_trait, + Self: 'async_trait, + { + Box::pin(async move { + tokio::time::sleep(self.next_list_latency()).await; + self.inner.list_with_delimiter(prefix).await + }) } - async fn copy_opts( - &self, - from: &Path, - to: &Path, + fn copy_opts<'life0, 'life1, 'life2, 'async_trait>( + &'life0 self, + from: &'life1 Path, + to: &'life2 Path, options: CopyOptions, - ) -> Result<()> { - self.inner.copy_opts(from, to, options).await + ) -> Pin> + Send + 'async_trait>> + where + 'life0: 'async_trait, + 'life1: 'async_trait, + 'life2: 'async_trait, + Self: 'async_trait, + { + Box::pin(async move { self.inner.copy_opts(from, to, options).await }) } } diff --git a/datafusion-cli/Cargo.toml b/datafusion-cli/Cargo.toml index 50b4002541f89..c7245d320cec5 100644 --- a/datafusion-cli/Cargo.toml +++ b/datafusion-cli/Cargo.toml @@ -36,7 +36,6 @@ backtrace = ["datafusion/backtrace"] [dependencies] arrow = { workspace = true } -async-trait = { workspace = true } aws-config = "1.8.18" aws-credential-types = "1.2.13" chrono = { workspace = true } diff --git a/datafusion-cli/examples/cli-session-context.rs b/datafusion-cli/examples/cli-session-context.rs index 6095072163870..bb872a52632c1 100644 --- a/datafusion-cli/examples/cli-session-context.rs +++ b/datafusion-cli/examples/cli-session-context.rs @@ -18,6 +18,7 @@ //! Shows an example of a custom session context that unions the input plan with itself. //! To run this example, use `cargo run --example cli-session-context` from within the `datafusion-cli` directory. +use futures::future::BoxFuture; use std::sync::Arc; use datafusion::{ @@ -47,7 +48,6 @@ impl Default for MyUnionerContext { } } -#[async_trait::async_trait] impl CliSessionContext for MyUnionerContext { fn task_ctx(&self) -> Arc { self.ctx.task_ctx() @@ -69,15 +69,17 @@ impl CliSessionContext for MyUnionerContext { unimplemented!() } - async fn execute_logical_plan( + fn execute_logical_plan( &self, plan: LogicalPlan, - ) -> Result { - let new_plan = LogicalPlanBuilder::from(plan.clone()) - .union(plan.clone())? - .build()?; + ) -> BoxFuture<'_, Result> { + Box::pin(async move { + let new_plan = LogicalPlanBuilder::from(plan.clone()) + .union(plan.clone())? + .build()?; - self.ctx.execute_logical_plan(new_plan).await + self.ctx.execute_logical_plan(new_plan).await + }) } } diff --git a/datafusion-cli/src/catalog.rs b/datafusion-cli/src/catalog.rs index ca24da7873bc1..64fb015457915 100644 --- a/datafusion-cli/src/catalog.rs +++ b/datafusion-cli/src/catalog.rs @@ -15,6 +15,7 @@ // specific language governing permissions and limitations // under the License. +use futures::future::BoxFuture; use std::sync::{Arc, Weak}; use crate::object_storage::{AwsOptions, GcpOptions, get_object_store}; @@ -28,7 +29,6 @@ use datafusion::error::Result; use datafusion::execution::context::SessionState; use datafusion::execution::session_state::SessionStateBuilder; -use async_trait::async_trait; use dirs::home_dir; use parking_lot::RwLock; @@ -122,8 +122,6 @@ impl DynamicObjectStoreSchemaProvider { Self { inner, state } } } - -#[async_trait] impl SchemaProvider for DynamicObjectStoreSchemaProvider { fn table_names(&self) -> Vec { self.inner.table_names() @@ -137,63 +135,69 @@ impl SchemaProvider for DynamicObjectStoreSchemaProvider { self.inner.register_table(name, table) } - async fn table(&self, name: &str) -> Result>> { - let inner_table = self.inner.table(name).await; - if inner_table.is_ok() - && let Some(inner_table) = inner_table? - { - return Ok(Some(inner_table)); - } + fn table<'a>( + &'a self, + name: &'a str, + ) -> BoxFuture<'a, Result>>> { + Box::pin(async move { + let inner_table = self.inner.table(name).await; + if inner_table.is_ok() + && let Some(inner_table) = inner_table? + { + return Ok(Some(inner_table)); + } - // if the inner schema provider didn't have a table by - // that name, try to treat it as a listing table - let mut state = self - .state - .upgrade() - .ok_or_else(|| plan_datafusion_err!("locking error"))? - .read() - .clone(); - let mut builder = SessionStateBuilder::from(state.clone()); - let optimized_name = substitute_tilde(name.to_owned()); - let table_url = ListingTableUrl::parse(optimized_name.as_str())?; - let scheme = table_url.scheme(); - let url = table_url.as_ref(); - - // If the store is already registered for this URL then `get_store` - // will return `Ok` which means we don't need to register it again. However, - // if `get_store` returns an `Err` then it means the corresponding store is - // not registered yet and we need to register it - match state.runtime_env().object_store_registry.get_store(url) { - Ok(_) => { /*Nothing to do here, store for this URL is already registered*/ } - Err(_) => { - // Register the store for this URL. Here we don't have access - // to any command options so the only choice is to use an empty collection - match scheme { - "s3" | "oss" | "cos" => { - if let Some(table_options) = builder.table_options() { - table_options.extensions.insert(AwsOptions::default()) + // if the inner schema provider didn't have a table by + // that name, try to treat it as a listing table + let mut state = self + .state + .upgrade() + .ok_or_else(|| plan_datafusion_err!("locking error"))? + .read() + .clone(); + let mut builder = SessionStateBuilder::from(state.clone()); + let optimized_name = substitute_tilde(name.to_owned()); + let table_url = ListingTableUrl::parse(optimized_name.as_str())?; + let scheme = table_url.scheme(); + let url = table_url.as_ref(); + + // If the store is already registered for this URL then `get_store` + // will return `Ok` which means we don't need to register it again. However, + // if `get_store` returns an `Err` then it means the corresponding store is + // not registered yet and we need to register it + match state.runtime_env().object_store_registry.get_store(url) { + Ok(_) => { /*Nothing to do here, store for this URL is already registered*/ + } + Err(_) => { + // Register the store for this URL. Here we don't have access + // to any command options so the only choice is to use an empty collection + match scheme { + "s3" | "oss" | "cos" => { + if let Some(table_options) = builder.table_options() { + table_options.extensions.insert(AwsOptions::default()) + } } - } - "gs" | "gcs" => { - if let Some(table_options) = builder.table_options() { - table_options.extensions.insert(GcpOptions::default()) + "gs" | "gcs" => { + if let Some(table_options) = builder.table_options() { + table_options.extensions.insert(GcpOptions::default()) + } } + _ => {} } - _ => {} + state = builder.build(); + let store = get_object_store( + &state, + table_url.scheme(), + url, + &state.default_table_options(), + false, + ) + .await?; + state.runtime_env().register_object_store(url, store); } - state = builder.build(); - let store = get_object_store( - &state, - table_url.scheme(), - url, - &state.default_table_options(), - false, - ) - .await?; - state.runtime_env().register_object_store(url, store); } - } - self.inner.table(name).await + self.inner.table(name).await + }) } fn deregister_table(&self, name: &str) -> Result>> { diff --git a/datafusion-cli/src/cli_context.rs b/datafusion-cli/src/cli_context.rs index a6320f03fe4de..bb01736a56f15 100644 --- a/datafusion-cli/src/cli_context.rs +++ b/datafusion-cli/src/cli_context.rs @@ -24,11 +24,11 @@ use datafusion::{ logical_expr::LogicalPlan, prelude::SessionContext, }; +use futures::future::BoxFuture; use object_store::ObjectStore; use crate::object_storage::{AwsOptions, GcpOptions}; -#[async_trait::async_trait] /// The CLI session context trait provides a way to have a session context that can be used with datafusion's CLI code. pub trait CliSessionContext { /// Get an atomic reference counted task context. @@ -48,13 +48,12 @@ pub trait CliSessionContext { fn register_table_options_extension_from_scheme(&self, scheme: &str); /// Execute a logical plan and return a DataFrame. - async fn execute_logical_plan( + fn execute_logical_plan( &self, plan: LogicalPlan, - ) -> Result; + ) -> BoxFuture<'_, Result>; } -#[async_trait::async_trait] impl CliSessionContext for SessionContext { fn task_ctx(&self) -> Arc { self.task_ctx() @@ -89,10 +88,10 @@ impl CliSessionContext for SessionContext { } } - async fn execute_logical_plan( + fn execute_logical_plan( &self, plan: LogicalPlan, - ) -> Result { - self.execute_logical_plan(plan).await + ) -> BoxFuture<'_, Result> { + Box::pin(async move { self.execute_logical_plan(plan).await }) } } diff --git a/datafusion-cli/src/functions.rs b/datafusion-cli/src/functions.rs index 0d7d8f33738fa..f595da0671f8b 100644 --- a/datafusion-cli/src/functions.rs +++ b/datafusion-cli/src/functions.rs @@ -18,6 +18,7 @@ //! Functions that are query-able and searchable via the `\h` command use datafusion_common::instant::Instant; +use futures::future::BoxFuture; use std::fmt; use std::fs::File; use std::str::FromStr; @@ -41,7 +42,6 @@ use datafusion::logical_expr::Expr; use datafusion::physical_plan::ExecutionPlan; use datafusion::scalar::ScalarValue; -use async_trait::async_trait; use datafusion_common::heap_size::{DFHeapSize, DFHeapSizeCtx}; use parquet::basic::ConvertedType; use parquet::data_type::{ByteArray, FixedLenByteArray}; @@ -227,8 +227,6 @@ struct ParquetMetadataTable { schema: SchemaRef, batch: RecordBatch, } - -#[async_trait] impl TableProvider for ParquetMetadataTable { fn schema(&self) -> SchemaRef { self.schema.clone() @@ -238,18 +236,20 @@ impl TableProvider for ParquetMetadataTable { datafusion::logical_expr::TableType::Base } - async fn scan( - &self, - _state: &dyn Session, - projection: Option<&[usize]>, - _filters: &[Expr], + fn scan<'a>( + &'a self, + _state: &'a dyn Session, + projection: Option<&'a [usize]>, + _filters: &'a [Expr], _limit: Option, - ) -> Result> { - Ok(MemorySourceConfig::try_new_exec( - &[vec![self.batch.clone()]], - TableProvider::schema(self), - projection.map(|p| p.to_vec()), - )?) + ) -> BoxFuture<'a, Result>> { + Box::pin(async move { + Ok(MemorySourceConfig::try_new_exec( + &[vec![self.batch.clone()]], + TableProvider::schema(self), + projection.map(|p| p.to_vec()), + )? as Arc) + }) } } @@ -469,8 +469,6 @@ struct MetadataCacheTable { schema: SchemaRef, batch: RecordBatch, } - -#[async_trait] impl TableProvider for MetadataCacheTable { fn schema(&self) -> SchemaRef { self.schema.clone() @@ -480,18 +478,20 @@ impl TableProvider for MetadataCacheTable { datafusion::logical_expr::TableType::Base } - async fn scan( - &self, - _state: &dyn Session, - projection: Option<&[usize]>, - _filters: &[Expr], + fn scan<'a>( + &'a self, + _state: &'a dyn Session, + projection: Option<&'a [usize]>, + _filters: &'a [Expr], _limit: Option, - ) -> Result> { - Ok(MemorySourceConfig::try_new_exec( - &[vec![self.batch.clone()]], - TableProvider::schema(self), - projection.map(|p| p.to_vec()), - )?) + ) -> BoxFuture<'a, Result>> { + Box::pin(async move { + Ok(MemorySourceConfig::try_new_exec( + &[vec![self.batch.clone()]], + TableProvider::schema(self), + projection.map(|p| p.to_vec()), + )? as Arc) + }) } } @@ -586,8 +586,6 @@ struct StatisticsCacheTable { schema: SchemaRef, batch: RecordBatch, } - -#[async_trait] impl TableProvider for StatisticsCacheTable { fn schema(&self) -> SchemaRef { self.schema.clone() @@ -597,18 +595,20 @@ impl TableProvider for StatisticsCacheTable { datafusion::logical_expr::TableType::Base } - async fn scan( - &self, - _state: &dyn Session, - projection: Option<&[usize]>, - _filters: &[Expr], + fn scan<'a>( + &'a self, + _state: &'a dyn Session, + projection: Option<&'a [usize]>, + _filters: &'a [Expr], _limit: Option, - ) -> Result> { - Ok(MemorySourceConfig::try_new_exec( - &[vec![self.batch.clone()]], - TableProvider::schema(self), - projection.map(|p| p.to_vec()), - )?) + ) -> BoxFuture<'a, Result>> { + Box::pin(async move { + Ok(MemorySourceConfig::try_new_exec( + &[vec![self.batch.clone()]], + TableProvider::schema(self), + projection.map(|p| p.to_vec()), + )? as Arc) + }) } } @@ -735,8 +735,6 @@ struct ListFilesCacheTable { schema: SchemaRef, batch: RecordBatch, } - -#[async_trait] impl TableProvider for ListFilesCacheTable { fn schema(&self) -> SchemaRef { self.schema.clone() @@ -746,18 +744,20 @@ impl TableProvider for ListFilesCacheTable { datafusion::logical_expr::TableType::Base } - async fn scan( - &self, - _state: &dyn Session, - projection: Option<&[usize]>, - _filters: &[Expr], + fn scan<'a>( + &'a self, + _state: &'a dyn Session, + projection: Option<&'a [usize]>, + _filters: &'a [Expr], _limit: Option, - ) -> Result> { - Ok(MemorySourceConfig::try_new_exec( - &[vec![self.batch.clone()]], - TableProvider::schema(self), - projection.map(|p| p.to_vec()), - )?) + ) -> BoxFuture<'a, Result>> { + Box::pin(async move { + Ok(MemorySourceConfig::try_new_exec( + &[vec![self.batch.clone()]], + TableProvider::schema(self), + projection.map(|p| p.to_vec()), + )? as Arc) + }) } } diff --git a/datafusion-cli/src/object_storage.rs b/datafusion-cli/src/object_storage.rs index 5e6337e303f6f..cdcf97fdd39fb 100644 --- a/datafusion-cli/src/object_storage.rs +++ b/datafusion-cli/src/object_storage.rs @@ -20,7 +20,6 @@ pub(crate) mod stdin; pub use stdin::{StdinCarriesCommands, is_stdin_location}; -use async_trait::async_trait; use aws_config::BehaviorVersion; use aws_credential_types::provider::{ ProvideCredentials, SharedCredentialsProvider, error::CredentialsError, @@ -43,6 +42,8 @@ use object_store::{ gcp::GoogleCloudStorageBuilder, http::HttpBuilder, }; +use std::future::Future; +use std::pin::Pin; use std::{ any::Any, error::Error, @@ -237,24 +238,37 @@ struct S3CredentialProvider { credentials: SharedCredentialsProvider, } -#[async_trait] impl CredentialProvider for S3CredentialProvider { type Credential = AwsCredential; - async fn get_credential(&self) -> object_store::Result> { - let creds = - self.credentials - .provide_credentials() - .await - .map_err(|e| Generic { - store: "S3", - source: Box::new(e), - })?; - Ok(Arc::new(AwsCredential { - key_id: creds.access_key_id().to_string(), - secret_key: creds.secret_access_key().to_string(), - token: creds.session_token().map(ToString::to_string), - })) + fn get_credential<'life0, 'async_trait>( + &'life0 self, + ) -> Pin< + Box< + dyn Future>> + + Send + + 'async_trait, + >, + > + where + 'life0: 'async_trait, + Self: 'async_trait, + { + Box::pin(async move { + let creds = + self.credentials + .provide_credentials() + .await + .map_err(|e| Generic { + store: "S3", + source: Box::new(e), + })?; + Ok(Arc::new(AwsCredential { + key_id: creds.access_key_id().to_string(), + secret_key: creds.secret_access_key().to_string(), + token: creds.session_token().map(ToString::to_string), + })) + }) } } diff --git a/datafusion-cli/src/object_storage/instrumented.rs b/datafusion-cli/src/object_storage/instrumented.rs index a0321cacb374b..abd8d5c2bdfa7 100644 --- a/datafusion-cli/src/object_storage/instrumented.rs +++ b/datafusion-cli/src/object_storage/instrumented.rs @@ -15,6 +15,8 @@ // specific language governing permissions and limitations // under the License. +use std::future::Future; +use std::pin::Pin; use std::{ cmp, fmt, ops::AddAssign, @@ -28,7 +30,6 @@ use std::{ use arrow::array::{ArrayRef, RecordBatch, StringArray}; use arrow::util::pretty::pretty_format_batches; -use async_trait::async_trait; use chrono::Utc; use datafusion::{ common::{HashMap, instant::Instant}, @@ -75,12 +76,12 @@ where type Item = Result; fn poll_next( - mut self: std::pin::Pin<&mut Self>, + mut self: Pin<&mut Self>, cx: &mut std::task::Context<'_>, ) -> std::task::Poll> { let start = *self.start.get_or_insert_with(Instant::now); - let poll_result = std::pin::Pin::new(&mut self.inner).poll_next(cx); + let poll_result = Pin::new(&mut self.inner).poll_next(cx); if !self.duration_recorded && poll_result.is_ready() { self.duration_recorded = true; @@ -377,39 +378,65 @@ impl fmt::Display for InstrumentedObjectStore { } } -#[async_trait] impl ObjectStore for InstrumentedObjectStore { - async fn put_opts( - &self, - location: &Path, + fn put_opts<'life0, 'life1, 'async_trait>( + &'life0 self, + location: &'life1 Path, payload: PutPayload, opts: PutOptions, - ) -> Result { - if self.enabled() { - return self.instrumented_put_opts(location, payload, opts).await; - } + ) -> Pin> + Send + 'async_trait>> + where + 'life0: 'async_trait, + 'life1: 'async_trait, + Self: 'async_trait, + { + Box::pin(async move { + if self.enabled() { + return self.instrumented_put_opts(location, payload, opts).await; + } - self.inner.put_opts(location, payload, opts).await + self.inner.put_opts(location, payload, opts).await + }) } - async fn put_multipart_opts( - &self, - location: &Path, + fn put_multipart_opts<'life0, 'life1, 'async_trait>( + &'life0 self, + location: &'life1 Path, opts: PutMultipartOptions, - ) -> Result> { - if self.enabled() { - return self.instrumented_put_multipart(location, opts).await; - } + ) -> Pin< + Box>> + Send + 'async_trait>, + > + where + 'life0: 'async_trait, + 'life1: 'async_trait, + Self: 'async_trait, + { + Box::pin(async move { + if self.enabled() { + return self.instrumented_put_multipart(location, opts).await; + } - self.inner.put_multipart_opts(location, opts).await + self.inner.put_multipart_opts(location, opts).await + }) } - async fn get_opts(&self, location: &Path, options: GetOptions) -> Result { - if self.enabled() { - return self.instrumented_get_opts(location, options).await; - } + fn get_opts<'life0, 'life1, 'async_trait>( + &'life0 self, + location: &'life1 Path, + options: GetOptions, + ) -> Pin> + Send + 'async_trait>> + where + 'life0: 'async_trait, + 'life1: 'async_trait, + Self: 'async_trait, + { + Box::pin(async move { + if self.enabled() { + return self.instrumented_get_opts(location, options).await; + } - self.inner.get_opts(location, options).await + self.inner.get_opts(location, options).await + }) } fn delete_stream( @@ -431,32 +458,50 @@ impl ObjectStore for InstrumentedObjectStore { self.inner.list(prefix) } - async fn list_with_delimiter(&self, prefix: Option<&Path>) -> Result { - if self.enabled() { - return self.instrumented_list_with_delimiter(prefix).await; - } + fn list_with_delimiter<'life0, 'life1, 'async_trait>( + &'life0 self, + prefix: Option<&'life1 Path>, + ) -> Pin> + Send + 'async_trait>> + where + 'life0: 'async_trait, + 'life1: 'async_trait, + Self: 'async_trait, + { + Box::pin(async move { + if self.enabled() { + return self.instrumented_list_with_delimiter(prefix).await; + } - self.inner.list_with_delimiter(prefix).await + self.inner.list_with_delimiter(prefix).await + }) } - async fn copy_opts( - &self, - from: &Path, - to: &Path, + fn copy_opts<'life0, 'life1, 'life2, 'async_trait>( + &'life0 self, + from: &'life1 Path, + to: &'life2 Path, options: CopyOptions, - ) -> Result<()> { - if self.enabled() { - 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 - } - }; - } + ) -> Pin> + Send + 'async_trait>> + where + 'life0: 'async_trait, + 'life1: 'async_trait, + 'life2: 'async_trait, + Self: 'async_trait, + { + Box::pin(async move { + if self.enabled() { + 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.copy_opts(from, to, options).await + self.inner.copy_opts(from, to, options).await + }) } } diff --git a/datafusion-cli/tests/cli_integration.rs b/datafusion-cli/tests/cli_integration.rs index 9bfeb65278a6d..5a2bcf9969ced 100644 --- a/datafusion-cli/tests/cli_integration.rs +++ b/datafusion-cli/tests/cli_integration.rs @@ -15,11 +15,11 @@ // specific language governing permissions and limitations // under the License. +use futures::future::BoxFuture; use std::process::Command; use rstest::rstest; -use async_trait::async_trait; use insta::internals::SettingsBindDropGuard; use insta::{Settings, glob}; use insta_cmd::{assert_cmd_snapshot, get_cargo_bin}; @@ -785,24 +785,25 @@ SELECT * from CARS LIMIT 1; } /// Extension trait to Add the minio connection information to a Command -#[async_trait] trait MinioCommandExt { - async fn with_minio(&mut self, container: &ContainerAsync) - -> &mut Self; + fn with_minio<'a>( + &'a mut self, + container: &'a ContainerAsync, + ) -> BoxFuture<'a, &'a mut Self>; } - -#[async_trait] impl MinioCommandExt for Command { - async fn with_minio( - &mut self, - container: &ContainerAsync, - ) -> &mut Self { - let port = container.get_host_port_ipv4(9000).await.unwrap(); - - self.env_clear() - .env("AWS_ACCESS_KEY_ID", "TEST-DataFusionLogin") - .env("AWS_SECRET_ACCESS_KEY", "TEST-DataFusionPassword") - .env("AWS_ENDPOINT", format!("http://localhost:{port}")) - .env("AWS_ALLOW_HTTP", "true") + fn with_minio<'a>( + &'a mut self, + container: &'a ContainerAsync, + ) -> BoxFuture<'a, &'a mut Self> { + Box::pin(async move { + let port = container.get_host_port_ipv4(9000).await.unwrap(); + + self.env_clear() + .env("AWS_ACCESS_KEY_ID", "TEST-DataFusionLogin") + .env("AWS_SECRET_ACCESS_KEY", "TEST-DataFusionPassword") + .env("AWS_ENDPOINT", format!("http://localhost:{port}")) + .env("AWS_ALLOW_HTTP", "true") + }) } } diff --git a/datafusion-examples/Cargo.toml b/datafusion-examples/Cargo.toml index 6d6d917ac46ec..42bbff587fc61 100644 --- a/datafusion-examples/Cargo.toml +++ b/datafusion-examples/Cargo.toml @@ -46,7 +46,6 @@ tokio = { workspace = true, features = ["rt-multi-thread", "parking_lot", "fs"] [dev-dependencies] arrow-flight = { workspace = true } -async-trait = { workspace = true } bytes = { workspace = true } dashmap = { workspace = true } # note only use main datafusion crate for examples diff --git a/datafusion-examples/examples/builtin_functions/function_factory.rs b/datafusion-examples/examples/builtin_functions/function_factory.rs index 3cc77371d44ce..830fc5237c47e 100644 --- a/datafusion-examples/examples/builtin_functions/function_factory.rs +++ b/datafusion-examples/examples/builtin_functions/function_factory.rs @@ -30,6 +30,7 @@ use datafusion::logical_expr::{ ColumnarValue, CreateFunction, Expr, ScalarFunctionArgs, ScalarUDF, ScalarUDFImpl, Signature, Volatility, }; +use futures::future::BoxFuture; use std::hash::Hash; use std::result::Result as RResult; use std::sync::Arc; @@ -92,18 +93,19 @@ pub async fn function_factory() -> Result<()> { #[derive(Debug, Default)] struct CustomFunctionFactory {} -#[async_trait::async_trait] impl FunctionFactory for CustomFunctionFactory { /// This function takes the parsed `CREATE FUNCTION` statement and returns /// the function instance. - async fn create( - &self, - _state: &SessionState, + fn create<'a>( + &'a self, + _state: &'a SessionState, statement: CreateFunction, - ) -> Result { - let f: ScalarFunctionWrapper = statement.try_into()?; + ) -> BoxFuture<'a, Result> { + Box::pin(async move { + let f: ScalarFunctionWrapper = statement.try_into()?; - Ok(RegisterFunction::Scalar(Arc::new(ScalarUDF::from(f)))) + Ok(RegisterFunction::Scalar(Arc::new(ScalarUDF::from(f)))) + }) } } diff --git a/datafusion-examples/examples/custom_data_source/custom_datasource.rs b/datafusion-examples/examples/custom_data_source/custom_datasource.rs index 91fb219f34736..0ab2bd61936a3 100644 --- a/datafusion-examples/examples/custom_data_source/custom_datasource.rs +++ b/datafusion-examples/examples/custom_data_source/custom_datasource.rs @@ -17,12 +17,12 @@ //! See `main.rs` for how to run it. +use futures::future::BoxFuture; use std::collections::{BTreeMap, HashMap}; use std::fmt::{self, Debug, Formatter}; use std::sync::{Arc, Mutex}; use std::time::Duration; -use async_trait::async_trait; use datafusion::arrow::array::{UInt8Builder, UInt64Builder}; use datafusion::arrow::datatypes::{DataType, Field, Schema, SchemaRef}; use datafusion::arrow::record_batch::RecordBatch; @@ -186,8 +186,6 @@ impl Default for CustomDataSource { } } } - -#[async_trait] impl TableProvider for CustomDataSource { fn schema(&self) -> SchemaRef { SchemaRef::new(Schema::new(vec![ @@ -200,15 +198,15 @@ impl TableProvider for CustomDataSource { TableType::Base } - async fn scan( - &self, - _state: &dyn Session, - projection: Option<&[usize]>, + fn scan<'a>( + &'a self, + _state: &'a dyn Session, + projection: Option<&'a [usize]>, // filters and limit can be used here to inject some push-down operations if needed - _filters: &[Expr], + _filters: &'a [Expr], _limit: Option, - ) -> Result> { - self.create_physical_plan(projection, self.schema()) + ) -> BoxFuture<'a, Result>> { + Box::pin(async move { self.create_physical_plan(projection, self.schema()) }) } } diff --git a/datafusion-examples/examples/custom_data_source/custom_file_format.rs b/datafusion-examples/examples/custom_data_source/custom_file_format.rs index 0cfbe11877e4d..2bc7fed73a1a6 100644 --- a/datafusion-examples/examples/custom_data_source/custom_file_format.rs +++ b/datafusion-examples/examples/custom_data_source/custom_file_format.rs @@ -17,6 +17,7 @@ //! See `main.rs` for how to run it. +use futures::future::BoxFuture; use std::sync::Arc; use arrow::{ @@ -102,7 +103,6 @@ impl TSVFileFormat { } } -#[async_trait::async_trait] impl FileFormat for TSVFileFormat { fn get_ext(&self) -> String { "tsv".to_string() @@ -120,47 +120,47 @@ impl FileFormat for TSVFileFormat { None } - async fn infer_schema( - &self, - state: &dyn Session, - store: &Arc, - objects: &[ObjectMeta], - ) -> Result { - self.csv_file_format - .infer_schema(state, store, objects) - .await + fn infer_schema<'a>( + &'a self, + state: &'a dyn Session, + store: &'a Arc, + objects: &'a [ObjectMeta], + ) -> BoxFuture<'a, Result> { + self.csv_file_format.infer_schema(state, store, objects) } - async fn infer_stats( - &self, - state: &dyn Session, - store: &Arc, + fn infer_stats<'a>( + &'a self, + state: &'a dyn Session, + store: &'a Arc, table_schema: SchemaRef, - object: &ObjectMeta, - ) -> Result { + object: &'a ObjectMeta, + ) -> BoxFuture<'a, Result> { self.csv_file_format .infer_stats(state, store, table_schema, object) - .await } - async fn create_physical_plan( - &self, - state: &dyn Session, + fn create_physical_plan<'a>( + &'a self, + state: &'a dyn Session, conf: FileScanConfig, - ) -> Result> { - self.csv_file_format.create_physical_plan(state, conf).await + ) -> BoxFuture<'a, Result>> { + self.csv_file_format.create_physical_plan(state, conf) } - async fn create_writer_physical_plan( - &self, + fn create_writer_physical_plan<'a>( + &'a self, input: Arc, - state: &dyn Session, + state: &'a dyn Session, conf: FileSinkConfig, order_requirements: Option, - ) -> Result> { - self.csv_file_format - .create_writer_physical_plan(input, state, conf, order_requirements) - .await + ) -> BoxFuture<'a, Result>> { + self.csv_file_format.create_writer_physical_plan( + input, + state, + conf, + order_requirements, + ) } fn file_source(&self, table_schema: TableSchema) -> Arc { diff --git a/datafusion-examples/examples/custom_data_source/default_column_values.rs b/datafusion-examples/examples/custom_data_source/default_column_values.rs index d2024621aad76..e2847c78cb63c 100644 --- a/datafusion-examples/examples/custom_data_source/default_column_values.rs +++ b/datafusion-examples/examples/custom_data_source/default_column_values.rs @@ -17,12 +17,12 @@ //! See `main.rs` for how to run it. +use futures::future::BoxFuture; use std::collections::HashMap; use std::sync::Arc; use arrow::array::RecordBatch; use arrow::datatypes::{DataType, Field, Schema, SchemaRef}; -use async_trait::async_trait; use datafusion::assert_batches_eq; use datafusion::catalog::memory::DataSourceExec; @@ -197,8 +197,6 @@ impl DefaultValueTableProvider { Self { schema } } } - -#[async_trait] impl TableProvider for DefaultValueTableProvider { fn schema(&self) -> SchemaRef { Arc::clone(&self.schema) @@ -215,52 +213,57 @@ impl TableProvider for DefaultValueTableProvider { Ok(vec![TableProviderFilterPushDown::Inexact; filters.len()]) } - async fn scan( - &self, - state: &dyn Session, - projection: Option<&[usize]>, - filters: &[Expr], + fn scan<'a>( + &'a self, + state: &'a dyn Session, + projection: Option<&'a [usize]>, + filters: &'a [Expr], limit: Option, - ) -> Result> { - let schema = Arc::clone(&self.schema); - let df_schema = DFSchema::try_from(schema.clone())?; - let filter = state.create_physical_expr( - conjunction(filters.iter().cloned()).unwrap_or_else(|| lit(true)), - &df_schema, - )?; - - let parquet_source = ParquetSource::new(schema.clone()) - .with_predicate(filter) - .with_pushdown_filters(true); - - let object_store_url = ObjectStoreUrl::parse("memory://")?; - let store = state.runtime_env().object_store(object_store_url)?; - - let mut files = vec![]; - let mut listing = store.list(None); - while let Some(file) = listing.next().await { - if let Ok(file) = file { - files.push(file); + ) -> BoxFuture<'a, Result>> { + Box::pin(async move { + let schema = Arc::clone(&self.schema); + let df_schema = DFSchema::try_from(schema.clone())?; + let filter = state.create_physical_expr( + conjunction(filters.iter().cloned()).unwrap_or_else(|| lit(true)), + &df_schema, + )?; + + let parquet_source = ParquetSource::new(schema.clone()) + .with_predicate(filter) + .with_pushdown_filters(true); + + let object_store_url = ObjectStoreUrl::parse("memory://")?; + let store = state.runtime_env().object_store(object_store_url)?; + + let mut files = vec![]; + let mut listing = store.list(None); + while let Some(file) = listing.next().await { + if let Ok(file) = file { + files.push(file); + } } - } - let file_group = files - .iter() - .map(|file| PartitionedFile::new(file.location.clone(), file.size)) - .collect(); - - let file_scan_config = FileScanConfigBuilder::new( - ObjectStoreUrl::parse("memory://")?, - Arc::new(parquet_source), - ) - .with_projection_indices(projection.map(|p| p.to_vec()))? - .with_limit(limit) - .with_file_group(file_group) - .with_expr_adapter(Some(Arc::new(DefaultValuePhysicalExprAdapterFactory) as _)); - - Ok(Arc::new(DataSourceExec::new(Arc::new( - file_scan_config.build(), - )))) + let file_group = files + .iter() + .map(|file| PartitionedFile::new(file.location.clone(), file.size)) + .collect(); + + let file_scan_config = FileScanConfigBuilder::new( + ObjectStoreUrl::parse("memory://")?, + Arc::new(parquet_source), + ) + .with_projection_indices(projection.map(|p| p.to_vec()))? + .with_limit(limit) + .with_file_group(file_group) + .with_expr_adapter(Some(Arc::new( + DefaultValuePhysicalExprAdapterFactory, + ) as _)); + + Ok( + Arc::new(DataSourceExec::new(Arc::new(file_scan_config.build()))) + as Arc, + ) + }) } } diff --git a/datafusion-examples/examples/data_io/catalog.rs b/datafusion-examples/examples/data_io/catalog.rs index 7e5cc5a4cfc05..047cef06f74b1 100644 --- a/datafusion-examples/examples/data_io/catalog.rs +++ b/datafusion-examples/examples/data_io/catalog.rs @@ -18,7 +18,6 @@ //! See `main.rs` for how to run it. //! //! Simple example of a catalog/schema implementation. -use async_trait::async_trait; use datafusion::{ arrow::util::pretty, catalog::{CatalogProvider, CatalogProviderList, SchemaProvider}, @@ -31,6 +30,7 @@ use datafusion::{ execution::context::SessionState, prelude::SessionContext, }; +use futures::future::BoxFuture; use std::sync::RwLock; use std::{collections::HashMap, path::Path, sync::Arc}; use std::{fs::File, io::Write}; @@ -175,17 +175,20 @@ impl DirSchema { })) } } - -#[async_trait] impl SchemaProvider for DirSchema { fn table_names(&self) -> Vec { let tables = self.tables.read().unwrap(); tables.keys().cloned().collect::>() } - async fn table(&self, name: &str) -> Result>> { - let tables = self.tables.read().unwrap(); - Ok(tables.get(name).cloned()) + fn table<'a>( + &'a self, + name: &'a str, + ) -> BoxFuture<'a, Result>>> { + Box::pin(async move { + let tables = self.tables.read().unwrap(); + Ok(tables.get(name).cloned()) + }) } fn table_exist(&self, name: &str) -> bool { diff --git a/datafusion-examples/examples/data_io/parquet_advanced_index.rs b/datafusion-examples/examples/data_io/parquet_advanced_index.rs index 43174bd76cf6b..ba80dc703a79f 100644 --- a/datafusion-examples/examples/data_io/parquet_advanced_index.rs +++ b/datafusion-examples/examples/data_io/parquet_advanced_index.rs @@ -55,7 +55,6 @@ use datafusion::prelude::*; use arrow::array::{ArrayRef, Int32Array, RecordBatch, StringArray}; use arrow::datatypes::SchemaRef; -use async_trait::async_trait; use bytes::Bytes; use datafusion::datasource::memory::DataSourceExec; use futures::FutureExt; @@ -452,7 +451,6 @@ impl IndexedFile { /// Implement the TableProvider trait for IndexTableProvider /// so that we can query it as a table. -#[async_trait] impl TableProvider for IndexTableProvider { fn schema(&self) -> SchemaRef { Arc::clone(&self.indexed_file.schema) @@ -462,51 +460,55 @@ impl TableProvider for IndexTableProvider { TableType::Base } - async fn scan( - &self, - state: &dyn Session, - projection: Option<&[usize]>, - filters: &[Expr], + fn scan<'a>( + &'a self, + state: &'a dyn Session, + projection: Option<&'a [usize]>, + filters: &'a [Expr], limit: Option, - ) -> Result> { - let indexed_file = &self.indexed_file; - let predicate = self.filters_to_predicate(state, filters)?; - - // Figure out which row groups to scan based on the predicate - let access_plan = self.create_plan(&predicate)?; - println!("{access_plan:?}"); - - let partitioned_file = indexed_file - .partitioned_file() - // provide the starting access plan to the DataSourceExec by - // storing it as "extensions" on PartitionedFile - .with_extension(access_plan); - - // Prepare for scanning - let schema = self.schema(); - let object_store_url = ObjectStoreUrl::parse("file://")?; - - // Configure a factory interface to avoid re-reading the metadata for each file - let reader_factory = - CachedParquetFileReaderFactory::new(Arc::clone(&self.object_store)) - .with_file(indexed_file); - - let file_source = Arc::new( - ParquetSource::new(schema.clone()) - // provide the predicate so the DataSourceExec can try and prune - // row groups internally - .with_predicate(predicate) - // provide the factory to create parquet reader without re-reading metadata - .with_parquet_file_reader_factory(Arc::new(reader_factory)), - ); - let file_scan_config = FileScanConfigBuilder::new(object_store_url, file_source) - .with_limit(limit) - .with_projection_indices(projection.map(|p| p.to_vec()))? - .with_file(partitioned_file) - .build(); - - // Finally, put it all together into a DataSourceExec - Ok(DataSourceExec::from_data_source(file_scan_config)) + ) -> BoxFuture<'a, Result>> { + Box::pin(async move { + let indexed_file = &self.indexed_file; + let predicate = self.filters_to_predicate(state, filters)?; + + // Figure out which row groups to scan based on the predicate + let access_plan = self.create_plan(&predicate)?; + println!("{access_plan:?}"); + + let partitioned_file = indexed_file + .partitioned_file() + // provide the starting access plan to the DataSourceExec by + // storing it as "extensions" on PartitionedFile + .with_extension(access_plan); + + // Prepare for scanning + let schema = self.schema(); + let object_store_url = ObjectStoreUrl::parse("file://")?; + + // Configure a factory interface to avoid re-reading the metadata for each file + let reader_factory = + CachedParquetFileReaderFactory::new(Arc::clone(&self.object_store)) + .with_file(indexed_file); + + let file_source = Arc::new( + ParquetSource::new(schema.clone()) + // provide the predicate so the DataSourceExec can try and prune + // row groups internally + .with_predicate(predicate) + // provide the factory to create parquet reader without re-reading metadata + .with_parquet_file_reader_factory(Arc::new(reader_factory)), + ); + let file_scan_config = + FileScanConfigBuilder::new(object_store_url, file_source) + .with_limit(limit) + .with_projection_indices(projection.map(|p| p.to_vec()))? + .with_file(partitioned_file) + .build(); + + // Finally, put it all together into a DataSourceExec + Ok(DataSourceExec::from_data_source(file_scan_config) + as Arc) + }) } /// Tell DataFusion to push filters down to the scan method diff --git a/datafusion-examples/examples/data_io/parquet_embedded_index.rs b/datafusion-examples/examples/data_io/parquet_embedded_index.rs index 9f205b8b8e306..1094aa3f2084d 100644 --- a/datafusion-examples/examples/data_io/parquet_embedded_index.rs +++ b/datafusion-examples/examples/data_io/parquet_embedded_index.rs @@ -116,7 +116,6 @@ use arrow::array::{ArrayRef, StringArray}; use arrow::record_batch::RecordBatch; use arrow_schema::{DataType, Field, Schema, SchemaRef}; -use async_trait::async_trait; use datafusion::catalog::{Session, TableProvider}; use datafusion::common::{HashMap, HashSet, Result, exec_err}; use datafusion::datasource::TableType; @@ -132,6 +131,7 @@ use datafusion::parquet::file::reader::{FileReader, SerializedFileReader}; use datafusion::physical_plan::ExecutionPlan; use datafusion::prelude::*; use datafusion::scalar::ScalarValue; +use futures::future::BoxFuture; use std::fs::{File, read_dir}; use std::io::{Read, Seek, SeekFrom, Write}; use std::path::{Path, PathBuf}; @@ -391,7 +391,6 @@ fn get_key_value<'a>(file_meta_data: &'a FileMetaData, key: &'_ str) -> Option<& } /// Implement TableProvider for DistinctIndexTable, using the distinct index to prune files -#[async_trait] impl TableProvider for DistinctIndexTable { fn schema(&self) -> SchemaRef { self.schema.clone() @@ -402,63 +401,66 @@ impl TableProvider for DistinctIndexTable { /// Prune files before reading: only keep files whose distinct set /// contains the filter value - async fn scan( - &self, - _ctx: &dyn Session, - _proj: Option<&[usize]>, - filters: &[Expr], + fn scan<'a>( + &'a self, + _ctx: &'a dyn Session, + _proj: Option<&'a [usize]>, + filters: &'a [Expr], _limit: Option, - ) -> Result> { - // This example only handles filters of the form - // `category = 'X'` where X is a string literal - // - // You can use `PruningPredicate` for much more general range and - // equality analysis or write your own custom logic. - let mut target: Option<&str> = None; - - if filters.len() == 1 - && let Expr::BinaryExpr(expr) = &filters[0] - && expr.op == Operator::Eq - && let (Expr::Column(c), Expr::Literal(ScalarValue::Utf8(Some(v)), _)) = - (&*expr.left, &*expr.right) - && c.name == "category" - { - println!("Filtering for category: {v}"); - target = Some(v); - } - // Determine which files to scan - let files_to_scan: Vec<_> = self - .files_and_index - .iter() - .filter_map(|(f, distinct_index)| { - // keep file if no target or target is in the distinct set - if target.is_none() || distinct_index.contains(target?) { - Some(f) - } else { - None - } - }) - .collect(); - - println!("Scanning only files: {files_to_scan:?}"); - - // Build ParquetSource to actually read the files - let url = ObjectStoreUrl::parse("file://")?; - let source = Arc::new( - ParquetSource::new(self.schema.clone()).with_enable_page_index(true), - ); - let mut builder = FileScanConfigBuilder::new(url, source); - for file in files_to_scan { - let path = self.dir.join(file); - let len = std::fs::metadata(&path)?.len(); - // If the index contained information about row groups or pages, - // you could also pass that information here to further prune - // the data read from the file. - let partitioned_file = - PartitionedFile::new(path.to_str().unwrap().to_string(), len); - builder = builder.with_file(partitioned_file); - } - Ok(DataSourceExec::from_data_source(builder.build())) + ) -> BoxFuture<'a, Result>> { + Box::pin(async move { + // This example only handles filters of the form + // `category = 'X'` where X is a string literal + // + // You can use `PruningPredicate` for much more general range and + // equality analysis or write your own custom logic. + let mut target: Option<&str> = None; + + if filters.len() == 1 + && let Expr::BinaryExpr(expr) = &filters[0] + && expr.op == Operator::Eq + && let (Expr::Column(c), Expr::Literal(ScalarValue::Utf8(Some(v)), _)) = + (&*expr.left, &*expr.right) + && c.name == "category" + { + println!("Filtering for category: {v}"); + target = Some(v); + } + // Determine which files to scan + let files_to_scan: Vec<_> = self + .files_and_index + .iter() + .filter_map(|(f, distinct_index)| { + // keep file if no target or target is in the distinct set + if target.is_none() || distinct_index.contains(target?) { + Some(f) + } else { + None + } + }) + .collect(); + + println!("Scanning only files: {files_to_scan:?}"); + + // Build ParquetSource to actually read the files + let url = ObjectStoreUrl::parse("file://")?; + let source = Arc::new( + ParquetSource::new(self.schema.clone()).with_enable_page_index(true), + ); + let mut builder = FileScanConfigBuilder::new(url, source); + for file in files_to_scan { + let path = self.dir.join(file); + let len = std::fs::metadata(&path)?.len(); + // If the index contained information about row groups or pages, + // you could also pass that information here to further prune + // the data read from the file. + let partitioned_file = + PartitionedFile::new(path.to_str().unwrap().to_string(), len); + builder = builder.with_file(partitioned_file); + } + Ok(DataSourceExec::from_data_source(builder.build()) + as Arc) + }) } /// Tell DataFusion that we can handle filters on the "category" column diff --git a/datafusion-examples/examples/data_io/parquet_encrypted_with_kms.rs b/datafusion-examples/examples/data_io/parquet_encrypted_with_kms.rs index 8e92f465eafe9..5efe0550d1613 100644 --- a/datafusion-examples/examples/data_io/parquet_encrypted_with_kms.rs +++ b/datafusion-examples/examples/data_io/parquet_encrypted_with_kms.rs @@ -19,7 +19,6 @@ use arrow::array::{ArrayRef, Int32Array, RecordBatch, StringArray}; use arrow_schema::SchemaRef; -use async_trait::async_trait; use base64::Engine; use datafusion::common::extensions_options; use datafusion::config::{EncryptionFactoryOptions, TableParquetOptions}; @@ -34,6 +33,7 @@ use datafusion::parquet::encryption::{ }; use datafusion::prelude::SessionContext; use futures::StreamExt; +use futures::future::BoxFuture; use object_store::path::Path; use rand::rand_core::{OsRng, TryRngCore}; use std::collections::HashSet; @@ -213,7 +213,6 @@ struct TestEncryptionFactory {} /// `EncryptionFactory` is a DataFusion trait for types that generate /// file encryption and decryption properties. -#[async_trait] impl EncryptionFactory for TestEncryptionFactory { /// Generate file encryption properties to use when writing a Parquet file. /// The `schema` is provided so that it may be used to dynamically configure @@ -222,58 +221,63 @@ impl EncryptionFactory for TestEncryptionFactory { /// but other implementations may want to use this to compute an /// AAD prefix for the file, or to allow use of external key material /// (where key metadata is stored in a JSON file alongside Parquet files). - async fn get_file_encryption_properties( - &self, - options: &EncryptionFactoryOptions, - schema: &SchemaRef, - _file_path: &Path, - ) -> Result>> { - let config: EncryptionConfig = options.to_extension_options()?; - - // Generate a random encryption key for this file. - let mut key = vec![0u8; 16]; - OsRng.try_fill_bytes(&mut key).unwrap(); - - // Generate the key metadata that allows retrieving the key when reading the file. - let key_metadata = wrap_key(&key); - - let mut builder = FileEncryptionProperties::builder(key.to_vec()) - .with_footer_key_metadata(key_metadata.clone()); - - let encrypted_columns: HashSet<&str> = - config.encrypted_columns.split(',').collect(); - if !encrypted_columns.is_empty() { - // Set up per-column encryption. - for field in schema.fields().iter() { - if encrypted_columns.contains(field.name().as_str()) { - // Here we re-use the same key for all encrypted columns, - // but new keys could also be generated per column. - builder = builder.with_column_key_and_metadata( - field.name().as_str(), - key.clone(), - key_metadata.clone(), - ); + fn get_file_encryption_properties<'a>( + &'a self, + options: &'a EncryptionFactoryOptions, + schema: &'a SchemaRef, + _file_path: &'a Path, + ) -> BoxFuture<'a, Result>>> { + Box::pin(async move { + let config: EncryptionConfig = options.to_extension_options()?; + + // Generate a random encryption key for this file. + let mut key = vec![0u8; 16]; + OsRng.try_fill_bytes(&mut key).unwrap(); + + // Generate the key metadata that allows retrieving the key when reading the file. + let key_metadata = wrap_key(&key); + + let mut builder = FileEncryptionProperties::builder(key.to_vec()) + .with_footer_key_metadata(key_metadata.clone()); + + let encrypted_columns: HashSet<&str> = + config.encrypted_columns.split(',').collect(); + if !encrypted_columns.is_empty() { + // Set up per-column encryption. + for field in schema.fields().iter() { + if encrypted_columns.contains(field.name().as_str()) { + // Here we re-use the same key for all encrypted columns, + // but new keys could also be generated per column. + builder = builder.with_column_key_and_metadata( + field.name().as_str(), + key.clone(), + key_metadata.clone(), + ); + } } } - } - let encryption_properties = builder.build()?; + let encryption_properties = builder.build()?; - Ok(Some(encryption_properties)) + Ok(Some(encryption_properties)) + }) } /// Generate file decryption properties to use when reading a Parquet file. /// Rather than provide the AES keys directly for decryption, we set a `KeyRetriever` /// that can determine the keys using the encryption metadata. - async fn get_file_decryption_properties( - &self, - _options: &EncryptionFactoryOptions, - _file_path: &Path, - ) -> Result>> { - let decryption_properties = - FileDecryptionProperties::with_key_retriever(Arc::new(TestKeyRetriever {})) - .build()?; - Ok(Some(decryption_properties)) + fn get_file_decryption_properties<'a>( + &'a self, + _options: &'a EncryptionFactoryOptions, + _file_path: &'a Path, + ) -> BoxFuture<'a, Result>>> { + Box::pin(async move { + let decryption_properties = FileDecryptionProperties::with_key_retriever( + Arc::new(TestKeyRetriever {}), + ) + .build()?; + Ok(Some(decryption_properties)) + }) } } diff --git a/datafusion-examples/examples/data_io/parquet_index.rs b/datafusion-examples/examples/data_io/parquet_index.rs index 8ca63516b2c08..17cade4e51630 100644 --- a/datafusion-examples/examples/data_io/parquet_index.rs +++ b/datafusion-examples/examples/data_io/parquet_index.rs @@ -23,7 +23,6 @@ use arrow::array::{ }; use arrow::datatypes::{Int32Type, SchemaRef}; use arrow::util::pretty::pretty_format_batches; -use async_trait::async_trait; use datafusion::catalog::Session; use datafusion::common::pruning::PruningStatistics; use datafusion::common::{ @@ -45,6 +44,7 @@ use datafusion::physical_expr::PhysicalExpr; use datafusion::physical_optimizer::pruning::PruningPredicateBuilder; use datafusion::physical_plan::ExecutionPlan; use datafusion::prelude::*; +use futures::future::BoxFuture; use std::collections::HashSet; use std::fmt::Display; use std::fs; @@ -204,8 +204,6 @@ impl IndexTableProvider { &self.index } } - -#[async_trait] impl TableProvider for IndexTableProvider { fn schema(&self) -> SchemaRef { self.index.schema().clone() @@ -215,48 +213,51 @@ impl TableProvider for IndexTableProvider { TableType::Base } - async fn scan( - &self, - state: &dyn Session, - projection: Option<&[usize]>, - filters: &[Expr], + fn scan<'a>( + &'a self, + state: &'a dyn Session, + projection: Option<&'a [usize]>, + filters: &'a [Expr], limit: Option, - ) -> Result> { - let df_schema = DFSchema::try_from(self.schema())?; - // convert filters like [`a = 1`, `b = 2`] to a single filter like `a = 1 AND b = 2` - let predicate = conjunction(filters.to_vec()); - let predicate = predicate - .map(|predicate| state.create_physical_expr(predicate, &df_schema)) - .transpose()? - // if there are no filters, use a literal true to have a predicate - // that always evaluates to true we can pass to the index - .unwrap_or_else(|| datafusion::physical_expr::expressions::lit(true)); - - // Use the index to find the files that might have data that matches the - // predicate. Any file that can not have data that matches the predicate - // will not be returned. - let files = self.index.get_files(predicate.clone())?; - - let object_store_url = ObjectStoreUrl::parse("file://")?; - let source = - Arc::new(ParquetSource::new(self.schema()).with_predicate(predicate)); - let mut file_scan_config_builder = - FileScanConfigBuilder::new(object_store_url, source) - .with_projection_indices(projection.map(|p| p.to_vec()))? - .with_limit(limit); - - // Transform to the format needed to pass to DataSourceExec - // Create one file group per file (default to scanning them all in parallel) - for (file_name, file_size) in files { - let path = self.dir.join(file_name); - let canonical_path = fs::canonicalize(path)?; - file_scan_config_builder = file_scan_config_builder.with_file( - PartitionedFile::new(canonical_path.display().to_string(), file_size), - ); - } - Ok(DataSourceExec::from_data_source( - file_scan_config_builder.build(), - )) + ) -> BoxFuture<'a, Result>> { + Box::pin(async move { + let df_schema = DFSchema::try_from(self.schema())?; + // convert filters like [`a = 1`, `b = 2`] to a single filter like `a = 1 AND b = 2` + let predicate = conjunction(filters.to_vec()); + let predicate = predicate + .map(|predicate| state.create_physical_expr(predicate, &df_schema)) + .transpose()? + // if there are no filters, use a literal true to have a predicate + // that always evaluates to true we can pass to the index + .unwrap_or_else(|| datafusion::physical_expr::expressions::lit(true)); + + // Use the index to find the files that might have data that matches the + // predicate. Any file that can not have data that matches the predicate + // will not be returned. + let files = self.index.get_files(predicate.clone())?; + + let object_store_url = ObjectStoreUrl::parse("file://")?; + let source = + Arc::new(ParquetSource::new(self.schema()).with_predicate(predicate)); + let mut file_scan_config_builder = + FileScanConfigBuilder::new(object_store_url, source) + .with_projection_indices(projection.map(|p| p.to_vec()))? + .with_limit(limit); + + // Transform to the format needed to pass to DataSourceExec + // Create one file group per file (default to scanning them all in parallel) + for (file_name, file_size) in files { + let path = self.dir.join(file_name); + let canonical_path = fs::canonicalize(path)?; + file_scan_config_builder = file_scan_config_builder.with_file( + PartitionedFile::new(canonical_path.display().to_string(), file_size), + ); + } + Ok( + DataSourceExec::from_data_source(file_scan_config_builder.build()) + as Arc, + ) + }) } /// Tell DataFusion to push filters down to the scan method diff --git a/datafusion-examples/examples/data_io/remote_catalog.rs b/datafusion-examples/examples/data_io/remote_catalog.rs index 49157c6b1f5b7..ae79b616e4732 100644 --- a/datafusion-examples/examples/data_io/remote_catalog.rs +++ b/datafusion-examples/examples/data_io/remote_catalog.rs @@ -33,7 +33,6 @@ /// [Hive]: https://hive.apache.org/ use arrow::array::record_batch; use arrow::datatypes::{Field, Fields, Schema, SchemaRef}; -use async_trait::async_trait; use datafusion::catalog::TableProvider; use datafusion::catalog::{AsyncSchemaProvider, Session}; use datafusion::common::Result; @@ -45,6 +44,7 @@ use datafusion::physical_plan::ExecutionPlan; use datafusion::physical_plan::stream::RecordBatchStreamAdapter; use datafusion::prelude::{DataFrame, SessionContext}; use futures::TryStreamExt; +use futures::future::BoxFuture; use std::sync::Arc; /// Interfacing with a remote catalog (e.g. over a network) @@ -184,19 +184,22 @@ impl RemoteCatalogInterface { /// Implements an async version of the DataFusion SchemaProvider API for tables /// stored in a remote catalog. struct RemoteCatalogDatafusionAdapter(Arc); - -#[async_trait] impl AsyncSchemaProvider for RemoteCatalogDatafusionAdapter { - async fn table(&self, name: &str) -> Result>> { - // Fetch information about the table from the remote catalog - // - // Note that a real remote catalog interface could return more - // information, but at the minimum, DataFusion requires the - // table's schema for planing. - Ok(self.0.table_info(name).await?.map(|schema| { - Arc::new(RemoteTable::new(Arc::clone(&self.0), name, schema)) - as Arc - })) + fn table<'a>( + &'a self, + name: &'a str, + ) -> BoxFuture<'a, Result>>> { + Box::pin(async move { + // Fetch information about the table from the remote catalog + // + // Note that a real remote catalog interface could return more + // information, but at the minimum, DataFusion requires the + // table's schema for planing. + Ok(self.0.table_info(name).await?.map(|schema| { + Arc::new(RemoteTable::new(Arc::clone(&self.0), name, schema)) + as Arc + })) + }) } } @@ -224,7 +227,6 @@ impl RemoteTable { } /// Implement the DataFusion Catalog API for [`RemoteTable`] -#[async_trait] impl TableProvider for RemoteTable { fn schema(&self) -> SchemaRef { self.schema.clone() @@ -234,31 +236,33 @@ impl TableProvider for RemoteTable { TableType::Base } - async fn scan( - &self, - _state: &dyn Session, - projection: Option<&[usize]>, - _filters: &[Expr], + fn scan<'a>( + &'a self, + _state: &'a dyn Session, + projection: Option<&'a [usize]>, + _filters: &'a [Expr], _limit: Option, - ) -> Result> { - // Note that `scan` is called once the plan begin execution, and thus is - // async. When interacting with remote data sources, this is the place - // to begin establishing the remote connections and interacting with the - // remote storage system. - // - // As this example is just modeling the catalog API interface, we buffer - // the results locally in memory for simplicity. - let batches = self - .remote_catalog_interface - .read_data(&self.name) - .await? - .try_collect() - .await?; - let exec = MemorySourceConfig::try_new_exec( - &[batches], - self.schema.clone(), - projection.map(|p| p.to_vec()), - )?; - Ok(exec) + ) -> BoxFuture<'a, Result>> { + Box::pin(async move { + // Note that `scan` is called once the plan begin execution, and thus is + // async. When interacting with remote data sources, this is the place + // to begin establishing the remote connections and interacting with the + // remote storage system. + // + // As this example is just modeling the catalog API interface, we buffer + // the results locally in memory for simplicity. + let batches = self + .remote_catalog_interface + .read_data(&self.name) + .await? + .try_collect() + .await?; + let exec = MemorySourceConfig::try_new_exec( + &[batches], + self.schema.clone(), + projection.map(|p| p.to_vec()), + )?; + Ok(exec as Arc) + }) } } diff --git a/datafusion-examples/examples/dataframe/cache_factory.rs b/datafusion-examples/examples/dataframe/cache_factory.rs index ffbce298b4f17..c10cfc8fbffc0 100644 --- a/datafusion-examples/examples/dataframe/cache_factory.rs +++ b/datafusion-examples/examples/dataframe/cache_factory.rs @@ -17,12 +17,12 @@ //! See `main.rs` for how to run it. +use futures::future::BoxFuture; use std::fmt::Debug; use std::hash::Hash; use std::sync::{Arc, RwLock}; use arrow::array::RecordBatch; -use async_trait::async_trait; use datafusion::catalog::Session; use datafusion::catalog::memory::MemorySourceConfig; use datafusion::common::DFSchemaRef; @@ -138,56 +138,57 @@ impl UserDefinedLogicalNodeCore for CacheNode { struct CacheNodePlanner { cache_manager: Arc>, } - -#[async_trait] impl ExtensionPlanner for CacheNodePlanner { - async fn plan_extension( - &self, - _planner: &dyn PhysicalPlanner, - node: &dyn UserDefinedLogicalNode, - logical_inputs: &[&LogicalPlan], - physical_inputs: &[Arc], - session_state: &dyn Session, - _planning_ctx: &PhysicalPlanningContext, - ) -> Result>> { - if let Some(cache_node) = node.as_any().downcast_ref::() { - assert_eq!(logical_inputs.len(), 1, "Inconsistent number of inputs"); - assert_eq!(physical_inputs.len(), 1, "Inconsistent number of inputs"); - if self - .cache_manager - .read() - .unwrap() - .get(&cache_node.input) - .is_none() - { - let ctx = session_state.task_ctx(); - println!("caching in memory"); - let batches = - collect_partitioned(physical_inputs[0].clone(), ctx).await?; - self.cache_manager - .write() + fn plan_extension<'a>( + &'a self, + _planner: &'a dyn PhysicalPlanner, + node: &'a dyn UserDefinedLogicalNode, + logical_inputs: &'a [&'a LogicalPlan], + physical_inputs: &'a [Arc], + session_state: &'a dyn Session, + _planning_ctx: &'a PhysicalPlanningContext, + ) -> BoxFuture<'a, Result>>> { + Box::pin(async move { + if let Some(cache_node) = node.as_any().downcast_ref::() { + assert_eq!(logical_inputs.len(), 1, "Inconsistent number of inputs"); + assert_eq!(physical_inputs.len(), 1, "Inconsistent number of inputs"); + if self + .cache_manager + .read() .unwrap() - .put(cache_node.input.clone(), batches); + .get(&cache_node.input) + .is_none() + { + let ctx = session_state.task_ctx(); + println!("caching in memory"); + let batches = + collect_partitioned(physical_inputs[0].clone(), ctx).await?; + self.cache_manager + .write() + .unwrap() + .put(cache_node.input.clone(), batches); + } else { + println!("fetching directly from cache manager"); + } + Ok(self + .cache_manager + .read() + .unwrap() + .get(&cache_node.input) + .map(|batches| { + let exec: Arc = + MemorySourceConfig::try_new_exec( + batches, + physical_inputs[0].schema(), + None, + ) + .unwrap(); + exec + })) } else { - println!("fetching directly from cache manager"); + Ok(None) } - Ok(self - .cache_manager - .read() - .unwrap() - .get(&cache_node.input) - .map(|batches| { - let exec: Arc = MemorySourceConfig::try_new_exec( - batches, - physical_inputs[0].schema(), - None, - ) - .unwrap(); - exec - })) - } else { - Ok(None) - } + }) } } @@ -195,23 +196,23 @@ impl ExtensionPlanner for CacheNodePlanner { struct CacheNodeQueryPlanner { cache_manager: Arc>, } - -#[async_trait] impl QueryPlanner for CacheNodeQueryPlanner { - async fn create_physical_plan( - &self, - logical_plan: &LogicalPlan, - session_state: &dyn Session, - ) -> Result> { - let physical_planner = - DefaultPhysicalPlanner::with_extension_planners(vec![Arc::new( - CacheNodePlanner { - cache_manager: Arc::clone(&self.cache_manager), - }, - )]); - physical_planner - .create_physical_plan(logical_plan, session_state) - .await + fn create_physical_plan<'a>( + &'a self, + logical_plan: &'a LogicalPlan, + session_state: &'a dyn Session, + ) -> BoxFuture<'a, Result>> { + Box::pin(async move { + let physical_planner = + DefaultPhysicalPlanner::with_extension_planners(vec![Arc::new( + CacheNodePlanner { + cache_manager: Arc::clone(&self.cache_manager), + }, + )]); + physical_planner + .create_physical_plan(logical_plan, session_state) + .await + }) } } diff --git a/datafusion-examples/examples/relation_planner/table_sample.rs b/datafusion-examples/examples/relation_planner/table_sample.rs index 7a8f533ac3a9b..988dd03267633 100644 --- a/datafusion-examples/examples/relation_planner/table_sample.rs +++ b/datafusion-examples/examples/relation_planner/table_sample.rs @@ -87,6 +87,8 @@ use std::{ task::{Context, Poll}, }; +use futures::future::BoxFuture; + use arrow::datatypes::{Float64Type, Int64Type}; use arrow::{ array::{ArrayRef, Int32Array, RecordBatch, StringArray, UInt32Array}, @@ -98,7 +100,6 @@ use futures::{ stream::{Stream, StreamExt}, }; use rand::{Rng, SeedableRng, rngs::StdRng}; -use tonic::async_trait; use datafusion::{ catalog::Session, @@ -563,49 +564,50 @@ impl Hash for HashableF64 { /// convert [`TableSamplePlanNode`] into [`SampleExec`]. #[derive(Debug)] struct TableSampleQueryPlanner; - -#[async_trait] impl QueryPlanner for TableSampleQueryPlanner { - async fn create_physical_plan( - &self, - logical_plan: &LogicalPlan, - session_state: &dyn Session, - ) -> Result> { - let planner = DefaultPhysicalPlanner::with_extension_planners(vec![Arc::new( - TableSampleExtensionPlanner, - )]); - planner - .create_physical_plan(logical_plan, session_state) - .await + fn create_physical_plan<'a>( + &'a self, + logical_plan: &'a LogicalPlan, + session_state: &'a dyn Session, + ) -> BoxFuture<'a, Result>> { + Box::pin(async move { + let planner = + DefaultPhysicalPlanner::with_extension_planners(vec![Arc::new( + TableSampleExtensionPlanner, + )]); + planner + .create_physical_plan(logical_plan, session_state) + .await + }) } } /// Extension planner that converts [`TableSamplePlanNode`] to [`SampleExec`]. struct TableSampleExtensionPlanner; - -#[async_trait] impl ExtensionPlanner for TableSampleExtensionPlanner { - async fn plan_extension( - &self, - _planner: &dyn PhysicalPlanner, - node: &dyn UserDefinedLogicalNode, - _logical_inputs: &[&LogicalPlan], - physical_inputs: &[Arc], - _session_state: &dyn Session, - _planning_ctx: &PhysicalPlanningContext, - ) -> Result>> { - let Some(sample_node) = node.as_any().downcast_ref::() - else { - return Ok(None); - }; - - let exec = SampleExec::try_new( - Arc::clone(&physical_inputs[0]), - sample_node.lower_bound.0, - sample_node.upper_bound.0, - sample_node.seed, - )?; - Ok(Some(Arc::new(exec))) + fn plan_extension<'a>( + &'a self, + _planner: &'a dyn PhysicalPlanner, + node: &'a dyn UserDefinedLogicalNode, + _logical_inputs: &'a [&'a LogicalPlan], + physical_inputs: &'a [Arc], + _session_state: &'a dyn Session, + _planning_ctx: &'a PhysicalPlanningContext, + ) -> BoxFuture<'a, Result>>> { + Box::pin(async move { + let Some(sample_node) = node.as_any().downcast_ref::() + else { + return Ok(None); + }; + + let exec = SampleExec::try_new( + Arc::clone(&physical_inputs[0]), + sample_node.lower_bound.0, + sample_node.upper_bound.0, + sample_node.seed, + )?; + Ok(Some(Arc::new(exec) as Arc)) + }) } } diff --git a/datafusion-examples/examples/udf/async_udf.rs b/datafusion-examples/examples/udf/async_udf.rs index 43b82c398c5c6..3500eecd3d1ac 100644 --- a/datafusion-examples/examples/udf/async_udf.rs +++ b/datafusion-examples/examples/udf/async_udf.rs @@ -23,11 +23,11 @@ //! making network requests. This can be used for tasks like fetching //! data from an external API such as a LLM service or an external database. +use futures::future::BoxFuture; use std::sync::Arc; use arrow::array::{ArrayRef, BooleanArray, Int64Array, RecordBatch, StringArray}; use arrow_schema::{DataType, Field, Schema}; -use async_trait::async_trait; use datafusion::assert_batches_eq; use datafusion::common::cast::as_string_view_array; use datafusion::common::error::Result; @@ -158,7 +158,6 @@ impl AskLLM { /// All async UDFs implement the `ScalarUDFImpl` trait, which provides the basic /// information for the function, such as its name, signature, and return type. -/// [async_trait] impl ScalarUDFImpl for AskLLM { fn name(&self) -> &str { "ask_llm" @@ -181,7 +180,6 @@ impl ScalarUDFImpl for AskLLM { /// In addition to [`ScalarUDFImpl`], we also need to implement the /// [`AsyncScalarUDFImpl`] trait. -#[async_trait] impl AsyncScalarUDFImpl for AskLLM { /// The `invoke_async_with_args` method is similar to `invoke_with_args`, /// but it returns a `Future` that resolves to the result. @@ -191,44 +189,47 @@ impl AsyncScalarUDFImpl for AskLLM { /// is processing the query, so you may wish to make actual network requests /// on a different `Runtime`, as explained in the `thread_pools.rs` example /// in this directory. - async fn invoke_async_with_args( + fn invoke_async_with_args( &self, args: ScalarFunctionArgs, - ) -> Result { - // in a real UDF you would likely want to special case constant - // arguments to improve performance, but this example converts the - // arguments to arrays for simplicity. - let args = ColumnarValue::values_to_arrays(&args.args)?; - let [content_column, question_column] = take_function_args(self.name(), args)?; - - // In a real function, you would use a library such as `reqwest` here to - // make an async HTTP request. Credentials and other configurations can - // be supplied via the `ConfigOptions` parameter. - - // In this example, we will simulate the LLM response by comparing the two - // input arguments using some static strings - let content_column = as_string_view_array(&content_column)?; - let question_column = as_string_view_array(&question_column)?; - - let result_array: BooleanArray = content_column - .iter() - .zip(question_column.iter()) - .map(|(a, b)| { - // If either value is null, return None - let a = a?; - let b = b?; - // Simulate an LLM response by checking the arguments to some - // hardcoded conditions. - if a.contains("cat") && b.contains("furry") - || a.contains("dog") && b.contains("furry") - { - Some(true) - } else { - Some(false) - } - }) - .collect(); - - Ok(ColumnarValue::from(Arc::new(result_array) as ArrayRef)) + ) -> BoxFuture<'_, Result> { + Box::pin(async move { + // in a real UDF you would likely want to special case constant + // arguments to improve performance, but this example converts the + // arguments to arrays for simplicity. + let args = ColumnarValue::values_to_arrays(&args.args)?; + let [content_column, question_column] = + take_function_args(self.name(), args)?; + + // In a real function, you would use a library such as `reqwest` here to + // make an async HTTP request. Credentials and other configurations can + // be supplied via the `ConfigOptions` parameter. + + // In this example, we will simulate the LLM response by comparing the two + // input arguments using some static strings + let content_column = as_string_view_array(&content_column)?; + let question_column = as_string_view_array(&question_column)?; + + let result_array: BooleanArray = content_column + .iter() + .zip(question_column.iter()) + .map(|(a, b)| { + // If either value is null, return None + let a = a?; + let b = b?; + // Simulate an LLM response by checking the arguments to some + // hardcoded conditions. + if a.contains("cat") && b.contains("furry") + || a.contains("dog") && b.contains("furry") + { + Some(true) + } else { + Some(false) + } + }) + .collect(); + + Ok(ColumnarValue::from(Arc::new(result_array) as ArrayRef)) + }) } } diff --git a/datafusion-examples/examples/udf/simple_udtf.rs b/datafusion-examples/examples/udf/simple_udtf.rs index 3b55a0456a0aa..670e4c9cf78bc 100644 --- a/datafusion-examples/examples/udf/simple_udtf.rs +++ b/datafusion-examples/examples/udf/simple_udtf.rs @@ -17,6 +17,7 @@ //! See `main.rs` for how to run it. +use futures::future::BoxFuture; use std::fs::File; use std::io::Seek; use std::path::Path; @@ -24,7 +25,6 @@ use std::sync::Arc; use arrow::csv::ReaderBuilder; use arrow::csv::reader::Format; -use async_trait::async_trait; use datafusion::arrow::datatypes::SchemaRef; use datafusion::arrow::record_batch::RecordBatch; use datafusion::catalog::{Session, TableFunctionArgs, TableFunctionImpl}; @@ -82,8 +82,6 @@ struct LocalCsvTable { limit: Option, batches: Vec, } - -#[async_trait] impl TableProvider for LocalCsvTable { fn schema(&self) -> SchemaRef { self.schema.clone() @@ -93,37 +91,39 @@ impl TableProvider for LocalCsvTable { TableType::Base } - async fn scan( - &self, - _state: &dyn Session, - projection: Option<&[usize]>, - _filters: &[Expr], + fn scan<'a>( + &'a self, + _state: &'a dyn Session, + projection: Option<&'a [usize]>, + _filters: &'a [Expr], _limit: Option, - ) -> Result> { - let batches = if let Some(max_return_lines) = self.limit { - // get max return rows from self.batches - let mut batches = vec![]; - let mut lines = 0; - for batch in &self.batches { - let batch_lines = batch.num_rows(); - if lines + batch_lines > max_return_lines { - let batch_lines = max_return_lines - lines; - batches.push(batch.slice(0, batch_lines)); - break; - } else { - batches.push(batch.clone()); - lines += batch_lines; + ) -> BoxFuture<'a, Result>> { + Box::pin(async move { + let batches = if let Some(max_return_lines) = self.limit { + // get max return rows from self.batches + let mut batches = vec![]; + let mut lines = 0; + for batch in &self.batches { + let batch_lines = batch.num_rows(); + if lines + batch_lines > max_return_lines { + let batch_lines = max_return_lines - lines; + batches.push(batch.slice(0, batch_lines)); + break; + } else { + batches.push(batch.clone()); + lines += batch_lines; + } } - } - batches - } else { - self.batches.clone() - }; - Ok(MemorySourceConfig::try_new_exec( - &[batches], - TableProvider::schema(self), - projection.map(|p| p.to_vec()), - )?) + batches + } else { + self.batches.clone() + }; + Ok(MemorySourceConfig::try_new_exec( + &[batches], + TableProvider::schema(self), + projection.map(|p| p.to_vec()), + )? as Arc) + }) } } diff --git a/datafusion/catalog-listing/Cargo.toml b/datafusion/catalog-listing/Cargo.toml index abe58f45994be..774d0a92f8565 100644 --- a/datafusion/catalog-listing/Cargo.toml +++ b/datafusion/catalog-listing/Cargo.toml @@ -32,7 +32,6 @@ all-features = true [dependencies] arrow = { workspace = true } -async-trait = { workspace = true } datafusion-catalog = { workspace = true } datafusion-common = { workspace = true, features = ["object_store"] } datafusion-datasource = { workspace = true } diff --git a/datafusion/catalog-listing/src/table.rs b/datafusion/catalog-listing/src/table.rs index 6c294fe077db4..ede51426bab99 100644 --- a/datafusion/catalog-listing/src/table.rs +++ b/datafusion/catalog-listing/src/table.rs @@ -21,7 +21,6 @@ use crate::helpers::{ }; use crate::{ListingOptions, ListingTableConfig}; use arrow::datatypes::{Field, Schema, SchemaBuilder, SchemaRef}; -use async_trait::async_trait; use datafusion_catalog::{ScanArgs, ScanResult, Session, TableProvider}; use datafusion_common::stats::{Precision, is_known_empty}; use datafusion_common::{ @@ -478,8 +477,6 @@ fn can_be_evaluated_for_partition_pruning( !partition_column_names.is_empty() && expr_applicable_for_cols(partition_column_names, expr) } - -#[async_trait] impl TableProvider for ListingTable { fn schema(&self) -> SchemaRef { Arc::clone(&self.table_schema) @@ -493,38 +490,21 @@ impl TableProvider for ListingTable { TableType::Base } - // Hand-written `#[async_trait]` expansion to reduce compile time. See - // - fn scan<'life0, 'life1, 'life2, 'life3, 'async_trait>( - &'life0 self, - state: &'life1 dyn Session, - projection: Option<&'life2 [usize]>, - filters: &'life3 [Expr], + fn scan<'a>( + &'a self, + state: &'a dyn Session, + projection: Option<&'a [usize]>, + filters: &'a [Expr], limit: Option, - ) -> BoxFuture<'async_trait, datafusion_common::Result>> - where - 'life0: 'async_trait, - 'life1: 'async_trait, - 'life2: 'async_trait, - 'life3: 'async_trait, - Self: 'async_trait, - { + ) -> BoxFuture<'a, datafusion_common::Result>> { self.scan_boxed(state, projection, filters, limit) } - // Hand-written `#[async_trait]` expansion to reduce compile time. See - // - fn scan_with_args<'a, 'life0, 'life1, 'async_trait>( - &'life0 self, - state: &'life1 dyn Session, + fn scan_with_args<'a>( + &'a self, + state: &'a dyn Session, args: ScanArgs<'a>, - ) -> BoxFuture<'async_trait, datafusion_common::Result> - where - 'a: 'async_trait, - 'life0: 'async_trait, - 'life1: 'async_trait, - Self: 'async_trait, - { + ) -> BoxFuture<'a, datafusion_common::Result> { self.scan_with_args_boxed(state, args) } @@ -556,19 +536,12 @@ impl TableProvider for ListingTable { self.definition.as_deref() } - // Hand-written `#[async_trait]` expansion to reduce compile time. See - // - fn insert_into<'life0, 'life1, 'async_trait>( - &'life0 self, - state: &'life1 dyn Session, + fn insert_into<'a>( + &'a self, + state: &'a dyn Session, input: Arc, insert_op: InsertOp, - ) -> BoxFuture<'async_trait, datafusion_common::Result>> - where - 'life0: 'async_trait, - 'life1: 'async_trait, - Self: 'async_trait, - { + ) -> BoxFuture<'a, datafusion_common::Result>> { self.insert_into_boxed(state, input, insert_op) } diff --git a/datafusion/catalog/Cargo.toml b/datafusion/catalog/Cargo.toml index 1009e9aee477b..ddc815c721170 100644 --- a/datafusion/catalog/Cargo.toml +++ b/datafusion/catalog/Cargo.toml @@ -32,7 +32,6 @@ all-features = true [dependencies] arrow = { workspace = true } -async-trait = { workspace = true } dashmap = { workspace = true } datafusion-common = { workspace = true } datafusion-common-runtime = { workspace = true } diff --git a/datafusion/catalog/src/async.rs b/datafusion/catalog/src/async.rs index 6f8e25cafd1c0..433253237268a 100644 --- a/datafusion/catalog/src/async.rs +++ b/datafusion/catalog/src/async.rs @@ -15,9 +15,9 @@ // specific language governing permissions and limitations // under the License. +use futures::future::BoxFuture; use std::sync::Arc; -use async_trait::async_trait; use datafusion_common::{HashMap, TableReference, error::Result, not_impl_err}; use datafusion_execution::config::SessionConfig; @@ -31,7 +31,6 @@ struct ResolvedSchemaProvider { owner_name: Option, cached_tables: HashMap>, } -#[async_trait] impl SchemaProvider for ResolvedSchemaProvider { fn owner_name(&self) -> Option<&str> { self.owner_name.as_deref() @@ -41,8 +40,11 @@ impl SchemaProvider for ResolvedSchemaProvider { self.cached_tables.keys().cloned().collect() } - async fn table(&self, name: &str) -> Result>> { - Ok(self.cached_tables.get(name).cloned()) + fn table<'a>( + &'a self, + name: &'a str, + ) -> BoxFuture<'a, Result>>> { + Box::pin(async move { Ok(self.cached_tables.get(name).cloned()) }) } fn register_table( @@ -184,10 +186,12 @@ impl CatalogProviderList for ResolvedCatalogProviderList { /// See the [remote_catalog.rs] for an end to end example /// /// [remote_catalog.rs]: https://github.com/apache/datafusion/blob/main/datafusion-examples/examples/data_io/remote_catalog.rs -#[async_trait] pub trait AsyncSchemaProvider: Send + Sync { /// Lookup a table in the schema provider - async fn table(&self, name: &str) -> Result>>; + fn table<'a>( + &'a self, + name: &'a str, + ) -> BoxFuture<'a, Result>>>; /// Creates a cached provider that can be used to execute a query containing given references /// /// This method will walk through the references and look them up once, creating a cache of table @@ -198,48 +202,51 @@ pub trait AsyncSchemaProvider: Send + Sync { /// for refresh or eviction of stale entries. /// /// See the [`AsyncSchemaProvider`] documentation for additional details - async fn resolve( - &self, - references: &[TableReference], - config: &SessionConfig, - catalog_name: &str, - schema_name: &str, - ) -> Result> { - let mut cached_tables = HashMap::>>::new(); - - for reference in references { - let ref_catalog_name = reference - .catalog() - .unwrap_or(&config.options().catalog.default_catalog); - - // Maybe this is a reference to some other catalog provided in another way - if ref_catalog_name != catalog_name { - continue; + fn resolve<'a>( + &'a self, + references: &'a [TableReference], + config: &'a SessionConfig, + catalog_name: &'a str, + schema_name: &'a str, + ) -> BoxFuture<'a, Result>> { + Box::pin(async move { + let mut cached_tables = + HashMap::>>::new(); + + for reference in references { + let ref_catalog_name = reference + .catalog() + .unwrap_or(&config.options().catalog.default_catalog); + + // Maybe this is a reference to some other catalog provided in another way + if ref_catalog_name != catalog_name { + continue; + } + + let ref_schema_name = reference + .schema() + .unwrap_or(&config.options().catalog.default_schema); + + if ref_schema_name != schema_name { + continue; + } + + if !cached_tables.contains_key(reference.table()) { + let resolved_table = self.table(reference.table()).await?; + cached_tables.insert(reference.table().to_string(), resolved_table); + } } - let ref_schema_name = reference - .schema() - .unwrap_or(&config.options().catalog.default_schema); - - if ref_schema_name != schema_name { - continue; - } - - if !cached_tables.contains_key(reference.table()) { - let resolved_table = self.table(reference.table()).await?; - cached_tables.insert(reference.table().to_string(), resolved_table); - } - } - - let cached_tables = cached_tables - .into_iter() - .filter_map(|(key, maybe_value)| maybe_value.map(|value| (key, value))) - .collect(); + let cached_tables = cached_tables + .into_iter() + .filter_map(|(key, maybe_value)| maybe_value.map(|value| (key, value))) + .collect(); - Ok(Arc::new(ResolvedSchemaProvider { - cached_tables, - owner_name: Some(catalog_name.to_string()), - })) + Ok(Arc::new(ResolvedSchemaProvider { + cached_tables, + owner_name: Some(catalog_name.to_string()), + }) as Arc) + }) } } @@ -248,11 +255,12 @@ pub trait AsyncSchemaProvider: Send + Sync { /// The [`CatalogProvider::schema`] method is synchronous because asynchronous operations should /// not be used during planning. This trait makes it easy to lookup schema references once and cache /// them for future planning use. See [`AsyncSchemaProvider`] for more details on motivation. - -#[async_trait] pub trait AsyncCatalogProvider: Send + Sync { /// Lookup a schema in the provider - async fn schema(&self, name: &str) -> Result>>; + fn schema<'a>( + &'a self, + name: &'a str, + ) -> BoxFuture<'a, Result>>>; /// Creates a cached provider that can be used to execute a query containing given references /// @@ -263,57 +271,60 @@ pub trait AsyncCatalogProvider: Send + Sync { /// /// This cache is intended to be short-lived for the execution of a single query. There is no mechanism /// for refresh or eviction of stale entries. - async fn resolve( - &self, - references: &[TableReference], - config: &SessionConfig, - catalog_name: &str, - ) -> Result> { - let mut cached_schemas = - HashMap::>::new(); - - for reference in references { - let ref_catalog_name = reference - .catalog() - .unwrap_or(&config.options().catalog.default_catalog); - - // Maybe this is a reference to some other catalog provided in another way - if ref_catalog_name != catalog_name { - continue; + fn resolve<'a>( + &'a self, + references: &'a [TableReference], + config: &'a SessionConfig, + catalog_name: &'a str, + ) -> BoxFuture<'a, Result>> { + Box::pin(async move { + let mut cached_schemas = + HashMap::>::new(); + + for reference in references { + let ref_catalog_name = reference + .catalog() + .unwrap_or(&config.options().catalog.default_catalog); + + // Maybe this is a reference to some other catalog provided in another way + if ref_catalog_name != catalog_name { + continue; + } + + let schema_name = reference + .schema() + .unwrap_or(&config.options().catalog.default_schema); + + let schema = if let Some(schema) = cached_schemas.get_mut(schema_name) { + schema + } else { + let resolved_schema = self.schema(schema_name).await?; + let resolved_schema = resolved_schema.map(|resolved_schema| { + ResolvedSchemaProviderBuilder::new( + catalog_name.to_string(), + resolved_schema, + ) + }); + cached_schemas.insert(schema_name.to_string(), resolved_schema); + cached_schemas.get_mut(schema_name).unwrap() + }; + + // If we can't find the catalog don't bother checking the table + let Some(schema) = schema else { continue }; + + schema.resolve_table(reference.table()).await?; } - let schema_name = reference - .schema() - .unwrap_or(&config.options().catalog.default_schema); - - let schema = if let Some(schema) = cached_schemas.get_mut(schema_name) { - schema - } else { - let resolved_schema = self.schema(schema_name).await?; - let resolved_schema = resolved_schema.map(|resolved_schema| { - ResolvedSchemaProviderBuilder::new( - catalog_name.to_string(), - resolved_schema, - ) - }); - cached_schemas.insert(schema_name.to_string(), resolved_schema); - cached_schemas.get_mut(schema_name).unwrap() - }; - - // If we can't find the catalog don't bother checking the table - let Some(schema) = schema else { continue }; - - schema.resolve_table(reference.table()).await?; - } - - let cached_schemas = cached_schemas - .into_iter() - .filter_map(|(key, maybe_builder)| { - maybe_builder.map(|schema_builder| (key, schema_builder.finish())) - }) - .collect::>(); + let cached_schemas = cached_schemas + .into_iter() + .filter_map(|(key, maybe_builder)| { + maybe_builder.map(|schema_builder| (key, schema_builder.finish())) + }) + .collect::>(); - Ok(Arc::new(ResolvedCatalogProvider { cached_schemas })) + Ok(Arc::new(ResolvedCatalogProvider { cached_schemas }) + as Arc) + }) } } @@ -322,10 +333,12 @@ pub trait AsyncCatalogProvider: Send + Sync { /// The [`CatalogProviderList::catalog`] method is synchronous because asynchronous operations should /// not be used during planning. This trait makes it easy to lookup catalog references once and cache /// them for future planning use. See [`AsyncSchemaProvider`] for more details on motivation. -#[async_trait] pub trait AsyncCatalogProviderList: Send + Sync { /// Lookup a catalog in the provider - async fn catalog(&self, name: &str) -> Result>>; + fn catalog<'a>( + &'a self, + name: &'a str, + ) -> BoxFuture<'a, Result>>>; /// Creates a cached provider that can be used to execute a query containing given references /// @@ -336,77 +349,82 @@ pub trait AsyncCatalogProviderList: Send + Sync { /// /// This cache is intended to be short-lived for the execution of a single query. There is no mechanism /// for refresh or eviction of stale entries. - async fn resolve( - &self, - references: &[TableReference], - config: &SessionConfig, - ) -> Result> { - let mut cached_catalogs = - HashMap::>::new(); - - for reference in references { - let catalog_name = reference - .catalog() - .unwrap_or(&config.options().catalog.default_catalog); - - // We will do three lookups here, one for the catalog, one for the schema, and one for the table - // We cache the result (both found results and not-found results) to speed up future lookups - // - // Note that a cache-miss is not an error at this point. We allow for the possibility that - // other providers may supply the reference. - // - // If this is the only provider then a not-found error will be raised during planning when it can't - // find the reference in the cache. - - let catalog = if let Some(catalog) = cached_catalogs.get_mut(catalog_name) { - catalog - } else { - let resolved_catalog = self.catalog(catalog_name).await?; - let resolved_catalog = - resolved_catalog.map(ResolvedCatalogProviderBuilder::new); - cached_catalogs.insert(catalog_name.to_string(), resolved_catalog); - cached_catalogs.get_mut(catalog_name).unwrap() - }; - - // If we can't find the catalog don't bother checking the schema / table - let Some(catalog) = catalog else { continue }; - - let schema_name = reference - .schema() - .unwrap_or(&config.options().catalog.default_schema); - - let schema = if let Some(schema) = catalog.cached_schemas.get_mut(schema_name) - { - schema - } else { - let resolved_schema = catalog.async_provider.schema(schema_name).await?; - let resolved_schema = resolved_schema.map(|async_schema| { - ResolvedSchemaProviderBuilder::new( - catalog_name.to_string(), - async_schema, - ) - }); - catalog - .cached_schemas - .insert(schema_name.to_string(), resolved_schema); - catalog.cached_schemas.get_mut(schema_name).unwrap() - }; - - // If we can't find the catalog don't bother checking the table - let Some(schema) = schema else { continue }; - - schema.resolve_table(reference.table()).await?; - } + fn resolve<'a>( + &'a self, + references: &'a [TableReference], + config: &'a SessionConfig, + ) -> BoxFuture<'a, Result>> { + Box::pin(async move { + let mut cached_catalogs = + HashMap::>::new(); + + for reference in references { + let catalog_name = reference + .catalog() + .unwrap_or(&config.options().catalog.default_catalog); + + // We will do three lookups here, one for the catalog, one for the schema, and one for the table + // We cache the result (both found results and not-found results) to speed up future lookups + // + // Note that a cache-miss is not an error at this point. We allow for the possibility that + // other providers may supply the reference. + // + // If this is the only provider then a not-found error will be raised during planning when it can't + // find the reference in the cache. + + let catalog = if let Some(catalog) = cached_catalogs.get_mut(catalog_name) + { + catalog + } else { + let resolved_catalog = self.catalog(catalog_name).await?; + let resolved_catalog = + resolved_catalog.map(ResolvedCatalogProviderBuilder::new); + cached_catalogs.insert(catalog_name.to_string(), resolved_catalog); + cached_catalogs.get_mut(catalog_name).unwrap() + }; + + // If we can't find the catalog don't bother checking the schema / table + let Some(catalog) = catalog else { continue }; + + let schema_name = reference + .schema() + .unwrap_or(&config.options().catalog.default_schema); + + let schema = + if let Some(schema) = catalog.cached_schemas.get_mut(schema_name) { + schema + } else { + let resolved_schema = + catalog.async_provider.schema(schema_name).await?; + let resolved_schema = resolved_schema.map(|async_schema| { + ResolvedSchemaProviderBuilder::new( + catalog_name.to_string(), + async_schema, + ) + }); + catalog + .cached_schemas + .insert(schema_name.to_string(), resolved_schema); + catalog.cached_schemas.get_mut(schema_name).unwrap() + }; + + // If we can't find the catalog don't bother checking the table + let Some(schema) = schema else { continue }; + + schema.resolve_table(reference.table()).await?; + } - // Build the cached catalog provider list - let cached_catalogs = cached_catalogs - .into_iter() - .filter_map(|(key, maybe_builder)| { - maybe_builder.map(|catalog_builder| (key, catalog_builder.finish())) - }) - .collect::>(); + // Build the cached catalog provider list + let cached_catalogs = cached_catalogs + .into_iter() + .filter_map(|(key, maybe_builder)| { + maybe_builder.map(|catalog_builder| (key, catalog_builder.finish())) + }) + .collect::>(); - Ok(Arc::new(ResolvedCatalogProviderList { cached_catalogs })) + Ok(Arc::new(ResolvedCatalogProviderList { cached_catalogs }) + as Arc) + }) } } @@ -418,11 +436,11 @@ mod tests { }; use arrow::datatypes::SchemaRef; - use async_trait::async_trait; use datafusion_common::{Statistics, TableReference, error::Result}; use datafusion_execution::config::SessionConfig; use datafusion_expr::{Expr, TableType}; use datafusion_physical_plan::ExecutionPlan; + use futures::future::BoxFuture; use crate::{Session, TableProvider}; @@ -430,7 +448,6 @@ mod tests { #[derive(Debug)] struct MockTableProvider {} - #[async_trait] impl TableProvider for MockTableProvider { /// Get a reference to the schema for this table fn schema(&self) -> SchemaRef { @@ -441,14 +458,14 @@ mod tests { unimplemented!() } - async fn scan( - &self, - _state: &dyn Session, - _projection: Option<&[usize]>, - _filters: &[Expr], + fn scan<'a>( + &'a self, + _state: &'a dyn Session, + _projection: Option<&'a [usize]>, + _filters: &'a [Expr], _limit: Option, - ) -> Result> { - unimplemented!() + ) -> BoxFuture<'a, Result>> { + Box::pin(async move { unimplemented!() }) } fn statistics(&self) -> Option { @@ -464,16 +481,21 @@ mod tests { const MOCK_CATALOG: &str = "mock_catalog"; const MOCK_SCHEMA: &str = "mock_schema"; const MOCK_TABLE: &str = "mock_table"; - - #[async_trait] impl AsyncSchemaProvider for MockAsyncSchemaProvider { - async fn table(&self, name: &str) -> Result>> { - self.lookup_count.fetch_add(1, Ordering::Release); - if name == MOCK_TABLE { - Ok(Some(Arc::new(MockTableProvider {}))) - } else { - Ok(None) - } + fn table<'a>( + &'a self, + name: &'a str, + ) -> BoxFuture<'a, Result>>> { + Box::pin(async move { + self.lookup_count.fetch_add(1, Ordering::Release); + if name == MOCK_TABLE { + Ok(Some( + Arc::new(MockTableProvider {}) as Arc + )) + } else { + Ok(None) + } + }) } } @@ -556,19 +578,20 @@ mod tests { struct MockAsyncCatalogProvider { lookup_count: AtomicU32, } - - #[async_trait] impl AsyncCatalogProvider for MockAsyncCatalogProvider { - async fn schema( - &self, - name: &str, - ) -> Result>> { - self.lookup_count.fetch_add(1, Ordering::Release); - if name == MOCK_SCHEMA { - Ok(Some(Arc::new(MockAsyncSchemaProvider::default()))) - } else { - Ok(None) - } + fn schema<'a>( + &'a self, + name: &'a str, + ) -> BoxFuture<'a, Result>>> { + Box::pin(async move { + self.lookup_count.fetch_add(1, Ordering::Release); + if name == MOCK_SCHEMA { + Ok(Some(Arc::new(MockAsyncSchemaProvider::default()) + as Arc)) + } else { + Ok(None) + } + }) } } @@ -641,19 +664,20 @@ mod tests { struct MockAsyncCatalogProviderList { lookup_count: AtomicU32, } - - #[async_trait] impl AsyncCatalogProviderList for MockAsyncCatalogProviderList { - async fn catalog( - &self, - name: &str, - ) -> Result>> { - self.lookup_count.fetch_add(1, Ordering::Release); - if name == MOCK_CATALOG { - Ok(Some(Arc::new(MockAsyncCatalogProvider::default()))) - } else { - Ok(None) - } + fn catalog<'a>( + &'a self, + name: &'a str, + ) -> BoxFuture<'a, Result>>> { + Box::pin(async move { + self.lookup_count.fetch_add(1, Ordering::Release); + if name == MOCK_CATALOG { + Ok(Some(Arc::new(MockAsyncCatalogProvider::default()) + as Arc)) + } else { + Ok(None) + } + }) } } diff --git a/datafusion/catalog/src/cte_worktable.rs b/datafusion/catalog/src/cte_worktable.rs index 9180e01043ed0..8debd6926e811 100644 --- a/datafusion/catalog/src/cte_worktable.rs +++ b/datafusion/catalog/src/cte_worktable.rs @@ -22,7 +22,6 @@ use std::future::ready; use std::sync::Arc; use arrow::datatypes::SchemaRef; -use async_trait::async_trait; use datafusion_common::error::Result; use datafusion_expr::{Expr, LogicalPlan, TableProviderFilterPushDown, TableType}; use datafusion_physical_plan::ExecutionPlan; @@ -65,8 +64,6 @@ impl CteWorkTable { Arc::clone(&self.table_schema) } } - -#[async_trait] impl TableProvider for CteWorkTable { fn get_logical_plan(&'_ self) -> Option> { None @@ -80,38 +77,21 @@ impl TableProvider for CteWorkTable { TableType::Temporary } - // Hand-written `#[async_trait]` expansion to reduce compile time. See - // - fn scan<'life0, 'life1, 'life2, 'life3, 'async_trait>( - &'life0 self, - state: &'life1 dyn Session, - projection: Option<&'life2 [usize]>, - filters: &'life3 [Expr], + fn scan<'a>( + &'a self, + state: &'a dyn Session, + projection: Option<&'a [usize]>, + filters: &'a [Expr], limit: Option, - ) -> BoxFuture<'async_trait, Result>> - where - 'life0: 'async_trait, - 'life1: 'async_trait, - 'life2: 'async_trait, - 'life3: 'async_trait, - Self: 'async_trait, - { + ) -> BoxFuture<'a, Result>> { self.scan_boxed(state, projection, filters, limit) } - // Hand-written `#[async_trait]` expansion to reduce compile time. See - // - fn scan_with_args<'a, 'life0, 'life1, 'async_trait>( - &'life0 self, - state: &'life1 dyn Session, + fn scan_with_args<'a>( + &'a self, + state: &'a dyn Session, args: ScanArgs<'a>, - ) -> BoxFuture<'async_trait, Result> - where - 'a: 'async_trait, - 'life0: 'async_trait, - 'life1: 'async_trait, - Self: 'async_trait, - { + ) -> BoxFuture<'a, Result> { Box::pin(ready(self.scan_with_args_inner(state, &args))) } diff --git a/datafusion/catalog/src/default_table_source.rs b/datafusion/catalog/src/default_table_source.rs index 3342db54de92f..5c9ee3907c8ba 100644 --- a/datafusion/catalog/src/default_table_source.rs +++ b/datafusion/catalog/src/default_table_source.rs @@ -99,13 +99,11 @@ pub fn source_as_provider( #[test] fn preserves_table_type() { - use async_trait::async_trait; use datafusion_common::DataFusionError; + use futures::future::BoxFuture; #[derive(Debug)] struct TestTempTable; - - #[async_trait] impl TableProvider for TestTempTable { fn table_type(&self) -> TableType { TableType::Temporary @@ -115,15 +113,17 @@ fn preserves_table_type() { unimplemented!() } - async fn scan( - &self, - _: &dyn crate::Session, - _: Option<&[usize]>, - _: &[Expr], + fn scan<'a>( + &'a self, + _: &'a dyn crate::Session, + _: Option<&'a [usize]>, + _: &'a [Expr], _: Option, - ) -> Result, DataFusionError> - { - unimplemented!() + ) -> BoxFuture< + 'a, + Result, DataFusionError>, + > { + Box::pin(async move { unimplemented!() }) } } diff --git a/datafusion/catalog/src/dynamic_file/catalog.rs b/datafusion/catalog/src/dynamic_file/catalog.rs index 4437d99667547..0b706819d7902 100644 --- a/datafusion/catalog/src/dynamic_file/catalog.rs +++ b/datafusion/catalog/src/dynamic_file/catalog.rs @@ -18,7 +18,7 @@ //! [`DynamicFileCatalog`] that creates tables from file paths use crate::{CatalogProvider, CatalogProviderList, SchemaProvider, TableProvider}; -use async_trait::async_trait; +use futures::future::BoxFuture; use std::fmt::Debug; use std::sync::Arc; @@ -125,22 +125,22 @@ impl DynamicFileSchemaProvider { Self { inner, factory } } } - -#[async_trait] impl SchemaProvider for DynamicFileSchemaProvider { fn table_names(&self) -> Vec { self.inner.table_names() } - async fn table( - &self, - name: &str, - ) -> datafusion_common::Result>> { - if let Some(table) = self.inner.table(name).await? { - return Ok(Some(table)); - } + fn table<'a>( + &'a self, + name: &'a str, + ) -> BoxFuture<'a, datafusion_common::Result>>> { + Box::pin(async move { + if let Some(table) = self.inner.table(name).await? { + return Ok(Some(table)); + } - self.factory.try_new(name).await + self.factory.try_new(name).await + }) } fn register_table( @@ -164,11 +164,10 @@ impl SchemaProvider for DynamicFileSchemaProvider { } /// [UrlTableFactory] is a factory that can create a table provider from the given url. -#[async_trait] pub trait UrlTableFactory: Debug + Sync + Send { /// create a new table provider from the provided url - async fn try_new( - &self, - url: &str, - ) -> datafusion_common::Result>>; + fn try_new<'a>( + &'a self, + url: &'a str, + ) -> BoxFuture<'a, datafusion_common::Result>>>; } diff --git a/datafusion/catalog/src/empty.rs b/datafusion/catalog/src/empty.rs index 66b62bfcec452..86ed0fd7fd40b 100644 --- a/datafusion/catalog/src/empty.rs +++ b/datafusion/catalog/src/empty.rs @@ -17,10 +17,10 @@ //! [`EmptyTable`] useful for testing. +use futures::future::BoxFuture; use std::sync::Arc; use arrow::datatypes::*; -use async_trait::async_trait; use datafusion_common::{Result, project_schema}; use datafusion_expr::{Expr, TableType}; use datafusion_physical_plan::ExecutionPlan; @@ -52,8 +52,6 @@ impl EmptyTable { self } } - -#[async_trait] impl TableProvider for EmptyTable { fn schema(&self) -> SchemaRef { Arc::clone(&self.schema) @@ -63,17 +61,21 @@ impl TableProvider for EmptyTable { TableType::Base } - async fn scan( - &self, - _state: &dyn Session, - projection: Option<&[usize]>, - _filters: &[Expr], + fn scan<'a>( + &'a self, + _state: &'a dyn Session, + projection: Option<&'a [usize]>, + _filters: &'a [Expr], _limit: Option, - ) -> Result> { - // even though there is no data, projections apply - let projected_schema = project_schema(&self.schema, projection)?; - Ok(Arc::new( - EmptyExec::new(projected_schema).with_partitions(self.partitions), - )) + ) -> BoxFuture<'a, Result>> { + Box::pin(async move { + // even though there is no data, projections apply + let projected_schema = project_schema(&self.schema, projection)?; + Ok( + Arc::new( + EmptyExec::new(projected_schema).with_partitions(self.partitions), + ) as Arc, + ) + }) } } diff --git a/datafusion/catalog/src/information_schema.rs b/datafusion/catalog/src/information_schema.rs index d9ad7791af67c..391f9b48dc550 100644 --- a/datafusion/catalog/src/information_schema.rs +++ b/datafusion/catalog/src/information_schema.rs @@ -28,7 +28,6 @@ use arrow::{ datatypes::{DataType, Field, FieldRef, Schema, SchemaRef}, record_batch::RecordBatch, }; -use async_trait::async_trait; use datafusion_common::DataFusionError; use datafusion_common::config::{ConfigEntry, ConfigOptions}; use datafusion_common::error::Result; @@ -43,6 +42,7 @@ use datafusion_expr::{TableType, Volatility}; use datafusion_physical_plan::SendableRecordBatchStream; use datafusion_physical_plan::stream::RecordBatchStreamAdapter; use datafusion_physical_plan::streaming::PartitionStream; +use futures::future::BoxFuture; use std::collections::{BTreeSet, HashMap, HashSet}; use std::fmt::Debug; use std::sync::Arc; @@ -572,8 +572,6 @@ fn get_udwf_args_and_return_types( fn remove_native_type_prefix(native_type: &NativeType) -> String { format!("{native_type}") } - -#[async_trait] impl SchemaProvider for InformationSchemaProvider { fn table_names(&self) -> Vec { INFORMATION_SCHEMA_TABLES @@ -582,25 +580,28 @@ impl SchemaProvider for InformationSchemaProvider { .collect() } - async fn table( - &self, - name: &str, - ) -> Result>, DataFusionError> { - let config = self.config.clone(); - let table: Arc = match name.to_ascii_lowercase().as_str() { - TABLES => Arc::new(InformationSchemaTables::new(config)), - COLUMNS => Arc::new(InformationSchemaColumns::new(config)), - VIEWS => Arc::new(InformationSchemaViews::new(config)), - DF_SETTINGS => Arc::new(InformationSchemaDfSettings::new(config)), - SCHEMATA => Arc::new(InformationSchemata::new(config)), - ROUTINES => Arc::new(InformationSchemaRoutines::new(config)), - PARAMETERS => Arc::new(InformationSchemaParameters::new(config)), - _ => return Ok(None), - }; - - Ok(Some(Arc::new( - StreamingTable::try_new(Arc::clone(table.schema()), vec![table]).unwrap(), - ))) + fn table<'a>( + &'a self, + name: &'a str, + ) -> BoxFuture<'a, Result>, DataFusionError>> { + Box::pin(async move { + let config = self.config.clone(); + let table: Arc = match name.to_ascii_lowercase().as_str() + { + TABLES => Arc::new(InformationSchemaTables::new(config)), + COLUMNS => Arc::new(InformationSchemaColumns::new(config)), + VIEWS => Arc::new(InformationSchemaViews::new(config)), + DF_SETTINGS => Arc::new(InformationSchemaDfSettings::new(config)), + SCHEMATA => Arc::new(InformationSchemata::new(config)), + ROUTINES => Arc::new(InformationSchemaRoutines::new(config)), + PARAMETERS => Arc::new(InformationSchemaParameters::new(config)), + _ => return Ok(None), + }; + + Ok(Some(Arc::new( + StreamingTable::try_new(Arc::clone(table.schema()), vec![table]).unwrap(), + ) as Arc)) + }) } fn table_exist(&self, name: &str) -> bool { @@ -1582,20 +1583,26 @@ mod tests { #[derive(Debug)] struct Fixture; - - #[async_trait] impl SchemaProvider for Fixture { // InformationSchemaConfig::make_tables should use this. - async fn table_type(&self, _: &str) -> Result> { - Ok(Some(TableType::Base)) + fn table_type<'a>( + &'a self, + _: &'a str, + ) -> BoxFuture<'a, Result>> { + Box::pin(async move { Ok(Some(TableType::Base)) }) } // InformationSchemaConfig::make_tables used this before `table_type` // existed but should not, as it may be expensive. - async fn table(&self, _: &str) -> Result>> { - panic!( - "InformationSchemaConfig::make_tables called SchemaProvider::table instead of table_type" - ) + fn table<'a>( + &'a self, + _: &'a str, + ) -> BoxFuture<'a, Result>>> { + Box::pin(async move { + panic!( + "InformationSchemaConfig::make_tables called SchemaProvider::table instead of table_type" + ) + }) } fn table_names(&self) -> Vec { diff --git a/datafusion/catalog/src/listing_schema.rs b/datafusion/catalog/src/listing_schema.rs index d38fe659aaa97..b1cfc11f9018d 100644 --- a/datafusion/catalog/src/listing_schema.rs +++ b/datafusion/catalog/src/listing_schema.rs @@ -17,6 +17,7 @@ //! [`ListingSchemaProvider`]: [`SchemaProvider`] that scans ObjectStores for tables automatically +use futures::future::BoxFuture; use std::collections::HashSet; use std::path::Path; use std::sync::{Arc, Mutex}; @@ -29,7 +30,6 @@ use datafusion_common::{ }; use datafusion_expr::CreateExternalTable; -use async_trait::async_trait; use futures::TryStreamExt; use itertools::Itertools; use object_store::ObjectStore; @@ -142,8 +142,6 @@ impl ListingSchemaProvider { Ok(()) } } - -#[async_trait] impl SchemaProvider for ListingSchemaProvider { fn table_names(&self) -> Vec { self.tables @@ -154,16 +152,18 @@ impl SchemaProvider for ListingSchemaProvider { .collect() } - async fn table( - &self, - name: &str, - ) -> Result>, DataFusionError> { - Ok(self - .tables - .lock() - .expect("Can't lock tables") - .get(name) - .cloned()) + fn table<'a>( + &'a self, + name: &'a str, + ) -> BoxFuture<'a, Result>, DataFusionError>> { + Box::pin(async move { + Ok(self + .tables + .lock() + .expect("Can't lock tables") + .get(name) + .cloned()) + }) } fn register_table( diff --git a/datafusion/catalog/src/memory/schema.rs b/datafusion/catalog/src/memory/schema.rs index 46b0beb440613..5803e5c9ef880 100644 --- a/datafusion/catalog/src/memory/schema.rs +++ b/datafusion/catalog/src/memory/schema.rs @@ -18,9 +18,9 @@ //! [`MemorySchemaProvider`]: In-memory implementations of [`SchemaProvider`]. use crate::{SchemaProvider, TableProvider}; -use async_trait::async_trait; use dashmap::DashMap; use datafusion_common::{DataFusionError, exec_err}; +use futures::future::BoxFuture; use std::sync::Arc; /// Simple in-memory implementation of a schema. @@ -43,8 +43,6 @@ impl Default for MemorySchemaProvider { Self::new() } } - -#[async_trait] impl SchemaProvider for MemorySchemaProvider { fn table_names(&self) -> Vec { self.tables @@ -53,11 +51,16 @@ impl SchemaProvider for MemorySchemaProvider { .collect() } - async fn table( - &self, - name: &str, - ) -> datafusion_common::Result>, DataFusionError> { - Ok(self.tables.get(name).map(|table| Arc::clone(table.value()))) + fn table<'a>( + &'a self, + name: &'a str, + ) -> BoxFuture< + 'a, + datafusion_common::Result>, DataFusionError>, + > { + Box::pin(async move { + Ok(self.tables.get(name).map(|table| Arc::clone(table.value()))) + }) } fn register_table( diff --git a/datafusion/catalog/src/memory/table.rs b/datafusion/catalog/src/memory/table.rs index d817aa8b7788a..5b81694960dff 100644 --- a/datafusion/catalog/src/memory/table.rs +++ b/datafusion/catalog/src/memory/table.rs @@ -51,7 +51,6 @@ use datafusion_physical_plan::{ }; use datafusion_session::Session; -use async_trait::async_trait; use futures::future::BoxFuture; use log::debug; use parking_lot::Mutex; @@ -171,8 +170,6 @@ impl MemTable { MemTable::try_new(schema, data).map(|table| table.with_constraints(constraints)) } } - -#[async_trait] impl TableProvider for MemTable { fn schema(&self) -> SchemaRef { Arc::clone(&self.schema) @@ -186,22 +183,13 @@ impl TableProvider for MemTable { TableType::Base } - // Hand-written `#[async_trait]` expansion to reduce compile time. See - // - fn scan<'life0, 'life1, 'life2, 'life3, 'async_trait>( - &'life0 self, - state: &'life1 dyn Session, - projection: Option<&'life2 [usize]>, - filters: &'life3 [Expr], + fn scan<'a>( + &'a self, + state: &'a dyn Session, + projection: Option<&'a [usize]>, + filters: &'a [Expr], limit: Option, - ) -> BoxFuture<'async_trait, Result>> - where - 'life0: 'async_trait, - 'life1: 'async_trait, - 'life2: 'async_trait, - 'life3: 'async_trait, - Self: 'async_trait, - { + ) -> BoxFuture<'a, Result>> { self.scan_boxed(state, projection, filters, limit) } @@ -219,19 +207,12 @@ impl TableProvider for MemTable { /// * A plan that returns the number of rows written. /// /// [`SessionState`]: https://docs.rs/datafusion/latest/datafusion/execution/session_state/struct.SessionState.html - // Hand-written `#[async_trait]` expansion to reduce compile time. See - // - fn insert_into<'life0, 'life1, 'async_trait>( - &'life0 self, - state: &'life1 dyn Session, + fn insert_into<'a>( + &'a self, + state: &'a dyn Session, input: Arc, insert_op: InsertOp, - ) -> BoxFuture<'async_trait, Result>> - where - 'life0: 'async_trait, - 'life1: 'async_trait, - Self: 'async_trait, - { + ) -> BoxFuture<'a, Result>> { self.insert_into_boxed(state, input, insert_op) } @@ -239,34 +220,20 @@ impl TableProvider for MemTable { self.column_defaults.get(column) } - // Hand-written `#[async_trait]` expansion to reduce compile time. See - // - fn delete_from<'life0, 'life1, 'async_trait>( - &'life0 self, - state: &'life1 dyn Session, + fn delete_from<'a>( + &'a self, + state: &'a dyn Session, filters: Vec, - ) -> BoxFuture<'async_trait, Result>> - where - 'life0: 'async_trait, - 'life1: 'async_trait, - Self: 'async_trait, - { + ) -> BoxFuture<'a, Result>> { self.delete_from_boxed(state, filters) } - // Hand-written `#[async_trait]` expansion to reduce compile time. See - // - fn update<'life0, 'life1, 'async_trait>( - &'life0 self, - state: &'life1 dyn Session, + fn update<'a>( + &'a self, + state: &'a dyn Session, assignments: Vec<(String, Expr)>, filters: Vec, - ) -> BoxFuture<'async_trait, Result>> - where - 'life0: 'async_trait, - 'life1: 'async_trait, - Self: 'async_trait, - { + ) -> BoxFuture<'a, Result>> { self.update_boxed(state, assignments, filters) } } diff --git a/datafusion/catalog/src/stream.rs b/datafusion/catalog/src/stream.rs index 97e5293765afb..96ef7841f264f 100644 --- a/datafusion/catalog/src/stream.rs +++ b/datafusion/catalog/src/stream.rs @@ -39,29 +39,18 @@ use datafusion_physical_plan::stream::RecordBatchReceiverStreamBuilder; use datafusion_physical_plan::streaming::{PartitionStream, StreamingTableExec}; use datafusion_physical_plan::{DisplayAs, DisplayFormatType, ExecutionPlan}; -use async_trait::async_trait; use futures::StreamExt; use futures::future::BoxFuture; /// A [`TableProviderFactory`] for [`StreamTable`] #[derive(Debug, Default)] pub struct StreamTableFactory {} - -#[async_trait] impl TableProviderFactory for StreamTableFactory { - // Hand-written `#[async_trait]` expansion to reduce compile time. See - // - fn create<'life0, 'life1, 'life2, 'async_trait>( - &'life0 self, - state: &'life1 dyn Session, - cmd: &'life2 CreateExternalTable, - ) -> BoxFuture<'async_trait, Result>> - where - 'life0: 'async_trait, - 'life1: 'async_trait, - 'life2: 'async_trait, - Self: 'async_trait, - { + fn create<'a>( + &'a self, + state: &'a dyn Session, + cmd: &'a CreateExternalTable, + ) -> BoxFuture<'a, Result>> { self.create_boxed(state, cmd) } } @@ -335,8 +324,6 @@ impl StreamTable { Self(config) } } - -#[async_trait] impl TableProvider for StreamTable { fn schema(&self) -> SchemaRef { Arc::clone(self.0.source.schema()) @@ -350,38 +337,22 @@ impl TableProvider for StreamTable { TableType::Base } - // Hand-written `#[async_trait]` expansion to reduce compile time. See - // - fn scan<'life0, 'life1, 'life2, 'life3, 'async_trait>( - &'life0 self, - state: &'life1 dyn Session, - projection: Option<&'life2 [usize]>, - filters: &'life3 [Expr], + fn scan<'a>( + &'a self, + state: &'a dyn Session, + projection: Option<&'a [usize]>, + filters: &'a [Expr], limit: Option, - ) -> BoxFuture<'async_trait, Result>> - where - 'life0: 'async_trait, - 'life1: 'async_trait, - 'life2: 'async_trait, - 'life3: 'async_trait, - Self: 'async_trait, - { + ) -> BoxFuture<'a, Result>> { self.scan_boxed(state, projection, filters, limit) } - // Hand-written `#[async_trait]` expansion to reduce compile time. See - // - fn insert_into<'life0, 'life1, 'async_trait>( - &'life0 self, - state: &'life1 dyn Session, + fn insert_into<'a>( + &'a self, + state: &'a dyn Session, input: Arc, insert_op: InsertOp, - ) -> BoxFuture<'async_trait, Result>> - where - 'life0: 'async_trait, - 'life1: 'async_trait, - Self: 'async_trait, - { + ) -> BoxFuture<'a, Result>> { self.insert_into_boxed(state, input, insert_op) } } @@ -488,25 +459,16 @@ impl DisplayAs for StreamWrite { self.0.source.stream_write_display(t, f) } } - -#[async_trait] impl DataSink for StreamWrite { fn schema(&self) -> &SchemaRef { self.0.source.schema() } - // Hand-written `#[async_trait]` expansion to reduce compile time. See - // - fn write_all<'life0, 'life1, 'async_trait>( - &'life0 self, + fn write_all<'a>( + &'a self, data: SendableRecordBatchStream, - context: &'life1 Arc, - ) -> BoxFuture<'async_trait, Result> - where - 'life0: 'async_trait, - 'life1: 'async_trait, - Self: 'async_trait, - { + context: &'a Arc, + ) -> BoxFuture<'a, Result> { self.write_all_boxed(data, context) } } diff --git a/datafusion/catalog/src/streaming.rs b/datafusion/catalog/src/streaming.rs index c2b5d691866b5..568a2996a430d 100644 --- a/datafusion/catalog/src/streaming.rs +++ b/datafusion/catalog/src/streaming.rs @@ -21,7 +21,6 @@ use std::future::ready; use std::sync::Arc; use arrow::datatypes::SchemaRef; -use async_trait::async_trait; use datafusion_common::{DFSchema, Result, plan_err}; use datafusion_expr::physical_planning_context::PhysicalPlanningContext; use datafusion_expr::{Expr, SortExpr, TableType}; @@ -109,8 +108,6 @@ impl StreamingTable { Ok(output_partitioning.project(&projection_mapping, &eq_properties)) } } - -#[async_trait] impl TableProvider for StreamingTable { fn schema(&self) -> SchemaRef { Arc::clone(&self.schema) @@ -120,22 +117,13 @@ impl TableProvider for StreamingTable { TableType::View } - // Hand-written `#[async_trait]` expansion to reduce compile time. See - // - fn scan<'life0, 'life1, 'life2, 'life3, 'async_trait>( - &'life0 self, - state: &'life1 dyn Session, - projection: Option<&'life2 [usize]>, - filters: &'life3 [Expr], + fn scan<'a>( + &'a self, + state: &'a dyn Session, + projection: Option<&'a [usize]>, + filters: &'a [Expr], limit: Option, - ) -> BoxFuture<'async_trait, Result>> - where - 'life0: 'async_trait, - 'life1: 'async_trait, - 'life2: 'async_trait, - 'life3: 'async_trait, - Self: 'async_trait, - { + ) -> BoxFuture<'a, Result>> { self.scan_boxed(state, projection, filters, limit) } } diff --git a/datafusion/catalog/src/view.rs b/datafusion/catalog/src/view.rs index 795d2985c8c44..d64b07d68d759 100644 --- a/datafusion/catalog/src/view.rs +++ b/datafusion/catalog/src/view.rs @@ -23,7 +23,6 @@ use crate::Session; use crate::TableProvider; use arrow::datatypes::SchemaRef; -use async_trait::async_trait; use datafusion_common::Column; use datafusion_common::error::Result; use datafusion_expr::TableType; @@ -70,8 +69,6 @@ impl ViewTable { &self.logical_plan } } - -#[async_trait] impl TableProvider for ViewTable { fn get_logical_plan(&'_ self) -> Option> { Some(Cow::Borrowed(&self.logical_plan)) @@ -96,22 +93,13 @@ impl TableProvider for ViewTable { Ok(vec![TableProviderFilterPushDown::Exact; filters.len()]) } - // Hand-written `#[async_trait]` expansion to reduce compile time. See - // - fn scan<'life0, 'life1, 'life2, 'life3, 'async_trait>( - &'life0 self, - state: &'life1 dyn Session, - projection: Option<&'life2 [usize]>, - filters: &'life3 [Expr], + fn scan<'a>( + &'a self, + state: &'a dyn Session, + projection: Option<&'a [usize]>, + filters: &'a [Expr], limit: Option, - ) -> BoxFuture<'async_trait, Result>> - where - 'life0: 'async_trait, - 'life1: 'async_trait, - 'life2: 'async_trait, - 'life3: 'async_trait, - Self: 'async_trait, - { + ) -> BoxFuture<'a, Result>> { self.scan_boxed(state, projection, filters, limit) } } diff --git a/datafusion/core/Cargo.toml b/datafusion/core/Cargo.toml index 222c0ec688b78..e97d409e86504 100644 --- a/datafusion/core/Cargo.toml +++ b/datafusion/core/Cargo.toml @@ -120,7 +120,6 @@ extended_tests = [] [dependencies] arrow = { workspace = true } arrow-schema = { workspace = true, features = ["canonical_extension_types"] } -async-trait = { workspace = true } bzip2 = { workspace = true, optional = true } chrono = { workspace = true } datafusion-catalog = { workspace = true } @@ -167,7 +166,6 @@ uuid = { workspace = true, features = ["v4", "js"] } zstd = { workspace = true, optional = true } [dev-dependencies] -async-trait = { workspace = true } criterion = { workspace = true, features = ["async_tokio", "async_futures"] } ctor = { workspace = true } dashmap = "6.2.1" diff --git a/datafusion/core/src/dataframe/mod.rs b/datafusion/core/src/dataframe/mod.rs index ed3dc5ea838b9..a1675d0b9c38f 100644 --- a/datafusion/core/src/dataframe/mod.rs +++ b/datafusion/core/src/dataframe/mod.rs @@ -67,7 +67,6 @@ use datafusion_functions_aggregate::expr_fn::{ avg, count, max, median, min, stddev, sum, }; -use async_trait::async_trait; use datafusion_catalog::Session; use datafusion_expr::extension_types::DFArrayFormatterFactory; use futures::future::BoxFuture; @@ -2709,8 +2708,6 @@ struct DataFrameTableProvider { plan: LogicalPlan, table_type: TableType, } - -#[async_trait] impl TableProvider for DataFrameTableProvider { fn get_logical_plan(&self) -> Option> { Some(Cow::Borrowed(&self.plan)) @@ -2732,22 +2729,13 @@ impl TableProvider for DataFrameTableProvider { self.table_type } - // Hand-written `#[async_trait]` expansion to reduce compile time. See - // - fn scan<'life0, 'life1, 'life2, 'life3, 'async_trait>( - &'life0 self, - state: &'life1 dyn Session, - projection: Option<&'life2 [usize]>, - filters: &'life3 [Expr], + fn scan<'a>( + &'a self, + state: &'a dyn Session, + projection: Option<&'a [usize]>, + filters: &'a [Expr], limit: Option, - ) -> BoxFuture<'async_trait, Result>> - where - 'life0: 'async_trait, - 'life1: 'async_trait, - 'life2: 'async_trait, - 'life3: 'async_trait, - Self: 'async_trait, - { + ) -> BoxFuture<'a, Result>> { self.scan_boxed(state, projection, filters, limit) } } diff --git a/datafusion/core/src/datasource/dynamic_file.rs b/datafusion/core/src/datasource/dynamic_file.rs index 0212222cdd1cb..eff0b553033ad 100644 --- a/datafusion/core/src/datasource/dynamic_file.rs +++ b/datafusion/core/src/datasource/dynamic_file.rs @@ -30,7 +30,6 @@ use datafusion_catalog::UrlTableFactory; use datafusion_common::plan_datafusion_err; use datafusion_session::SessionStore; -use async_trait::async_trait; use futures::future::BoxFuture; /// [DynamicListTableFactory] is a factory that can create a [ListingTable] from the given url. @@ -51,20 +50,11 @@ impl DynamicListTableFactory { &self.session_store } } - -#[async_trait] impl UrlTableFactory for DynamicListTableFactory { - // Hand-written `#[async_trait]` expansion to reduce compile time. See - // - fn try_new<'life0, 'life1, 'async_trait>( - &'life0 self, - url: &'life1 str, - ) -> BoxFuture<'async_trait, Result>>> - where - 'life0: 'async_trait, - 'life1: 'async_trait, - Self: 'async_trait, - { + fn try_new<'a>( + &'a self, + url: &'a str, + ) -> BoxFuture<'a, Result>>> { self.try_new_boxed(url) } } diff --git a/datafusion/core/src/datasource/file_format/csv.rs b/datafusion/core/src/datasource/file_format/csv.rs index 2fb64fd6486e6..4efa2fd179e1c 100644 --- a/datafusion/core/src/datasource/file_format/csv.rs +++ b/datafusion/core/src/datasource/file_format/csv.rs @@ -21,7 +21,9 @@ pub use datafusion_datasource_csv::file_format::*; #[cfg(test)] mod tests { use std::fmt::{self, Display}; + use std::future::Future; use std::ops::Range; + use std::pin::Pin; use std::sync::{Arc, Mutex}; use super::*; @@ -54,7 +56,6 @@ mod tests { use arrow::compute::concat_batches; use arrow::csv::ReaderBuilder; use arrow::util::pretty::pretty_format_batches; - use async_trait::async_trait; use bytes::Bytes; use chrono::DateTime; use datafusion_common::parsers::CompressionTypeVariant; @@ -86,65 +87,110 @@ mod tests { write!(f, "VariableStream") } } - - #[async_trait] impl ObjectStore for VariableStream { - async fn put_opts( - &self, - _location: &Path, + fn put_opts<'life0, 'life1, 'async_trait>( + &'life0 self, + _location: &'life1 Path, _payload: PutPayload, _opts: PutOptions, - ) -> object_store::Result { - unimplemented!() + ) -> Pin< + Box< + dyn Future> + + Send + + 'async_trait, + >, + > + where + 'life0: 'async_trait, + 'life1: 'async_trait, + Self: 'async_trait, + { + Box::pin(async move { unimplemented!() }) } - async fn put_multipart_opts( - &self, - _location: &Path, + fn put_multipart_opts<'life0, 'life1, 'async_trait>( + &'life0 self, + _location: &'life1 Path, _opts: PutMultipartOptions, - ) -> object_store::Result> { - unimplemented!() + ) -> Pin< + Box< + dyn Future>> + + Send + + 'async_trait, + >, + > + where + 'life0: 'async_trait, + 'life1: 'async_trait, + Self: 'async_trait, + { + Box::pin(async move { unimplemented!() }) } - async fn get_opts( - &self, - location: &Path, + fn get_opts<'life0, 'life1, 'async_trait>( + &'life0 self, + location: &'life1 Path, _opts: GetOptions, - ) -> object_store::Result { - let bytes = self.bytes_to_repeat.clone(); - let len = bytes.len() as u64; - let range = 0..len * self.max_iterations; - let arc = self.iterations_detected.clone(); - #[expect(clippy::result_large_err)] - // closure only ever returns Ok; Err type is never constructed - let stream = futures::stream::repeat_with(move || { - let arc_inner = arc.clone(); - *arc_inner.lock().unwrap() += 1; - Ok(bytes.clone()) - }) - .take(self.max_iterations as usize) - .boxed(); - - Ok(GetResult { - payload: GetResultPayload::Stream(stream), - meta: ObjectMeta { - location: location.clone(), - last_modified: Default::default(), - size: range.end, - e_tag: None, - version: None, - }, - range: Default::default(), - attributes: Attributes::default(), + ) -> Pin< + Box< + dyn Future> + + Send + + 'async_trait, + >, + > + where + 'life0: 'async_trait, + 'life1: 'async_trait, + Self: 'async_trait, + { + Box::pin(async move { + let bytes = self.bytes_to_repeat.clone(); + let len = bytes.len() as u64; + let range = 0..len * self.max_iterations; + let arc = self.iterations_detected.clone(); + #[expect(clippy::result_large_err)] + // closure only ever returns Ok; Err type is never constructed + let stream = futures::stream::repeat_with(move || { + let arc_inner = arc.clone(); + *arc_inner.lock().unwrap() += 1; + Ok(bytes.clone()) + }) + .take(self.max_iterations as usize) + .boxed(); + + Ok(GetResult { + payload: GetResultPayload::Stream(stream), + meta: ObjectMeta { + location: location.clone(), + last_modified: Default::default(), + size: range.end, + e_tag: None, + version: None, + }, + range: Default::default(), + attributes: Attributes::default(), + }) }) } - async fn get_ranges( - &self, - _location: &Path, - _ranges: &[Range], - ) -> object_store::Result> { - unimplemented!() + fn get_ranges<'life0, 'life1, 'life2, 'async_trait>( + &'life0 self, + _location: &'life1 Path, + _ranges: &'life2 [Range], + ) -> Pin< + Box< + dyn Future>> + + Send + + 'async_trait, + >, + > + where + 'life0: 'async_trait, + 'life1: 'async_trait, + 'life2: 'async_trait, + Self: 'async_trait, + { + Box::pin(async move { unimplemented!() }) } fn list( @@ -154,20 +200,37 @@ mod tests { unimplemented!() } - async fn list_with_delimiter( - &self, - _prefix: Option<&Path>, - ) -> object_store::Result { - unimplemented!() + fn list_with_delimiter<'life0, 'life1, 'async_trait>( + &'life0 self, + _prefix: Option<&'life1 Path>, + ) -> Pin< + Box< + dyn Future> + + Send + + 'async_trait, + >, + > + where + 'life0: 'async_trait, + 'life1: 'async_trait, + Self: 'async_trait, + { + Box::pin(async move { unimplemented!() }) } - async fn copy_opts( - &self, - _from: &Path, - _to: &Path, + fn copy_opts<'life0, 'life1, 'life2, 'async_trait>( + &'life0 self, + _from: &'life1 Path, + _to: &'life2 Path, _options: object_store::CopyOptions, - ) -> object_store::Result<()> { - unimplemented!() + ) -> Pin> + Send + 'async_trait>> + where + 'life0: 'async_trait, + 'life1: 'async_trait, + 'life2: 'async_trait, + Self: 'async_trait, + { + Box::pin(async move { unimplemented!() }) } fn delete_stream( diff --git a/datafusion/core/src/datasource/file_format/options.rs b/datafusion/core/src/datasource/file_format/options.rs index 8250ef45be2f7..0f3ce944ebc56 100644 --- a/datafusion/core/src/datasource/file_format/options.rs +++ b/datafusion/core/src/datasource/file_format/options.rs @@ -42,7 +42,6 @@ use datafusion_common::{ DEFAULT_JSON_EXTENSION, DEFAULT_PARQUET_EXTENSION, }; -use async_trait::async_trait; use datafusion_datasource_json::file_format::JsonFormat; use datafusion_expr::SortExpr; use futures::future::BoxFuture; @@ -579,8 +578,6 @@ impl<'a> JsonReadOptions<'a> { self } } - -#[async_trait] /// [`ReadOptions`] is implemented by Options like [`CsvReadOptions`] that control the reading of respective files/sources. pub trait ReadOptions<'a> { /// Helper to convert these user facing options to `ListingTable` options @@ -591,12 +588,12 @@ pub trait ReadOptions<'a> { ) -> ListingOptions; /// Infer and resolve the schema from the files/sources provided. - async fn get_resolved_schema( - &self, - config: &SessionConfig, + fn get_resolved_schema<'s>( + &'s self, + config: &'s SessionConfig, state: SessionState, table_path: ListingTableUrl, - ) -> Result; + ) -> BoxFuture<'s, Result>; /// Returns whether the read schema was inferred or specified. fn schema_source(&self) -> SchemaSource { @@ -604,20 +601,13 @@ pub trait ReadOptions<'a> { } /// helper function to reduce repetitive code. Infers the schema from sources if not provided. Infinite data sources not supported through this function. - // Hand-written `#[async_trait]` expansion to reduce compile time. See - // - fn _get_resolved_schema<'life0, 'async_trait>( - &'a self, - config: &'life0 SessionConfig, + fn _get_resolved_schema<'s>( + &'s self, + config: &'s SessionConfig, state: SessionState, table_path: ListingTableUrl, - schema: Option<&'a Schema>, - ) -> BoxFuture<'async_trait, Result> - where - 'a: 'async_trait, - 'life0: 'async_trait, - Self: 'async_trait, - { + schema: Option<&'s Schema>, + ) -> BoxFuture<'s, Result> { if let Some(s) = schema { return Box::pin(ready(Ok(Arc::new(s.to_owned())))); } @@ -636,8 +626,6 @@ fn infer_schema_boxed( ) -> BoxFuture<'static, Result> { Box::pin(async move { listing_options.infer_schema(&state, &table_path).await }) } - -#[async_trait] impl ReadOptions<'_> for CsvReadOptions<'_> { fn to_listing_options( &self, @@ -664,19 +652,12 @@ impl ReadOptions<'_> for CsvReadOptions<'_> { .with_file_sort_order(self.file_sort_order.clone()) } - // Hand-written `#[async_trait]` expansion to reduce compile time. See - // - fn get_resolved_schema<'life0, 'life1, 'async_trait>( - &'life0 self, - config: &'life1 SessionConfig, + fn get_resolved_schema<'s>( + &'s self, + config: &'s SessionConfig, state: SessionState, table_path: ListingTableUrl, - ) -> BoxFuture<'async_trait, Result> - where - 'life0: 'async_trait, - 'life1: 'async_trait, - Self: 'async_trait, - { + ) -> BoxFuture<'s, Result> { self._get_resolved_schema(config, state, table_path, self.schema) } @@ -686,7 +667,6 @@ impl ReadOptions<'_> for CsvReadOptions<'_> { } #[cfg(feature = "parquet")] -#[async_trait] impl ReadOptions<'_> for ParquetReadOptions<'_> { fn to_listing_options( &self, @@ -717,19 +697,12 @@ impl ReadOptions<'_> for ParquetReadOptions<'_> { .with_file_sort_order(self.file_sort_order.clone()) } - // Hand-written `#[async_trait]` expansion to reduce compile time. See - // - fn get_resolved_schema<'life0, 'life1, 'async_trait>( - &'life0 self, - config: &'life1 SessionConfig, + fn get_resolved_schema<'s>( + &'s self, + config: &'s SessionConfig, state: SessionState, table_path: ListingTableUrl, - ) -> BoxFuture<'async_trait, Result> - where - 'life0: 'async_trait, - 'life1: 'async_trait, - Self: 'async_trait, - { + ) -> BoxFuture<'s, Result> { self._get_resolved_schema(config, state, table_path, self.schema) } @@ -737,8 +710,6 @@ impl ReadOptions<'_> for ParquetReadOptions<'_> { schema_source_from_option(self.schema) } } - -#[async_trait] impl ReadOptions<'_> for JsonReadOptions<'_> { fn to_listing_options( &self, @@ -757,19 +728,12 @@ impl ReadOptions<'_> for JsonReadOptions<'_> { .with_file_sort_order(self.file_sort_order.clone()) } - // Hand-written `#[async_trait]` expansion to reduce compile time. See - // - fn get_resolved_schema<'life0, 'life1, 'async_trait>( - &'life0 self, - config: &'life1 SessionConfig, + fn get_resolved_schema<'s>( + &'s self, + config: &'s SessionConfig, state: SessionState, table_path: ListingTableUrl, - ) -> BoxFuture<'async_trait, Result> - where - 'life0: 'async_trait, - 'life1: 'async_trait, - Self: 'async_trait, - { + ) -> BoxFuture<'s, Result> { self._get_resolved_schema(config, state, table_path, self.schema) } @@ -779,7 +743,6 @@ impl ReadOptions<'_> for JsonReadOptions<'_> { } #[cfg(feature = "avro")] -#[async_trait] impl ReadOptions<'_> for AvroReadOptions<'_> { fn to_listing_options( &self, @@ -793,19 +756,12 @@ impl ReadOptions<'_> for AvroReadOptions<'_> { .with_table_partition_cols(self.table_partition_cols.clone()) } - // Hand-written `#[async_trait]` expansion to reduce compile time. See - // - fn get_resolved_schema<'life0, 'life1, 'async_trait>( - &'life0 self, - config: &'life1 SessionConfig, + fn get_resolved_schema<'s>( + &'s self, + config: &'s SessionConfig, state: SessionState, table_path: ListingTableUrl, - ) -> BoxFuture<'async_trait, Result> - where - 'life0: 'async_trait, - 'life1: 'async_trait, - Self: 'async_trait, - { + ) -> BoxFuture<'s, Result> { self._get_resolved_schema(config, state, table_path, self.schema) } @@ -813,8 +769,6 @@ impl ReadOptions<'_> for AvroReadOptions<'_> { schema_source_from_option(self.schema) } } - -#[async_trait] impl ReadOptions<'_> for ArrowReadOptions<'_> { fn to_listing_options( &self, @@ -828,19 +782,12 @@ impl ReadOptions<'_> for ArrowReadOptions<'_> { .with_table_partition_cols(self.table_partition_cols.clone()) } - // Hand-written `#[async_trait]` expansion to reduce compile time. See - // - fn get_resolved_schema<'life0, 'life1, 'async_trait>( - &'life0 self, - config: &'life1 SessionConfig, + fn get_resolved_schema<'s>( + &'s self, + config: &'s SessionConfig, state: SessionState, table_path: ListingTableUrl, - ) -> BoxFuture<'async_trait, Result> - where - 'life0: 'async_trait, - 'life1: 'async_trait, - Self: 'async_trait, - { + ) -> BoxFuture<'s, Result> { self._get_resolved_schema(config, state, table_path, self.schema) } diff --git a/datafusion/core/src/datasource/file_format/parquet.rs b/datafusion/core/src/datasource/file_format/parquet.rs index bfcfb74848861..324e7b1567bc7 100644 --- a/datafusion/core/src/datasource/file_format/parquet.rs +++ b/datafusion/core/src/datasource/file_format/parquet.rs @@ -106,6 +106,10 @@ pub(crate) mod test_util { #[cfg(test)] mod tests { + use std::future::Future; + + use std::pin::Pin; + use std::fmt::{self, Display, Formatter}; use std::sync::Arc; use std::sync::atomic::{AtomicUsize, Ordering}; @@ -151,7 +155,6 @@ mod tests { types::Int32Type, }; use arrow::datatypes::{DataType, Field}; - use async_trait::async_trait; use datafusion_datasource::file_groups::FileGroup; use datafusion_datasource_parquet::metadata::DFParquetMetadata; use futures::StreamExt; @@ -303,33 +306,66 @@ mod tests { self.clone() } } - - #[async_trait] impl ObjectStore for RequestCountingObjectStore { - async fn put_opts( - &self, - _location: &Path, + fn put_opts<'life0, 'life1, 'async_trait>( + &'life0 self, + _location: &'life1 Path, _payload: PutPayload, _opts: PutOptions, - ) -> object_store::Result { - unimplemented!() + ) -> Pin< + Box< + dyn Future> + + Send + + 'async_trait, + >, + > + where + 'life0: 'async_trait, + 'life1: 'async_trait, + Self: 'async_trait, + { + Box::pin(async move { unimplemented!() }) } - async fn put_multipart_opts( - &self, - _location: &Path, + fn put_multipart_opts<'life0, 'life1, 'async_trait>( + &'life0 self, + _location: &'life1 Path, _opts: PutMultipartOptions, - ) -> object_store::Result> { - unimplemented!() + ) -> Pin< + Box< + dyn Future>> + + Send + + 'async_trait, + >, + > + where + 'life0: 'async_trait, + 'life1: 'async_trait, + Self: 'async_trait, + { + Box::pin(async move { unimplemented!() }) } - async fn get_opts( - &self, - location: &Path, + fn get_opts<'life0, 'life1, 'async_trait>( + &'life0 self, + location: &'life1 Path, options: GetOptions, - ) -> object_store::Result { - self.request_count.fetch_add(1, Ordering::SeqCst); - self.inner.get_opts(location, options).await + ) -> Pin< + Box< + dyn Future> + + Send + + 'async_trait, + >, + > + where + 'life0: 'async_trait, + 'life1: 'async_trait, + Self: 'async_trait, + { + Box::pin(async move { + self.request_count.fetch_add(1, Ordering::SeqCst); + self.inner.get_opts(location, options).await + }) } fn delete_stream( @@ -346,20 +382,37 @@ mod tests { unimplemented!() } - async fn list_with_delimiter( - &self, - _prefix: Option<&Path>, - ) -> object_store::Result { - unimplemented!() + fn list_with_delimiter<'life0, 'life1, 'async_trait>( + &'life0 self, + _prefix: Option<&'life1 Path>, + ) -> Pin< + Box< + dyn Future> + + Send + + 'async_trait, + >, + > + where + 'life0: 'async_trait, + 'life1: 'async_trait, + Self: 'async_trait, + { + Box::pin(async move { unimplemented!() }) } - async fn copy_opts( - &self, - _from: &Path, - _to: &Path, + fn copy_opts<'life0, 'life1, 'life2, 'async_trait>( + &'life0 self, + _from: &'life1 Path, + _to: &'life2 Path, _options: CopyOptions, - ) -> object_store::Result<()> { - unimplemented!() + ) -> Pin> + Send + 'async_trait>> + where + 'life0: 'async_trait, + 'life1: 'async_trait, + 'life2: 'async_trait, + Self: 'async_trait, + { + Box::pin(async move { unimplemented!() }) } } diff --git a/datafusion/core/src/datasource/listing/table.rs b/datafusion/core/src/datasource/listing/table.rs index 982766dc88519..f1d2bbb873d8e 100644 --- a/datafusion/core/src/datasource/listing/table.rs +++ b/datafusion/core/src/datasource/listing/table.rs @@ -16,7 +16,6 @@ // under the License. use crate::execution::SessionState; -use async_trait::async_trait; use datafusion_catalog_listing::{ListingOptions, ListingTableConfig}; use datafusion_common::{config_datafusion_err, internal_datafusion_err}; use datafusion_session::Session; @@ -29,48 +28,33 @@ use std::collections::HashMap; /// This trait exists because the following inference methods only /// work for [`SessionState`] implementations of [`Session`]. /// See [`ListingTableConfig`] for the remaining inference methods. -#[async_trait] pub trait ListingTableConfigExt { /// Infer `ListingOptions` based on `table_path` and file suffix. /// /// The format is inferred based on the first `table_path`. - async fn infer_options( + fn infer_options( self, state: &dyn Session, - ) -> datafusion_common::Result; + ) -> BoxFuture<'_, datafusion_common::Result>; /// Convenience method to call both [`Self::infer_options`] and [`ListingTableConfig::infer_schema`] - async fn infer( + fn infer( self, state: &dyn Session, - ) -> datafusion_common::Result; + ) -> BoxFuture<'_, datafusion_common::Result>; } - -#[async_trait] impl ListingTableConfigExt for ListingTableConfig { - // Hand-written `#[async_trait]` expansion to reduce compile time. See - // - fn infer_options<'life0, 'async_trait>( + fn infer_options( self, - state: &'life0 dyn Session, - ) -> BoxFuture<'async_trait, datafusion_common::Result> - where - 'life0: 'async_trait, - Self: 'async_trait, - { + state: &dyn Session, + ) -> BoxFuture<'_, datafusion_common::Result> { infer_options_boxed(self, state) } - // Hand-written `#[async_trait]` expansion to reduce compile time. See - // - fn infer<'life0, 'async_trait>( + fn infer( self, - state: &'life0 dyn Session, - ) -> BoxFuture<'async_trait, datafusion_common::Result> - where - 'life0: 'async_trait, - Self: 'async_trait, - { + state: &dyn Session, + ) -> BoxFuture<'_, datafusion_common::Result> { infer_boxed(self, state) } } diff --git a/datafusion/core/src/datasource/listing_table_factory.rs b/datafusion/core/src/datasource/listing_table_factory.rs index 1e597e38fb5b1..55531cf9d630a 100644 --- a/datafusion/core/src/datasource/listing_table_factory.rs +++ b/datafusion/core/src/datasource/listing_table_factory.rs @@ -34,7 +34,6 @@ use datafusion_common::{ }; use datafusion_expr::CreateExternalTable; -use async_trait::async_trait; use datafusion_catalog::Session; use futures::future::BoxFuture; @@ -48,22 +47,12 @@ impl ListingTableFactory { Self::default() } } - -#[async_trait] impl TableProviderFactory for ListingTableFactory { - // Hand-written `#[async_trait]` expansion to reduce compile time. See - // - fn create<'life0, 'life1, 'life2, 'async_trait>( - &'life0 self, - state: &'life1 dyn Session, - cmd: &'life2 CreateExternalTable, - ) -> BoxFuture<'async_trait, Result>> - where - 'life0: 'async_trait, - 'life1: 'async_trait, - 'life2: 'async_trait, - Self: 'async_trait, - { + fn create<'a>( + &'a self, + state: &'a dyn Session, + cmd: &'a CreateExternalTable, + ) -> BoxFuture<'a, Result>> { self.create_boxed(state, cmd) } } @@ -872,8 +861,6 @@ mod tests { // A mock Session that is NOT SessionState #[derive(Debug)] struct MockSession; - - #[async_trait] impl Session for MockSession { fn session_id(&self) -> &str { "mock_session" @@ -884,12 +871,13 @@ mod tests { fn catalog_list(&self) -> Arc { Arc::new(EmptyCatalogProviderList) } - async fn create_physical_plan( - &self, + fn create_physical_plan<'a>( + &'a self, _logical_plan: &datafusion_expr::LogicalPlan, - ) -> Result> { - unimplemented!() + ) -> BoxFuture<'a, Result>> { + Box::pin(async move { unimplemented!() }) } + fn create_physical_expr( &self, _expr: datafusion_expr::Expr, diff --git a/datafusion/core/src/datasource/provider.rs b/datafusion/core/src/datasource/provider.rs index e574042813a7b..0f3f1f5bcaed1 100644 --- a/datafusion/core/src/datasource/provider.rs +++ b/datafusion/core/src/datasource/provider.rs @@ -19,7 +19,6 @@ use std::sync::Arc; -use async_trait::async_trait; use datafusion_catalog::Session; use datafusion_expr::CreateExternalTable; pub use datafusion_expr::{TableProviderFilterPushDown, TableType}; @@ -46,22 +45,12 @@ impl DefaultTableFactory { Self::default() } } - -#[async_trait] impl TableProviderFactory for DefaultTableFactory { - // Hand-written `#[async_trait]` expansion to reduce compile time. See - // - fn create<'life0, 'life1, 'life2, 'async_trait>( - &'life0 self, - state: &'life1 dyn Session, - cmd: &'life2 CreateExternalTable, - ) -> BoxFuture<'async_trait, Result>> - where - 'life0: 'async_trait, - 'life1: 'async_trait, - 'life2: 'async_trait, - Self: 'async_trait, - { + fn create<'a>( + &'a self, + state: &'a dyn Session, + cmd: &'a CreateExternalTable, + ) -> BoxFuture<'a, Result>> { self.create_boxed(state, cmd) } } diff --git a/datafusion/core/src/execution/context/mod.rs b/datafusion/core/src/execution/context/mod.rs index ff1ad25811440..78a53efc53aae 100644 --- a/datafusion/core/src/execution/context/mod.rs +++ b/datafusion/core/src/execution/context/mod.rs @@ -17,6 +17,7 @@ //! [`SessionContext`] API for registering data sources and executing queries +use futures::future::BoxFuture; use std::collections::HashSet; use std::fmt::Debug; use std::sync::{Arc, Weak}; @@ -101,7 +102,6 @@ use datafusion_optimizer::{Analyzer, OptimizerContext}; use datafusion_optimizer::{AnalyzerRule, OptimizerRule}; use datafusion_session::SessionStore; -use async_trait::async_trait; use chrono::{DateTime, Utc}; use object_store::ObjectStore; use parking_lot::RwLock; @@ -2216,15 +2216,13 @@ pub use datafusion_session::{QueryPlanner, UnsupportedQueryPlanner}; /// return pc.multiply(km_data, conversation_rate_multiplier) /// ' /// ``` - -#[async_trait] pub trait FunctionFactory: Debug + Sync + Send { /// Creates a new dynamic function from the SQL in the [CreateFunction] statement - async fn create( - &self, - state: &SessionState, + fn create<'a>( + &'a self, + state: &'a SessionState, statement: CreateFunction, - ) -> Result; + ) -> BoxFuture<'a, Result>; } /// The result of processing a [`CreateFunction`] statement with [`FunctionFactory`]. @@ -2376,7 +2374,6 @@ mod tests { use crate::catalog::SchemaProvider; use crate::execution::session_state::SessionStateBuilder; use crate::physical_planner::PhysicalPlanner; - use async_trait::async_trait; use datafusion_expr::planner::TypePlanner; use datafusion_session::Session; use sqlparser::ast; @@ -2819,15 +2816,13 @@ mod tests { } struct MyPhysicalPlanner {} - - #[async_trait] impl PhysicalPlanner for MyPhysicalPlanner { - async fn create_physical_plan( - &self, - _logical_plan: &LogicalPlan, - _session_state: &dyn Session, - ) -> Result> { - not_impl_err!("query not supported") + fn create_physical_plan<'a>( + &'a self, + _logical_plan: &'a LogicalPlan, + _session_state: &'a dyn Session, + ) -> BoxFuture<'a, Result>> { + Box::pin(async move { not_impl_err!("query not supported") }) } fn create_physical_expr( @@ -2843,18 +2838,18 @@ mod tests { #[derive(Debug)] struct MyQueryPlanner {} - - #[async_trait] impl QueryPlanner for MyQueryPlanner { - async fn create_physical_plan( - &self, - logical_plan: &LogicalPlan, - session_state: &dyn Session, - ) -> Result> { - let physical_planner = MyPhysicalPlanner {}; - physical_planner - .create_physical_plan(logical_plan, session_state) - .await + fn create_physical_plan<'a>( + &'a self, + logical_plan: &'a LogicalPlan, + session_state: &'a dyn Session, + ) -> BoxFuture<'a, Result>> { + Box::pin(async move { + let physical_planner = MyPhysicalPlanner {}; + physical_planner + .create_physical_plan(logical_plan, session_state) + .await + }) } } diff --git a/datafusion/core/src/execution/session_state.rs b/datafusion/core/src/execution/session_state.rs index aa8ba4c3b733b..f6f23fba532bc 100644 --- a/datafusion/core/src/execution/session_state.rs +++ b/datafusion/core/src/execution/session_state.rs @@ -83,7 +83,6 @@ use datafusion_sql::{ planner::{ContextProvider, ParserOptions, PlannerContext, SqlToRel}, }; -use async_trait::async_trait; use chrono::{DateTime, Utc}; use futures::future::BoxFuture; use itertools::Itertools; @@ -269,8 +268,6 @@ impl Debug for SessionState { .finish() } } - -#[async_trait] impl Session for SessionState { fn session_id(&self) -> &str { self.session_id() @@ -303,17 +300,10 @@ impl Session for SessionState { SessionState::statistics_registry(self) } - // Hand-written `#[async_trait]` expansion to reduce compile time. See - // - fn create_physical_plan<'life0, 'life1, 'async_trait>( - &'life0 self, - logical_plan: &'life1 LogicalPlan, - ) -> BoxFuture<'async_trait, datafusion_common::Result>> - where - 'life0: 'async_trait, - 'life1: 'async_trait, - Self: 'async_trait, - { + fn create_physical_plan<'a>( + &'a self, + logical_plan: &'a LogicalPlan, + ) -> BoxFuture<'a, datafusion_common::Result>> { self.create_physical_plan_boxed(logical_plan) } @@ -2413,23 +2403,13 @@ impl From<&SessionState> for TaskContext { /// The query planner used if no user defined planner is provided #[derive(Debug)] struct DefaultQueryPlanner {} - -#[async_trait] impl QueryPlanner for DefaultQueryPlanner { /// Given a `LogicalPlan`, create an [`ExecutionPlan`] suitable for execution - // Hand-written `#[async_trait]` expansion to reduce compile time. See - // - fn create_physical_plan<'life0, 'life1, 'life2, 'async_trait>( - &'life0 self, - logical_plan: &'life1 LogicalPlan, - session_state: &'life2 dyn Session, - ) -> BoxFuture<'async_trait, datafusion_common::Result>> - where - 'life0: 'async_trait, - 'life1: 'async_trait, - 'life2: 'async_trait, - Self: 'async_trait, - { + fn create_physical_plan<'a>( + &'a self, + logical_plan: &'a LogicalPlan, + session_state: &'a dyn Session, + ) -> BoxFuture<'a, datafusion_common::Result>> { self.create_physical_plan_boxed(logical_plan, session_state) } } diff --git a/datafusion/core/src/physical_planner.rs b/datafusion/core/src/physical_planner.rs index 882956a63c114..12cf51577409c 100644 --- a/datafusion/core/src/physical_planner.rs +++ b/datafusion/core/src/physical_planner.rs @@ -111,7 +111,6 @@ use datafusion_physical_plan::scalar_subquery::{ScalarSubqueryExec, ScalarSubque use datafusion_physical_plan::unnest::ListUnnest; use datafusion_session::{PhysicalOptimizerContext, PhysicalOptimizerRule, Session}; -use async_trait::async_trait; use datafusion_physical_plan::async_func::{AsyncFuncExec, AsyncMapper}; use futures::future::BoxFuture; use futures::{StreamExt, TryStreamExt}; @@ -151,23 +150,13 @@ impl PhysicalOptimizerContext for SessionOptimizerContext<'_> { pub struct DefaultPhysicalPlanner { extension_planners: Vec>, } - -#[async_trait] impl PhysicalPlanner for DefaultPhysicalPlanner { /// Create a physical plan from a logical plan - // Hand-written `#[async_trait]` expansion to reduce compile time. See - // - fn create_physical_plan<'life0, 'life1, 'life2, 'async_trait>( - &'life0 self, - logical_plan: &'life1 LogicalPlan, - session_state: &'life2 dyn Session, - ) -> BoxFuture<'async_trait, Result>> - where - 'life0: 'async_trait, - 'life1: 'async_trait, - 'life2: 'async_trait, - Self: 'async_trait, - { + fn create_physical_plan<'a>( + &'a self, + logical_plan: &'a LogicalPlan, + session_state: &'a dyn Session, + ) -> BoxFuture<'a, Result>> { self.create_physical_plan_boxed(logical_plan, session_state) } @@ -3444,18 +3433,18 @@ mod tests { struct TestQueryPlanner { invoked: Arc, } - - #[async_trait] impl QueryPlanner for TestQueryPlanner { - async fn create_physical_plan( - &self, - logical_plan: &LogicalPlan, - session: &dyn Session, - ) -> Result> { - self.invoked.store(true, AtomicOrdering::Relaxed); - DefaultPhysicalPlanner::default() - .create_physical_plan(logical_plan, session) - .await + fn create_physical_plan<'a>( + &'a self, + logical_plan: &'a LogicalPlan, + session: &'a dyn Session, + ) -> BoxFuture<'a, Result>> { + Box::pin(async move { + self.invoked.store(true, AtomicOrdering::Relaxed); + DefaultPhysicalPlanner::default() + .create_physical_plan(logical_plan, session) + .await + }) } } @@ -3463,8 +3452,6 @@ mod tests { inner: SessionState, query_planner: Arc, } - - #[async_trait] impl Session for TestSession { fn session_id(&self) -> &str { self.inner.session_id() @@ -3497,14 +3484,16 @@ mod tests { self.inner.statistics_registry() } - async fn create_physical_plan( - &self, - logical_plan: &LogicalPlan, - ) -> Result> { - let logical_plan = self.optimize(logical_plan)?; - self.query_planner() - .create_physical_plan(&logical_plan, self) - .await + fn create_physical_plan<'a>( + &'a self, + logical_plan: &'a LogicalPlan, + ) -> BoxFuture<'a, Result>> { + Box::pin(async move { + let logical_plan = self.optimize(logical_plan)?; + self.query_planner() + .create_physical_plan(&logical_plan, self) + .await + }) } fn create_physical_expr( @@ -3576,8 +3565,6 @@ mod tests { schema: SchemaRef, captured: Mutex>, } - - #[async_trait] impl TableProvider for CaptureMergeProvider { fn schema(&self) -> SchemaRef { Arc::clone(&self.schema) @@ -3587,28 +3574,33 @@ mod tests { TableType::Base } - async fn scan( - &self, - _state: &dyn Session, - _projection: Option<&[usize]>, - _filters: &[Expr], + fn scan<'a>( + &'a self, + _state: &'a dyn Session, + _projection: Option<&'a [usize]>, + _filters: &'a [Expr], _limit: Option, - ) -> Result> { - Ok(Arc::new(EmptyExec::new(Arc::clone(&self.schema)))) + ) -> BoxFuture<'a, Result>> { + Box::pin(async move { + Ok(Arc::new(EmptyExec::new(Arc::clone(&self.schema))) + as Arc) + }) } - async fn merge_into( - &self, - state: &dyn Session, + fn merge_into<'a>( + &'a self, + state: &'a dyn Session, source: Arc, merge_schema: DFSchemaRef, on: Expr, clauses: Vec, - ) -> Result> { - let physical_on = state.create_physical_expr(on, &merge_schema)?; - *self.captured.lock().await = - Some((merge_schema, format!("{physical_on:?}"), clauses.len())); - Ok(source) + ) -> BoxFuture<'a, Result>> { + Box::pin(async move { + let physical_on = state.create_physical_expr(on, &merge_schema)?; + *self.captured.lock().await = + Some((merge_schema, format!("{physical_on:?}"), clauses.len())); + Ok(source) + }) } } @@ -4927,20 +4919,18 @@ mod tests { } struct ErrorExtensionPlanner {} - - #[async_trait] impl ExtensionPlanner for ErrorExtensionPlanner { /// Create a physical plan for an extension node - async fn plan_extension( - &self, - _planner: &dyn PhysicalPlanner, - _node: &dyn UserDefinedLogicalNode, - _logical_inputs: &[&LogicalPlan], - _physical_inputs: &[Arc], - _session_state: &dyn Session, - _planning_ctx: &PhysicalPlanningContext, - ) -> Result>> { - internal_err!("BOOM") + fn plan_extension<'a>( + &'a self, + _planner: &'a dyn PhysicalPlanner, + _node: &'a dyn UserDefinedLogicalNode, + _logical_inputs: &'a [&'a LogicalPlan], + _physical_inputs: &'a [Arc], + _session_state: &'a dyn Session, + _planning_ctx: &'a PhysicalPlanningContext, + ) -> BoxFuture<'a, Result>>> { + Box::pin(async move { internal_err!("BOOM") }) } } /// An example extension node that doesn't do anything @@ -5107,51 +5097,53 @@ mod tests { } struct ExpressionExtensionPlanner; - - #[async_trait] impl ExtensionPlanner for ExpressionExtensionPlanner { - async fn plan_extension( - &self, - planner: &dyn PhysicalPlanner, - node: &dyn UserDefinedLogicalNode, - _logical_inputs: &[&LogicalPlan], - _physical_inputs: &[Arc], - session_state: &dyn Session, - planning_ctx: &PhysicalPlanningContext, - ) -> Result>> { - for expr in node.expressions() { - planner.create_physical_expr( - &expr, - node.schema(), - session_state, - planning_ctx, - )?; - } - Ok(Some(Arc::new(NoOpExecutionPlan::new(Arc::clone( - node.schema().inner(), - ))))) + fn plan_extension<'a>( + &'a self, + planner: &'a dyn PhysicalPlanner, + node: &'a dyn UserDefinedLogicalNode, + _logical_inputs: &'a [&'a LogicalPlan], + _physical_inputs: &'a [Arc], + session_state: &'a dyn Session, + planning_ctx: &'a PhysicalPlanningContext, + ) -> BoxFuture<'a, Result>>> { + Box::pin(async move { + for expr in node.expressions() { + planner.create_physical_expr( + &expr, + node.schema(), + session_state, + planning_ctx, + )?; + } + Ok(Some(Arc::new(NoOpExecutionPlan::new(Arc::clone( + node.schema().inner(), + ))) as Arc)) + }) } } // Produces an execution plan where the schema is mismatched from // the logical plan node. struct BadExtensionPlanner {} - - #[async_trait] impl ExtensionPlanner for BadExtensionPlanner { /// Create a physical plan for an extension node - async fn plan_extension( - &self, - _planner: &dyn PhysicalPlanner, - _node: &dyn UserDefinedLogicalNode, - _logical_inputs: &[&LogicalPlan], - _physical_inputs: &[Arc], - _session_state: &dyn Session, - _planning_ctx: &PhysicalPlanningContext, - ) -> Result>> { - Ok(Some(Arc::new(NoOpExecutionPlan::new(SchemaRef::new( - Schema::new(vec![Field::new("b", DataType::Int32, false)]), - ))))) + fn plan_extension<'a>( + &'a self, + _planner: &'a dyn PhysicalPlanner, + _node: &'a dyn UserDefinedLogicalNode, + _logical_inputs: &'a [&'a LogicalPlan], + _physical_inputs: &'a [Arc], + _session_state: &'a dyn Session, + _planning_ctx: &'a PhysicalPlanningContext, + ) -> BoxFuture<'a, Result>>> { + Box::pin(async move { + Ok(Some( + Arc::new(NoOpExecutionPlan::new(SchemaRef::new(Schema::new(vec![ + Field::new("b", DataType::Int32, false), + ])))) as Arc, + )) + }) } } @@ -5618,8 +5610,6 @@ digraph { logical_schema: SchemaRef, physical_schema: SchemaRef, } - - #[async_trait] impl TableProvider for MockSchemaTableProvider { fn schema(&self) -> SchemaRef { Arc::clone(&self.logical_schema) @@ -5629,16 +5619,19 @@ digraph { TableType::Base } - async fn scan( - &self, - _state: &dyn Session, - _projection: Option<&[usize]>, - _filters: &[Expr], + fn scan<'a>( + &'a self, + _state: &'a dyn Session, + _projection: Option<&'a [usize]>, + _filters: &'a [Expr], _limit: Option, - ) -> Result> { - Ok(Arc::new(NoOpExecutionPlan::new(Arc::clone( - &self.physical_schema, - )))) + ) -> BoxFuture<'a, Result>> { + Box::pin(async move { + Ok( + Arc::new(NoOpExecutionPlan::new(Arc::clone(&self.physical_schema))) + as Arc, + ) + }) } } @@ -5837,35 +5830,35 @@ digraph { } struct MockTableScanExtensionPlanner; - - #[async_trait] impl ExtensionPlanner for MockTableScanExtensionPlanner { - async fn plan_extension( - &self, - _planner: &dyn PhysicalPlanner, - _node: &dyn UserDefinedLogicalNode, - _logical_inputs: &[&LogicalPlan], - _physical_inputs: &[Arc], - _session_state: &dyn Session, - _planning_ctx: &PhysicalPlanningContext, - ) -> Result>> { - Ok(None) - } - - async fn plan_table_scan( - &self, - _planner: &dyn PhysicalPlanner, - scan: &TableScan, - _session_state: &dyn Session, - _planning_ctx: &PhysicalPlanningContext, - ) -> Result>> { - if scan.source.is::() { - Ok(Some(Arc::new(EmptyExec::new(Arc::clone( - scan.projected_schema.inner(), - ))))) - } else { - Ok(None) - } + fn plan_extension<'a>( + &'a self, + _planner: &'a dyn PhysicalPlanner, + _node: &'a dyn UserDefinedLogicalNode, + _logical_inputs: &'a [&'a LogicalPlan], + _physical_inputs: &'a [Arc], + _session_state: &'a dyn Session, + _planning_ctx: &'a PhysicalPlanningContext, + ) -> BoxFuture<'a, Result>>> { + Box::pin(async move { Ok(None) }) + } + + fn plan_table_scan<'a>( + &'a self, + _planner: &'a dyn PhysicalPlanner, + scan: &'a TableScan, + _session_state: &'a dyn Session, + _planning_ctx: &'a PhysicalPlanningContext, + ) -> BoxFuture<'a, Result>>> { + Box::pin(async move { + if scan.source.is::() { + Ok(Some(Arc::new(EmptyExec::new(Arc::clone( + scan.projected_schema.inner(), + ))) as Arc)) + } else { + Ok(None) + } + }) } } diff --git a/datafusion/core/src/test/object_store.rs b/datafusion/core/src/test/object_store.rs index 62c6699f8fcd1..918e5608b152c 100644 --- a/datafusion/core/src/test/object_store.rs +++ b/datafusion/core/src/test/object_store.rs @@ -28,6 +28,8 @@ use crate::{ }; use futures::{FutureExt, stream::BoxStream}; use object_store::{CopyOptions, ObjectStoreExt}; +use std::future::Future; +use std::pin::Pin; use std::{ fmt::{Debug, Display, Formatter}, sync::Arc, @@ -108,57 +110,84 @@ impl Display for BlockingObjectStore { /// All trait methods are forwarded to the inner object store, except for /// the `head` method which waits until the expected number of concurrent calls is reached. -#[async_trait::async_trait] impl ObjectStore for BlockingObjectStore { - async fn put_opts( - &self, - location: &Path, + fn put_opts<'life0, 'life1, 'async_trait>( + &'life0 self, + location: &'life1 Path, payload: PutPayload, opts: PutOptions, - ) -> object_store::Result { - self.inner.put_opts(location, payload, opts).await + ) -> Pin< + Box> + Send + 'async_trait>, + > + where + 'life0: 'async_trait, + 'life1: 'async_trait, + Self: 'async_trait, + { + Box::pin(async move { self.inner.put_opts(location, payload, opts).await }) } - async fn put_multipart_opts( - &self, - location: &Path, + fn put_multipart_opts<'life0, 'life1, 'async_trait>( + &'life0 self, + location: &'life1 Path, opts: PutMultipartOptions, - ) -> object_store::Result> { - self.inner.put_multipart_opts(location, opts).await + ) -> Pin< + Box< + dyn Future>> + + Send + + 'async_trait, + >, + > + where + 'life0: 'async_trait, + 'life1: 'async_trait, + Self: 'async_trait, + { + Box::pin(async move { self.inner.put_multipart_opts(location, opts).await }) } - async fn get_opts( - &self, - location: &Path, + fn get_opts<'life0, 'life1, 'async_trait>( + &'life0 self, + location: &'life1 Path, options: GetOptions, - ) -> object_store::Result { - if options.head { - println!( - "{} received head call for {location}", - BlockingObjectStore::NAME - ); - // Wait until the expected number of concurrent calls is reached, but timeout after 1 second to avoid hanging failing tests. - let wait_result = timeout(Duration::from_secs(1), self.barrier.wait()).await; - match wait_result { - Ok(_) => println!( - "{} barrier reached for {location}", + ) -> Pin< + Box> + Send + 'async_trait>, + > + where + 'life0: 'async_trait, + 'life1: 'async_trait, + Self: 'async_trait, + { + Box::pin(async move { + if options.head { + println!( + "{} received head call for {location}", BlockingObjectStore::NAME - ), - Err(_) => { - let error_message = format!( - "{} barrier wait timed out for {location}", + ); + // Wait until the expected number of concurrent calls is reached, but timeout after 1 second to avoid hanging failing tests. + let wait_result = + timeout(Duration::from_secs(1), self.barrier.wait()).await; + match wait_result { + Ok(_) => println!( + "{} barrier reached for {location}", BlockingObjectStore::NAME - ); - log::error!("{error_message}"); - return Err(Error::Generic { - store: BlockingObjectStore::NAME, - source: error_message.into(), - }); + ), + Err(_) => { + let error_message = format!( + "{} barrier wait timed out for {location}", + BlockingObjectStore::NAME + ); + log::error!("{error_message}"); + return Err(Error::Generic { + store: BlockingObjectStore::NAME, + source: error_message.into(), + }); + } } } - } - // Forward the call to the inner object store. - self.inner.get_opts(location, options).await + // Forward the call to the inner object store. + self.inner.get_opts(location, options).await + }) } fn delete_stream( &self, @@ -174,19 +203,32 @@ impl ObjectStore for BlockingObjectStore { self.inner.list(prefix) } - async fn list_with_delimiter( - &self, - prefix: Option<&Path>, - ) -> object_store::Result { - self.inner.list_with_delimiter(prefix).await + fn list_with_delimiter<'life0, 'life1, 'async_trait>( + &'life0 self, + prefix: Option<&'life1 Path>, + ) -> Pin< + Box> + Send + 'async_trait>, + > + where + 'life0: 'async_trait, + 'life1: 'async_trait, + Self: 'async_trait, + { + Box::pin(async move { self.inner.list_with_delimiter(prefix).await }) } - async fn copy_opts( - &self, - from: &Path, - to: &Path, + fn copy_opts<'life0, 'life1, 'life2, 'async_trait>( + &'life0 self, + from: &'life1 Path, + to: &'life2 Path, options: CopyOptions, - ) -> object_store::Result<()> { - self.inner.copy_opts(from, to, options).await + ) -> Pin> + Send + 'async_trait>> + where + 'life0: 'async_trait, + 'life1: 'async_trait, + 'life2: 'async_trait, + Self: 'async_trait, + { + Box::pin(async move { self.inner.copy_opts(from, to, options).await }) } } diff --git a/datafusion/core/src/test_util/mod.rs b/datafusion/core/src/test_util/mod.rs index ab7edef885b40..a83e8c9218087 100644 --- a/datafusion/core/src/test_util/mod.rs +++ b/datafusion/core/src/test_util/mod.rs @@ -53,7 +53,6 @@ use datafusion_expr::{ }; use std::pin::Pin; -use async_trait::async_trait; use futures::future::BoxFuture; use tempfile::TempDir; @@ -181,22 +180,12 @@ pub fn populate_csv_partitions( /// TableFactory for tests #[derive(Default, Debug)] pub struct TestTableFactory {} - -#[async_trait] impl TableProviderFactory for TestTableFactory { - // Hand-written `#[async_trait]` expansion to reduce compile time. See - // - fn create<'life0, 'life1, 'life2, 'async_trait>( - &'life0 self, - session: &'life1 dyn Session, - cmd: &'life2 CreateExternalTable, - ) -> BoxFuture<'async_trait, Result>> - where - 'life0: 'async_trait, - 'life1: 'async_trait, - 'life2: 'async_trait, - Self: 'async_trait, - { + fn create<'a>( + &'a self, + session: &'a dyn Session, + cmd: &'a CreateExternalTable, + ) -> BoxFuture<'a, Result>> { self.create_boxed(session, cmd) } } @@ -236,8 +225,6 @@ pub struct TestTableProvider { } impl TestTableProvider {} - -#[async_trait] impl TableProvider for TestTableProvider { fn schema(&self) -> SchemaRef { Arc::clone(&self.schema) @@ -247,14 +234,16 @@ impl TableProvider for TestTableProvider { unimplemented!("TestTableProvider is a stub for testing.") } - async fn scan( - &self, - _state: &dyn Session, - _projection: Option<&[usize]>, - _filters: &[Expr], + fn scan<'a>( + &'a self, + _state: &'a dyn Session, + _projection: Option<&'a [usize]>, + _filters: &'a [Expr], _limit: Option, - ) -> Result> { - unimplemented!("TestTableProvider is a stub for testing.") + ) -> BoxFuture<'a, Result>> { + Box::pin( + async move { unimplemented!("TestTableProvider is a stub for testing.") }, + ) } } diff --git a/datafusion/core/tests/custom_sources_cases/dml_planning.rs b/datafusion/core/tests/custom_sources_cases/dml_planning.rs index cb5b134fab04a..7cfbf07a5c496 100644 --- a/datafusion/core/tests/custom_sources_cases/dml_planning.rs +++ b/datafusion/core/tests/custom_sources_cases/dml_planning.rs @@ -17,10 +17,10 @@ //! Tests for DELETE, UPDATE, and TRUNCATE planning to verify filter and assignment extraction. +use futures::future::BoxFuture; use std::sync::{Arc, Mutex}; use arrow::datatypes::{DataType, Field, Schema, SchemaRef}; -use async_trait::async_trait; use datafusion::datasource::{TableProvider, TableType}; use datafusion::error::Result; use datafusion::execution::context::{SessionConfig, SessionContext}; @@ -87,8 +87,6 @@ impl std::fmt::Debug for CaptureDeleteProvider { .finish() } } - -#[async_trait] impl TableProvider for CaptureDeleteProvider { fn schema(&self) -> SchemaRef { Arc::clone(&self.schema) @@ -98,25 +96,34 @@ impl TableProvider for CaptureDeleteProvider { TableType::Base } - async fn scan( - &self, - _state: &dyn Session, - _projection: Option<&[usize]>, - _filters: &[Expr], + fn scan<'a>( + &'a self, + _state: &'a dyn Session, + _projection: Option<&'a [usize]>, + _filters: &'a [Expr], _limit: Option, - ) -> Result> { - Ok(Arc::new(EmptyExec::new(Arc::clone(&self.schema)))) + ) -> BoxFuture<'a, Result>> { + Box::pin(async move { + Ok(Arc::new(EmptyExec::new(Arc::clone(&self.schema))) + as Arc) + }) } - async fn delete_from( - &self, - _state: &dyn Session, + fn delete_from<'a>( + &'a self, + _state: &'a dyn Session, filters: Vec, - ) -> Result> { - *self.received_filters.lock().unwrap() = Some(filters); - Ok(Arc::new(EmptyExec::new(Arc::new(Schema::new(vec![ - Field::new("count", DataType::UInt64, false), - ]))))) + ) -> BoxFuture<'a, Result>> { + Box::pin(async move { + *self.received_filters.lock().unwrap() = Some(filters); + Ok( + Arc::new(EmptyExec::new(Arc::new(Schema::new(vec![Field::new( + "count", + DataType::UInt64, + false, + )])))) as Arc, + ) + }) } fn supports_filters_pushdown( @@ -183,8 +190,6 @@ impl std::fmt::Debug for CaptureUpdateProvider { .finish() } } - -#[async_trait] impl TableProvider for CaptureUpdateProvider { fn schema(&self) -> SchemaRef { Arc::clone(&self.schema) @@ -194,27 +199,36 @@ impl TableProvider for CaptureUpdateProvider { TableType::Base } - async fn scan( - &self, - _state: &dyn Session, - _projection: Option<&[usize]>, - _filters: &[Expr], + fn scan<'a>( + &'a self, + _state: &'a dyn Session, + _projection: Option<&'a [usize]>, + _filters: &'a [Expr], _limit: Option, - ) -> Result> { - Ok(Arc::new(EmptyExec::new(Arc::clone(&self.schema)))) + ) -> BoxFuture<'a, Result>> { + Box::pin(async move { + Ok(Arc::new(EmptyExec::new(Arc::clone(&self.schema))) + as Arc) + }) } - async fn update( - &self, - _state: &dyn Session, + fn update<'a>( + &'a self, + _state: &'a dyn Session, assignments: Vec<(String, Expr)>, filters: Vec, - ) -> Result> { - *self.received_filters.lock().unwrap() = Some(filters); - *self.received_assignments.lock().unwrap() = Some(assignments); - Ok(Arc::new(EmptyExec::new(Arc::new(Schema::new(vec![ - Field::new("count", DataType::UInt64, false), - ]))))) + ) -> BoxFuture<'a, Result>> { + Box::pin(async move { + *self.received_filters.lock().unwrap() = Some(filters); + *self.received_assignments.lock().unwrap() = Some(assignments); + Ok( + Arc::new(EmptyExec::new(Arc::new(Schema::new(vec![Field::new( + "count", + DataType::UInt64, + false, + )])))) as Arc, + ) + }) } fn supports_filters_pushdown( @@ -257,8 +271,6 @@ impl std::fmt::Debug for CaptureTruncateProvider { .finish() } } - -#[async_trait] impl TableProvider for CaptureTruncateProvider { fn schema(&self) -> SchemaRef { Arc::clone(&self.schema) @@ -268,22 +280,34 @@ impl TableProvider for CaptureTruncateProvider { TableType::Base } - async fn scan( - &self, - _state: &dyn Session, - _projection: Option<&[usize]>, - _filters: &[Expr], + fn scan<'a>( + &'a self, + _state: &'a dyn Session, + _projection: Option<&'a [usize]>, + _filters: &'a [Expr], _limit: Option, - ) -> Result> { - Ok(Arc::new(EmptyExec::new(Arc::clone(&self.schema)))) + ) -> BoxFuture<'a, Result>> { + Box::pin(async move { + Ok(Arc::new(EmptyExec::new(Arc::clone(&self.schema))) + as Arc) + }) } - async fn truncate(&self, _state: &dyn Session) -> Result> { - *self.truncate_called.lock().unwrap() = true; - - Ok(Arc::new(EmptyExec::new(Arc::new(Schema::new(vec![ - Field::new("count", DataType::UInt64, false), - ]))))) + fn truncate<'a>( + &'a self, + _state: &'a dyn Session, + ) -> BoxFuture<'a, Result>> { + Box::pin(async move { + *self.truncate_called.lock().unwrap() = true; + + Ok( + Arc::new(EmptyExec::new(Arc::new(Schema::new(vec![Field::new( + "count", + DataType::UInt64, + false, + )])))) as Arc, + ) + }) } } diff --git a/datafusion/core/tests/custom_sources_cases/mod.rs b/datafusion/core/tests/custom_sources_cases/mod.rs index ea3cdaa34df54..6419697fdbc1d 100644 --- a/datafusion/core/tests/custom_sources_cases/mod.rs +++ b/datafusion/core/tests/custom_sources_cases/mod.rs @@ -15,6 +15,7 @@ // specific language governing permissions and limitations // under the License. +use futures::future::BoxFuture; use std::pin::Pin; use std::sync::Arc; use std::task::{Context, Poll}; @@ -47,7 +48,6 @@ use datafusion_physical_plan::{ ChildrenPropertiesMode, PlanProperties, ReplaceChildrenOptions, }; -use async_trait::async_trait; use futures::stream::Stream; mod dml_planning; @@ -233,8 +233,6 @@ impl ExecutionPlan for CustomExecutionPlan { Ok(TreeNodeRecursion::Continue) } } - -#[async_trait] impl TableProvider for CustomTableProvider { fn schema(&self) -> SchemaRef { TEST_CUSTOM_SCHEMA_REF!() @@ -244,16 +242,19 @@ impl TableProvider for CustomTableProvider { TableType::Base } - async fn scan( - &self, - _state: &dyn Session, - projection: Option<&[usize]>, - _filters: &[Expr], + fn scan<'a>( + &'a self, + _state: &'a dyn Session, + projection: Option<&'a [usize]>, + _filters: &'a [Expr], _limit: Option, - ) -> Result> { - Ok(Arc::new(CustomExecutionPlan::new( - projection.map(|p| p.to_vec()), - ))) + ) -> BoxFuture<'a, Result>> { + Box::pin(async move { + Ok( + Arc::new(CustomExecutionPlan::new(projection.map(|p| p.to_vec()))) + as Arc, + ) + }) } } diff --git a/datafusion/core/tests/custom_sources_cases/provider_filter_pushdown.rs b/datafusion/core/tests/custom_sources_cases/provider_filter_pushdown.rs index 4ec9747058140..6c81c5480a2e1 100644 --- a/datafusion/core/tests/custom_sources_cases/provider_filter_pushdown.rs +++ b/datafusion/core/tests/custom_sources_cases/provider_filter_pushdown.rs @@ -15,6 +15,7 @@ // specific language governing permissions and limitations // under the License. +use futures::future::BoxFuture; use std::sync::Arc; use arrow::array::{Int32Builder, Int64Array}; @@ -42,8 +43,6 @@ use datafusion_physical_expr::EquivalenceProperties; use datafusion_physical_plan::execution_plan::{Boundedness, EmissionType}; use datafusion_physical_plan::{ChildrenPropertiesMode, ReplaceChildrenOptions}; -use async_trait::async_trait; - fn create_batch(value: i32, num_rows: usize) -> Result { let mut builder = Int32Builder::with_capacity(num_rows); for _ in 0..num_rows { @@ -176,8 +175,6 @@ struct CustomProvider { zero_batch: RecordBatch, one_batch: RecordBatch, } - -#[async_trait] impl TableProvider for CustomProvider { fn schema(&self) -> SchemaRef { self.zero_batch.schema() @@ -187,63 +184,67 @@ impl TableProvider for CustomProvider { TableType::Base } - async fn scan( - &self, - _state: &dyn Session, - projection: Option<&[usize]>, - filters: &[Expr], + fn scan<'a>( + &'a self, + _state: &'a dyn Session, + projection: Option<&'a [usize]>, + filters: &'a [Expr], _: Option, - ) -> Result> { - let empty = Vec::new(); - let projection = projection.unwrap_or(&empty); - match &filters[0] { - Expr::BinaryExpr(BinaryExpr { right, .. }) => { - let int_value = match &**right { - Expr::Literal(ScalarValue::Int8(Some(i)), _) => *i as i64, - Expr::Literal(ScalarValue::Int16(Some(i)), _) => *i as i64, - Expr::Literal(ScalarValue::Int32(Some(i)), _) => *i as i64, - Expr::Literal(ScalarValue::Int64(Some(i)), _) => *i, - Expr::Cast(Cast { expr, field: _ }) => match &**expr { - Expr::Literal(lit_value, _) => match lit_value { - ScalarValue::Int8(Some(v)) => *v as i64, - ScalarValue::Int16(Some(v)) => *v as i64, - ScalarValue::Int32(Some(v)) => *v as i64, - ScalarValue::Int64(Some(v)) => *v, - other_value => { + ) -> BoxFuture<'a, Result>> { + Box::pin(async move { + let empty = Vec::new(); + let projection = projection.unwrap_or(&empty); + match &filters[0] { + Expr::BinaryExpr(BinaryExpr { right, .. }) => { + let int_value = match &**right { + Expr::Literal(ScalarValue::Int8(Some(i)), _) => *i as i64, + Expr::Literal(ScalarValue::Int16(Some(i)), _) => *i as i64, + Expr::Literal(ScalarValue::Int32(Some(i)), _) => *i as i64, + Expr::Literal(ScalarValue::Int64(Some(i)), _) => *i, + Expr::Cast(Cast { expr, field: _ }) => match &**expr { + Expr::Literal(lit_value, _) => match lit_value { + ScalarValue::Int8(Some(v)) => *v as i64, + ScalarValue::Int16(Some(v)) => *v as i64, + ScalarValue::Int32(Some(v)) => *v as i64, + ScalarValue::Int64(Some(v)) => *v, + other_value => { + return not_impl_err!( + "Do not support value {other_value:?}" + ); + } + }, + other_expr => { return not_impl_err!( - "Do not support value {other_value:?}" + "Do not support expr {other_expr:?}" ); } }, other_expr => { return not_impl_err!("Do not support expr {other_expr:?}"); } - }, - other_expr => { - return not_impl_err!("Do not support expr {other_expr:?}"); - } - }; + }; - Ok(Arc::new(CustomPlan::new( + Ok(Arc::new(CustomPlan::new( + match projection.is_empty() { + true => Arc::new(Schema::empty()), + false => self.zero_batch.schema(), + }, + match int_value { + 0 => vec![self.zero_batch.clone()], + 1 => vec![self.one_batch.clone()], + _ => vec![], + }, + )) as Arc) + } + _ => Ok(Arc::new(CustomPlan::new( match projection.is_empty() { true => Arc::new(Schema::empty()), false => self.zero_batch.schema(), }, - match int_value { - 0 => vec![self.zero_batch.clone()], - 1 => vec![self.one_batch.clone()], - _ => vec![], - }, - ))) + vec![], + )) as Arc), } - _ => Ok(Arc::new(CustomPlan::new( - match projection.is_empty() { - true => Arc::new(Schema::empty()), - false => self.zero_batch.schema(), - }, - vec![], - ))), - } + }) } fn supports_filters_pushdown( diff --git a/datafusion/core/tests/custom_sources_cases/statistics.rs b/datafusion/core/tests/custom_sources_cases/statistics.rs index 7701ca0279125..43c59613a18bc 100644 --- a/datafusion/core/tests/custom_sources_cases/statistics.rs +++ b/datafusion/core/tests/custom_sources_cases/statistics.rs @@ -17,6 +17,7 @@ //! This module contains end to end tests of statistics propagation +use futures::future::BoxFuture; use std::sync::Arc; use arrow::datatypes::{DataType, Field, Schema, SchemaRef}; @@ -41,8 +42,6 @@ use datafusion_physical_plan::{ ChildrenPropertiesMode, ReplaceChildrenOptions, StatisticsArgs, StatisticsContext, }; -use async_trait::async_trait; - /// This is a testing structure for statistics /// It will act both as a table provider and execution plan #[derive(Debug, Clone)] @@ -77,8 +76,6 @@ impl StatisticsValidation { ) } } - -#[async_trait] impl TableProvider for StatisticsValidation { fn schema(&self) -> SchemaRef { Arc::clone(&self.schema) @@ -88,41 +85,43 @@ impl TableProvider for StatisticsValidation { TableType::Base } - async fn scan( - &self, - _state: &dyn Session, - projection: Option<&[usize]>, - filters: &[Expr], + fn scan<'a>( + &'a self, + _state: &'a dyn Session, + projection: Option<&'a [usize]>, + filters: &'a [Expr], // limit is ignored because it is not mandatory for a `TableProvider` to honor it _limit: Option, - ) -> Result> { - // Filters should not be pushed down as they are marked as unsupported by default. - assert_eq!( - 0, - filters.len(), - "Unsupported expressions should not be pushed down" - ); - let projection = match projection.map(|p| p.to_vec()) { - Some(p) => p, - None => (0..self.schema.fields().len()).collect(), - }; - let projected_schema = project_schema(&self.schema, Some(&projection))?; - - let current_stat = self.stats.clone(); - - let proj_col_stats = projection - .iter() - .map(|i| current_stat.column_statistics[*i].clone()) - .collect(); - Ok(Arc::new(Self::new( - Statistics { - num_rows: current_stat.num_rows, - column_statistics: proj_col_stats, - // TODO stats: knowing the type of the new columns we can guess the output size - total_byte_size: Precision::Absent, - }, - projected_schema, - ))) + ) -> BoxFuture<'a, Result>> { + Box::pin(async move { + // Filters should not be pushed down as they are marked as unsupported by default. + assert_eq!( + 0, + filters.len(), + "Unsupported expressions should not be pushed down" + ); + let projection = match projection.map(|p| p.to_vec()) { + Some(p) => p, + None => (0..self.schema.fields().len()).collect(), + }; + let projected_schema = project_schema(&self.schema, Some(&projection))?; + + let current_stat = self.stats.clone(); + + let proj_col_stats = projection + .iter() + .map(|i| current_stat.column_statistics[*i].clone()) + .collect(); + Ok(Arc::new(Self::new( + Statistics { + num_rows: current_stat.num_rows, + column_statistics: proj_col_stats, + // TODO stats: knowing the type of the new columns we can guess the output size + total_byte_size: Precision::Absent, + }, + projected_schema, + )) as Arc) + }) } } diff --git a/datafusion/core/tests/datasource/object_store_access.rs b/datafusion/core/tests/datasource/object_store_access.rs index 16d894bde1303..cbf14463b320f 100644 --- a/datafusion/core/tests/datasource/object_store_access.rs +++ b/datafusion/core/tests/datasource/object_store_access.rs @@ -25,7 +25,6 @@ //! [`ListingTable`]: datafusion::datasource::listing::ListingTable use arrow::array::{ArrayRef, Int32Array, RecordBatch}; -use async_trait::async_trait; use bytes::Bytes; use datafusion::prelude::{ CsvReadOptions, JsonReadOptions, ParquetReadOptions, SessionContext, @@ -46,7 +45,9 @@ use object_store::{ use parking_lot::Mutex; use std::fmt; use std::fmt::{Display, Formatter}; +use std::future::Future; use std::ops::Range; +use std::pin::Pin; use std::sync::Arc; use url::Url; @@ -1375,49 +1376,85 @@ impl RequestCountingObjectStore { } } -#[async_trait] impl ObjectStore for RequestCountingObjectStore { - async fn put_opts( - &self, - _location: &Path, + fn put_opts<'life0, 'life1, 'async_trait>( + &'life0 self, + _location: &'life1 Path, _payload: PutPayload, _opts: PutOptions, - ) -> object_store::Result { - unimplemented!() + ) -> Pin< + Box> + Send + 'async_trait>, + > + where + 'life0: 'async_trait, + 'life1: 'async_trait, + Self: 'async_trait, + { + Box::pin(async move { unimplemented!() }) } - async fn put_multipart_opts( - &self, - _location: &Path, + fn put_multipart_opts<'life0, 'life1, 'async_trait>( + &'life0 self, + _location: &'life1 Path, _opts: PutMultipartOptions, - ) -> object_store::Result> { - unimplemented!() + ) -> Pin< + Box< + dyn Future>> + + Send + + 'async_trait, + >, + > + where + 'life0: 'async_trait, + 'life1: 'async_trait, + Self: 'async_trait, + { + Box::pin(async move { unimplemented!() }) } - async fn get_opts( - &self, - location: &Path, + fn get_opts<'life0, 'life1, 'async_trait>( + &'life0 self, + location: &'life1 Path, options: GetOptions, - ) -> object_store::Result { - let result = self.inner.get_opts(location, options.clone()).await?; - self.requests.lock().push(RequestDetails::GetOpts { - path: location.to_owned(), - get_options: options, - }); - Ok(result) + ) -> Pin< + Box> + Send + 'async_trait>, + > + where + 'life0: 'async_trait, + 'life1: 'async_trait, + Self: 'async_trait, + { + Box::pin(async move { + let result = self.inner.get_opts(location, options.clone()).await?; + self.requests.lock().push(RequestDetails::GetOpts { + path: location.to_owned(), + get_options: options, + }); + Ok(result) + }) } - async fn get_ranges( - &self, - location: &Path, - ranges: &[Range], - ) -> object_store::Result> { - let result = self.inner.get_ranges(location, ranges).await?; - self.requests.lock().push(RequestDetails::GetRanges { - path: location.to_owned(), - ranges: ranges.to_vec(), - }); - Ok(result) + fn get_ranges<'life0, 'life1, 'life2, 'async_trait>( + &'life0 self, + location: &'life1 Path, + ranges: &'life2 [Range], + ) -> Pin< + Box>> + Send + 'async_trait>, + > + where + 'life0: 'async_trait, + 'life1: 'async_trait, + 'life2: 'async_trait, + Self: 'async_trait, + { + Box::pin(async move { + let result = self.inner.get_ranges(location, ranges).await?; + self.requests.lock().push(RequestDetails::GetRanges { + path: location.to_owned(), + ranges: ranges.to_vec(), + }); + Ok(result) + }) } fn list( @@ -1443,16 +1480,25 @@ impl ObjectStore for RequestCountingObjectStore { self.inner.list_with_offset(prefix, offset) } - async fn list_with_delimiter( - &self, - prefix: Option<&Path>, - ) -> object_store::Result { - self.requests - .lock() - .push(RequestDetails::ListWithDelimiter { - prefix: prefix.map(|p| p.to_owned()), - }); - self.inner.list_with_delimiter(prefix).await + fn list_with_delimiter<'life0, 'life1, 'async_trait>( + &'life0 self, + prefix: Option<&'life1 Path>, + ) -> Pin< + Box> + Send + 'async_trait>, + > + where + 'life0: 'async_trait, + 'life1: 'async_trait, + Self: 'async_trait, + { + Box::pin(async move { + self.requests + .lock() + .push(RequestDetails::ListWithDelimiter { + prefix: prefix.map(|p| p.to_owned()), + }); + self.inner.list_with_delimiter(prefix).await + }) } fn delete_stream( @@ -1462,12 +1508,18 @@ impl ObjectStore for RequestCountingObjectStore { unimplemented!() } - async fn copy_opts( - &self, - _from: &Path, - _to: &Path, + fn copy_opts<'life0, 'life1, 'life2, 'async_trait>( + &'life0 self, + _from: &'life1 Path, + _to: &'life2 Path, _options: CopyOptions, - ) -> object_store::Result<()> { - unimplemented!() + ) -> Pin> + Send + 'async_trait>> + where + 'life0: 'async_trait, + 'life1: 'async_trait, + 'life2: 'async_trait, + Self: 'async_trait, + { + Box::pin(async move { unimplemented!() }) } } diff --git a/datafusion/core/tests/memory_limit/mod.rs b/datafusion/core/tests/memory_limit/mod.rs index 1e86c1fb7f89b..3bc3877eb2298 100644 --- a/datafusion/core/tests/memory_limit/mod.rs +++ b/datafusion/core/tests/memory_limit/mod.rs @@ -17,6 +17,7 @@ //! This module contains tests for limiting memory at runtime in DataFusion +use futures::future::BoxFuture; use std::num::NonZeroUsize; use std::sync::{Arc, LazyLock}; @@ -65,7 +66,6 @@ use datafusion_physical_plan::spill::get_record_batch_memory_size; use rand::Rng; use test_utils::AccessLogGenerator; -use async_trait::async_trait; use futures::StreamExt; use tokio::fs::File; @@ -1659,8 +1659,6 @@ impl SortedTableProvider { } } } - -#[async_trait] impl TableProvider for SortedTableProvider { fn schema(&self) -> SchemaRef { self.schema.clone() @@ -1670,20 +1668,22 @@ impl TableProvider for SortedTableProvider { TableType::Base } - async fn scan( - &self, - _state: &dyn Session, - projection: Option<&[usize]>, - _filters: &[Expr], + fn scan<'a>( + &'a self, + _state: &'a dyn Session, + projection: Option<&'a [usize]>, + _filters: &'a [Expr], _limit: Option, - ) -> Result> { - let mem_conf = MemorySourceConfig::try_new( - &self.batches, - self.schema(), - projection.map(|p| p.to_vec()), - )? - .try_with_sort_information(self.sort_information.clone())?; - - Ok(DataSourceExec::from_data_source(mem_conf)) + ) -> BoxFuture<'a, Result>> { + Box::pin(async move { + let mem_conf = MemorySourceConfig::try_new( + &self.batches, + self.schema(), + projection.map(|p| p.to_vec()), + )? + .try_with_sort_information(self.sort_information.clone())?; + + Ok(DataSourceExec::from_data_source(mem_conf) as Arc) + }) } } diff --git a/datafusion/core/tests/parquet/encryption.rs b/datafusion/core/tests/parquet/encryption.rs index b7bfcefa30f34..898bd3ea63e0a 100644 --- a/datafusion/core/tests/parquet/encryption.rs +++ b/datafusion/core/tests/parquet/encryption.rs @@ -20,7 +20,6 @@ use arrow::array::{ArrayRef, Int32Array, StringArray}; use arrow::record_batch::RecordBatch; use arrow_schema::{DataType, SchemaRef}; -use async_trait::async_trait; use datafusion::dataframe::DataFrameWriteOptions; use datafusion::datasource::listing::ListingOptions; use datafusion::prelude::{ParquetReadOptions, SessionContext}; @@ -28,6 +27,7 @@ use datafusion_common::config::{EncryptionFactoryOptions, TableParquetOptions}; use datafusion_common::{DataFusionError, assert_batches_sorted_eq, exec_datafusion_err}; use datafusion_datasource_parquet::ParquetFormat; use datafusion_execution::parquet_encryption::EncryptionFactory; +use futures::future::BoxFuture; use parquet::arrow::ArrowWriter; use parquet::arrow::arrow_reader::{ArrowReaderMetadata, ArrowReaderOptions}; use parquet::encryption::decrypt::FileDecryptionProperties; @@ -329,42 +329,46 @@ struct MockEncryptionFactory { pub encryption_keys: Mutex>>, pub counter: AtomicU8, } - -#[async_trait] impl EncryptionFactory for MockEncryptionFactory { - async fn get_file_encryption_properties( - &self, - config: &EncryptionFactoryOptions, - _schema: &SchemaRef, - file_path: &object_store::path::Path, - ) -> datafusion_common::Result>> { - assert_eq!( - config.options.get("test_key"), - Some(&"test value".to_string()) - ); - let file_idx = self.counter.fetch_add(1, Ordering::Relaxed); - let key = vec![file_idx, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0]; - let mut keys = self.encryption_keys.lock().unwrap(); - keys.insert(file_path.clone(), key.clone()); - let encryption_properties = FileEncryptionProperties::builder(key).build()?; - Ok(Some(encryption_properties)) + fn get_file_encryption_properties<'a>( + &'a self, + config: &'a EncryptionFactoryOptions, + _schema: &'a SchemaRef, + file_path: &'a object_store::path::Path, + ) -> BoxFuture<'a, datafusion_common::Result>>> + { + Box::pin(async move { + assert_eq!( + config.options.get("test_key"), + Some(&"test value".to_string()) + ); + let file_idx = self.counter.fetch_add(1, Ordering::Relaxed); + let key = vec![file_idx, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0]; + let mut keys = self.encryption_keys.lock().unwrap(); + keys.insert(file_path.clone(), key.clone()); + let encryption_properties = FileEncryptionProperties::builder(key).build()?; + Ok(Some(encryption_properties)) + }) } - async fn get_file_decryption_properties( - &self, - config: &EncryptionFactoryOptions, - file_path: &object_store::path::Path, - ) -> datafusion_common::Result>> { - assert_eq!( - config.options.get("test_key"), - Some(&"test value".to_string()) - ); - let keys = self.encryption_keys.lock().unwrap(); - let key = keys - .get(file_path) - .ok_or_else(|| exec_datafusion_err!("No key for file {file_path:?}"))?; - let decryption_properties = - FileDecryptionProperties::builder(key.clone()).build()?; - Ok(Some(decryption_properties)) + fn get_file_decryption_properties<'a>( + &'a self, + config: &'a EncryptionFactoryOptions, + file_path: &'a object_store::path::Path, + ) -> BoxFuture<'a, datafusion_common::Result>>> + { + Box::pin(async move { + assert_eq!( + config.options.get("test_key"), + Some(&"test value".to_string()) + ); + let keys = self.encryption_keys.lock().unwrap(); + let key = keys + .get(file_path) + .ok_or_else(|| exec_datafusion_err!("No key for file {file_path:?}"))?; + let decryption_properties = + FileDecryptionProperties::builder(key.clone()).build()?; + Ok(Some(decryption_properties)) + }) } } diff --git a/datafusion/core/tests/physical_optimizer/sanity_checker.rs b/datafusion/core/tests/physical_optimizer/sanity_checker.rs index 32f02d0f6a2f5..9eb0db811a4d0 100644 --- a/datafusion/core/tests/physical_optimizer/sanity_checker.rs +++ b/datafusion/core/tests/physical_optimizer/sanity_checker.rs @@ -15,6 +15,7 @@ // specific language governing permissions and limitations // under the License. +use futures::future::BoxFuture; use insta::assert_snapshot; use std::sync::Arc; @@ -40,8 +41,6 @@ use datafusion_physical_plan::joins::{StreamJoinPartitionMode, SymmetricHashJoin use datafusion_physical_plan::repartition::RepartitionExec; use datafusion_physical_plan::{ExecutionPlan, displayable}; -use async_trait::async_trait; - async fn register_current_csv( ctx: &SessionContext, table_name: &str, @@ -71,10 +70,9 @@ pub enum SourceType { Unbounded, Bounded, } - -#[async_trait] pub trait SqlTestCase { - async fn register_table(&self, ctx: &SessionContext) -> Result<()>; + fn register_table<'a>(&'a self, ctx: &'a SessionContext) + -> BoxFuture<'a, Result<()>>; fn expect_fail(&self) -> bool; } @@ -83,13 +81,16 @@ pub struct UnaryTestCase { pub source_type: SourceType, pub expect_fail: bool, } - -#[async_trait] impl SqlTestCase for UnaryTestCase { - async fn register_table(&self, ctx: &SessionContext) -> Result<()> { - let table_is_infinite = self.source_type == SourceType::Unbounded; - register_current_csv(ctx, "test", table_is_infinite).await?; - Ok(()) + fn register_table<'a>( + &'a self, + ctx: &'a SessionContext, + ) -> BoxFuture<'a, Result<()>> { + Box::pin(async move { + let table_is_infinite = self.source_type == SourceType::Unbounded; + register_current_csv(ctx, "test", table_is_infinite).await?; + Ok(()) + }) } fn expect_fail(&self) -> bool { @@ -102,15 +103,18 @@ pub struct BinaryTestCase { pub source_types: (SourceType, SourceType), pub expect_fail: bool, } - -#[async_trait] impl SqlTestCase for BinaryTestCase { - async fn register_table(&self, ctx: &SessionContext) -> Result<()> { - let left_table_is_infinite = self.source_types.0 == SourceType::Unbounded; - let right_table_is_infinite = self.source_types.1 == SourceType::Unbounded; - register_current_csv(ctx, "left", left_table_is_infinite).await?; - register_current_csv(ctx, "right", right_table_is_infinite).await?; - Ok(()) + fn register_table<'a>( + &'a self, + ctx: &'a SessionContext, + ) -> BoxFuture<'a, Result<()>> { + Box::pin(async move { + let left_table_is_infinite = self.source_types.0 == SourceType::Unbounded; + let right_table_is_infinite = self.source_types.1 == SourceType::Unbounded; + register_current_csv(ctx, "left", left_table_is_infinite).await?; + register_current_csv(ctx, "right", right_table_is_infinite).await?; + Ok(()) + }) } fn expect_fail(&self) -> bool { diff --git a/datafusion/core/tests/sql/path_partition.rs b/datafusion/core/tests/sql/path_partition.rs index f8701ccc568a3..e56f4d3af4ed9 100644 --- a/datafusion/core/tests/sql/path_partition.rs +++ b/datafusion/core/tests/sql/path_partition.rs @@ -19,7 +19,9 @@ use std::collections::BTreeSet; use std::fs::File; +use std::future::Future; use std::io::{Read, Seek, SeekFrom}; +use std::pin::Pin; use std::sync::Arc; use arrow::datatypes::DataType; @@ -40,7 +42,6 @@ use datafusion_common::test_util::batches_to_sort_string; use datafusion_execution::config::SessionConfig; use datafusion_physical_plan::statistics::{StatisticsArgs, StatisticsContext}; -use async_trait::async_trait; use bytes::Bytes; use chrono::{TimeZone, Utc}; use futures::StreamExt; @@ -655,69 +656,97 @@ impl MirroringObjectStore { } } -#[async_trait] impl ObjectStore for MirroringObjectStore { - async fn put_opts( - &self, - _location: &Path, + fn put_opts<'life0, 'life1, 'async_trait>( + &'life0 self, + _location: &'life1 Path, _put_payload: PutPayload, _opts: PutOptions, - ) -> object_store::Result { - unimplemented!() + ) -> Pin< + Box> + Send + 'async_trait>, + > + where + 'life0: 'async_trait, + 'life1: 'async_trait, + Self: 'async_trait, + { + Box::pin(async move { unimplemented!() }) } - async fn put_multipart_opts( - &self, - _location: &Path, + fn put_multipart_opts<'life0, 'life1, 'async_trait>( + &'life0 self, + _location: &'life1 Path, _opts: PutMultipartOptions, - ) -> object_store::Result> { - unimplemented!() + ) -> Pin< + Box< + dyn Future>> + + Send + + 'async_trait, + >, + > + where + 'life0: 'async_trait, + 'life1: 'async_trait, + Self: 'async_trait, + { + Box::pin(async move { unimplemented!() }) } - async fn get_opts( - &self, - location: &Path, + fn get_opts<'life0, 'life1, 'async_trait>( + &'life0 self, + location: &'life1 Path, options: GetOptions, - ) -> object_store::Result { - self.files.iter().find(|x| *x == location).unwrap(); - let path = std::path::PathBuf::from(&self.mirrored_file); - let file = File::open(&path).unwrap(); - let metadata = file.metadata().unwrap(); - - let meta = ObjectMeta { - location: location.clone(), - last_modified: metadata.modified().map(chrono::DateTime::from).unwrap(), - size: metadata.len(), - e_tag: None, - version: None, - }; + ) -> Pin< + Box> + Send + 'async_trait>, + > + where + 'life0: 'async_trait, + 'life1: 'async_trait, + Self: 'async_trait, + { + Box::pin(async move { + self.files.iter().find(|x| *x == location).unwrap(); + let path = std::path::PathBuf::from(&self.mirrored_file); + let file = File::open(&path).unwrap(); + let metadata = file.metadata().unwrap(); + + let meta = ObjectMeta { + location: location.clone(), + last_modified: metadata.modified().map(chrono::DateTime::from).unwrap(), + size: metadata.len(), + e_tag: None, + version: None, + }; - let payload = if options.head { - // no content for head requests - GetResultPayload::Stream(stream::empty().boxed()) - } else if let Some(range) = options.range { - let GetRange::Bounded(range) = range else { - unimplemented!("Unbounded range not supported in MirroringObjectStore"); + let payload = if options.head { + // no content for head requests + GetResultPayload::Stream(stream::empty().boxed()) + } else if let Some(range) = options.range { + let GetRange::Bounded(range) = range else { + unimplemented!( + "Unbounded range not supported in MirroringObjectStore" + ); + }; + let mut file = File::open(path).unwrap(); + file.seek(SeekFrom::Start(range.start)).unwrap(); + + let to_read = range.end - range.start; + let to_read: usize = to_read.try_into().unwrap(); + let mut data = Vec::with_capacity(to_read); + let read = file.take(to_read as u64).read_to_end(&mut data).unwrap(); + assert_eq!(read, to_read); + let stream = stream::once(async move { Ok(Bytes::from(data)) }).boxed(); + GetResultPayload::Stream(stream) + } else { + GetResultPayload::File(file, path) }; - let mut file = File::open(path).unwrap(); - file.seek(SeekFrom::Start(range.start)).unwrap(); - - let to_read = range.end - range.start; - let to_read: usize = to_read.try_into().unwrap(); - let mut data = Vec::with_capacity(to_read); - let read = file.take(to_read as u64).read_to_end(&mut data).unwrap(); - assert_eq!(read, to_read); - let stream = stream::once(async move { Ok(Bytes::from(data)) }).boxed(); - GetResultPayload::Stream(stream) - } else { - GetResultPayload::File(file, path) - }; - Ok(GetResult { - range: 0..meta.size, - payload, - meta, - attributes: Attributes::default(), + Ok(GetResult { + range: 0..meta.size, + payload, + meta, + attributes: Attributes::default(), + }) }) } @@ -750,42 +779,51 @@ impl ObjectStore for MirroringObjectStore { ))) } - async fn list_with_delimiter( - &self, - prefix: Option<&Path>, - ) -> object_store::Result { - let root = Path::default(); - let prefix = prefix.unwrap_or(&root); - - let mut common_prefixes = BTreeSet::new(); - let mut objects = vec![]; - - for k in &self.files { - let Some(mut parts) = k.prefix_match(prefix) else { - continue; - }; - - // Pop first element - let Some(common_prefix) = parts.next() else { - continue; - }; + fn list_with_delimiter<'life0, 'life1, 'async_trait>( + &'life0 self, + prefix: Option<&'life1 Path>, + ) -> Pin< + Box> + Send + 'async_trait>, + > + where + 'life0: 'async_trait, + 'life1: 'async_trait, + Self: 'async_trait, + { + Box::pin(async move { + let root = Path::default(); + let prefix = prefix.unwrap_or(&root); + + let mut common_prefixes = BTreeSet::new(); + let mut objects = vec![]; + + for k in &self.files { + let Some(mut parts) = k.prefix_match(prefix) else { + continue; + }; - if parts.next().is_some() { - common_prefixes.insert(prefix.clone().join(common_prefix)); - } else { - let object = ObjectMeta { - location: k.clone(), - last_modified: Utc.timestamp_nanos(0), - size: self.file_size, - e_tag: None, - version: None, + // Pop first element + let Some(common_prefix) = parts.next() else { + continue; }; - objects.push(object); + + if parts.next().is_some() { + common_prefixes.insert(prefix.clone().join(common_prefix)); + } else { + let object = ObjectMeta { + location: k.clone(), + last_modified: Utc.timestamp_nanos(0), + size: self.file_size, + e_tag: None, + version: None, + }; + objects.push(object); + } } - } - Ok(ListResult { - common_prefixes: common_prefixes.into_iter().collect(), - objects, + Ok(ListResult { + common_prefixes: common_prefixes.into_iter().collect(), + objects, + }) }) } @@ -796,12 +834,18 @@ impl ObjectStore for MirroringObjectStore { unimplemented!() } - async fn copy_opts( - &self, - _from: &Path, - _to: &Path, + fn copy_opts<'life0, 'life1, 'life2, 'async_trait>( + &'life0 self, + _from: &'life1 Path, + _to: &'life2 Path, _options: CopyOptions, - ) -> object_store::Result<()> { - unimplemented!() + ) -> Pin> + Send + 'async_trait>> + where + 'life0: 'async_trait, + 'life1: 'async_trait, + 'life2: 'async_trait, + Self: 'async_trait, + { + Box::pin(async move { unimplemented!() }) } } diff --git a/datafusion/core/tests/tracing/traceable_object_store.rs b/datafusion/core/tests/tracing/traceable_object_store.rs index 71a61dbf8772a..6b9b33a31ee38 100644 --- a/datafusion/core/tests/tracing/traceable_object_store.rs +++ b/datafusion/core/tests/tracing/traceable_object_store.rs @@ -25,6 +25,8 @@ use object_store::{ ObjectStore, PutMultipartOptions, PutOptions, PutPayload, PutResult, path::Path, }; use std::fmt::{Debug, Display, Formatter}; +use std::future::Future; +use std::pin::Pin; use std::sync::Arc; /// Returns an `ObjectStore` that asserts it can trace its calls back to the root tokio task. @@ -54,34 +56,64 @@ impl Display for TraceableObjectStore { /// All trait methods are forwarded to the inner object store, /// after asserting they can trace their calls back to the root tokio task. -#[async_trait::async_trait] impl ObjectStore for TraceableObjectStore { - async fn put_opts( - &self, - location: &Path, + fn put_opts<'life0, 'life1, 'async_trait>( + &'life0 self, + location: &'life1 Path, payload: PutPayload, opts: PutOptions, - ) -> object_store::Result { - assert_traceability().await; - self.inner.put_opts(location, payload, opts).await + ) -> Pin< + Box> + Send + 'async_trait>, + > + where + 'life0: 'async_trait, + 'life1: 'async_trait, + Self: 'async_trait, + { + Box::pin(async move { + assert_traceability().await; + self.inner.put_opts(location, payload, opts).await + }) } - async fn put_multipart_opts( - &self, - location: &Path, + fn put_multipart_opts<'life0, 'life1, 'async_trait>( + &'life0 self, + location: &'life1 Path, opts: PutMultipartOptions, - ) -> object_store::Result> { - assert_traceability().await; - self.inner.put_multipart_opts(location, opts).await + ) -> Pin< + Box< + dyn Future>> + + Send + + 'async_trait, + >, + > + where + 'life0: 'async_trait, + 'life1: 'async_trait, + Self: 'async_trait, + { + Box::pin(async move { + assert_traceability().await; + self.inner.put_multipart_opts(location, opts).await + }) } - async fn get_opts( - &self, - location: &Path, + fn get_opts<'life0, 'life1, 'async_trait>( + &'life0 self, + location: &'life1 Path, options: GetOptions, - ) -> object_store::Result { - assert_traceability().await; - self.inner.get_opts(location, options).await + ) -> Pin< + Box> + Send + 'async_trait>, + > + where + 'life0: 'async_trait, + 'life1: 'async_trait, + Self: 'async_trait, + { + Box::pin(async move { + assert_traceability().await; + self.inner.get_opts(location, options).await + }) } fn delete_stream( @@ -105,21 +137,38 @@ impl ObjectStore for TraceableObjectStore { self.inner.list(prefix) } - async fn list_with_delimiter( - &self, - prefix: Option<&Path>, - ) -> object_store::Result { - assert_traceability().await; - self.inner.list_with_delimiter(prefix).await + fn list_with_delimiter<'life0, 'life1, 'async_trait>( + &'life0 self, + prefix: Option<&'life1 Path>, + ) -> Pin< + Box> + Send + 'async_trait>, + > + where + 'life0: 'async_trait, + 'life1: 'async_trait, + Self: 'async_trait, + { + Box::pin(async move { + assert_traceability().await; + self.inner.list_with_delimiter(prefix).await + }) } - async fn copy_opts( - &self, - from: &Path, - to: &Path, + fn copy_opts<'life0, 'life1, 'life2, 'async_trait>( + &'life0 self, + from: &'life1 Path, + to: &'life2 Path, options: CopyOptions, - ) -> object_store::Result<()> { - assert_traceability().await; - self.inner.copy_opts(from, to, options).await + ) -> Pin> + Send + 'async_trait>> + where + 'life0: 'async_trait, + 'life1: 'async_trait, + 'life2: 'async_trait, + Self: 'async_trait, + { + Box::pin(async move { + assert_traceability().await; + self.inner.copy_opts(from, to, options).await + }) } } diff --git a/datafusion/core/tests/user_defined/insert_operation.rs b/datafusion/core/tests/user_defined/insert_operation.rs index 49d0ecc1aeb23..cc515f5d6512a 100644 --- a/datafusion/core/tests/user_defined/insert_operation.rs +++ b/datafusion/core/tests/user_defined/insert_operation.rs @@ -15,10 +15,10 @@ // specific language governing permissions and limitations // under the License. +use futures::future::BoxFuture; use std::{str::FromStr, sync::Arc}; use arrow::datatypes::{DataType, Field, Schema, SchemaRef}; -use async_trait::async_trait; use datafusion::{ error::Result, prelude::{SessionConfig, SessionContext}, @@ -87,8 +87,6 @@ impl TestInsertTableProvider { } } } - -#[async_trait] impl TableProvider for TestInsertTableProvider { fn schema(&self) -> SchemaRef { self.schema.clone() @@ -98,23 +96,27 @@ impl TableProvider for TestInsertTableProvider { TableType::Base } - async fn scan( - &self, - _state: &dyn Session, - _projection: Option<&[usize]>, - _filters: &[Expr], + fn scan<'a>( + &'a self, + _state: &'a dyn Session, + _projection: Option<&'a [usize]>, + _filters: &'a [Expr], _limit: Option, - ) -> Result> { - unimplemented!("TestInsertTableProvider is a stub for testing.") + ) -> BoxFuture<'a, Result>> { + Box::pin(async move { + unimplemented!("TestInsertTableProvider is a stub for testing.") + }) } - async fn insert_into( - &self, - _state: &dyn Session, + fn insert_into<'a>( + &'a self, + _state: &'a dyn Session, _input: Arc, insert_op: InsertOp, - ) -> Result> { - Ok(Arc::new(TestInsertExec::new(insert_op))) + ) -> BoxFuture<'a, Result>> { + Box::pin(async move { + Ok(Arc::new(TestInsertExec::new(insert_op)) as Arc) + }) } } diff --git a/datafusion/core/tests/user_defined/statistics_requests.rs b/datafusion/core/tests/user_defined/statistics_requests.rs index 079639baf1db3..56db059dcdaa7 100644 --- a/datafusion/core/tests/user_defined/statistics_requests.rs +++ b/datafusion/core/tests/user_defined/statistics_requests.rs @@ -24,11 +24,11 @@ //! plays both roles, demonstrating that the request-side hooks are //! sufficient to build the whole feature outside of DataFusion. +use futures::future::BoxFuture; use std::sync::{Arc, Mutex}; use arrow::array::{Int64Array, RecordBatch}; use arrow::datatypes::{DataType, Field, Schema, SchemaRef}; -use async_trait::async_trait; use datafusion::catalog::{ScanArgs, ScanResult, Session, TableProvider}; use datafusion::common::tree_node::Transformed; use datafusion::common::{Column, Result}; @@ -99,8 +99,6 @@ struct RecordingTable { batch: RecordBatch, last_requests: Arc>>, } - -#[async_trait] impl TableProvider for RecordingTable { fn schema(&self) -> SchemaRef { Arc::clone(&self.schema) @@ -110,36 +108,40 @@ impl TableProvider for RecordingTable { TableType::Base } - async fn scan( - &self, - _state: &dyn Session, - projection: Option<&[usize]>, - _filters: &[Expr], + fn scan<'a>( + &'a self, + _state: &'a dyn Session, + projection: Option<&'a [usize]>, + _filters: &'a [Expr], _limit: Option, - ) -> Result> { - Ok(MemorySourceConfig::try_new_exec( - &[vec![self.batch.clone()]], - Arc::clone(&self.schema), - projection.map(|p| p.to_vec()), - )?) + ) -> BoxFuture<'a, Result>> { + Box::pin(async move { + Ok(MemorySourceConfig::try_new_exec( + &[vec![self.batch.clone()]], + Arc::clone(&self.schema), + projection.map(|p| p.to_vec()), + )? as Arc) + }) } - async fn scan_with_args<'a>( - &self, - state: &dyn Session, + fn scan_with_args<'a>( + &'a self, + state: &'a dyn Session, args: ScanArgs<'a>, - ) -> Result { - // Record what reached us, then delegate to `scan`. - *self.last_requests.lock().unwrap() = args.statistics_requests().to_vec(); - let plan = self - .scan( - state, - args.projection(), - args.filters().unwrap_or(&[]), - args.limit(), - ) - .await?; - Ok(ScanResult::new(plan)) + ) -> BoxFuture<'a, Result> { + Box::pin(async move { + // Record what reached us, then delegate to `scan`. + *self.last_requests.lock().unwrap() = args.statistics_requests().to_vec(); + let plan = self + .scan( + state, + args.projection(), + args.filters().unwrap_or(&[]), + args.limit(), + ) + .await?; + Ok(ScanResult::new(plan)) + }) } } diff --git a/datafusion/core/tests/user_defined/user_defined_async_scalar_functions.rs b/datafusion/core/tests/user_defined/user_defined_async_scalar_functions.rs index 150801e38f49a..40bb9365b0c6d 100644 --- a/datafusion/core/tests/user_defined/user_defined_async_scalar_functions.rs +++ b/datafusion/core/tests/user_defined/user_defined_async_scalar_functions.rs @@ -15,11 +15,11 @@ // specific language governing permissions and limitations // under the License. +use futures::future::BoxFuture; use std::sync::Arc; use arrow::array::{Int32Array, RecordBatch, StringArray}; use arrow::datatypes::{DataType, Field, FieldRef, Schema}; -use async_trait::async_trait; use datafusion::prelude::*; use datafusion_common::test_util::format_batches; use datafusion_common::{Result, assert_batches_eq}; @@ -158,14 +158,12 @@ async fn test_async_udf_preserves_result_field_metadata() -> Result<()> { panic!("Call invoke_async_with_args instead") } } - - #[async_trait] impl AsyncScalarUDFImpl for AsyncExtensionUDF { - async fn invoke_async_with_args( + fn invoke_async_with_args( &self, args: ScalarFunctionArgs, - ) -> Result { - Ok(args.args[0].clone()) + ) -> BoxFuture<'_, Result> { + Box::pin(async move { Ok(args.args[0].clone()) }) } } @@ -250,19 +248,19 @@ impl ScalarUDFImpl for TestAsyncUDFImpl { panic!("Call invoke_async_with_args instead") } } - -#[async_trait] impl AsyncScalarUDFImpl for TestAsyncUDFImpl { fn ideal_batch_size(&self) -> Option { Some(self.batch_size) } - async fn invoke_async_with_args( + fn invoke_async_with_args( &self, args: ScalarFunctionArgs, - ) -> Result { - let arg1 = &args.args[0]; - let results = call_external_service(arg1.clone()).await?; - Ok(results) + ) -> BoxFuture<'_, Result> { + Box::pin(async move { + let arg1 = &args.args[0]; + let results = call_external_service(arg1.clone()).await?; + Ok(results) + }) } } diff --git a/datafusion/core/tests/user_defined/user_defined_plan.rs b/datafusion/core/tests/user_defined/user_defined_plan.rs index a1a8e3aa148cd..ad46026c9bfe0 100644 --- a/datafusion/core/tests/user_defined/user_defined_plan.rs +++ b/datafusion/core/tests/user_defined/user_defined_plan.rs @@ -57,6 +57,7 @@ //! The same answer can be produced by simply keeping track of the top //! N elements, reducing the total amount of required buffer memory. +use futures::future::BoxFuture; use std::fmt::Debug; use std::hash::Hash; use std::task::{Context, Poll}; @@ -101,7 +102,6 @@ use datafusion_optimizer::optimizer::ApplyOrder; use datafusion_physical_plan::execution_plan::{Boundedness, EmissionType}; use datafusion_physical_plan::{ChildrenPropertiesMode, ReplaceChildrenOptions}; -use async_trait::async_trait; use datafusion_common::cast::as_string_view_array; use datafusion_expr::physical_planning_context::PhysicalPlanningContext; use futures::{Stream, StreamExt}; @@ -463,25 +463,25 @@ impl OptimizerRule for OptimizerMakeExtensionNodeInvalid { #[derive(Debug)] struct TopKQueryPlanner {} - -#[async_trait] impl QueryPlanner for TopKQueryPlanner { /// Given a `LogicalPlan` created from above, create an /// `ExecutionPlan` suitable for execution - async fn create_physical_plan( - &self, - logical_plan: &LogicalPlan, - session_state: &dyn Session, - ) -> Result> { - // Teach the default physical planner how to plan TopK nodes. - let physical_planner = - DefaultPhysicalPlanner::with_extension_planners(vec![Arc::new( - TopKPlanner {}, - )]); - // Delegate most work of physical planning to the default physical planner - physical_planner - .create_physical_plan(logical_plan, session_state) - .await + fn create_physical_plan<'a>( + &'a self, + logical_plan: &'a LogicalPlan, + session_state: &'a dyn Session, + ) -> BoxFuture<'a, Result>> { + Box::pin(async move { + // Teach the default physical planner how to plan TopK nodes. + let physical_planner = + DefaultPhysicalPlanner::with_extension_planners(vec![Arc::new( + TopKPlanner {}, + )]); + // Delegate most work of physical planning to the default physical planner + physical_planner + .create_physical_plan(logical_plan, session_state) + .await + }) } } @@ -624,32 +624,32 @@ impl UserDefinedLogicalNodeCore for TopKPlanNode { /// Physical planner for TopK nodes struct TopKPlanner {} - -#[async_trait] impl ExtensionPlanner for TopKPlanner { /// Create a physical plan for an extension node - async fn plan_extension( - &self, - _planner: &dyn PhysicalPlanner, - node: &dyn UserDefinedLogicalNode, - logical_inputs: &[&LogicalPlan], - physical_inputs: &[Arc], - _session_state: &dyn Session, - _planning_ctx: &PhysicalPlanningContext, - ) -> Result>> { - Ok( - if let Some(topk_node) = node.as_any().downcast_ref::() { - assert_eq!(logical_inputs.len(), 1, "Inconsistent number of inputs"); - assert_eq!(physical_inputs.len(), 1, "Inconsistent number of inputs"); - // figure out input name - Some(Arc::new(TopKExec::new( - physical_inputs[0].clone(), - topk_node.k, - ))) - } else { - None - }, - ) + fn plan_extension<'a>( + &'a self, + _planner: &'a dyn PhysicalPlanner, + node: &'a dyn UserDefinedLogicalNode, + logical_inputs: &'a [&'a LogicalPlan], + physical_inputs: &'a [Arc], + _session_state: &'a dyn Session, + _planning_ctx: &'a PhysicalPlanningContext, + ) -> BoxFuture<'a, Result>>> { + Box::pin(async move { + Ok( + if let Some(topk_node) = node.as_any().downcast_ref::() { + assert_eq!(logical_inputs.len(), 1, "Inconsistent number of inputs"); + assert_eq!(physical_inputs.len(), 1, "Inconsistent number of inputs"); + // figure out input name + Some( + Arc::new(TopKExec::new(physical_inputs[0].clone(), topk_node.k)) + as Arc, + ) + } else { + None + }, + ) + }) } } @@ -702,8 +702,6 @@ impl DisplayAs for TopKExec { } } } - -#[async_trait] impl ExecutionPlan for TopKExec { fn name(&self) -> &'static str { Self::static_name() diff --git a/datafusion/core/tests/user_defined/user_defined_scalar_functions.rs b/datafusion/core/tests/user_defined/user_defined_scalar_functions.rs index 36d094cd0360b..2448d73051244 100644 --- a/datafusion/core/tests/user_defined/user_defined_scalar_functions.rs +++ b/datafusion/core/tests/user_defined/user_defined_scalar_functions.rs @@ -51,6 +51,7 @@ use datafusion_expr::{ use datafusion_expr_common::signature::Coercion; use datafusion_expr_common::signature::TypeSignature; use datafusion_functions_nested::range::range_udf; +use futures::future::BoxFuture; use parking_lot::Mutex; use regex::Regex; use sqlparser::ast::Ident; @@ -909,16 +910,17 @@ async fn verify_udf_return_type() -> Result<()> { #[derive(Debug, Default)] struct CustomFunctionFactory {} -#[async_trait::async_trait] impl FunctionFactory for CustomFunctionFactory { - async fn create( - &self, - _state: &SessionState, + fn create<'a>( + &'a self, + _state: &'a SessionState, statement: CreateFunction, - ) -> Result { - let f: ScalarFunctionWrapper = statement.try_into()?; + ) -> BoxFuture<'a, Result> { + Box::pin(async move { + let f: ScalarFunctionWrapper = statement.try_into()?; - Ok(RegisterFunction::Scalar(Arc::new(ScalarUDF::from(f)))) + Ok(RegisterFunction::Scalar(Arc::new(ScalarUDF::from(f)))) + }) } } // a wrapper type to be used to register @@ -1314,17 +1316,18 @@ impl RecordingFunctionFactory { } } -#[async_trait::async_trait] impl FunctionFactory for RecordingFunctionFactory { - async fn create( - &self, - _state: &SessionState, + fn create<'a>( + &'a self, + _state: &'a SessionState, statement: CreateFunction, - ) -> Result { - self.calls.lock().push(statement); + ) -> BoxFuture<'a, Result> { + Box::pin(async move { + self.calls.lock().push(statement); - let udf = range_udf(); - Ok(RegisterFunction::Scalar(udf)) + let udf = range_udf(); + Ok(RegisterFunction::Scalar(udf)) + }) } } diff --git a/datafusion/core/tests/user_defined/user_defined_table_functions.rs b/datafusion/core/tests/user_defined/user_defined_table_functions.rs index 24205cf8c4010..15a8b4fbc54d4 100644 --- a/datafusion/core/tests/user_defined/user_defined_table_functions.rs +++ b/datafusion/core/tests/user_defined/user_defined_table_functions.rs @@ -15,6 +15,7 @@ // specific language governing permissions and limitations // under the License. +use futures::future::BoxFuture; use std::fs::File; use std::io::Seek; use std::path::Path; @@ -38,8 +39,6 @@ use datafusion_catalog::{Session, TableFunctionArgs}; use datafusion_common::{DFSchema, ScalarValue}; use datafusion_expr::{EmptyRelation, Expr, LogicalPlan, Projection, TableType}; -use async_trait::async_trait; - /// test simple udtf with define read_csv with parameters #[tokio::test] async fn test_simple_read_csv_udtf() -> Result<()> { @@ -115,8 +114,6 @@ struct SimpleCsvTable { exprs: Vec, batches: Vec, } - -#[async_trait] impl TableProvider for SimpleCsvTable { fn schema(&self) -> SchemaRef { self.schema.clone() @@ -126,38 +123,40 @@ impl TableProvider for SimpleCsvTable { TableType::Base } - async fn scan( - &self, - state: &dyn Session, - projection: Option<&[usize]>, - _filters: &[Expr], + fn scan<'a>( + &'a self, + state: &'a dyn Session, + projection: Option<&'a [usize]>, + _filters: &'a [Expr], _limit: Option, - ) -> Result> { - let batches = if !self.exprs.is_empty() { - let max_return_lines = self.interpreter_expr(state).await?; - // get max return rows from self.batches - let mut batches = vec![]; - let mut lines = 0; - for batch in &self.batches { - let batch_lines = batch.num_rows(); - if lines + batch_lines > max_return_lines as usize { - let batch_lines = max_return_lines as usize - lines; - batches.push(batch.slice(0, batch_lines)); - break; - } else { - batches.push(batch.clone()); - lines += batch_lines; + ) -> BoxFuture<'a, Result>> { + Box::pin(async move { + let batches = if !self.exprs.is_empty() { + let max_return_lines = self.interpreter_expr(state).await?; + // get max return rows from self.batches + let mut batches = vec![]; + let mut lines = 0; + for batch in &self.batches { + let batch_lines = batch.num_rows(); + if lines + batch_lines > max_return_lines as usize { + let batch_lines = max_return_lines as usize - lines; + batches.push(batch.slice(0, batch_lines)); + break; + } else { + batches.push(batch.clone()); + lines += batch_lines; + } } - } - batches - } else { - self.batches.clone() - }; - Ok(MemorySourceConfig::try_new_exec( - &[batches], - TableProvider::schema(self), - projection.map(|p| p.to_vec()), - )?) + batches + } else { + self.batches.clone() + }; + Ok(MemorySourceConfig::try_new_exec( + &[batches], + TableProvider::schema(self), + projection.map(|p| p.to_vec()), + )? as Arc) + }) } } diff --git a/datafusion/datasource-arrow/Cargo.toml b/datafusion/datasource-arrow/Cargo.toml index 6f50135403d69..fe3e42e517662 100644 --- a/datafusion/datasource-arrow/Cargo.toml +++ b/datafusion/datasource-arrow/Cargo.toml @@ -33,7 +33,6 @@ all-features = true [dependencies] arrow = { workspace = true } arrow-ipc = { workspace = true } -async-trait = { workspace = true } bytes = { workspace = true } datafusion-common = { workspace = true, features = ["object_store"] } datafusion-common-runtime = { workspace = true } diff --git a/datafusion/datasource-arrow/src/file_format.rs b/datafusion/datasource-arrow/src/file_format.rs index 2bee57ef17581..9a9dc64e8b8a2 100644 --- a/datafusion/datasource-arrow/src/file_format.rs +++ b/datafusion/datasource-arrow/src/file_format.rs @@ -19,6 +19,7 @@ //! //! Works with files following the [Arrow IPC format](https://arrow.apache.org/docs/format/Columnar.html#ipc-file-format) +use futures::future::BoxFuture; use std::collections::HashMap; use std::fmt::{self, Debug}; use std::io::{Seek, SeekFrom}; @@ -50,7 +51,6 @@ use datafusion_expr::dml::InsertOp; use datafusion_physical_expr_common::sort_expr::LexRequirement; use crate::source::ArrowSource; -use async_trait::async_trait; use bytes::Bytes; use datafusion_datasource::file_compression_type::FileCompressionType; use datafusion_datasource::file_format::{FileFormat, FileFormatFactory}; @@ -109,8 +109,6 @@ impl GetExt for ArrowFormatFactory { /// Arrow [`FileFormat`] implementation. #[derive(Default, Debug)] pub struct ArrowFormat; - -#[async_trait] impl FileFormat for ArrowFormat { fn get_ext(&self) -> String { ArrowFormatFactory::new().get_ext() @@ -133,110 +131,124 @@ impl FileFormat for ArrowFormat { None } - async fn infer_schema( - &self, - _state: &dyn Session, - store: &Arc, - objects: &[ObjectMeta], - ) -> Result { - let mut schemas = vec![]; - for object in objects { - let r = store.as_ref().get(&object.location).await?; - let schema = match r.payload { - #[cfg(not(target_arch = "wasm32"))] - GetResultPayload::File(mut file, _) => { - match FileReader::try_new(&mut file, None) { - Ok(reader) => reader.schema(), - Err(file_error) => { - // not in the file format, but FileReader read some bytes - // while trying to parse the file and so we need to rewind - // it to the beginning of the file - file.seek(SeekFrom::Start(0))?; - match StreamReader::try_new(&mut file, None) { - Ok(reader) => reader.schema(), - Err(stream_error) => { - return Err(internal_datafusion_err!( - "Failed to parse Arrow file as either file format or stream format. File format error: {file_error}. Stream format error: {stream_error}" - )); + fn infer_schema<'a>( + &'a self, + _state: &'a dyn Session, + store: &'a Arc, + objects: &'a [ObjectMeta], + ) -> BoxFuture<'a, Result> { + Box::pin(async move { + let mut schemas = vec![]; + for object in objects { + let r = store.as_ref().get(&object.location).await?; + let schema = match r.payload { + #[cfg(not(target_arch = "wasm32"))] + GetResultPayload::File(mut file, _) => { + match FileReader::try_new(&mut file, None) { + Ok(reader) => reader.schema(), + Err(file_error) => { + // not in the file format, but FileReader read some bytes + // while trying to parse the file and so we need to rewind + // it to the beginning of the file + file.seek(SeekFrom::Start(0))?; + match StreamReader::try_new(&mut file, None) { + Ok(reader) => reader.schema(), + Err(stream_error) => { + return Err(internal_datafusion_err!( + "Failed to parse Arrow file as either file format or stream format. File format error: {file_error}. Stream format error: {stream_error}" + )); + } } } } } - } - GetResultPayload::Stream(stream) => infer_stream_schema(stream).await?, - }; - schemas.push(Arc::unwrap_or_clone(schema)); - } - let merged_schema = Schema::try_merge(schemas)?; - Ok(Arc::new(merged_schema)) + GetResultPayload::Stream(stream) => { + infer_stream_schema(stream).await? + } + }; + schemas.push(Arc::unwrap_or_clone(schema)); + } + let merged_schema = Schema::try_merge(schemas)?; + Ok(Arc::new(merged_schema)) + }) } - async fn infer_stats( - &self, - _state: &dyn Session, - _store: &Arc, + fn infer_stats<'a>( + &'a self, + _state: &'a dyn Session, + _store: &'a Arc, table_schema: SchemaRef, - _object: &ObjectMeta, - ) -> Result { - Ok(Statistics::new_unknown(&table_schema)) + _object: &'a ObjectMeta, + ) -> BoxFuture<'a, Result> { + Box::pin(async move { Ok(Statistics::new_unknown(&table_schema)) }) } - async fn create_physical_plan( - &self, - state: &dyn Session, + fn create_physical_plan<'a>( + &'a self, + state: &'a dyn Session, conf: FileScanConfig, - ) -> Result> { - let object_store = state.runtime_env().object_store(&conf.object_store_url)?; - let object_location = &conf - .file_groups - .first() - .ok_or_else(|| internal_datafusion_err!("No files found in file group"))? - .files() - .first() - .ok_or_else(|| internal_datafusion_err!("No files found in file group"))? - .object_meta - .location; - - let table_schema = TableSchemaBuilder::from(conf.file_schema()) - .with_table_partition_cols(conf.table_partition_cols().clone()) - .build(); - - let mut source: Arc = - match is_object_in_arrow_ipc_file_format(object_store, object_location).await - { - Ok(true) => Arc::new(ArrowSource::new_file_source(table_schema)), - Ok(false) => Arc::new(ArrowSource::new_stream_file_source(table_schema)), - Err(e) => Err(e)?, - }; + ) -> BoxFuture<'a, Result>> { + Box::pin(async move { + let object_store = + state.runtime_env().object_store(&conf.object_store_url)?; + let object_location = &conf + .file_groups + .first() + .ok_or_else(|| internal_datafusion_err!("No files found in file group"))? + .files() + .first() + .ok_or_else(|| internal_datafusion_err!("No files found in file group"))? + .object_meta + .location; + + let table_schema = TableSchemaBuilder::from(conf.file_schema()) + .with_table_partition_cols(conf.table_partition_cols().clone()) + .build(); + + let mut source: Arc = + match is_object_in_arrow_ipc_file_format(object_store, object_location) + .await + { + Ok(true) => Arc::new(ArrowSource::new_file_source(table_schema)), + Ok(false) => { + Arc::new(ArrowSource::new_stream_file_source(table_schema)) + } + Err(e) => Err(e)?, + }; - // Preserve projection from the original file source - if let Some(projection) = conf.file_source.projection() - && let Some(new_source) = source.try_pushdown_projection(projection)? - { - source = new_source; - } + // Preserve projection from the original file source + if let Some(projection) = conf.file_source.projection() + && let Some(new_source) = source.try_pushdown_projection(projection)? + { + source = new_source; + } - let config = FileScanConfigBuilder::from(conf) - .with_source(source) - .build(); + let config = FileScanConfigBuilder::from(conf) + .with_source(source) + .build(); - Ok(DataSourceExec::from_data_source(config)) + Ok(DataSourceExec::from_data_source(config) as Arc) + }) } - async fn create_writer_physical_plan( - &self, + fn create_writer_physical_plan<'a>( + &'a self, input: Arc, - _state: &dyn Session, + _state: &'a dyn Session, conf: FileSinkConfig, order_requirements: Option, - ) -> Result> { - if conf.insert_op != InsertOp::Append { - return not_impl_err!("Overwrites are not implemented yet for Arrow format"); - } + ) -> BoxFuture<'a, Result>> { + Box::pin(async move { + if conf.insert_op != InsertOp::Append { + return not_impl_err!( + "Overwrites are not implemented yet for Arrow format" + ); + } - let sink = Arc::new(ArrowFileSink::new(conf)); + let sink = Arc::new(ArrowFileSink::new(conf)); - Ok(Arc::new(DataSinkExec::new(input, sink, order_requirements)) as _) + Ok(Arc::new(DataSinkExec::new(input, sink, order_requirements)) as _) + }) } fn file_source(&self, table_schema: TableSchema) -> Arc { @@ -254,89 +266,91 @@ impl ArrowFileSink { Self { config } } } - -#[async_trait] impl FileSink for ArrowFileSink { fn config(&self) -> &FileSinkConfig { &self.config } - async fn spawn_writer_tasks_and_join( - &self, - context: &Arc, + fn spawn_writer_tasks_and_join<'a>( + &'a self, + context: &'a Arc, demux_task: SpawnedTask>, mut file_stream_rx: DemuxedStreamReceiver, object_store: Arc, - ) -> Result { - let mut file_write_tasks: JoinSet> = - JoinSet::new(); - - let ipc_options = - IpcWriteOptions::try_new(64, false, arrow_ipc::MetadataVersion::V5)? - .try_with_compression(Some(CompressionType::LZ4_FRAME))?; - while let Some((path, mut rx)) = file_stream_rx.recv().await { - let shared_buffer = SharedBuffer::new(INITIAL_BUFFER_BYTES); - let mut arrow_writer = arrow_ipc::writer::FileWriter::try_new_with_options( - shared_buffer.clone(), - &get_writer_schema(&self.config), - ipc_options.clone(), - )?; - let mut object_store_writer = ObjectWriterBuilder::new( - FileCompressionType::UNCOMPRESSED, - &path, - Arc::clone(&object_store), - ) - .with_buffer_size(Some( - context - .session_config() - .options() - .execution - .objectstore_writer_buffer_size, - )) - .build()?; - file_write_tasks.spawn(async move { - let mut row_count = 0; - while let Some(batch) = rx.recv().await { - row_count += batch.num_rows(); - arrow_writer.write(&batch)?; - let mut buff_to_flush = shared_buffer.buffer.try_lock().unwrap(); - if buff_to_flush.len() > BUFFER_FLUSH_BYTES { - object_store_writer - .write_all(buff_to_flush.as_slice()) - .await?; - buff_to_flush.clear(); + ) -> BoxFuture<'a, Result> { + Box::pin(async move { + let mut file_write_tasks: JoinSet< + std::result::Result, + > = JoinSet::new(); + + let ipc_options = + IpcWriteOptions::try_new(64, false, arrow_ipc::MetadataVersion::V5)? + .try_with_compression(Some(CompressionType::LZ4_FRAME))?; + while let Some((path, mut rx)) = file_stream_rx.recv().await { + let shared_buffer = SharedBuffer::new(INITIAL_BUFFER_BYTES); + let mut arrow_writer = + arrow_ipc::writer::FileWriter::try_new_with_options( + shared_buffer.clone(), + &get_writer_schema(&self.config), + ipc_options.clone(), + )?; + let mut object_store_writer = ObjectWriterBuilder::new( + FileCompressionType::UNCOMPRESSED, + &path, + Arc::clone(&object_store), + ) + .with_buffer_size(Some( + context + .session_config() + .options() + .execution + .objectstore_writer_buffer_size, + )) + .build()?; + file_write_tasks.spawn(async move { + let mut row_count = 0; + while let Some(batch) = rx.recv().await { + row_count += batch.num_rows(); + arrow_writer.write(&batch)?; + let mut buff_to_flush = shared_buffer.buffer.try_lock().unwrap(); + if buff_to_flush.len() > BUFFER_FLUSH_BYTES { + object_store_writer + .write_all(buff_to_flush.as_slice()) + .await?; + buff_to_flush.clear(); + } } - } - arrow_writer.finish()?; - let final_buff = shared_buffer.buffer.try_lock().unwrap(); + arrow_writer.finish()?; + let final_buff = shared_buffer.buffer.try_lock().unwrap(); - object_store_writer.write_all(final_buff.as_slice()).await?; - object_store_writer.shutdown().await?; - Ok(row_count) - }); - } + object_store_writer.write_all(final_buff.as_slice()).await?; + object_store_writer.shutdown().await?; + Ok(row_count) + }); + } - let mut row_count = 0; - while let Some(result) = file_write_tasks.join_next().await { - match result { - Ok(r) => { - row_count += r?; - } - Err(e) => { - if e.is_panic() { - std::panic::resume_unwind(e.into_panic()); - } else { - unreachable!(); + let mut row_count = 0; + while let Some(result) = file_write_tasks.join_next().await { + match result { + Ok(r) => { + row_count += r?; + } + Err(e) => { + if e.is_panic() { + std::panic::resume_unwind(e.into_panic()); + } else { + unreachable!(); + } } } } - } - demux_task - .join_unwind() - .await - .map_err(|e| DataFusionError::ExecutionJoin(Box::new(e)))??; - Ok(row_count as u64) + demux_task + .join_unwind() + .await + .map_err(|e| DataFusionError::ExecutionJoin(Box::new(e)))??; + Ok(row_count as u64) + }) } } @@ -361,19 +375,17 @@ impl DisplayAs for ArrowFileSink { } } } - -#[async_trait] impl DataSink for ArrowFileSink { fn schema(&self) -> &SchemaRef { self.config.output_schema() } - async fn write_all( - &self, + fn write_all<'a>( + &'a self, data: SendableRecordBatchStream, - context: &Arc, - ) -> Result { - FileSink::write_all(self, data, context).await + context: &'a Arc, + ) -> BoxFuture<'a, Result> { + Box::pin(async move { FileSink::write_all(self, data, context).await }) } } @@ -565,7 +577,6 @@ mod tests { } } - #[async_trait::async_trait] impl Session for MockSession { fn session_id(&self) -> &str { unimplemented!() @@ -579,10 +590,10 @@ mod tests { Arc::new(EmptyCatalogProviderList) } - async fn create_physical_plan( - &self, - _logical_plan: &LogicalPlan, - ) -> Result> { + fn create_physical_plan<'a>( + &'a self, + _logical_plan: &'a LogicalPlan, + ) -> BoxFuture<'a, Result>> { unimplemented!() } diff --git a/datafusion/datasource-avro/Cargo.toml b/datafusion/datasource-avro/Cargo.toml index 63e949e21def9..beee9d263dddb 100644 --- a/datafusion/datasource-avro/Cargo.toml +++ b/datafusion/datasource-avro/Cargo.toml @@ -42,7 +42,6 @@ proto = [ [dependencies] arrow = { workspace = true } arrow-avro = { workspace = true } -async-trait = { workspace = true } bytes = { workspace = true } datafusion-common = { workspace = true, features = ["object_store"] } datafusion-common-runtime = { workspace = true } diff --git a/datafusion/datasource-avro/src/file_format.rs b/datafusion/datasource-avro/src/file_format.rs index 93dad800c6f0d..d85a8a624a8cd 100644 --- a/datafusion/datasource-avro/src/file_format.rs +++ b/datafusion/datasource-avro/src/file_format.rs @@ -16,6 +16,7 @@ // under the License. //! Apache Avro [`FileFormat`] abstractions +use futures::future::BoxFuture; use std::collections::HashMap; use std::fmt::Debug; use std::sync::Arc; @@ -54,7 +55,6 @@ use datafusion_physical_expr_common::sort_expr::LexRequirement; use datafusion_physical_plan::{DisplayAs, DisplayFormatType, ExecutionPlan}; use datafusion_session::Session; -use async_trait::async_trait; use object_store::{GetResultPayload, ObjectMeta, ObjectStore, ObjectStoreExt}; use tokio::io::AsyncWriteExt; @@ -106,8 +106,6 @@ impl GetExt for AvroFormatFactory { /// Avro [`FileFormat`] implementation. #[derive(Default, Debug)] pub struct AvroFormat; - -#[async_trait] impl FileFormat for AvroFormat { fn get_ext(&self) -> String { AvroFormatFactory::new().get_ext() @@ -128,63 +126,71 @@ impl FileFormat for AvroFormat { None } - async fn infer_schema( - &self, - _state: &dyn Session, - store: &Arc, - objects: &[ObjectMeta], - ) -> Result { - let mut schemas = vec![]; - for object in objects { - let r = store.as_ref().get(&object.location).await?; - let schema = match r.payload { - GetResultPayload::File(mut file, _) => { - read_avro_schema_from_reader(&mut file)? - } - GetResultPayload::Stream(_) => { - // TODO: Fetching entire file to get schema is potentially wasteful - let data = r.bytes().await?; - read_avro_schema_from_reader(&mut data.as_ref())? - } - }; - schemas.push(schema); - } - let merged_schema = Schema::try_merge(schemas)?; - Ok(Arc::new(merged_schema)) + fn infer_schema<'a>( + &'a self, + _state: &'a dyn Session, + store: &'a Arc, + objects: &'a [ObjectMeta], + ) -> BoxFuture<'a, Result> { + Box::pin(async move { + let mut schemas = vec![]; + for object in objects { + let r = store.as_ref().get(&object.location).await?; + let schema = match r.payload { + GetResultPayload::File(mut file, _) => { + read_avro_schema_from_reader(&mut file)? + } + GetResultPayload::Stream(_) => { + // TODO: Fetching entire file to get schema is potentially wasteful + let data = r.bytes().await?; + read_avro_schema_from_reader(&mut data.as_ref())? + } + }; + schemas.push(schema); + } + let merged_schema = Schema::try_merge(schemas)?; + Ok(Arc::new(merged_schema)) + }) } - async fn infer_stats( - &self, - _state: &dyn Session, - _store: &Arc, + fn infer_stats<'a>( + &'a self, + _state: &'a dyn Session, + _store: &'a Arc, table_schema: SchemaRef, - _object: &ObjectMeta, - ) -> Result { - Ok(Statistics::new_unknown(&table_schema)) + _object: &'a ObjectMeta, + ) -> BoxFuture<'a, Result> { + Box::pin(async move { Ok(Statistics::new_unknown(&table_schema)) }) } - async fn create_physical_plan( - &self, - _state: &dyn Session, + fn create_physical_plan<'a>( + &'a self, + _state: &'a dyn Session, conf: FileScanConfig, - ) -> Result> { - Ok(DataSourceExec::from_data_source(conf)) + ) -> BoxFuture<'a, Result>> { + Box::pin(async move { + Ok(DataSourceExec::from_data_source(conf) as Arc) + }) } - async fn create_writer_physical_plan( - &self, + fn create_writer_physical_plan<'a>( + &'a self, input: Arc, - _state: &dyn Session, + _state: &'a dyn Session, conf: FileSinkConfig, order_requirements: Option, - ) -> Result> { - if conf.insert_op != InsertOp::Append { - return not_impl_err!("Overwrites are not implemented yet for Avro format"); - } + ) -> BoxFuture<'a, Result>> { + Box::pin(async move { + if conf.insert_op != InsertOp::Append { + return not_impl_err!( + "Overwrites are not implemented yet for Avro format" + ); + } - let sink = Arc::new(AvroFileSink::new(conf)); + let sink = Arc::new(AvroFileSink::new(conf)); - Ok(Arc::new(DataSinkExec::new(input, sink, order_requirements)) as _) + Ok(Arc::new(DataSinkExec::new(input, sink, order_requirements)) as _) + }) } fn file_source( @@ -206,102 +212,105 @@ impl AvroFileSink { } } -#[async_trait] impl FileSink for AvroFileSink { fn config(&self) -> &FileSinkConfig { &self.config } - async fn spawn_writer_tasks_and_join( - &self, - context: &Arc, + fn spawn_writer_tasks_and_join<'a>( + &'a self, + context: &'a Arc, demux_task: SpawnedTask>, mut file_stream_rx: DemuxedStreamReceiver, object_store: Arc, - ) -> Result { - let mut file_write_tasks: JoinSet> = - JoinSet::new(); - - let writer_schema = get_writer_schema(&self.config); - while let Some((path, mut rx)) = file_stream_rx.recv().await { - let shared_buffer = SharedBuffer::new(INITIAL_BUFFER_BYTES); - let mut avro_writer: AvroWriter = - WriterBuilder::new(writer_schema.as_ref().clone()) - .build::<_, AvroOcfFormat>(shared_buffer.clone()) - .map_err(|e| { - internal_datafusion_err!("Failed to create Avro writer: {e}") - })?; - let mut object_store_writer = ObjectWriterBuilder::new( - FileCompressionType::UNCOMPRESSED, - &path, - Arc::clone(&object_store), - ) - .with_buffer_size(Some( - context - .session_config() - .options() - .execution - .objectstore_writer_buffer_size, - )) - .build()?; - file_write_tasks.spawn(async move { - let mut row_count = 0; - while let Some(batch) = rx.recv().await { - row_count += batch.num_rows(); - avro_writer - .write(&batch) - .map_err(|e| internal_datafusion_err!("{e}"))?; - let mut buff_to_flush = shared_buffer.buffer.try_lock().unwrap(); - if buff_to_flush.len() > BUFFER_FLUSH_BYTES { - object_store_writer - .write_all(buff_to_flush.as_slice()) - .await?; - buff_to_flush.clear(); + ) -> BoxFuture<'a, Result> { + Box::pin(async move { + let mut file_write_tasks: JoinSet< + std::result::Result, + > = JoinSet::new(); + + let writer_schema = get_writer_schema(&self.config); + while let Some((path, mut rx)) = file_stream_rx.recv().await { + let shared_buffer = SharedBuffer::new(INITIAL_BUFFER_BYTES); + let mut avro_writer: AvroWriter = + WriterBuilder::new(writer_schema.as_ref().clone()) + .build::<_, AvroOcfFormat>(shared_buffer.clone()) + .map_err(|e| { + internal_datafusion_err!("Failed to create Avro writer: {e}") + })?; + let mut object_store_writer = ObjectWriterBuilder::new( + FileCompressionType::UNCOMPRESSED, + &path, + Arc::clone(&object_store), + ) + .with_buffer_size(Some( + context + .session_config() + .options() + .execution + .objectstore_writer_buffer_size, + )) + .build()?; + file_write_tasks.spawn(async move { + let mut row_count = 0; + while let Some(batch) = rx.recv().await { + row_count += batch.num_rows(); + avro_writer + .write(&batch) + .map_err(|e| internal_datafusion_err!("{e}"))?; + let mut buff_to_flush = shared_buffer.buffer.try_lock().unwrap(); + if buff_to_flush.len() > BUFFER_FLUSH_BYTES { + object_store_writer + .write_all(buff_to_flush.as_slice()) + .await?; + buff_to_flush.clear(); + } } - } - if let Err(e) = avro_writer.finish() { - return Err(match e { - AvroError::NYI(e) => DataFusionError::NotImplemented(e), - AvroError::EOF(e) => DataFusionError::IoError(io::Error::new( - io::ErrorKind::UnexpectedEof, - e, - )), - AvroError::ArrowError(e) => DataFusionError::ArrowError(e, None), - AvroError::External(e) => DataFusionError::External(e), - AvroError::IoError(msg, e) => DataFusionError::IoError(e) - .with_diagnostic(Diagnostic::new_error(msg, None)), - _ => internal_datafusion_err!("{e}"), - }); - } - let final_buff = shared_buffer.buffer.try_lock().unwrap(); + if let Err(e) = avro_writer.finish() { + return Err(match e { + AvroError::NYI(e) => DataFusionError::NotImplemented(e), + AvroError::EOF(e) => DataFusionError::IoError( + io::Error::new(io::ErrorKind::UnexpectedEof, e), + ), + AvroError::ArrowError(e) => { + DataFusionError::ArrowError(e, None) + } + AvroError::External(e) => DataFusionError::External(e), + AvroError::IoError(msg, e) => DataFusionError::IoError(e) + .with_diagnostic(Diagnostic::new_error(msg, None)), + _ => internal_datafusion_err!("{e}"), + }); + } + let final_buff = shared_buffer.buffer.try_lock().unwrap(); - object_store_writer.write_all(final_buff.as_slice()).await?; - object_store_writer.shutdown().await?; - Ok(row_count) - }); - } + object_store_writer.write_all(final_buff.as_slice()).await?; + object_store_writer.shutdown().await?; + Ok(row_count) + }); + } - let mut row_count = 0; - while let Some(result) = file_write_tasks.join_next().await { - match result { - Ok(r) => { - row_count += r?; - } - Err(e) => { - if e.is_panic() { - std::panic::resume_unwind(e.into_panic()); - } else { - unreachable!(); + let mut row_count = 0; + while let Some(result) = file_write_tasks.join_next().await { + match result { + Ok(r) => { + row_count += r?; + } + Err(e) => { + if e.is_panic() { + std::panic::resume_unwind(e.into_panic()); + } else { + unreachable!(); + } } } } - } - demux_task - .join_unwind() - .await - .map_err(|e| DataFusionError::ExecutionJoin(Box::new(e)))??; - Ok(row_count as u64) + demux_task + .join_unwind() + .await + .map_err(|e| DataFusionError::ExecutionJoin(Box::new(e)))??; + Ok(row_count as u64) + }) } } @@ -327,17 +336,16 @@ impl DisplayAs for AvroFileSink { } } -#[async_trait] impl DataSink for AvroFileSink { fn schema(&self) -> &SchemaRef { self.config.output_schema() } - async fn write_all( - &self, + fn write_all<'a>( + &'a self, data: SendableRecordBatchStream, - context: &Arc, - ) -> Result { - FileSink::write_all(self, data, context).await + context: &'a Arc, + ) -> BoxFuture<'a, Result> { + Box::pin(async move { FileSink::write_all(self, data, context).await }) } } diff --git a/datafusion/datasource-csv/Cargo.toml b/datafusion/datasource-csv/Cargo.toml index 7e7195dfda9d5..1696b136bf768 100644 --- a/datafusion/datasource-csv/Cargo.toml +++ b/datafusion/datasource-csv/Cargo.toml @@ -40,7 +40,6 @@ proto = [ [dependencies] arrow = { workspace = true } -async-trait = { workspace = true } bytes = { workspace = true } datafusion-common = { workspace = true, features = ["object_store"] } datafusion-common-runtime = { workspace = true } diff --git a/datafusion/datasource-csv/src/file_format.rs b/datafusion/datasource-csv/src/file_format.rs index a8c27369cb9c1..191775082f194 100644 --- a/datafusion/datasource-csv/src/file_format.rs +++ b/datafusion/datasource-csv/src/file_format.rs @@ -17,6 +17,7 @@ //! [`CsvFormat`], Comma Separated Value (CSV) [`FileFormat`] abstractions +use futures::future::BoxFuture; use std::collections::{HashMap, HashSet}; use std::fmt::{self, Debug}; use std::sync::Arc; @@ -55,7 +56,6 @@ use datafusion_physical_expr_common::sort_expr::LexRequirement; use datafusion_physical_plan::{DisplayAs, DisplayFormatType, ExecutionPlan}; use datafusion_session::Session; -use async_trait::async_trait; use bytes::{Buf, Bytes}; use datafusion_datasource::source::DataSourceExec; use futures::stream::BoxStream; @@ -354,8 +354,6 @@ impl Debug for CsvSerializer { .finish() } } - -#[async_trait] impl FileFormat for CsvFormat { fn get_ext(&self) -> String { CsvFormatFactory::new().get_ext() @@ -373,133 +371,139 @@ impl FileFormat for CsvFormat { Some(self.options.compression.into()) } - async fn infer_schema( - &self, - state: &dyn Session, - store: &Arc, - objects: &[ObjectMeta], - ) -> Result { - let mut schemas = vec![]; - - let mut records_to_read = self - .options - .schema_infer_max_rec - .unwrap_or(DEFAULT_SCHEMA_INFER_MAX_RECORD); - - for object in objects { - let stream = self.read_to_delimited_chunks(store, object).await; - let (schema, records_read) = self - .infer_schema_from_stream(state, records_to_read, stream) - .await - .map_err(|err| { + fn infer_schema<'a>( + &'a self, + state: &'a dyn Session, + store: &'a Arc, + objects: &'a [ObjectMeta], + ) -> BoxFuture<'a, Result> { + Box::pin(async move { + let mut schemas = vec![]; + + let mut records_to_read = self + .options + .schema_infer_max_rec + .unwrap_or(DEFAULT_SCHEMA_INFER_MAX_RECORD); + + for object in objects { + let stream = self.read_to_delimited_chunks(store, object).await; + let (schema, records_read) = self + .infer_schema_from_stream(state, records_to_read, stream) + .await + .map_err(|err| { + DataFusionError::Context( + format!("Error when processing CSV file {}", object.location), + Box::new(err), + ) + })?; + records_to_read -= records_read; + schemas.push((&object.location, schema)); + if records_to_read == 0 { + break; + } + } + + let mut seen = HashSet::new(); + for (location, schema) in &schemas { + ensure_unique_field_names(schema, &mut seen).map_err(|err| { DataFusionError::Context( - format!("Error when processing CSV file {}", object.location), + format!("Error when processing CSV file {location}"), Box::new(err), ) })?; - records_to_read -= records_read; - schemas.push((&object.location, schema)); - if records_to_read == 0 { - break; } - } - - let mut seen = HashSet::new(); - for (location, schema) in &schemas { - ensure_unique_field_names(schema, &mut seen).map_err(|err| { - DataFusionError::Context( - format!("Error when processing CSV file {location}"), - Box::new(err), - ) - })?; - } - drop(seen); + drop(seen); - let schemas = schemas.into_iter().map(|(_, schema)| schema); - let merged_schema = Schema::try_merge(schemas)?; - Ok(Arc::new(merged_schema)) + let schemas = schemas.into_iter().map(|(_, schema)| schema); + let merged_schema = Schema::try_merge(schemas)?; + Ok(Arc::new(merged_schema)) + }) } - async fn infer_stats( - &self, - _state: &dyn Session, - _store: &Arc, + fn infer_stats<'a>( + &'a self, + _state: &'a dyn Session, + _store: &'a Arc, table_schema: SchemaRef, - _object: &ObjectMeta, - ) -> Result { - Ok(Statistics::new_unknown(&table_schema)) + _object: &'a ObjectMeta, + ) -> BoxFuture<'a, Result> { + Box::pin(async move { Ok(Statistics::new_unknown(&table_schema)) }) } - async fn create_physical_plan( - &self, - state: &dyn Session, + fn create_physical_plan<'a>( + &'a self, + state: &'a dyn Session, conf: FileScanConfig, - ) -> Result> { - // Consult configuration options for default values - let has_header = self - .options - .has_header - .unwrap_or_else(|| state.config_options().catalog.has_header); - let newlines_in_values = self - .options - .newlines_in_values - .unwrap_or_else(|| state.config_options().catalog.newlines_in_values); - - let mut csv_options = self.options.clone(); - csv_options.has_header = Some(has_header); - csv_options.newlines_in_values = Some(newlines_in_values); - - // Get the existing CsvSource and update its options - // We need to preserve the table_schema from the original source (which includes partition columns) - let csv_source = conf - .file_source - .downcast_ref::() - .expect("file_source should be a CsvSource"); - let source = Arc::new(csv_source.clone().with_csv_options(csv_options)); - - let config = FileScanConfigBuilder::from(conf) - .with_file_compression_type(self.options.compression.into()) - .with_source(source) - .build(); - - Ok(DataSourceExec::from_data_source(config)) + ) -> BoxFuture<'a, Result>> { + Box::pin(async move { + // Consult configuration options for default values + let has_header = self + .options + .has_header + .unwrap_or_else(|| state.config_options().catalog.has_header); + let newlines_in_values = self + .options + .newlines_in_values + .unwrap_or_else(|| state.config_options().catalog.newlines_in_values); + + let mut csv_options = self.options.clone(); + csv_options.has_header = Some(has_header); + csv_options.newlines_in_values = Some(newlines_in_values); + + // Get the existing CsvSource and update its options + // We need to preserve the table_schema from the original source (which includes partition columns) + let csv_source = conf + .file_source + .downcast_ref::() + .expect("file_source should be a CsvSource"); + let source = Arc::new(csv_source.clone().with_csv_options(csv_options)); + + let config = FileScanConfigBuilder::from(conf) + .with_file_compression_type(self.options.compression.into()) + .with_source(source) + .build(); + + Ok(DataSourceExec::from_data_source(config) as Arc) + }) } - async fn create_writer_physical_plan( - &self, + fn create_writer_physical_plan<'a>( + &'a self, input: Arc, - state: &dyn Session, + state: &'a dyn Session, conf: FileSinkConfig, order_requirements: Option, - ) -> Result> { - if conf.insert_op != InsertOp::Append { - return not_impl_err!("Overwrites are not implemented yet for CSV"); - } + ) -> BoxFuture<'a, Result>> { + Box::pin(async move { + if conf.insert_op != InsertOp::Append { + return not_impl_err!("Overwrites are not implemented yet for CSV"); + } - // `has_header` and `newlines_in_values` fields of CsvOptions may inherit - // their values from session from configuration settings. To support - // this logic, writer options are built from the copy of `self.options` - // with updated values of these special fields. - let has_header = self - .options() - .has_header - .unwrap_or_else(|| state.config_options().catalog.has_header); - let newlines_in_values = self - .options() - .newlines_in_values - .unwrap_or_else(|| state.config_options().catalog.newlines_in_values); - - let options = self - .options() - .clone() - .with_has_header(has_header) - .with_newlines_in_values(newlines_in_values); - - let writer_options = CsvWriterOptions::try_from(&options)?; - - let sink = Arc::new(CsvSink::new(conf, writer_options)); - - Ok(Arc::new(DataSinkExec::new(input, sink, order_requirements)) as _) + // `has_header` and `newlines_in_values` fields of CsvOptions may inherit + // their values from session from configuration settings. To support + // this logic, writer options are built from the copy of `self.options` + // with updated values of these special fields. + let has_header = self + .options() + .has_header + .unwrap_or_else(|| state.config_options().catalog.has_header); + let newlines_in_values = self + .options() + .newlines_in_values + .unwrap_or_else(|| state.config_options().catalog.newlines_in_values); + + let options = self + .options() + .clone() + .with_has_header(has_header) + .with_newlines_in_values(newlines_in_values); + + let writer_options = CsvWriterOptions::try_from(&options)?; + + let sink = Arc::new(CsvSink::new(conf, writer_options)); + + Ok(Arc::new(DataSinkExec::new(input, sink, order_requirements)) as _) + }) } fn file_source(&self, table_schema: TableSchema) -> Arc { @@ -791,52 +795,50 @@ impl CsvSink { &self.writer_options } } - -#[async_trait] impl FileSink for CsvSink { fn config(&self) -> &FileSinkConfig { &self.config } - async fn spawn_writer_tasks_and_join( - &self, - context: &Arc, + fn spawn_writer_tasks_and_join<'a>( + &'a self, + context: &'a Arc, demux_task: SpawnedTask>, file_stream_rx: DemuxedStreamReceiver, object_store: Arc, - ) -> Result { - let builder = self.writer_options.writer_options.clone(); - let header = builder.header(); - let serializer = Arc::new( - CsvSerializer::new() - .with_builder(builder) - .with_header(header), - ) as _; - spawn_writer_tasks_and_join( - context, - serializer, - self.writer_options.compression.into(), - self.writer_options.compression_level, - object_store, - demux_task, - file_stream_rx, - ) - .await + ) -> BoxFuture<'a, Result> { + Box::pin(async move { + let builder = self.writer_options.writer_options.clone(); + let header = builder.header(); + let serializer = Arc::new( + CsvSerializer::new() + .with_builder(builder) + .with_header(header), + ) as _; + spawn_writer_tasks_and_join( + context, + serializer, + self.writer_options.compression.into(), + self.writer_options.compression_level, + object_store, + demux_task, + file_stream_rx, + ) + .await + }) } } - -#[async_trait] impl DataSink for CsvSink { fn schema(&self) -> &SchemaRef { self.config.output_schema() } - async fn write_all( - &self, + fn write_all<'a>( + &'a self, data: SendableRecordBatchStream, - context: &Arc, - ) -> Result { - FileSink::write_all(self, data, context).await + context: &'a Arc, + ) -> BoxFuture<'a, Result> { + Box::pin(async move { FileSink::write_all(self, data, context).await }) } #[cfg(feature = "proto")] diff --git a/datafusion/datasource-json/Cargo.toml b/datafusion/datasource-json/Cargo.toml index 04192083f583a..61dd5d84d7bc8 100644 --- a/datafusion/datasource-json/Cargo.toml +++ b/datafusion/datasource-json/Cargo.toml @@ -40,7 +40,6 @@ proto = [ [dependencies] arrow = { workspace = true } -async-trait = { workspace = true } bytes = { workspace = true } datafusion-common = { workspace = true, features = ["object_store"] } datafusion-common-runtime = { workspace = true } diff --git a/datafusion/datasource-json/src/file_format.rs b/datafusion/datasource-json/src/file_format.rs index aef89560a3b5f..26c5cdf90c20c 100644 --- a/datafusion/datasource-json/src/file_format.rs +++ b/datafusion/datasource-json/src/file_format.rs @@ -17,6 +17,7 @@ //! [`JsonFormat`]: Line delimited and array JSON [`FileFormat`] abstractions +use futures::future::BoxFuture; use std::collections::HashMap; use std::fmt; use std::fmt::Debug; @@ -59,7 +60,6 @@ use datafusion_physical_plan::{DisplayAs, DisplayFormatType, ExecutionPlan}; use datafusion_session::Session; use crate::utils::JsonArrayToNdjsonReader; -use async_trait::async_trait; use object_store::{GetResultPayload, ObjectMeta, ObjectStore, ObjectStoreExt}; #[derive(Default)] @@ -233,8 +233,6 @@ fn infer_schema_from_json_array( Ok((schema, count)) } - -#[async_trait] impl FileFormat for JsonFormat { fn get_ext(&self) -> String { JsonFormatFactory::new().get_ext() @@ -252,124 +250,132 @@ impl FileFormat for JsonFormat { Some(self.options.compression.into()) } - async fn infer_schema( - &self, - _state: &dyn Session, - store: &Arc, - objects: &[ObjectMeta], - ) -> Result { - let mut schemas = Vec::new(); - let mut records_to_read = self - .options - .schema_infer_max_rec - .unwrap_or(DEFAULT_SCHEMA_INFER_MAX_RECORD); - let file_compression_type = FileCompressionType::from(self.options.compression); - let newline_delimited = self.options.newline_delimited; - - for object in objects { - // Early exit if we've read enough records - if records_to_read == 0 { - break; - } + fn infer_schema<'a>( + &'a self, + _state: &'a dyn Session, + store: &'a Arc, + objects: &'a [ObjectMeta], + ) -> BoxFuture<'a, Result> { + Box::pin(async move { + let mut schemas = Vec::new(); + let mut records_to_read = self + .options + .schema_infer_max_rec + .unwrap_or(DEFAULT_SCHEMA_INFER_MAX_RECORD); + let file_compression_type = + FileCompressionType::from(self.options.compression); + let newline_delimited = self.options.newline_delimited; + + for object in objects { + // Early exit if we've read enough records + if records_to_read == 0 { + break; + } - let r = store.as_ref().get(&object.location).await?; - - let (schema, records_consumed) = match r.payload { - #[cfg(not(target_arch = "wasm32"))] - GetResultPayload::File(file, _) => { - let decoder = file_compression_type.convert_read(file)?; - let reader = BufReader::new(decoder); - - if newline_delimited { - // NDJSON: use ValueIter directly - let iter = ValueIter::new(reader, None); - let mut count = 0; - let schema = - infer_json_schema_from_iterator(iter.take_while(|_| { - let should_take = count < records_to_read; - if should_take { - count += 1; - } - should_take - }))?; - (schema, count) - } else { - // JSON array format: use streaming converter - infer_schema_from_json_array(reader, records_to_read)? + let r = store.as_ref().get(&object.location).await?; + + let (schema, records_consumed) = match r.payload { + #[cfg(not(target_arch = "wasm32"))] + GetResultPayload::File(file, _) => { + let decoder = file_compression_type.convert_read(file)?; + let reader = BufReader::new(decoder); + + if newline_delimited { + // NDJSON: use ValueIter directly + let iter = ValueIter::new(reader, None); + let mut count = 0; + let schema = + infer_json_schema_from_iterator(iter.take_while(|_| { + let should_take = count < records_to_read; + if should_take { + count += 1; + } + should_take + }))?; + (schema, count) + } else { + // JSON array format: use streaming converter + infer_schema_from_json_array(reader, records_to_read)? + } } - } - GetResultPayload::Stream(_) => { - let data = r.bytes().await?; - let decoder = file_compression_type.convert_read(data.reader())?; - let reader = BufReader::new(decoder); - - if newline_delimited { - let iter = ValueIter::new(reader, None); - let mut count = 0; - let schema = - infer_json_schema_from_iterator(iter.take_while(|_| { - let should_take = count < records_to_read; - if should_take { - count += 1; - } - should_take - }))?; - (schema, count) - } else { - // JSON array format: use streaming converter - infer_schema_from_json_array(reader, records_to_read)? + GetResultPayload::Stream(_) => { + let data = r.bytes().await?; + let decoder = + file_compression_type.convert_read(data.reader())?; + let reader = BufReader::new(decoder); + + if newline_delimited { + let iter = ValueIter::new(reader, None); + let mut count = 0; + let schema = + infer_json_schema_from_iterator(iter.take_while(|_| { + let should_take = count < records_to_read; + if should_take { + count += 1; + } + should_take + }))?; + (schema, count) + } else { + // JSON array format: use streaming converter + infer_schema_from_json_array(reader, records_to_read)? + } } - } - }; + }; - schemas.push(schema); - // Correctly decrement records_to_read - records_to_read = records_to_read.saturating_sub(records_consumed); - } + schemas.push(schema); + // Correctly decrement records_to_read + records_to_read = records_to_read.saturating_sub(records_consumed); + } - let schema = Schema::try_merge(schemas)?; - Ok(Arc::new(schema)) + let schema = Schema::try_merge(schemas)?; + Ok(Arc::new(schema)) + }) } - async fn infer_stats( - &self, - _state: &dyn Session, - _store: &Arc, + fn infer_stats<'a>( + &'a self, + _state: &'a dyn Session, + _store: &'a Arc, table_schema: SchemaRef, - _object: &ObjectMeta, - ) -> Result { - Ok(Statistics::new_unknown(&table_schema)) + _object: &'a ObjectMeta, + ) -> BoxFuture<'a, Result> { + Box::pin(async move { Ok(Statistics::new_unknown(&table_schema)) }) } - async fn create_physical_plan( - &self, - _state: &dyn Session, + fn create_physical_plan<'a>( + &'a self, + _state: &'a dyn Session, conf: FileScanConfig, - ) -> Result> { - let conf = FileScanConfigBuilder::from(conf) - .with_file_compression_type(FileCompressionType::from( - self.options.compression, - )) - .build(); - Ok(DataSourceExec::from_data_source(conf)) + ) -> BoxFuture<'a, Result>> { + Box::pin(async move { + let conf = FileScanConfigBuilder::from(conf) + .with_file_compression_type(FileCompressionType::from( + self.options.compression, + )) + .build(); + Ok(DataSourceExec::from_data_source(conf) as Arc) + }) } - async fn create_writer_physical_plan( - &self, + fn create_writer_physical_plan<'a>( + &'a self, input: Arc, - _state: &dyn Session, + _state: &'a dyn Session, conf: FileSinkConfig, order_requirements: Option, - ) -> Result> { - if conf.insert_op != InsertOp::Append { - return not_impl_err!("Overwrites are not implemented yet for Json"); - } + ) -> BoxFuture<'a, Result>> { + Box::pin(async move { + if conf.insert_op != InsertOp::Append { + return not_impl_err!("Overwrites are not implemented yet for Json"); + } - let writer_options = JsonWriterOptions::try_from(&self.options)?; + let writer_options = JsonWriterOptions::try_from(&self.options)?; - let sink = Arc::new(JsonSink::new(conf, writer_options)); + let sink = Arc::new(JsonSink::new(conf, writer_options)); - Ok(Arc::new(DataSinkExec::new(input, sink, order_requirements)) as _) + Ok(Arc::new(DataSinkExec::new(input, sink, order_requirements)) as _) + }) } fn file_source(&self, table_schema: TableSchema) -> Arc { @@ -449,46 +455,44 @@ impl JsonSink { &self.writer_options } } - -#[async_trait] impl FileSink for JsonSink { fn config(&self) -> &FileSinkConfig { &self.config } - async fn spawn_writer_tasks_and_join( - &self, - context: &Arc, + fn spawn_writer_tasks_and_join<'a>( + &'a self, + context: &'a Arc, demux_task: SpawnedTask>, file_stream_rx: DemuxedStreamReceiver, object_store: Arc, - ) -> Result { - let serializer = Arc::new(JsonSerializer::new()) as _; - spawn_writer_tasks_and_join( - context, - serializer, - self.writer_options.compression.into(), - self.writer_options.compression_level, - object_store, - demux_task, - file_stream_rx, - ) - .await + ) -> BoxFuture<'a, Result> { + Box::pin(async move { + let serializer = Arc::new(JsonSerializer::new()) as _; + spawn_writer_tasks_and_join( + context, + serializer, + self.writer_options.compression.into(), + self.writer_options.compression_level, + object_store, + demux_task, + file_stream_rx, + ) + .await + }) } } - -#[async_trait] impl DataSink for JsonSink { fn schema(&self) -> &SchemaRef { self.config.output_schema() } - async fn write_all( - &self, + fn write_all<'a>( + &'a self, data: SendableRecordBatchStream, - context: &Arc, - ) -> Result { - FileSink::write_all(self, data, context).await + context: &'a Arc, + ) -> BoxFuture<'a, Result> { + Box::pin(async move { FileSink::write_all(self, data, context).await }) } #[cfg(feature = "proto")] diff --git a/datafusion/datasource-parquet/Cargo.toml b/datafusion/datasource-parquet/Cargo.toml index a2589af19a6ee..6ecfa95648cea 100644 --- a/datafusion/datasource-parquet/Cargo.toml +++ b/datafusion/datasource-parquet/Cargo.toml @@ -33,7 +33,6 @@ all-features = true [dependencies] arrow = { workspace = true } arrow-schema = { workspace = true } -async-trait = { workspace = true } bytes = { workspace = true } datafusion-common = { workspace = true, features = ["object_store", "parquet"] } datafusion-common-runtime = { workspace = true } diff --git a/datafusion/datasource-parquet/src/file_format.rs b/datafusion/datasource-parquet/src/file_format.rs index 18f2b5a650c8d..143a2ecbf74f5 100644 --- a/datafusion/datasource-parquet/src/file_format.rs +++ b/datafusion/datasource-parquet/src/file_format.rs @@ -64,7 +64,6 @@ use crate::reader::CachedParquetFileReaderFactory; use crate::source::{ ParquetSource, parse_coerce_int96_string, parse_coerce_int96_tz_string, }; -use async_trait::async_trait; use bytes::Bytes; use datafusion_datasource::source::DataSourceExec; use datafusion_execution::cache::cache_manager::FileMetadataCache; @@ -307,8 +306,6 @@ async fn get_file_decryption_properties( ) -> Result>> { Ok(None) } - -#[async_trait] impl FileFormat for ParquetFormat { fn get_ext(&self) -> String { ParquetFormatFactory::new().get_ext() @@ -329,246 +326,262 @@ impl FileFormat for ParquetFormat { None } - async fn infer_schema( - &self, - state: &dyn Session, - store: &Arc, - objects: &[ObjectMeta], - ) -> Result { - let coerce_int96 = match self.coerce_int96() { - Some(time_unit) => Some(parse_coerce_int96_string(time_unit.as_str())?), - None => None, - }; - let coerce_int96_tz = self - .options - .global - .coerce_int96_tz - .as_ref() - .map(|tz| parse_coerce_int96_tz_string(tz)) - .transpose()?; - - let file_metadata_cache = - state.runtime_env().cache_manager.get_file_metadata_cache(); - - let mut schemas: Vec<_> = futures::stream::iter(objects) - .map(|object| async { - let file_decryption_properties = get_file_decryption_properties( - state, - &self.options, - &object.location, - ) - .await?; - let result = DFParquetMetadata::new(store.as_ref(), object) - .with_metadata_size_hint(self.metadata_size_hint()) - .with_decryption_properties(file_decryption_properties) - .with_file_metadata_cache(Some(Arc::clone(&file_metadata_cache))) - .with_coerce_int96(coerce_int96) - .with_coerce_int96_tz(coerce_int96_tz.clone()) - .fetch_schema_with_location() + fn infer_schema<'a>( + &'a self, + state: &'a dyn Session, + store: &'a Arc, + objects: &'a [ObjectMeta], + ) -> BoxFuture<'a, Result> { + Box::pin(async move { + let coerce_int96 = match self.coerce_int96() { + Some(time_unit) => Some(parse_coerce_int96_string(time_unit.as_str())?), + None => None, + }; + let coerce_int96_tz = self + .options + .global + .coerce_int96_tz + .as_ref() + .map(|tz| parse_coerce_int96_tz_string(tz)) + .transpose()?; + + let file_metadata_cache = + state.runtime_env().cache_manager.get_file_metadata_cache(); + + let mut schemas: Vec<_> = futures::stream::iter(objects) + .map(|object| async { + let file_decryption_properties = get_file_decryption_properties( + state, + &self.options, + &object.location, + ) .await?; - Ok::<_, DataFusionError>(result) - }) - .boxed() // Workaround https://github.com/rust-lang/rust/issues/64552 - // fetch schemas concurrently, if requested - .buffer_unordered( - state - .config_options() - .execution - .meta_fetch_concurrency - .get(), - ) - .try_collect() - .await?; - - // Schema inference adds fields based the order they are seen - // which depends on the order the files are processed. For some - // object stores (like local file systems) the order returned from list - // is not deterministic. Thus, to ensure deterministic schema inference - // sort the files first. - // https://github.com/apache/datafusion/pull/6629 - schemas - .sort_unstable_by(|(location1, _), (location2, _)| location1.cmp(location2)); - - let mut seen = HashSet::new(); - for (location, schema) in &schemas { - ensure_unique_field_names(schema, &mut seen).map_err(|err| { - DataFusionError::Context( - format!("Error when processing Parquet file {location}"), - Box::new(err), + let result = DFParquetMetadata::new(store.as_ref(), object) + .with_metadata_size_hint(self.metadata_size_hint()) + .with_decryption_properties(file_decryption_properties) + .with_file_metadata_cache(Some(Arc::clone(&file_metadata_cache))) + .with_coerce_int96(coerce_int96) + .with_coerce_int96_tz(coerce_int96_tz.clone()) + .fetch_schema_with_location() + .await?; + Ok::<_, DataFusionError>(result) + }) + .boxed() // Workaround https://github.com/rust-lang/rust/issues/64552 + // fetch schemas concurrently, if requested + .buffer_unordered( + state + .config_options() + .execution + .meta_fetch_concurrency + .get(), ) - })?; - } - drop(seen); + .try_collect() + .await?; - let schemas = schemas.into_iter().map(|(_, schema)| schema); + // Schema inference adds fields based the order they are seen + // which depends on the order the files are processed. For some + // object stores (like local file systems) the order returned from list + // is not deterministic. Thus, to ensure deterministic schema inference + // sort the files first. + // https://github.com/apache/datafusion/pull/6629 + schemas.sort_unstable_by(|(location1, _), (location2, _)| { + location1.cmp(location2) + }); + + let mut seen = HashSet::new(); + for (location, schema) in &schemas { + ensure_unique_field_names(schema, &mut seen).map_err(|err| { + DataFusionError::Context( + format!("Error when processing Parquet file {location}"), + Box::new(err), + ) + })?; + } + drop(seen); - let schema = if self.skip_metadata() { - Schema::try_merge(clear_metadata(schemas)) - } else { - Schema::try_merge(schemas) - }?; + let schemas = schemas.into_iter().map(|(_, schema)| schema); - let schema = if self.binary_as_string() { - transform_binary_to_string(&schema) - } else { - schema - }; + let schema = if self.skip_metadata() { + Schema::try_merge(clear_metadata(schemas)) + } else { + Schema::try_merge(schemas) + }?; - let schema = if self.force_view_types() { - transform_schema_to_view(&schema) - } else { - schema - }; + let schema = if self.binary_as_string() { + transform_binary_to_string(&schema) + } else { + schema + }; - Ok(Arc::new(schema)) + let schema = if self.force_view_types() { + transform_schema_to_view(&schema) + } else { + schema + }; + + Ok(Arc::new(schema)) + }) } - async fn infer_stats( - &self, - state: &dyn Session, - store: &Arc, + fn infer_stats<'a>( + &'a self, + state: &'a dyn Session, + store: &'a Arc, table_schema: SchemaRef, - object: &ObjectMeta, - ) -> Result { - let file_decryption_properties = - get_file_decryption_properties(state, &self.options, &object.location) - .await?; - let file_metadata_cache = - state.runtime_env().cache_manager.get_file_metadata_cache(); - DFParquetMetadata::new(store, object) - .with_metadata_size_hint(self.metadata_size_hint()) - .with_decryption_properties(file_decryption_properties) - .with_file_metadata_cache(Some(file_metadata_cache)) - .fetch_statistics(&table_schema) - .await + object: &'a ObjectMeta, + ) -> BoxFuture<'a, Result> { + Box::pin(async move { + let file_decryption_properties = + get_file_decryption_properties(state, &self.options, &object.location) + .await?; + let file_metadata_cache = + state.runtime_env().cache_manager.get_file_metadata_cache(); + DFParquetMetadata::new(store, object) + .with_metadata_size_hint(self.metadata_size_hint()) + .with_decryption_properties(file_decryption_properties) + .with_file_metadata_cache(Some(file_metadata_cache)) + .fetch_statistics(&table_schema) + .await + }) } - async fn infer_ordering( - &self, - state: &dyn Session, - store: &Arc, + fn infer_ordering<'a>( + &'a self, + state: &'a dyn Session, + store: &'a Arc, table_schema: SchemaRef, - object: &ObjectMeta, - ) -> Result> { - let file_decryption_properties = - get_file_decryption_properties(state, &self.options, &object.location) + object: &'a ObjectMeta, + ) -> BoxFuture<'a, Result>> { + Box::pin(async move { + let file_decryption_properties = + get_file_decryption_properties(state, &self.options, &object.location) + .await?; + let file_metadata_cache = + state.runtime_env().cache_manager.get_file_metadata_cache(); + let metadata = DFParquetMetadata::new(store, object) + .with_metadata_size_hint(self.metadata_size_hint()) + .with_decryption_properties(file_decryption_properties) + .with_file_metadata_cache(Some(file_metadata_cache)) + .fetch_metadata() .await?; - let file_metadata_cache = - state.runtime_env().cache_manager.get_file_metadata_cache(); - let metadata = DFParquetMetadata::new(store, object) - .with_metadata_size_hint(self.metadata_size_hint()) - .with_decryption_properties(file_decryption_properties) - .with_file_metadata_cache(Some(file_metadata_cache)) - .fetch_metadata() - .await?; - crate::metadata::ordering_from_parquet_metadata(&metadata, &table_schema) - } - - async fn infer_stats_and_ordering( - &self, - state: &dyn Session, - store: &Arc, + crate::metadata::ordering_from_parquet_metadata(&metadata, &table_schema) + }) + } + + fn infer_stats_and_ordering<'a>( + &'a self, + state: &'a dyn Session, + store: &'a Arc, table_schema: SchemaRef, - object: &ObjectMeta, - ) -> Result { - let file_decryption_properties = - get_file_decryption_properties(state, &self.options, &object.location) + object: &'a ObjectMeta, + ) -> BoxFuture<'a, Result> { + Box::pin(async move { + let file_decryption_properties = + get_file_decryption_properties(state, &self.options, &object.location) + .await?; + let file_metadata_cache = + state.runtime_env().cache_manager.get_file_metadata_cache(); + let metadata = DFParquetMetadata::new(store, object) + .with_metadata_size_hint(self.metadata_size_hint()) + .with_decryption_properties(file_decryption_properties) + .with_file_metadata_cache(Some(file_metadata_cache)) + .fetch_metadata() .await?; - let file_metadata_cache = - state.runtime_env().cache_manager.get_file_metadata_cache(); - let metadata = DFParquetMetadata::new(store, object) - .with_metadata_size_hint(self.metadata_size_hint()) - .with_decryption_properties(file_decryption_properties) - .with_file_metadata_cache(Some(file_metadata_cache)) - .fetch_metadata() - .await?; - let statistics = DFParquetMetadata::statistics_from_parquet_metadata( - &metadata, - &table_schema, - )?; - let ordering = - crate::metadata::ordering_from_parquet_metadata(&metadata, &table_schema)?; - Ok( - datafusion_datasource::file_format::FileMeta::new(statistics) - .with_ordering(ordering), - ) + let statistics = DFParquetMetadata::statistics_from_parquet_metadata( + &metadata, + &table_schema, + )?; + let ordering = crate::metadata::ordering_from_parquet_metadata( + &metadata, + &table_schema, + )?; + Ok( + datafusion_datasource::file_format::FileMeta::new(statistics) + .with_ordering(ordering), + ) + }) } - async fn create_physical_plan( - &self, - state: &dyn Session, + fn create_physical_plan<'a>( + &'a self, + state: &'a dyn Session, conf: FileScanConfig, - ) -> Result> { - let mut metadata_size_hint = None; + ) -> BoxFuture<'a, Result>> { + Box::pin(async move { + let mut metadata_size_hint = None; - if let Some(metadata) = self.metadata_size_hint() { - metadata_size_hint = Some(metadata); - } + if let Some(metadata) = self.metadata_size_hint() { + metadata_size_hint = Some(metadata); + } - let mut source = conf - .file_source() - .downcast_ref::() - .cloned() - .ok_or_else(|| internal_datafusion_err!("Expected ParquetSource"))?; - source = source.with_table_parquet_options(self.options.clone()); - - // Use the CachedParquetFileReaderFactory - let metadata_cache = state.runtime_env().cache_manager.get_file_metadata_cache(); - let store = state - .runtime_env() - .object_store(conf.object_store_url.clone())?; - let cached_parquet_read_factory = - Arc::new(CachedParquetFileReaderFactory::new(store, metadata_cache)); - source = source.with_parquet_file_reader_factory(cached_parquet_read_factory); - - if let Some(metadata_size_hint) = metadata_size_hint { - source = source.with_metadata_size_hint(metadata_size_hint) - } + let mut source = conf + .file_source() + .downcast_ref::() + .cloned() + .ok_or_else(|| internal_datafusion_err!("Expected ParquetSource"))?; + source = source.with_table_parquet_options(self.options.clone()); + + // Use the CachedParquetFileReaderFactory + let metadata_cache = + state.runtime_env().cache_manager.get_file_metadata_cache(); + let store = state + .runtime_env() + .object_store(conf.object_store_url.clone())?; + let cached_parquet_read_factory = + Arc::new(CachedParquetFileReaderFactory::new(store, metadata_cache)); + source = source.with_parquet_file_reader_factory(cached_parquet_read_factory); + + if let Some(metadata_size_hint) = metadata_size_hint { + source = source.with_metadata_size_hint(metadata_size_hint) + } - source = self.set_source_encryption_factory(source, state)?; + source = self.set_source_encryption_factory(source, state)?; - let conf = FileScanConfigBuilder::from(conf) - .with_source(Arc::new(source)) - .build(); - Ok(DataSourceExec::from_data_source(conf)) + let conf = FileScanConfigBuilder::from(conf) + .with_source(Arc::new(source)) + .build(); + Ok(DataSourceExec::from_data_source(conf) as Arc) + }) } - async fn create_writer_physical_plan( - &self, + fn create_writer_physical_plan<'a>( + &'a self, input: Arc, - _state: &dyn Session, + _state: &'a dyn Session, conf: FileSinkConfig, order_requirements: Option, - ) -> Result> { - if conf.insert_op != InsertOp::Append { - return not_impl_err!("Overwrites are not implemented yet for Parquet"); - } - - // Convert ordering requirements to Parquet SortingColumns for file metadata - let sorting_columns = if let Some(ref requirements) = order_requirements { - let ordering: LexOrdering = requirements.clone().into(); - let writer_schema = get_writer_schema(&conf); - // In cases like `COPY (... ORDER BY ...) TO ...` the ORDER BY clause - // may not be compatible with Parquet sorting columns (e.g. ordering on `random()`). - // So if we cannot create a Parquet sorting column from the ordering requirement, - // we skip setting sorting columns on the Parquet sink. - lex_ordering_to_sorting_columns( - &ordering, - conf.output_schema(), - &writer_schema, - ) - .ok() - .filter(|columns| !columns.is_empty()) - } else { - None - }; - - let sink = Arc::new( - ParquetSink::new(conf, self.options.clone()) - .with_sorting_columns(sorting_columns), - ); + ) -> BoxFuture<'a, Result>> { + Box::pin(async move { + if conf.insert_op != InsertOp::Append { + return not_impl_err!("Overwrites are not implemented yet for Parquet"); + } - Ok(Arc::new(DataSinkExec::new(input, sink, order_requirements)) as _) + // Convert ordering requirements to Parquet SortingColumns for file metadata + let sorting_columns = if let Some(ref requirements) = order_requirements { + let ordering: LexOrdering = requirements.clone().into(); + let writer_schema = get_writer_schema(&conf); + // In cases like `COPY (... ORDER BY ...) TO ...` the ORDER BY clause + // may not be compatible with Parquet sorting columns (e.g. ordering on `random()`). + // So if we cannot create a Parquet sorting column from the ordering requirement, + // we skip setting sorting columns on the Parquet sink. + lex_ordering_to_sorting_columns( + &ordering, + conf.output_schema(), + &writer_schema, + ) + .ok() + .filter(|columns| !columns.is_empty()) + } else { + None + }; + + let sink = Arc::new( + ParquetSink::new(conf, self.options.clone()) + .with_sorting_columns(sorting_columns), + ); + + Ok(Arc::new(DataSinkExec::new(input, sink, order_requirements)) as _) + }) } fn file_source(&self, table_schema: TableSchema) -> Arc { diff --git a/datafusion/datasource-parquet/src/sink.rs b/datafusion/datasource-parquet/src/sink.rs index 3c66d4dcd74fb..d79b725909733 100644 --- a/datafusion/datasource-parquet/src/sink.rs +++ b/datafusion/datasource-parquet/src/sink.rs @@ -19,13 +19,13 @@ //! or more Parquet files to an [`ObjectStore`], optionally with parallel //! per-column and per-row-group serialization. +use futures::future::BoxFuture; use std::fmt; use std::fmt::Debug; use std::sync::Arc; use arrow::array::RecordBatch; use arrow::datatypes::{Schema, SchemaRef}; -use async_trait::async_trait; use datafusion_common::config::TableParquetOptions; use datafusion_common::{DataFusionError, HashMap, Result, internal_datafusion_err}; use datafusion_common_runtime::{JoinSet, SpawnedTask}; @@ -253,150 +253,152 @@ async fn set_writer_encryption_properties( ) -> Result { Ok(builder) } - -#[async_trait] impl FileSink for ParquetSink { fn config(&self) -> &FileSinkConfig { &self.config } - async fn spawn_writer_tasks_and_join( - &self, - context: &Arc, + fn spawn_writer_tasks_and_join<'a>( + &'a self, + context: &'a Arc, demux_task: SpawnedTask>, mut file_stream_rx: DemuxedStreamReceiver, object_store: Arc, - ) -> Result { - let rows_written_counter = MetricBuilder::new(&self.metrics) - .with_category(MetricCategory::Rows) - .global_counter("rows_written"); - // Note: bytes_written is the sum of compressed row group sizes, which - // may differ slightly from the actual on-disk file size (excludes footer, - // page indexes, and other Parquet metadata overhead). - let bytes_written_counter = - MetricBuilder::new(&self.metrics).global_bytes_counter("bytes_written"); - let elapsed_compute = MetricBuilder::new(&self.metrics).elapsed_compute(0); - - let parquet_opts = &self.parquet_options; - - let mut file_write_tasks: JoinSet< - std::result::Result<(Path, ParquetMetaData), DataFusionError>, - > = JoinSet::new(); - - let runtime = context.runtime_env(); - let parallel_options = ParallelParquetWriterOptions { - max_parallel_row_groups: parquet_opts - .global - .maximum_parallel_row_group_writers, - max_buffered_record_batches_per_stream: parquet_opts - .global - .maximum_buffered_record_batches_per_stream, - }; - - while let Some((path, mut rx)) = file_stream_rx.recv().await { - let parquet_props = self.create_writer_props(&runtime, &path).await?; - // CDC requires the sequential writer: the chunker state lives in ArrowWriter - // and persists across row groups. The parallel path bypasses ArrowWriter entirely. - if !parquet_opts.global.allow_single_file_parallelism - || parquet_opts.global.content_defined_chunking.enabled - { - let mut writer = self.create_async_arrow_writer( - &path, - Arc::clone(&object_store), - context, - parquet_props.clone(), - )?; - let reservation = MemoryConsumer::new(format!("ParquetSink[{path}]")) - .register(context.memory_pool()); - file_write_tasks.spawn( - async move { - while let Some(batch) = rx.recv().await { - writer.write(&batch).await?; - reservation.try_resize(writer.memory_size())?; + ) -> BoxFuture<'a, Result> { + Box::pin(async move { + let rows_written_counter = MetricBuilder::new(&self.metrics) + .with_category(MetricCategory::Rows) + .global_counter("rows_written"); + // Note: bytes_written is the sum of compressed row group sizes, which + // may differ slightly from the actual on-disk file size (excludes footer, + // page indexes, and other Parquet metadata overhead). + let bytes_written_counter = + MetricBuilder::new(&self.metrics).global_bytes_counter("bytes_written"); + let elapsed_compute = MetricBuilder::new(&self.metrics).elapsed_compute(0); + + let parquet_opts = &self.parquet_options; + + let mut file_write_tasks: JoinSet< + std::result::Result<(Path, ParquetMetaData), DataFusionError>, + > = JoinSet::new(); + + let runtime = context.runtime_env(); + let parallel_options = ParallelParquetWriterOptions { + max_parallel_row_groups: parquet_opts + .global + .maximum_parallel_row_group_writers, + max_buffered_record_batches_per_stream: parquet_opts + .global + .maximum_buffered_record_batches_per_stream, + }; + + while let Some((path, mut rx)) = file_stream_rx.recv().await { + let parquet_props = self.create_writer_props(&runtime, &path).await?; + // CDC requires the sequential writer: the chunker state lives in ArrowWriter + // and persists across row groups. The parallel path bypasses ArrowWriter entirely. + if !parquet_opts.global.allow_single_file_parallelism + || parquet_opts.global.content_defined_chunking.enabled + { + let mut writer = self.create_async_arrow_writer( + &path, + Arc::clone(&object_store), + context, + parquet_props.clone(), + )?; + let reservation = MemoryConsumer::new(format!("ParquetSink[{path}]")) + .register(context.memory_pool()); + file_write_tasks.spawn( + async move { + while let Some(batch) = rx.recv().await { + writer.write(&batch).await?; + reservation.try_resize(writer.memory_size())?; + } + let parquet_meta_data = + writer.close().await.map_err(|e| { + DataFusionError::ParquetError(Box::new(e)) + })?; + Ok((path, parquet_meta_data)) } - let parquet_meta_data = writer - .close() - .await - .map_err(|e| DataFusionError::ParquetError(Box::new(e)))?; - Ok((path, parquet_meta_data)) - } - .with_elapsed_compute(elapsed_compute.clone()), - ); - } else { - let writer = ObjectWriterBuilder::new( - // Parquet files as a whole are never compressed, since they - // manage compressed blocks themselves. - FileCompressionType::UNCOMPRESSED, - &path, - Arc::clone(&object_store), - ) - .with_buffer_size(Some( - context - .session_config() - .options() - .execution - .objectstore_writer_buffer_size, - )) - .build()?; - let ctx = ParquetFileWriteContext { - schema: get_writer_schema(&self.config), - props: Arc::new(parquet_props), - skip_arrow_metadata: self.parquet_options.global.skip_arrow_metadata, - parallel_options: Arc::new(parallel_options.clone()), - pool: Arc::clone(context.memory_pool()), - }; - let encoding_time = elapsed_compute.clone(); - file_write_tasks.spawn(async move { - let parquet_meta_data = output_single_parquet_file_parallelized( - writer, - rx, - ctx, - encoding_time, + .with_elapsed_compute(elapsed_compute.clone()), + ); + } else { + let writer = ObjectWriterBuilder::new( + // Parquet files as a whole are never compressed, since they + // manage compressed blocks themselves. + FileCompressionType::UNCOMPRESSED, + &path, + Arc::clone(&object_store), ) - .await?; - Ok((path, parquet_meta_data)) - }); + .with_buffer_size(Some( + context + .session_config() + .options() + .execution + .objectstore_writer_buffer_size, + )) + .build()?; + let ctx = ParquetFileWriteContext { + schema: get_writer_schema(&self.config), + props: Arc::new(parquet_props), + skip_arrow_metadata: self + .parquet_options + .global + .skip_arrow_metadata, + parallel_options: Arc::new(parallel_options.clone()), + pool: Arc::clone(context.memory_pool()), + }; + let encoding_time = elapsed_compute.clone(); + file_write_tasks.spawn(async move { + let parquet_meta_data = output_single_parquet_file_parallelized( + writer, + rx, + ctx, + encoding_time, + ) + .await?; + Ok((path, parquet_meta_data)) + }); + } } - } - while let Some(result) = file_write_tasks.join_next().await { - match result { - Ok(r) => { - let (path, parquet_meta_data) = r?; - let file_rows = parquet_meta_data.file_metadata().num_rows() as usize; - let file_bytes: usize = parquet_meta_data - .row_groups() - .iter() - .map(|rg| rg.compressed_size() as usize) - .sum(); - rows_written_counter.add(file_rows); - bytes_written_counter.add(file_bytes); - let mut written_files = self.written.lock(); - written_files - .try_insert(path.clone(), parquet_meta_data) - .map_err(|e| internal_datafusion_err!("duplicate entry detected for partitioned file {path}: {e}"))?; - drop(written_files); - } - Err(e) => { - if e.is_panic() { - std::panic::resume_unwind(e.into_panic()); - } else { - unreachable!(); + while let Some(result) = file_write_tasks.join_next().await { + match result { + Ok(r) => { + let (path, parquet_meta_data) = r?; + let file_rows = + parquet_meta_data.file_metadata().num_rows() as usize; + let file_bytes: usize = parquet_meta_data + .row_groups() + .iter() + .map(|rg| rg.compressed_size() as usize) + .sum(); + rows_written_counter.add(file_rows); + bytes_written_counter.add(file_bytes); + let mut written_files = self.written.lock(); + written_files + .try_insert(path.clone(), parquet_meta_data) + .map_err(|e| internal_datafusion_err!("duplicate entry detected for partitioned file {path}: {e}"))?; + drop(written_files); + } + Err(e) => { + if e.is_panic() { + std::panic::resume_unwind(e.into_panic()); + } else { + unreachable!(); + } } } } - } - demux_task - .join_unwind() - .await - .map_err(|e| DataFusionError::ExecutionJoin(Box::new(e)))??; + demux_task + .join_unwind() + .await + .map_err(|e| DataFusionError::ExecutionJoin(Box::new(e)))??; - Ok(rows_written_counter.value() as u64) + Ok(rows_written_counter.value() as u64) + }) } } - -#[async_trait] impl DataSink for ParquetSink { fn metrics(&self) -> Option { Some(self.metrics.clone_inner()) @@ -406,12 +408,12 @@ impl DataSink for ParquetSink { self.config.output_schema() } - async fn write_all( - &self, + fn write_all<'a>( + &'a self, data: SendableRecordBatchStream, - context: &Arc, - ) -> Result { - FileSink::write_all(self, data, context).await + context: &'a Arc, + ) -> BoxFuture<'a, Result> { + Box::pin(async move { FileSink::write_all(self, data, context).await }) } #[cfg(feature = "proto")] diff --git a/datafusion/datasource/Cargo.toml b/datafusion/datasource/Cargo.toml index f09447b694f52..a508bbcd20ee9 100644 --- a/datafusion/datasource/Cargo.toml +++ b/datafusion/datasource/Cargo.toml @@ -52,7 +52,6 @@ async-compression = { version = "0.4.40", features = [ "zstd", "tokio", ], optional = true } -async-trait = { workspace = true } bytes = { workspace = true } bzip2 = { workspace = true, optional = true } chrono = { workspace = true } diff --git a/datafusion/datasource/src/file_format.rs b/datafusion/datasource/src/file_format.rs index 2c2a94e8cc4b4..25b0af4a580ca 100644 --- a/datafusion/datasource/src/file_format.rs +++ b/datafusion/datasource/src/file_format.rs @@ -38,7 +38,7 @@ use datafusion_physical_expr_common::sort_expr::LexOrdering; use datafusion_physical_plan::ExecutionPlan; use datafusion_session::Session; -use async_trait::async_trait; +use futures::future::BoxFuture; use object_store::{ObjectMeta, ObjectStore}; /// Default max records to scan to infer the schema @@ -99,7 +99,6 @@ impl FileMeta { /// providers that support the same file formats. /// /// [`TableProvider`]: https://docs.rs/datafusion/latest/datafusion/catalog/trait.TableProvider.html -#[async_trait] pub trait FileFormat: Any + Send + Sync + fmt::Debug { /// Returns the extension for this FileFormat, e.g. "file.csv" -> csv fn get_ext(&self) -> String; @@ -117,12 +116,12 @@ pub trait FileFormat: Any + Send + Sync + fmt::Debug { /// be analysed up to a given number of records or files (as specified in the /// format config) then give the estimated common schema. This might fail if /// the files have schemas that cannot be merged. - async fn infer_schema( - &self, - state: &dyn Session, - store: &Arc, - objects: &[ObjectMeta], - ) -> Result; + fn infer_schema<'a>( + &'a self, + state: &'a dyn Session, + store: &'a Arc, + objects: &'a [ObjectMeta], + ) -> BoxFuture<'a, Result>; /// Infer the statistics for the provided object. The cost and accuracy of the /// estimated statistics might vary greatly between file formats. @@ -131,13 +130,13 @@ pub trait FileFormat: Any + Send + Sync + fmt::Debug { /// and may be a superset of the schema contained in this file. /// /// TODO: should the file source return statistics for only columns referred to in the table schema? - async fn infer_stats( - &self, - state: &dyn Session, - store: &Arc, + fn infer_stats<'a>( + &'a self, + state: &'a dyn Session, + store: &'a Arc, table_schema: SchemaRef, - object: &ObjectMeta, - ) -> Result; + object: &'a ObjectMeta, + ) -> BoxFuture<'a, Result>; /// Infer the ordering (sort order) for the provided object from file metadata. /// @@ -148,14 +147,14 @@ pub trait FileFormat: Any + Send + Sync + fmt::Debug { /// and may be a superset of the schema contained in this file. /// /// The default implementation returns `Ok(None)`. - async fn infer_ordering( - &self, - _state: &dyn Session, - _store: &Arc, + fn infer_ordering<'a>( + &'a self, + _state: &'a dyn Session, + _store: &'a Arc, _table_schema: SchemaRef, - _object: &ObjectMeta, - ) -> Result> { - Ok(None) + _object: &'a ObjectMeta, + ) -> BoxFuture<'a, Result>> { + Box::pin(async { Ok(None) }) } /// Infer both statistics and ordering from a single metadata read. @@ -166,43 +165,45 @@ pub trait FileFormat: Any + Send + Sync + fmt::Debug { /// /// The default implementation calls both methods separately. File formats /// that can extract both from a single read should override this method. - async fn infer_stats_and_ordering( - &self, - state: &dyn Session, - store: &Arc, + fn infer_stats_and_ordering<'a>( + &'a self, + state: &'a dyn Session, + store: &'a Arc, table_schema: SchemaRef, - object: &ObjectMeta, - ) -> Result { - let statistics = self - .infer_stats(state, store, Arc::clone(&table_schema), object) - .await?; - let ordering = self - .infer_ordering(state, store, table_schema, object) - .await?; - Ok(FileMeta { - statistics, - ordering, + object: &'a ObjectMeta, + ) -> BoxFuture<'a, Result> { + Box::pin(async move { + let statistics = self + .infer_stats(state, store, Arc::clone(&table_schema), object) + .await?; + let ordering = self + .infer_ordering(state, store, table_schema, object) + .await?; + Ok(FileMeta { + statistics, + ordering, + }) }) } /// Take a list of files and convert it to the appropriate executor /// according to this file format. - async fn create_physical_plan( - &self, - state: &dyn Session, + fn create_physical_plan<'a>( + &'a self, + state: &'a dyn Session, conf: FileScanConfig, - ) -> Result>; + ) -> BoxFuture<'a, Result>>; /// Take a list of files and the configuration to convert it to the /// appropriate writer executor according to this file format. - async fn create_writer_physical_plan( - &self, + fn create_writer_physical_plan<'a>( + &'a self, _input: Arc, - _state: &dyn Session, + _state: &'a dyn Session, _conf: FileSinkConfig, _order_requirements: Option, - ) -> Result> { - not_impl_err!("Writer not implemented for this format") + ) -> BoxFuture<'a, Result>> { + Box::pin(async { not_impl_err!("Writer not implemented for this format") }) } /// Return the related FileSource such as `CsvSource`, `JsonSource`, etc. diff --git a/datafusion/datasource/src/file_sink_config.rs b/datafusion/datasource/src/file_sink_config.rs index 48dce9a0cdb3e..c84c58eec4f72 100644 --- a/datafusion/datasource/src/file_sink_config.rs +++ b/datafusion/datasource/src/file_sink_config.rs @@ -29,7 +29,6 @@ use datafusion_execution::object_store::ObjectStoreUrl; use datafusion_execution::{SendableRecordBatchStream, TaskContext}; use datafusion_expr::dml::InsertOp; -use async_trait::async_trait; use object_store::ObjectStore; #[cfg(feature = "proto")] @@ -81,8 +80,9 @@ impl From for Option { } } +use futures::future::BoxFuture; + /// General behaviors for files that do `DataSink` operations -#[async_trait] pub trait FileSink: DataSink { /// Retrieves the file sink configuration. fn config(&self) -> &FileSinkConfig; @@ -107,32 +107,34 @@ pub trait FileSink: DataSink { /// /// # Returns /// - `Result`: Returns the total number of rows written across all files. - async fn spawn_writer_tasks_and_join( - &self, - context: &Arc, + fn spawn_writer_tasks_and_join<'a>( + &'a self, + context: &'a Arc, demux_task: SpawnedTask>, file_stream_rx: DemuxedStreamReceiver, object_store: Arc, - ) -> Result; + ) -> BoxFuture<'a, Result>; /// File sink implementation of the [`DataSink::write_all`] method. - async fn write_all( - &self, + fn write_all<'a>( + &'a self, data: SendableRecordBatchStream, - context: &Arc, - ) -> Result { - let config = self.config(); - let object_store = context - .runtime_env() - .object_store(&config.object_store_url)?; - let (demux_task, file_stream_rx) = start_demuxer_task(config, data, context); - self.spawn_writer_tasks_and_join( - context, - demux_task, - file_stream_rx, - object_store, - ) - .await + context: &'a Arc, + ) -> BoxFuture<'a, Result> { + Box::pin(async move { + let config = self.config(); + let object_store = context + .runtime_env() + .object_store(&config.object_store_url)?; + let (demux_task, file_stream_rx) = start_demuxer_task(config, data, context); + self.spawn_writer_tasks_and_join( + context, + demux_task, + file_stream_rx, + object_store, + ) + .await + }) } } diff --git a/datafusion/datasource/src/memory.rs b/datafusion/datasource/src/memory.rs index 4c79cf4a9851d..ada7df520f30b 100644 --- a/datafusion/datasource/src/memory.rs +++ b/datafusion/datasource/src/memory.rs @@ -15,6 +15,7 @@ // specific language governing permissions and limitations // under the License. +use futures::future::BoxFuture; use std::cmp::Ordering; use std::collections::BinaryHeap; use std::fmt; @@ -46,7 +47,6 @@ use datafusion_physical_plan::{ SendableRecordBatchStream, Statistics, common, }; -use async_trait::async_trait; use datafusion_physical_plan::coop::cooperative; use datafusion_physical_plan::execution_plan::SchedulingType; use futures::StreamExt; @@ -956,38 +956,38 @@ impl MemSink { Ok(Self { batches, schema }) } } - -#[async_trait] impl DataSink for MemSink { fn schema(&self) -> &SchemaRef { &self.schema } - async fn write_all( - &self, + fn write_all<'a>( + &'a self, mut data: SendableRecordBatchStream, - _context: &Arc, - ) -> Result { - let num_partitions = self.batches.len(); - - // buffer up the data round robin style into num_partitions - - let mut new_batches = vec![vec![]; num_partitions]; - let mut i = 0; - let mut row_count = 0; - while let Some(batch) = data.next().await.transpose()? { - row_count += batch.num_rows(); - new_batches[i].push(batch); - i = (i + 1) % num_partitions; - } + _context: &'a Arc, + ) -> BoxFuture<'a, Result> { + Box::pin(async move { + let num_partitions = self.batches.len(); + + // buffer up the data round robin style into num_partitions + + let mut new_batches = vec![vec![]; num_partitions]; + let mut i = 0; + let mut row_count = 0; + while let Some(batch) = data.next().await.transpose()? { + row_count += batch.num_rows(); + new_batches[i].push(batch); + i = (i + 1) % num_partitions; + } - // write the outputs into the batches - for (target, mut batches) in self.batches.iter().zip(new_batches) { - // Append all the new batches in one go to minimize locking overhead - target.write().await.append(&mut batches); - } + // write the outputs into the batches + for (target, mut batches) in self.batches.iter().zip(new_batches) { + // Append all the new batches in one go to minimize locking overhead + target.write().await.append(&mut batches); + } - Ok(row_count as u64) + Ok(row_count as u64) + }) } } diff --git a/datafusion/datasource/src/sink.rs b/datafusion/datasource/src/sink.rs index 1557e41531ab7..843da180c2562 100644 --- a/datafusion/datasource/src/sink.rs +++ b/datafusion/datasource/src/sink.rs @@ -37,16 +37,15 @@ use datafusion_physical_plan::{ ReplaceChildrenOptions, SendableRecordBatchStream, execute_input_stream, }; -use async_trait::async_trait; use datafusion_physical_plan::execution_plan::{EvaluationType, SchedulingType}; use futures::StreamExt; +use futures::future::BoxFuture; /// `DataSink` implements writing streams of [`RecordBatch`]es to /// user defined destinations. /// /// The `Display` impl is used to format the sink for explain plan /// output. -#[async_trait] pub trait DataSink: Any + DisplayAs + Debug + Send + Sync { /// Return a snapshot of the [MetricsSet] for this /// [DataSink]. @@ -67,11 +66,11 @@ pub trait DataSink: Any + DisplayAs + Debug + Send + Sync { /// This method will be called exactly once during each DML /// statement. Thus prior to return, the sink should do any commit /// or rollback required. - async fn write_all( - &self, + fn write_all<'a>( + &'a self, data: SendableRecordBatchStream, - context: &Arc, - ) -> Result; + context: &'a Arc, + ) -> BoxFuture<'a, Result>; /// Serialize this sink into a full protobuf plan node, if it knows how. /// diff --git a/datafusion/datasource/src/url.rs b/datafusion/datasource/src/url.rs index cfb6608ca0a78..0906a56e4dc34 100644 --- a/datafusion/datasource/src/url.rs +++ b/datafusion/datasource/src/url.rs @@ -509,7 +509,6 @@ fn split_glob_expression(path: &str) -> Option<(&str, &str)> { #[cfg(test)] mod tests { use super::*; - use async_trait::async_trait; use bytes::Bytes; use datafusion_common::DFSchema; use datafusion_common::config::TableOptions; @@ -530,7 +529,9 @@ mod tests { }; use std::any::Any; use std::collections::HashMap; + use std::future::Future; use std::ops::Range; + use std::pin::Pin; use tempfile::tempdir; #[test] @@ -1086,47 +1087,92 @@ mod tests { self.in_mem.fmt(f) } } - - #[async_trait] impl ObjectStore for MockObjectStore { - async fn put_opts( - &self, - location: &Path, + fn put_opts<'life0, 'life1, 'async_trait>( + &'life0 self, + location: &'life1 Path, payload: PutPayload, opts: object_store::PutOptions, - ) -> object_store::Result { - self.in_mem.put_opts(location, payload, opts).await + ) -> Pin< + Box< + dyn Future> + + Send + + 'async_trait, + >, + > + where + 'life0: 'async_trait, + 'life1: 'async_trait, + Self: 'async_trait, + { + Box::pin(async move { self.in_mem.put_opts(location, payload, opts).await }) } - async fn put_multipart_opts( - &self, - location: &Path, + fn put_multipart_opts<'life0, 'life1, 'async_trait>( + &'life0 self, + location: &'life1 Path, opts: PutMultipartOptions, - ) -> object_store::Result> { - self.in_mem.put_multipart_opts(location, opts).await + ) -> Pin< + Box< + dyn Future>> + + Send + + 'async_trait, + >, + > + where + 'life0: 'async_trait, + 'life1: 'async_trait, + Self: 'async_trait, + { + Box::pin(async move { self.in_mem.put_multipart_opts(location, opts).await }) } - async fn get_opts( - &self, - location: &Path, + fn get_opts<'life0, 'life1, 'async_trait>( + &'life0 self, + location: &'life1 Path, options: GetOptions, - ) -> object_store::Result { - if options.head && self.forbidden_paths.contains(location) { - Err(object_store::Error::PermissionDenied { - path: location.to_string(), - source: "forbidden".into(), - }) - } else { - self.in_mem.get_opts(location, options).await - } + ) -> Pin< + Box< + dyn Future> + + Send + + 'async_trait, + >, + > + where + 'life0: 'async_trait, + 'life1: 'async_trait, + Self: 'async_trait, + { + Box::pin(async move { + if options.head && self.forbidden_paths.contains(location) { + Err(object_store::Error::PermissionDenied { + path: location.to_string(), + source: "forbidden".into(), + }) + } else { + self.in_mem.get_opts(location, options).await + } + }) } - async fn get_ranges( - &self, - location: &Path, - ranges: &[Range], - ) -> object_store::Result> { - self.in_mem.get_ranges(location, ranges).await + fn get_ranges<'life0, 'life1, 'life2, 'async_trait>( + &'life0 self, + location: &'life1 Path, + ranges: &'life2 [Range], + ) -> Pin< + Box< + dyn Future>> + + Send + + 'async_trait, + >, + > + where + 'life0: 'async_trait, + 'life1: 'async_trait, + 'life2: 'async_trait, + Self: 'async_trait, + { + Box::pin(async move { self.in_mem.get_ranges(location, ranges).await }) } fn delete_stream( @@ -1143,20 +1189,37 @@ mod tests { self.in_mem.list(prefix) } - async fn list_with_delimiter( - &self, - prefix: Option<&Path>, - ) -> object_store::Result { - self.in_mem.list_with_delimiter(prefix).await + fn list_with_delimiter<'life0, 'life1, 'async_trait>( + &'life0 self, + prefix: Option<&'life1 Path>, + ) -> Pin< + Box< + dyn Future> + + Send + + 'async_trait, + >, + > + where + 'life0: 'async_trait, + 'life1: 'async_trait, + Self: 'async_trait, + { + Box::pin(async move { self.in_mem.list_with_delimiter(prefix).await }) } - async fn copy_opts( - &self, - from: &Path, - to: &Path, + fn copy_opts<'life0, 'life1, 'life2, 'async_trait>( + &'life0 self, + from: &'life1 Path, + to: &'life2 Path, options: CopyOptions, - ) -> object_store::Result<()> { - self.in_mem.copy_opts(from, to, options).await + ) -> Pin> + Send + 'async_trait>> + where + 'life0: 'async_trait, + 'life1: 'async_trait, + 'life2: 'async_trait, + Self: 'async_trait, + { + Box::pin(async move { self.in_mem.copy_opts(from, to, options).await }) } } @@ -1182,7 +1245,6 @@ mod tests { } } - #[async_trait::async_trait] impl Session for MockSession { fn session_id(&self) -> &str { unimplemented!() @@ -1196,10 +1258,10 @@ mod tests { Arc::new(EmptyCatalogProviderList) } - async fn create_physical_plan( - &self, - _logical_plan: &LogicalPlan, - ) -> Result> { + fn create_physical_plan<'a>( + &'a self, + _logical_plan: &'a LogicalPlan, + ) -> futures::future::BoxFuture<'a, Result>> { unimplemented!() } diff --git a/datafusion/execution/Cargo.toml b/datafusion/execution/Cargo.toml index c9d4acd3644ba..3ef73790db2ae 100644 --- a/datafusion/execution/Cargo.toml +++ b/datafusion/execution/Cargo.toml @@ -54,7 +54,6 @@ sql = [] [dependencies] arrow = { workspace = true } arrow-buffer = { workspace = true } -async-trait = { workspace = true } bytes = { workspace = true } dashmap = { workspace = true } datafusion-common = { workspace = true, default-features = false } diff --git a/datafusion/execution/src/parquet_encryption.rs b/datafusion/execution/src/parquet_encryption.rs index 45eac10264e88..502d50f3d61d4 100644 --- a/datafusion/execution/src/parquet_encryption.rs +++ b/datafusion/execution/src/parquet_encryption.rs @@ -16,11 +16,11 @@ // under the License. use arrow::datatypes::SchemaRef; -use async_trait::async_trait; use dashmap::DashMap; use datafusion_common::config::EncryptionFactoryOptions; use datafusion_common::error::Result; use datafusion_common::internal_datafusion_err; +use futures::future::BoxFuture; use object_store::path::Path; use parquet::encryption::decrypt::FileDecryptionProperties; use parquet::encryption::encrypt::FileEncryptionProperties; @@ -33,22 +33,21 @@ use std::sync::Arc; /// For example usage, see the [`parquet_encrypted_with_kms` example]. /// /// [`parquet_encrypted_with_kms` example]: https://github.com/apache/datafusion/blob/main/datafusion-examples/examples/data_io/parquet_encrypted_with_kms.rs -#[async_trait] pub trait EncryptionFactory: Send + Sync + std::fmt::Debug + 'static { /// Generate file encryption properties to use when writing a Parquet file. - async fn get_file_encryption_properties( - &self, - config: &EncryptionFactoryOptions, - schema: &SchemaRef, - file_path: &Path, - ) -> Result>>; + fn get_file_encryption_properties<'a>( + &'a self, + config: &'a EncryptionFactoryOptions, + schema: &'a SchemaRef, + file_path: &'a Path, + ) -> BoxFuture<'a, Result>>>; /// Generate file decryption properties to use when reading a Parquet file. - async fn get_file_decryption_properties( - &self, - config: &EncryptionFactoryOptions, - file_path: &Path, - ) -> Result>>; + fn get_file_decryption_properties<'a>( + &'a self, + config: &'a EncryptionFactoryOptions, + file_path: &'a Path, + ) -> BoxFuture<'a, Result>>>; } /// Stores [`EncryptionFactory`] implementations that can be retrieved by a unique string identifier diff --git a/datafusion/expr/Cargo.toml b/datafusion/expr/Cargo.toml index 4fe7b65f6d05f..62118cb57a866 100644 --- a/datafusion/expr/Cargo.toml +++ b/datafusion/expr/Cargo.toml @@ -52,7 +52,7 @@ sql = ["sqlparser"] [dependencies] arrow = { workspace = true, features = ["canonical_extension_types"] } arrow-schema = { workspace = true, features = ["canonical_extension_types"] } -async-trait = { workspace = true } +futures = { workspace = true } chrono = { workspace = true } datafusion-common = { workspace = true, default-features = false } datafusion-doc = { workspace = true } diff --git a/datafusion/expr/src/async_udf.rs b/datafusion/expr/src/async_udf.rs index 02a6d2ece8cdb..78eb39957fe12 100644 --- a/datafusion/expr/src/async_udf.rs +++ b/datafusion/expr/src/async_udf.rs @@ -17,11 +17,11 @@ use crate::{ReturnFieldArgs, ScalarFunctionArgs, ScalarUDF, ScalarUDFImpl}; use arrow::datatypes::{DataType, FieldRef}; -use async_trait::async_trait; use datafusion_common::error::Result; use datafusion_common::internal_err; use datafusion_expr_common::columnar_value::ColumnarValue; use datafusion_expr_common::signature::Signature; +use futures::future::BoxFuture; use std::any::Any; use std::fmt::{Debug, Display}; use std::hash::{Hash, Hasher}; @@ -33,7 +33,6 @@ use std::sync::Arc; /// to register remote functions in the context. /// /// The name is chosen to mirror ScalarUDFImpl -#[async_trait] pub trait AsyncScalarUDFImpl: ScalarUDFImpl { /// The ideal batch size for this function. /// @@ -44,10 +43,10 @@ pub trait AsyncScalarUDFImpl: ScalarUDFImpl { } /// Invoke the function asynchronously with the async arguments - async fn invoke_async_with_args( + fn invoke_async_with_args( &self, args: ScalarFunctionArgs, - ) -> Result; + ) -> BoxFuture<'_, Result>; } /// A scalar UDF that must be invoked using async methods @@ -137,9 +136,9 @@ mod tests { }; use arrow::datatypes::DataType; - use async_trait::async_trait; use datafusion_common::error::Result; use datafusion_expr_common::{columnar_value::ColumnarValue, signature::Signature}; + use futures::future::BoxFuture; use crate::{ ScalarFunctionArgs, ScalarUDFImpl, @@ -168,14 +167,12 @@ mod tests { todo!() } } - - #[async_trait] impl AsyncScalarUDFImpl for TestAsyncUDFImpl1 { - async fn invoke_async_with_args( + fn invoke_async_with_args( &self, _args: ScalarFunctionArgs, - ) -> Result { - todo!() + ) -> BoxFuture<'_, Result> { + Box::pin(async move { todo!() }) } } @@ -201,14 +198,12 @@ mod tests { todo!() } } - - #[async_trait] impl AsyncScalarUDFImpl for TestAsyncUDFImpl2 { - async fn invoke_async_with_args( + fn invoke_async_with_args( &self, _args: ScalarFunctionArgs, - ) -> Result { - todo!() + ) -> BoxFuture<'_, Result> { + Box::pin(async move { todo!() }) } } diff --git a/datafusion/ffi/Cargo.toml b/datafusion/ffi/Cargo.toml index affcff3dbdcd9..1ada6c80db2c8 100644 --- a/datafusion/ffi/Cargo.toml +++ b/datafusion/ffi/Cargo.toml @@ -47,7 +47,6 @@ crate-type = ["cdylib", "rlib"] arrow = { workspace = true, features = ["ffi"] } arrow-schema = { workspace = true } async-ffi = { version = "0.5.0" } -async-trait = { workspace = true } chrono = { workspace = true } datafusion-catalog = { workspace = true } datafusion-common = { workspace = true } diff --git a/datafusion/ffi/src/physical_optimizer.rs b/datafusion/ffi/src/physical_optimizer.rs index 87fe3314af9d6..e266dae75f157 100644 --- a/datafusion/ffi/src/physical_optimizer.rs +++ b/datafusion/ffi/src/physical_optimizer.rs @@ -18,7 +18,6 @@ use std::ffi::c_void; use std::sync::Arc; -use async_trait::async_trait; use datafusion_common::config::ConfigOptions; use datafusion_common::error::Result; use datafusion_physical_optimizer::{PhysicalOptimizerContext, PhysicalOptimizerRule}; @@ -306,8 +305,6 @@ impl Clone for FFI_PhysicalOptimizerRule { unsafe { (self.clone)(self) } } } - -#[async_trait] impl PhysicalOptimizerRule for ForeignPhysicalOptimizerRule { fn optimize( &self, diff --git a/datafusion/ffi/src/query_planner.rs b/datafusion/ffi/src/query_planner.rs index 90705a447fbd9..8d1d4e251b04b 100644 --- a/datafusion/ffi/src/query_planner.rs +++ b/datafusion/ffi/src/query_planner.rs @@ -54,11 +54,11 @@ //! reference-counted planner, so it outlives A's original session, whereas //! `FFI_SessionRef` borrows its session with the lifetime erased. +use futures::future::BoxFuture; use std::ffi::c_void; use std::sync::Arc; use async_ffi::{FfiFuture, FutureExt}; -use async_trait::async_trait; use datafusion_common::error::{DataFusionError, Result}; use datafusion_expr::LogicalPlan; use datafusion_physical_plan::ExecutionPlan; @@ -340,17 +340,17 @@ impl From<&FFI_QueryPlanner> for Arc { } } } - -#[async_trait] impl QueryPlanner for ForeignQueryPlanner { - async fn create_physical_plan( - &self, - logical_plan: &LogicalPlan, - session: &dyn Session, - ) -> Result> { - self.0 - .create_physical_plan_with_session_runtime(logical_plan, session, None) - .await + fn create_physical_plan<'a>( + &'a self, + logical_plan: &'a LogicalPlan, + session: &'a dyn Session, + ) -> BoxFuture<'a, Result>> { + Box::pin(async move { + self.0 + .create_physical_plan_with_session_runtime(logical_plan, session, None) + .await + }) } } @@ -371,17 +371,17 @@ mod tests { #[derive(Debug)] struct EmptyQueryPlanner; - - #[async_trait] impl QueryPlanner for EmptyQueryPlanner { - async fn create_physical_plan( - &self, - _logical_plan: &LogicalPlan, - _session: &dyn Session, - ) -> Result> { - let schema = - Arc::new(Schema::new(vec![Field::new("a", DataType::Int32, true)])); - Ok(Arc::new(EmptyExec::new(schema))) + fn create_physical_plan<'a>( + &'a self, + _logical_plan: &'a LogicalPlan, + _session: &'a dyn Session, + ) -> BoxFuture<'a, Result>> { + Box::pin(async move { + let schema = + Arc::new(Schema::new(vec![Field::new("a", DataType::Int32, true)])); + Ok(Arc::new(EmptyExec::new(schema)) as Arc) + }) } } diff --git a/datafusion/ffi/src/schema_provider.rs b/datafusion/ffi/src/schema_provider.rs index 322d3bbfcdf86..04aa2708aed89 100644 --- a/datafusion/ffi/src/schema_provider.rs +++ b/datafusion/ffi/src/schema_provider.rs @@ -15,11 +15,11 @@ // specific language governing permissions and limitations // under the License. +use futures::future::BoxFuture; use std::ffi::c_void; use std::sync::Arc; use async_ffi::{FfiFuture, FutureExt}; -use async_trait::async_trait; use datafusion_catalog::{SchemaProvider, TableProvider}; use datafusion_common::error::{DataFusionError, Result}; use datafusion_proto::logical_plan::{ @@ -309,8 +309,6 @@ impl Clone for FFI_SchemaProvider { unsafe { (self.clone)(self) } } } - -#[async_trait] impl SchemaProvider for ForeignSchemaProvider { fn owner_name(&self) -> Option<&str> { let name: Option<&SString> = self.0.owner_name.as_ref(); @@ -326,18 +324,20 @@ impl SchemaProvider for ForeignSchemaProvider { } } - async fn table( - &self, - name: &str, - ) -> Result>, DataFusionError> { - unsafe { - let table: Option = - df_result!((self.0.table)(&self.0, name.into()).await)?.into(); + fn table<'a>( + &'a self, + name: &'a str, + ) -> BoxFuture<'a, Result>, DataFusionError>> { + Box::pin(async move { + unsafe { + let table: Option = + df_result!((self.0.table)(&self.0, name.into()).await)?.into(); - let table = table.as_ref().map(>::from); + let table = table.as_ref().map(>::from); - Ok(table) - } + Ok(table) + } + }) } fn register_table( diff --git a/datafusion/ffi/src/session/mod.rs b/datafusion/ffi/src/session/mod.rs index c3a7519cf584e..5981b9885039d 100644 --- a/datafusion/ffi/src/session/mod.rs +++ b/datafusion/ffi/src/session/mod.rs @@ -30,6 +30,7 @@ //! planner, and C must retain and invoke that planner directly. See the //! [`crate::query_planner`] module for details. +use futures::future::BoxFuture; use std::any::Any; use std::collections::HashMap; use std::ffi::c_void; @@ -38,7 +39,6 @@ use std::sync::{Arc, OnceLock}; use arrow_schema::SchemaRef; use arrow_schema::ffi::FFI_ArrowSchema; use async_ffi::{FfiFuture, FutureExt}; -use async_trait::async_trait; use datafusion_common::config::{ConfigFileType, ConfigOptions, TableOptions}; use datafusion_common::{DFSchema, DataFusionError, not_impl_err}; use datafusion_execution::TaskContext; @@ -688,8 +688,6 @@ fn table_options_from_rhashmap(options: SVec<(SString, SString)>) -> TableOption }); table_options } - -#[async_trait] impl Session for ForeignSession { fn session_id(&self) -> &str { unsafe { (self.session.session_id)(&self.session).as_str() } @@ -732,13 +730,15 @@ impl Session for ForeignSession { } } - async fn create_physical_plan( - &self, - _logical_plan: &LogicalPlan, - ) -> datafusion_common::Result> { - not_impl_err!( - "ForeignSession::create_physical_plan is unsupported; export and invoke an FFI_QueryPlanner captured before installing a foreign planner" - ) + fn create_physical_plan<'a>( + &'a self, + _logical_plan: &'a LogicalPlan, + ) -> BoxFuture<'a, datafusion_common::Result>> { + Box::pin(async move { + not_impl_err!( + "ForeignSession::create_physical_plan is unsupported; export and invoke an FFI_QueryPlanner captured before installing a foreign planner" + ) + }) } fn create_physical_expr( @@ -848,19 +848,19 @@ mod tests { #[derive(Debug)] struct ReenteringQueryPlanner; - - #[async_trait] impl QueryPlanner for ReenteringQueryPlanner { - async fn create_physical_plan( - &self, - logical_plan: &LogicalPlan, - session: &dyn Session, - ) -> Result> { - if REENTERING_PLANNER_CALLS.fetch_add(1, Ordering::Relaxed) == 0 { - session.create_physical_plan(logical_plan).await - } else { - exec_err!("query planner was re-entered through the session") - } + fn create_physical_plan<'a>( + &'a self, + logical_plan: &'a LogicalPlan, + session: &'a dyn Session, + ) -> BoxFuture<'a, Result>> { + Box::pin(async move { + if REENTERING_PLANNER_CALLS.fetch_add(1, Ordering::Relaxed) == 0 { + session.create_physical_plan(logical_plan).await + } else { + exec_err!("query planner was re-entered through the session") + } + }) } } diff --git a/datafusion/ffi/src/table_provider.rs b/datafusion/ffi/src/table_provider.rs index 5103297fcba57..4d5b699a2e07b 100644 --- a/datafusion/ffi/src/table_provider.rs +++ b/datafusion/ffi/src/table_provider.rs @@ -15,12 +15,12 @@ // specific language governing permissions and limitations // under the License. +use futures::future::BoxFuture; use std::ffi::c_void; use std::sync::Arc; use arrow::datatypes::SchemaRef; use async_ffi::{FfiFuture, FutureExt}; -use async_trait::async_trait; use datafusion_catalog::{Session, TableProvider}; use datafusion_common::Statistics; use datafusion_common::error::{DataFusionError, Result}; @@ -636,8 +636,6 @@ impl Clone for FFI_TableProvider { unsafe { (self.clone)(self) } } } - -#[async_trait] impl TableProvider for ForeignTableProvider { fn schema(&self) -> SchemaRef { let wrapped_schema = unsafe { (self.0.schema)(&self.0) }; @@ -663,36 +661,39 @@ impl TableProvider for ForeignTableProvider { } } - async fn scan( - &self, - session: &dyn Session, - projection: Option<&[usize]>, - filters: &[Expr], + fn scan<'a>( + &'a self, + session: &'a dyn Session, + projection: Option<&'a [usize]>, + filters: &'a [Expr], limit: Option, - ) -> Result> { - let session = FFI_SessionRef::new(session, None, self.0.logical_codec.clone()); - - let projections: FFI_Option> = projection - .map(|p| p.iter().map(|v| v.to_owned()).collect()) - .into(); + ) -> BoxFuture<'a, Result>> { + Box::pin(async move { + let session = + FFI_SessionRef::new(session, None, self.0.logical_codec.clone()); - let codec: Arc = (&self.0.logical_codec).into(); - let filters_serialized = serialize_expr_list(filters.iter(), codec.as_ref())?; - - let plan = unsafe { - let maybe_plan = (self.0.scan)( - &self.0, - session, - projections, - filters_serialized, - limit.into(), - ) - .await; + let projections: FFI_Option> = projection + .map(|p| p.iter().map(|v| v.to_owned()).collect()) + .into(); - >::try_from(&df_result!(maybe_plan)?)? - }; + let codec: Arc = (&self.0.logical_codec).into(); + let filters_serialized = serialize_expr_list(filters.iter(), codec.as_ref())?; + + let plan = unsafe { + let maybe_plan = (self.0.scan)( + &self.0, + session, + projections, + filters_serialized, + limit.into(), + ) + .await; + + >::try_from(&df_result!(maybe_plan)?)? + }; - Ok(plan) + Ok(plan) + }) } /// Tests whether the table provider can make use of a filter expression @@ -723,92 +724,108 @@ impl TableProvider for ForeignTableProvider { } } - async fn insert_into( - &self, - session: &dyn Session, + fn insert_into<'a>( + &'a self, + session: &'a dyn Session, input: Arc, insert_op: InsertOp, - ) -> Result> { - let session = FFI_SessionRef::new(session, None, self.0.logical_codec.clone()); + ) -> BoxFuture<'a, Result>> { + Box::pin(async move { + let session = + FFI_SessionRef::new(session, None, self.0.logical_codec.clone()); - let rc = Handle::try_current().ok(); - let input = FFI_ExecutionPlan::new(input, rc); - let insert_op: FFI_InsertOp = insert_op.into(); + let rc = Handle::try_current().ok(); + let input = FFI_ExecutionPlan::new(input, rc); + let insert_op: FFI_InsertOp = insert_op.into(); - let plan = unsafe { - let maybe_plan = - (self.0.insert_into)(&self.0, session, &input, insert_op).await; + let plan = unsafe { + let maybe_plan = + (self.0.insert_into)(&self.0, session, &input, insert_op).await; - >::try_from(&df_result!(maybe_plan)?)? - }; + >::try_from(&df_result!(maybe_plan)?)? + }; - Ok(plan) + Ok(plan) + }) } - async fn delete_from( - &self, - session: &dyn Session, + fn delete_from<'a>( + &'a self, + session: &'a dyn Session, filters: Vec, - ) -> Result> { - let session = FFI_SessionRef::new(session, None, self.0.logical_codec.clone()); - let codec: Arc = (&self.0.logical_codec).into(); - let filters_serialized = serialize_expr_list(filters.iter(), codec.as_ref())?; + ) -> BoxFuture<'a, Result>> { + Box::pin(async move { + let session = + FFI_SessionRef::new(session, None, self.0.logical_codec.clone()); + let codec: Arc = (&self.0.logical_codec).into(); + let filters_serialized = serialize_expr_list(filters.iter(), codec.as_ref())?; - let plan = unsafe { - let maybe_plan = - (self.0.delete_from)(&self.0, session, filters_serialized).await; + let plan = unsafe { + let maybe_plan = + (self.0.delete_from)(&self.0, session, filters_serialized).await; - >::try_from(&df_result!(maybe_plan)?)? - }; + >::try_from(&df_result!(maybe_plan)?)? + }; - Ok(plan) + Ok(plan) + }) } - async fn update( - &self, - session: &dyn Session, + fn update<'a>( + &'a self, + session: &'a dyn Session, assignments: Vec<(String, Expr)>, filters: Vec, - ) -> Result> { - let session = FFI_SessionRef::new(session, None, self.0.logical_codec.clone()); - let codec: Arc = (&self.0.logical_codec).into(); - - let assignments: SVec<_> = assignments - .iter() - .map(|(column, expr)| { - Ok(FFI_TableProviderUpdateAssignment { - column: SString::from(column.as_str()), - expr_serialized: serialize_expr_list( - std::iter::once(expr), - codec.as_ref(), - )?, + ) -> BoxFuture<'a, Result>> { + Box::pin(async move { + let session = + FFI_SessionRef::new(session, None, self.0.logical_codec.clone()); + let codec: Arc = (&self.0.logical_codec).into(); + + let assignments: SVec<_> = assignments + .iter() + .map(|(column, expr)| { + Ok(FFI_TableProviderUpdateAssignment { + column: SString::from(column.as_str()), + expr_serialized: serialize_expr_list( + std::iter::once(expr), + codec.as_ref(), + )?, + }) }) - }) - .collect::>>()? - .into_iter() - .collect(); - let filters_serialized = serialize_expr_list(filters.iter(), codec.as_ref())?; + .collect::>>()? + .into_iter() + .collect(); + let filters_serialized = serialize_expr_list(filters.iter(), codec.as_ref())?; - let plan = unsafe { - let maybe_plan = - (self.0.update)(&self.0, session, assignments, filters_serialized).await; + let plan = unsafe { + let maybe_plan = + (self.0.update)(&self.0, session, assignments, filters_serialized) + .await; - >::try_from(&df_result!(maybe_plan)?)? - }; + >::try_from(&df_result!(maybe_plan)?)? + }; - Ok(plan) + Ok(plan) + }) } - async fn truncate(&self, session: &dyn Session) -> Result> { - let session = FFI_SessionRef::new(session, None, self.0.logical_codec.clone()); + fn truncate<'a>( + &'a self, + session: &'a dyn Session, + ) -> BoxFuture<'a, Result>> { + Box::pin(async move { + let session = + FFI_SessionRef::new(session, None, self.0.logical_codec.clone()); - let plan = unsafe { - let maybe_plan = (self.0.truncate)(&self.0, session).await; + let plan = unsafe { + let maybe_plan = (self.0.truncate)(&self.0, session).await; - >::try_from(&df_result!(maybe_plan)?)? - }; + >::try_from(&df_result!(maybe_plan)?)? + }; - Ok(plan) + Ok(plan) + }) } } @@ -866,8 +883,6 @@ mod tests { } } } - - #[async_trait] impl TableProvider for TableWithStats { fn schema(&self) -> SchemaRef { self.inner.schema() @@ -881,69 +896,84 @@ mod tests { self.stats.clone() } - async fn scan( - &self, - session: &dyn Session, - projection: Option<&[usize]>, - filters: &[Expr], + fn scan<'a>( + &'a self, + session: &'a dyn Session, + projection: Option<&'a [usize]>, + filters: &'a [Expr], limit: Option, - ) -> Result> { - self.inner.scan(session, projection, filters, limit).await + ) -> BoxFuture<'a, Result>> { + Box::pin( + async move { self.inner.scan(session, projection, filters, limit).await }, + ) } - async fn delete_from( - &self, - _state: &dyn Session, + fn delete_from<'a>( + &'a self, + _state: &'a dyn Session, filters: Vec, - ) -> Result> { - let call = self.delete_calls.fetch_add(1, Ordering::Relaxed); - let valid = match call { - 0 => filters == vec![col("a").gt(lit(10_i64)), col("b").lt(lit(2.5_f64))], - 1 => filters.is_empty(), - _ => false, - }; - - if !valid { - return Err(DataFusionError::Internal(format!( - "Unexpected DELETE filters for call {call}" - ))); - } - Ok(dml_count_plan()) + ) -> BoxFuture<'a, Result>> { + Box::pin(async move { + let call = self.delete_calls.fetch_add(1, Ordering::Relaxed); + let valid = match call { + 0 => { + filters + == vec![col("a").gt(lit(10_i64)), col("b").lt(lit(2.5_f64))] + } + 1 => filters.is_empty(), + _ => false, + }; + + if !valid { + return Err(DataFusionError::Internal(format!( + "Unexpected DELETE filters for call {call}" + ))); + } + Ok(dml_count_plan()) + }) } - async fn update( - &self, - _state: &dyn Session, + fn update<'a>( + &'a self, + _state: &'a dyn Session, assignments: Vec<(String, Expr)>, filters: Vec, - ) -> Result> { - if assignments - != vec![ - ("b".to_string(), lit(42_f64)), - ("a".to_string(), lit(7_i64)), - ] - { - return Err(DataFusionError::Internal( - "Unexpected UPDATE assignments".to_string(), - )); - } - let call = self.update_calls.fetch_add(1, Ordering::Relaxed); - let valid = match call { - 0 => filters == vec![col("a").eq(lit(7_i64)), col("b").gt(lit(1.5_f64))], - 1 => filters.is_empty(), - _ => false, - }; - - if !valid { - return Err(DataFusionError::Internal(format!( - "Unexpected UPDATE filters for call {call}" - ))); - } - Ok(dml_count_plan()) + ) -> BoxFuture<'a, Result>> { + Box::pin(async move { + if assignments + != vec![ + ("b".to_string(), lit(42_f64)), + ("a".to_string(), lit(7_i64)), + ] + { + return Err(DataFusionError::Internal( + "Unexpected UPDATE assignments".to_string(), + )); + } + let call = self.update_calls.fetch_add(1, Ordering::Relaxed); + let valid = match call { + 0 => { + filters + == vec![col("a").eq(lit(7_i64)), col("b").gt(lit(1.5_f64))] + } + 1 => filters.is_empty(), + _ => false, + }; + + if !valid { + return Err(DataFusionError::Internal(format!( + "Unexpected UPDATE filters for call {call}" + ))); + } + Ok(dml_count_plan()) + }) } - async fn truncate(&self, _state: &dyn Session) -> Result> { - Ok(dml_count_plan()) + fn truncate<'a>( + &'a self, + _state: &'a dyn Session, + ) -> BoxFuture<'a, Result>> { + Box::pin(async move { Ok(dml_count_plan()) }) } } diff --git a/datafusion/ffi/src/table_provider_factory.rs b/datafusion/ffi/src/table_provider_factory.rs index 70cab63ede2cb..b3a65c06c688c 100644 --- a/datafusion/ffi/src/table_provider_factory.rs +++ b/datafusion/ffi/src/table_provider_factory.rs @@ -15,10 +15,10 @@ // specific language governing permissions and limitations // under the License. +use futures::future::BoxFuture; use std::{ffi::c_void, sync::Arc}; use async_ffi::{FfiFuture, FutureExt}; -use async_trait::async_trait; use datafusion_catalog::{Session, TableProvider, TableProviderFactory}; use datafusion_common::error::{DataFusionError, Result}; use datafusion_execution::TaskContext; @@ -287,25 +287,26 @@ impl ForeignTableProviderFactory { unsafe impl Send for ForeignTableProviderFactory {} unsafe impl Sync for ForeignTableProviderFactory {} - -#[async_trait] impl TableProviderFactory for ForeignTableProviderFactory { - async fn create( - &self, - session: &dyn Session, - cmd: &CreateExternalTable, - ) -> Result> { - let session = FFI_SessionRef::new(session, None, self.0.logical_codec.clone()); - let cmd = self.serialize_cmd(cmd.clone())?; - - let provider = unsafe { - let maybe_provider = (self.0.create)(&self.0, session, cmd).await; - - let ffi_provider = df_result!(maybe_provider)?; - ForeignTableProvider(ffi_provider) - }; - - Ok(Arc::new(provider)) + fn create<'a>( + &'a self, + session: &'a dyn Session, + cmd: &'a CreateExternalTable, + ) -> BoxFuture<'a, Result>> { + Box::pin(async move { + let session = + FFI_SessionRef::new(session, None, self.0.logical_codec.clone()); + let cmd = self.serialize_cmd(cmd.clone())?; + + let provider = unsafe { + let maybe_provider = (self.0.create)(&self.0, session, cmd).await; + + let ffi_provider = df_result!(maybe_provider)?; + ForeignTableProvider(ffi_provider) + }; + + Ok(Arc::new(provider) as Arc) + }) } } @@ -321,36 +322,39 @@ mod tests { #[derive(Debug)] struct TestTableProviderFactory {} - - #[async_trait] impl TableProviderFactory for TestTableProviderFactory { - async fn create( - &self, - _session: &dyn Session, - _cmd: &CreateExternalTable, - ) -> Result> { - use arrow::datatypes::Field; - use datafusion::arrow::array::Float32Array; - use datafusion::arrow::datatypes::DataType; - use datafusion::arrow::record_batch::RecordBatch; - use datafusion::datasource::MemTable; - - let schema = - Arc::new(Schema::new(vec![Field::new("a", DataType::Float32, false)])); - - let batch1 = RecordBatch::try_new( - Arc::clone(&schema), - vec![Arc::new(Float32Array::from(vec![2.0, 4.0, 8.0]))], - )?; - let batch2 = RecordBatch::try_new( - Arc::clone(&schema), - vec![Arc::new(Float32Array::from(vec![64.0]))], - )?; - - Ok(Arc::new(MemTable::try_new( - schema, - vec![vec![batch1], vec![batch2]], - )?)) + fn create<'a>( + &'a self, + _session: &'a dyn Session, + _cmd: &'a CreateExternalTable, + ) -> BoxFuture<'a, Result>> { + Box::pin(async move { + use arrow::datatypes::Field; + use datafusion::arrow::array::Float32Array; + use datafusion::arrow::datatypes::DataType; + use datafusion::arrow::record_batch::RecordBatch; + use datafusion::datasource::MemTable; + + let schema = Arc::new(Schema::new(vec![Field::new( + "a", + DataType::Float32, + false, + )])); + + let batch1 = RecordBatch::try_new( + Arc::clone(&schema), + vec![Arc::new(Float32Array::from(vec![2.0, 4.0, 8.0]))], + )?; + let batch2 = RecordBatch::try_new( + Arc::clone(&schema), + vec![Arc::new(Float32Array::from(vec![64.0]))], + )?; + + Ok( + Arc::new(MemTable::try_new(schema, vec![vec![batch1], vec![batch2]])?) + as Arc, + ) + }) } } diff --git a/datafusion/ffi/src/tests/async_provider.rs b/datafusion/ffi/src/tests/async_provider.rs index 8217f52fa29f7..418843cb7f1a3 100644 --- a/datafusion/ffi/src/tests/async_provider.rs +++ b/datafusion/ffi/src/tests/async_provider.rs @@ -25,12 +25,12 @@ //! access the runtime, then you will get a panic when trying to do operations //! such as spawning a tokio task. +use futures::future::BoxFuture; use std::fmt::Debug; use std::sync::Arc; use arrow::array::RecordBatch; use arrow::datatypes::Schema; -use async_trait::async_trait; use datafusion_catalog::{MemoryCatalogProvider, TableProvider}; use datafusion_common::tree_node::TreeNodeRecursion; use datafusion_common::{Result, exec_err}; @@ -125,8 +125,6 @@ pub fn start_async_provider() -> (AsyncTableProvider, Handle) { (table_provider, tokio_rt) } - -#[async_trait] impl TableProvider for AsyncTableProvider { fn schema(&self) -> Arc { super::create_test_schema() @@ -136,34 +134,37 @@ impl TableProvider for AsyncTableProvider { datafusion_expr::TableType::Base } - async fn scan( - &self, - state: &dyn Session, - _projection: Option<&[usize]>, - _filters: &[Expr], + fn scan<'a>( + &'a self, + state: &'a dyn Session, + _projection: Option<&'a [usize]>, + _filters: &'a [Expr], _limit: Option, - ) -> Result> { - let catalog = state.catalog_list().catalog("datafusion").ok_or_else(|| { - datafusion_common::exec_datafusion_err!("missing datafusion catalog") - })?; - let schema = catalog.schema("public").ok_or_else(|| { - datafusion_common::exec_datafusion_err!("missing public schema") - })?; - if schema.table("external_table").await?.is_none() { - return exec_err!("missing external_table"); - } + ) -> BoxFuture<'a, Result>> { + Box::pin(async move { + let catalog = + state.catalog_list().catalog("datafusion").ok_or_else(|| { + datafusion_common::exec_datafusion_err!("missing datafusion catalog") + })?; + let schema = catalog.schema("public").ok_or_else(|| { + datafusion_common::exec_datafusion_err!("missing public schema") + })?; + if schema.table("external_table").await?.is_none() { + return exec_err!("missing external_table"); + } - // Register a catalog from the dynamically loaded library so the host - // can verify that catalog mutations cross the FFI boundary as well. - state.catalog_list().register_catalog( - "ffi_registered".to_owned(), - Arc::new(MemoryCatalogProvider::new()), - ); - - Ok(Arc::new(AsyncTestExecutionPlan::new( - self.batch_request.clone(), - self.batch_receiver.resubscribe(), - ))) + // Register a catalog from the dynamically loaded library so the host + // can verify that catalog mutations cross the FFI boundary as well. + state.catalog_list().register_catalog( + "ffi_registered".to_owned(), + Arc::new(MemoryCatalogProvider::new()), + ); + + Ok(Arc::new(AsyncTestExecutionPlan::new( + self.batch_request.clone(), + self.batch_receiver.resubscribe(), + ))) + }) } } diff --git a/datafusion/ffi/src/tests/catalog.rs b/datafusion/ffi/src/tests/catalog.rs index b0b0858a8a3d7..74541ac0f48ef 100644 --- a/datafusion/ffi/src/tests/catalog.rs +++ b/datafusion/ffi/src/tests/catalog.rs @@ -25,11 +25,11 @@ //! access the runtime, then you will get a panic when trying to do operations //! such as spawning a tokio task. +use futures::future::BoxFuture; use std::fmt::Debug; use std::sync::Arc; use arrow::datatypes::Schema; -use async_trait::async_trait; use datafusion_catalog::{ CatalogProvider, CatalogProviderList, MemTable, MemoryCatalogProvider, MemoryCatalogProviderList, MemorySchemaProvider, SchemaProvider, TableProvider, @@ -85,15 +85,16 @@ impl Default for FixedSchemaProvider { Self { inner } } } - -#[async_trait] impl SchemaProvider for FixedSchemaProvider { fn table_names(&self) -> Vec { self.inner.table_names() } - async fn table(&self, name: &str) -> Result>> { - self.inner.table(name).await + fn table<'a>( + &'a self, + name: &'a str, + ) -> BoxFuture<'a, Result>>> { + Box::pin(async move { self.inner.table(name).await }) } fn table_exist(&self, name: &str) -> bool { diff --git a/datafusion/ffi/src/tests/mod.rs b/datafusion/ffi/src/tests/mod.rs index 7e583b6c5d5bf..9bdd2a6e22b4f 100644 --- a/datafusion/ffi/src/tests/mod.rs +++ b/datafusion/ffi/src/tests/mod.rs @@ -15,13 +15,13 @@ // specific language governing permissions and limitations // under the License. +use futures::future::BoxFuture; use std::sync::Arc; use std::sync::atomic::{AtomicUsize, Ordering}; use arrow::array::{RecordBatch, record_batch}; use arrow_schema::{DataType, Field, Schema}; use async_provider::create_async_table_provider; -use async_trait::async_trait; use catalog::create_catalog_provider; use datafusion_catalog::MemTable; use datafusion_catalog::{Session, TableProvider}; @@ -282,8 +282,6 @@ struct TableWithStats { delete_calls: AtomicUsize, update_calls: AtomicUsize, } - -#[async_trait] impl TableProvider for TableWithStats { fn schema(&self) -> arrow_schema::SchemaRef { self.inner.schema() @@ -297,64 +295,73 @@ impl TableProvider for TableWithStats { Some(self.stats.clone()) } - async fn scan( - &self, - session: &dyn Session, - projection: Option<&[usize]>, - filters: &[Expr], + fn scan<'a>( + &'a self, + session: &'a dyn Session, + projection: Option<&'a [usize]>, + filters: &'a [Expr], limit: Option, - ) -> Result> { - self.inner.scan(session, projection, filters, limit).await + ) -> BoxFuture<'a, Result>> { + Box::pin( + async move { self.inner.scan(session, projection, filters, limit).await }, + ) } - async fn delete_from( - &self, - _state: &dyn Session, + fn delete_from<'a>( + &'a self, + _state: &'a dyn Session, filters: Vec, - ) -> Result> { - let call = self.delete_calls.fetch_add(1, Ordering::Relaxed); - let valid = match call { - 0 => filters == vec![col("a").gt(lit(10_i32)), col("b").lt(lit(2.5_f64))], - 1 => filters.is_empty(), - _ => false, - }; - if !valid { - return exec_err!("Unexpected DELETE filters for call {call}"); - } - Ok(dml_count_plan()) + ) -> BoxFuture<'a, Result>> { + Box::pin(async move { + let call = self.delete_calls.fetch_add(1, Ordering::Relaxed); + let valid = match call { + 0 => filters == vec![col("a").gt(lit(10_i32)), col("b").lt(lit(2.5_f64))], + 1 => filters.is_empty(), + _ => false, + }; + if !valid { + return exec_err!("Unexpected DELETE filters for call {call}"); + } + Ok(dml_count_plan()) + }) } - async fn update( - &self, - _state: &dyn Session, + fn update<'a>( + &'a self, + _state: &'a dyn Session, assignments: Vec<(String, Expr)>, filters: Vec, - ) -> Result> { - if assignments - != vec![ - ("b".to_string(), lit(42_f64)), - ("a".to_string(), lit(7_i32)), - ] - { - return exec_err!("Unexpected UPDATE assignments"); - } - - let call = self.update_calls.fetch_add(1, Ordering::Relaxed); - let valid = match call { - 0 => filters == vec![col("a").eq(lit(7_i32)), col("b").gt(lit(1.5_f64))], - 1 => filters.is_empty(), - _ => false, - }; - - if !valid { - return exec_err!("Unexpected UPDATE filters for call {call}"); - } - - Ok(dml_count_plan()) + ) -> BoxFuture<'a, Result>> { + Box::pin(async move { + if assignments + != vec![ + ("b".to_string(), lit(42_f64)), + ("a".to_string(), lit(7_i32)), + ] + { + return exec_err!("Unexpected UPDATE assignments"); + } + + let call = self.update_calls.fetch_add(1, Ordering::Relaxed); + let valid = match call { + 0 => filters == vec![col("a").eq(lit(7_i32)), col("b").gt(lit(1.5_f64))], + 1 => filters.is_empty(), + _ => false, + }; + + if !valid { + return exec_err!("Unexpected UPDATE filters for call {call}"); + } + + Ok(dml_count_plan()) + }) } - async fn truncate(&self, _state: &dyn Session) -> Result> { - Ok(dml_count_plan()) + fn truncate<'a>( + &'a self, + _state: &'a dyn Session, + ) -> BoxFuture<'a, Result>> { + Box::pin(async move { Ok(dml_count_plan()) }) } } diff --git a/datafusion/ffi/src/tests/query_planner.rs b/datafusion/ffi/src/tests/query_planner.rs index 9b5d6a07800bb..0e90c02c9adb4 100644 --- a/datafusion/ffi/src/tests/query_planner.rs +++ b/datafusion/ffi/src/tests/query_planner.rs @@ -15,11 +15,11 @@ // specific language governing permissions and limitations // under the License. +use futures::future::BoxFuture; use std::any::Any; use std::sync::Arc; use arrow::datatypes::{DataType, Field, Schema}; -use async_trait::async_trait; use datafusion_catalog::default_table_source::source_as_provider; use datafusion_common::{DataFusionError, Result, exec_err}; use datafusion_expr::LogicalPlan; @@ -39,58 +39,63 @@ use crate::util::FFI_Option; #[derive(Debug)] struct TestQueryPlanner; - -#[async_trait] impl QueryPlanner for TestQueryPlanner { - async fn create_physical_plan( - &self, - logical_plan: &LogicalPlan, - session: &dyn Session, - ) -> Result> { - if let LogicalPlan::TableScan(scan) = logical_plan { - if session.as_any().downcast_ref::().is_none() { - return exec_err!("library A's session was not foreign to library C"); + fn create_physical_plan<'a>( + &'a self, + logical_plan: &'a LogicalPlan, + session: &'a dyn Session, + ) -> BoxFuture<'a, Result>> { + Box::pin(async move { + if let LogicalPlan::TableScan(scan) = logical_plan { + if session.as_any().downcast_ref::().is_none() { + return exec_err!("library A's session was not foreign to library C"); + } + + let provider = source_as_provider(&scan.source)?; + if provider.downcast_ref::().is_none() { + return exec_err!( + "library B's provider was not foreign to library C" + ); + } + let library_b_plan = provider + .scan( + session, + scan.projection.as_deref(), + &scan.filters, + scan.fetch, + ) + .await?; + + if !library_b_plan.is::() { + return exec_err!( + "library B's plan unexpectedly downcast as C-local" + ); + } + + let plan = UnionExec::try_new(vec![ + Arc::clone(&library_b_plan), + Arc::clone(&library_b_plan), + ])?; + if !plan.is::() { + return exec_err!("library C could not downcast its local UnionExec"); + } + return Ok(plan); } - let provider = source_as_provider(&scan.source)?; - if provider.downcast_ref::().is_none() { - return exec_err!("library B's provider was not foreign to library C"); + let query_planner = session.query_planner(); + let planner_any: &dyn Any = query_planner.as_ref(); + if planner_any.downcast_ref::().is_none() { + return exec_err!("query planner did not cross the FFI boundary"); } - let library_b_plan = provider - .scan( - session, - scan.projection.as_deref(), - &scan.filters, - scan.fetch, - ) - .await?; - - if !library_b_plan.is::() { - return exec_err!("library B's plan unexpectedly downcast as C-local"); + session.optimize(logical_plan)?; + if session.physical_optimizers().is_empty() { + return exec_err!("physical optimizers did not cross the FFI boundary"); } - let plan = UnionExec::try_new(vec![ - Arc::clone(&library_b_plan), - Arc::clone(&library_b_plan), - ])?; - if !plan.is::() { - return exec_err!("library C could not downcast its local UnionExec"); - } - return Ok(plan); - } - - let query_planner = session.query_planner(); - let planner_any: &dyn Any = query_planner.as_ref(); - if planner_any.downcast_ref::().is_none() { - return exec_err!("query planner did not cross the FFI boundary"); - } - session.optimize(logical_plan)?; - if session.physical_optimizers().is_empty() { - return exec_err!("physical optimizers did not cross the FFI boundary"); - } - - let schema = Arc::new(Schema::new(vec![Field::new("a", DataType::Int32, true)])); - Ok(Arc::new(EmptyExec::new(schema))) + let schema = + Arc::new(Schema::new(vec![Field::new("a", DataType::Int32, true)])); + Ok(Arc::new(EmptyExec::new(schema))) + }) } } @@ -102,68 +107,68 @@ impl QueryPlanner for TestQueryPlanner { struct SwappedQueryPlanner { library_a_planner: Arc, } - -#[async_trait] impl QueryPlanner for SwappedQueryPlanner { - async fn create_physical_plan( - &self, - logical_plan: &LogicalPlan, - session: &dyn Session, - ) -> Result> { - if session.as_any().downcast_ref::().is_none() { - return exec_err!("library A's session was not foreign to library C"); - } - - // After the swap, the planner installed on library A's session is this - // planner, so `session.query_planner()` is a self-reference. - let installed = session.query_planner(); - let installed: &dyn Any = installed.as_ref(); - if installed.downcast_ref::().is_none() { - return exec_err!( - "expected the swapped session to report library C's own planner" - ); - } - - // Direct session delegation used to re-enter this planner recursively. - // A foreign session must reject it before dispatching to the installed - // planner, while the captured planner path below remains usable. - let direct_error = session - .create_physical_plan(logical_plan) - .await - .expect_err("direct foreign-session planning should be unsupported"); - if !matches!(direct_error, DataFusionError::NotImplemented(_)) { - return exec_err!( - "expected direct foreign-session planning to return NotImplemented; got {direct_error}" - ); - } - - // Delegate to library A. The result crosses the FFI boundary as - // serialized bytes, so library C receives nodes carrying its own local - // Rust type identities. - let plan = self - .library_a_planner - .create_physical_plan(logical_plan, session) - .await?; - - if plan.is::() { - return exec_err!("library A's plan was opaque to library C"); - } - let Some(sort) = plan.downcast_ref::() else { - return exec_err!( - "library C could not downcast library A's SortExec; got {}", - plan.name() - ); - }; - // Library B's scan is still foreign to library C. Only a codec boundary - // reconstructs it, and library A's codec hands back an A-local node. - if !sort.input().is::() { - return exec_err!("library B's scan unexpectedly downcast as C-local"); - } - - Ok(UnionExec::try_new(vec![ - Arc::clone(&plan), - Arc::clone(&plan), - ])?) + fn create_physical_plan<'a>( + &'a self, + logical_plan: &'a LogicalPlan, + session: &'a dyn Session, + ) -> BoxFuture<'a, Result>> { + Box::pin(async move { + if session.as_any().downcast_ref::().is_none() { + return exec_err!("library A's session was not foreign to library C"); + } + + // After the swap, the planner installed on library A's session is this + // planner, so `session.query_planner()` is a self-reference. + let installed = session.query_planner(); + let installed: &dyn Any = installed.as_ref(); + if installed.downcast_ref::().is_none() { + return exec_err!( + "expected the swapped session to report library C's own planner" + ); + } + + // Direct session delegation used to re-enter this planner recursively. + // A foreign session must reject it before dispatching to the installed + // planner, while the captured planner path below remains usable. + let direct_error = session + .create_physical_plan(logical_plan) + .await + .expect_err("direct foreign-session planning should be unsupported"); + if !matches!(direct_error, DataFusionError::NotImplemented(_)) { + return exec_err!( + "expected direct foreign-session planning to return NotImplemented; got {direct_error}" + ); + } + + // Delegate to library A. The result crosses the FFI boundary as + // serialized bytes, so library C receives nodes carrying its own local + // Rust type identities. + let plan = self + .library_a_planner + .create_physical_plan(logical_plan, session) + .await?; + + if plan.is::() { + return exec_err!("library A's plan was opaque to library C"); + } + let Some(sort) = plan.downcast_ref::() else { + return exec_err!( + "library C could not downcast library A's SortExec; got {}", + plan.name() + ); + }; + // Library B's scan is still foreign to library C. Only a codec boundary + // reconstructs it, and library A's codec hands back an A-local node. + if !sort.input().is::() { + return exec_err!("library B's scan unexpectedly downcast as C-local"); + } + + Ok(UnionExec::try_new(vec![ + Arc::clone(&plan), + Arc::clone(&plan), + ])?) + }) } } diff --git a/datafusion/ffi/src/tests/table_provider_factory.rs b/datafusion/ffi/src/tests/table_provider_factory.rs index 29af6aacf6484..18ba4166e03c7 100644 --- a/datafusion/ffi/src/tests/table_provider_factory.rs +++ b/datafusion/ffi/src/tests/table_provider_factory.rs @@ -15,9 +15,9 @@ // specific language governing permissions and limitations // under the License. +use futures::future::BoxFuture; use std::sync::Arc; -use async_trait::async_trait; use datafusion_catalog::{MemTable, Session, TableProvider, TableProviderFactory}; use datafusion_common::Result; use datafusion_expr::CreateExternalTable; @@ -28,27 +28,27 @@ use crate::table_provider_factory::FFI_TableProviderFactory; #[derive(Debug)] pub struct TestTableProviderFactory {} - -#[async_trait] impl TableProviderFactory for TestTableProviderFactory { - async fn create( - &self, - _session: &dyn Session, - _cmd: &CreateExternalTable, - ) -> Result> { - let schema = create_test_schema(); - - // It is useful to create these as multiple record batches - // so that we can demonstrate the FFI stream. - let batches = vec![ - create_record_batch(1, 5), - create_record_batch(6, 1), - create_record_batch(7, 5), - ]; - - let table_provider = MemTable::try_new(schema, vec![batches]).unwrap(); - - Ok(Arc::new(table_provider)) + fn create<'a>( + &'a self, + _session: &'a dyn Session, + _cmd: &'a CreateExternalTable, + ) -> BoxFuture<'a, Result>> { + Box::pin(async move { + let schema = create_test_schema(); + + // It is useful to create these as multiple record batches + // so that we can demonstrate the FFI stream. + let batches = vec![ + create_record_batch(1, 5), + create_record_batch(6, 1), + create_record_batch(7, 5), + ]; + + let table_provider = MemTable::try_new(schema, vec![batches]).unwrap(); + + Ok(Arc::new(table_provider)) + }) } } diff --git a/datafusion/functions-table/Cargo.toml b/datafusion/functions-table/Cargo.toml index fb02c2c5e2cb3..95b486c491f18 100644 --- a/datafusion/functions-table/Cargo.toml +++ b/datafusion/functions-table/Cargo.toml @@ -42,9 +42,9 @@ name = "datafusion_functions_table" [dependencies] arrow = { workspace = true } -async-trait = { workspace = true } datafusion-catalog = { workspace = true } datafusion-common = { workspace = true } +futures = { workspace = true } datafusion-expr = { workspace = true } datafusion-physical-expr = { workspace = true } datafusion-physical-plan = { workspace = true } diff --git a/datafusion/functions-table/src/generate_series.rs b/datafusion/functions-table/src/generate_series.rs index cf5e026581584..99a18ab63f51e 100644 --- a/datafusion/functions-table/src/generate_series.rs +++ b/datafusion/functions-table/src/generate_series.rs @@ -23,7 +23,6 @@ use arrow::datatypes::{ DataType, Field, IntervalMonthDayNano, Schema, SchemaRef, TimeUnit, }; use arrow::record_batch::RecordBatch; -use async_trait::async_trait; use datafusion_catalog::TableFunctionImpl; use datafusion_catalog::TableProvider; use datafusion_catalog::{Session, TableFunctionArgs}; @@ -33,6 +32,7 @@ use datafusion_physical_expr::PhysicalSortExpr; use datafusion_physical_expr::expressions::Column; use datafusion_physical_plan::ExecutionPlan; use datafusion_physical_plan::memory::{LazyBatchGenerator, LazyMemoryExec}; +use futures::future::BoxFuture; use parking_lot::RwLock; use std::any::Any; use std::fmt; @@ -541,8 +541,6 @@ fn validate_interval_step(step: IntervalMonthDayNano) -> Result<()> { Ok(()) } - -#[async_trait] impl TableProvider for GenerateSeriesTable { fn schema(&self) -> SchemaRef { Arc::clone(&self.schema) @@ -552,23 +550,25 @@ impl TableProvider for GenerateSeriesTable { TableType::Base } - async fn scan( - &self, - state: &dyn Session, - projection: Option<&[usize]>, - _filters: &[Expr], + fn scan<'a>( + &'a self, + state: &'a dyn Session, + projection: Option<&'a [usize]>, + _filters: &'a [Expr], _limit: Option, - ) -> Result> { - let batch_size = state.config_options().execution.batch_size.get(); - let generator = self.as_generator(batch_size)?; - let mut exec = LazyMemoryExec::try_new(self.schema(), vec![generator])? - .with_projection(projection.map(|p| p.to_vec())); - - if let Some(ordering) = self.output_ordering(exec.schema().as_ref()) { - exec.add_ordering([ordering]); - } + ) -> BoxFuture<'a, Result>> { + Box::pin(async move { + let batch_size = state.config_options().execution.batch_size.get(); + let generator = self.as_generator(batch_size)?; + let mut exec = LazyMemoryExec::try_new(self.schema(), vec![generator])? + .with_projection(projection.map(|p| p.to_vec())); + + if let Some(ordering) = self.output_ordering(exec.schema().as_ref()) { + exec.add_ordering([ordering]); + } - Ok(Arc::new(exec)) + Ok(Arc::new(exec) as Arc) + }) } } diff --git a/datafusion/optimizer/Cargo.toml b/datafusion/optimizer/Cargo.toml index 31a73062e119e..b04b115a9b1a6 100644 --- a/datafusion/optimizer/Cargo.toml +++ b/datafusion/optimizer/Cargo.toml @@ -66,7 +66,6 @@ regex = { workspace = true } regex-syntax = "0.8.9" [dev-dependencies] -async-trait = { workspace = true } criterion = { workspace = true } ctor = { workspace = true } datafusion-functions-aggregate = { workspace = true } diff --git a/datafusion/optimizer/src/push_down_filter.rs b/datafusion/optimizer/src/push_down_filter.rs index 8a1dcc12ef874..a3e34a3274044 100644 --- a/datafusion/optimizer/src/push_down_filter.rs +++ b/datafusion/optimizer/src/push_down_filter.rs @@ -1438,7 +1438,6 @@ mod tests { use std::fmt::{Debug, Formatter}; use arrow::datatypes::{Field, Schema, SchemaRef}; - use async_trait::async_trait; use datafusion_common::{DFSchemaRef, DataFusionError, ScalarValue}; use datafusion_expr::expr::ScalarFunction; @@ -3100,8 +3099,6 @@ mod tests { struct PushDownProvider { pub filter_support: TableProviderFilterPushDown, } - - #[async_trait] impl TableSource for PushDownProvider { fn schema(&self) -> SchemaRef { Arc::new(Schema::new(vec![ diff --git a/datafusion/physical-plan/Cargo.toml b/datafusion/physical-plan/Cargo.toml index 534cea8ea9cbb..d1b983b3dc180 100644 --- a/datafusion/physical-plan/Cargo.toml +++ b/datafusion/physical-plan/Cargo.toml @@ -64,7 +64,6 @@ arrow-data = { workspace = true } arrow-ipc = { workspace = true, features = ["lz4", "zstd"] } arrow-ord = { workspace = true } arrow-schema = { workspace = true } -async-trait = { workspace = true } bytes = { workspace = true } datafusion-common = { workspace = true } datafusion-common-runtime = { workspace = true, default-features = true } diff --git a/datafusion/physical-plan/src/streaming.rs b/datafusion/physical-plan/src/streaming.rs index dd57f958e83ec..7093ba2a76353 100644 --- a/datafusion/physical-plan/src/streaming.rs +++ b/datafusion/physical-plan/src/streaming.rs @@ -43,7 +43,6 @@ use datafusion_physical_expr::PhysicalExpr; use datafusion_physical_expr::projection::ProjectionMapping; use datafusion_physical_expr::{EquivalenceProperties, LexOrdering}; -use async_trait::async_trait; use futures::stream::StreamExt; use log::debug; @@ -257,8 +256,6 @@ impl DisplayAs for StreamingTableExec { } } } - -#[async_trait] impl ExecutionPlan for StreamingTableExec { fn name(&self) -> &'static str { "StreamingTableExec" diff --git a/datafusion/physical-plan/src/unnest.rs b/datafusion/physical-plan/src/unnest.rs index d93e0280515c6..105369454eeff 100644 --- a/datafusion/physical-plan/src/unnest.rs +++ b/datafusion/physical-plan/src/unnest.rs @@ -44,7 +44,6 @@ use arrow::compute::{cast, is_not_null, kernels, sum}; use arrow::datatypes::{DataType, Int64Type, Schema, SchemaRef}; use arrow::record_batch::RecordBatch; use arrow_ord::cmp::lt; -use async_trait::async_trait; use datafusion_common::tree_node::TreeNodeRecursion; use datafusion_common::{ Constraints, HashMap, HashSet, Result, UnnestOptions, exec_datafusion_err, exec_err, @@ -605,8 +604,6 @@ impl RecordBatchStream for UnnestStream { Arc::clone(&self.schema) } } - -#[async_trait] impl Stream for UnnestStream { type Item = Result; diff --git a/datafusion/proto/Cargo.toml b/datafusion/proto/Cargo.toml index 4b11bea01103f..1619f7c7822ef 100644 --- a/datafusion/proto/Cargo.toml +++ b/datafusion/proto/Cargo.toml @@ -80,7 +80,7 @@ recursive = { workspace = true, optional = true } serde_json = { workspace = true, optional = true } [dev-dependencies] -async-trait = { workspace = true } +futures = { workspace = true } datafusion = { workspace = true, default-features = false, features = [ "sql", "compression", diff --git a/datafusion/proto/tests/cases/plans/sinks.rs b/datafusion/proto/tests/cases/plans/sinks.rs index 63bf11dd30f23..73ebd11a3e425 100644 --- a/datafusion/proto/tests/cases/plans/sinks.rs +++ b/datafusion/proto/tests/cases/plans/sinks.rs @@ -19,7 +19,6 @@ use super::{roundtrip_test, roundtrip_test_and_return}; use arrow::csv::WriterBuilder; -use async_trait::async_trait; use datafusion::arrow::compute::kernels::sort::SortOptions; use datafusion::arrow::datatypes::{DataType, Field, Schema, SchemaRef}; use datafusion::datasource::file_format::csv::CsvSink; @@ -51,6 +50,7 @@ use datafusion_proto::physical_plan::{ }; use datafusion_proto::protobuf; use datafusion_proto::protobuf::PhysicalPlanNode; +use futures::future::BoxFuture; use std::fmt::Formatter; use std::sync::Arc; use std::vec; @@ -65,19 +65,19 @@ impl DisplayAs for ProtoHookSink { write!(f, "ProtoHookSink") } } - -#[async_trait] impl DataSink for ProtoHookSink { fn schema(&self) -> &SchemaRef { &self.schema } - async fn write_all( - &self, + fn write_all<'a>( + &'a self, _data: SendableRecordBatchStream, - _context: &Arc, - ) -> Result { - unreachable!("serialization test does not execute the sink") + _context: &'a Arc, + ) -> BoxFuture<'a, Result> { + Box::pin( + async move { unreachable!("serialization test does not execute the sink") }, + ) } fn try_to_proto( diff --git a/datafusion/proto/tests/cases/plans/udfs.rs b/datafusion/proto/tests/cases/plans/udfs.rs index 08d00030e04e1..db06eaa475717 100644 --- a/datafusion/proto/tests/cases/plans/udfs.rs +++ b/datafusion/proto/tests/cases/plans/udfs.rs @@ -547,13 +547,12 @@ async fn roundtrip_async_func_exec() -> Result<()> { } } - #[async_trait::async_trait] impl AsyncScalarUDFImpl for TestAsyncUDF { - async fn invoke_async_with_args( + fn invoke_async_with_args( &self, args: ScalarFunctionArgs, - ) -> Result { - Ok(args.args[0].clone()) + ) -> futures::future::BoxFuture<'_, Result> { + Box::pin(async move { Ok(args.args[0].clone()) }) } } diff --git a/datafusion/session/Cargo.toml b/datafusion/session/Cargo.toml index e6c277a8b8493..6af0195b893f0 100644 --- a/datafusion/session/Cargo.toml +++ b/datafusion/session/Cargo.toml @@ -32,7 +32,6 @@ all-features = true [dependencies] arrow-schema = { workspace = true } -async-trait = { workspace = true } datafusion-common = { workspace = true } datafusion-execution = { workspace = true } datafusion-expr = { workspace = true } diff --git a/datafusion/session/src/planner.rs b/datafusion/session/src/planner.rs index 37726009f0f4d..13729aae4cb4e 100644 --- a/datafusion/session/src/planner.rs +++ b/datafusion/session/src/planner.rs @@ -21,23 +21,22 @@ use std::any::Any; use std::fmt::Debug; use std::sync::Arc; -use async_trait::async_trait; use datafusion_common::{DFSchema, Result, not_impl_err}; use datafusion_expr::physical_planning_context::PhysicalPlanningContext; use datafusion_expr::{Expr, LogicalPlan, TableScan, UserDefinedLogicalNode}; use datafusion_physical_plan::{ExecutionPlan, PhysicalExpr}; +use futures::future::BoxFuture; use crate::Session; /// A planner that creates a physical plan for a query. -#[async_trait] pub trait QueryPlanner: Any + Debug { /// Given a [`LogicalPlan`], create an [`ExecutionPlan`] suitable for execution - async fn create_physical_plan( - &self, - logical_plan: &LogicalPlan, - session: &dyn Session, - ) -> Result>; + fn create_physical_plan<'a>( + &'a self, + logical_plan: &'a LogicalPlan, + session: &'a dyn Session, + ) -> BoxFuture<'a, Result>>; } /// A query planner that reports that planning is not implemented. @@ -47,27 +46,27 @@ pub trait QueryPlanner: Any + Debug { #[derive(Debug, Default)] pub struct UnsupportedQueryPlanner; -#[async_trait] impl QueryPlanner for UnsupportedQueryPlanner { - async fn create_physical_plan( - &self, - _logical_plan: &LogicalPlan, - _session: &dyn Session, - ) -> Result> { - not_impl_err!("This session does not expose its query planner") + fn create_physical_plan<'a>( + &'a self, + _logical_plan: &'a LogicalPlan, + _session: &'a dyn Session, + ) -> BoxFuture<'a, Result>> { + Box::pin(async { + not_impl_err!("This session does not expose its query planner") + }) } } /// Physical query planner that converts a [`LogicalPlan`] to an /// [`ExecutionPlan`] suitable for execution. -#[async_trait] pub trait PhysicalPlanner: Send + Sync { /// Create a physical plan from a logical plan - async fn create_physical_plan( - &self, - logical_plan: &LogicalPlan, - session: &dyn Session, - ) -> Result>; + fn create_physical_plan<'a>( + &'a self, + logical_plan: &'a LogicalPlan, + session: &'a dyn Session, + ) -> BoxFuture<'a, Result>>; /// Create a physical expression from a logical expression /// suitable for evaluation @@ -92,7 +91,6 @@ pub trait PhysicalPlanner: Send + Sync { } /// This trait exposes the ability to plan an [`ExecutionPlan`] out of a [`LogicalPlan`]. -#[async_trait] pub trait ExtensionPlanner { /// Create a physical plan for a [`UserDefinedLogicalNode`]. /// @@ -110,15 +108,15 @@ pub trait ExtensionPlanner { /// [`PhysicalPlanner::create_physical_expr`] when creating this node's /// physical expressions so that scalar subqueries resolve against the same /// subquery state as the rest of the plan. - async fn plan_extension( - &self, - planner: &dyn PhysicalPlanner, - node: &dyn UserDefinedLogicalNode, - logical_inputs: &[&LogicalPlan], - physical_inputs: &[Arc], - session: &dyn Session, - planning_ctx: &PhysicalPlanningContext, - ) -> Result>>; + fn plan_extension<'a>( + &'a self, + planner: &'a dyn PhysicalPlanner, + node: &'a dyn UserDefinedLogicalNode, + logical_inputs: &'a [&'a LogicalPlan], + physical_inputs: &'a [Arc], + session: &'a dyn Session, + planning_ctx: &'a PhysicalPlanningContext, + ) -> BoxFuture<'a, Result>>>; /// Create a physical plan for a [`LogicalPlan::TableScan`]. /// @@ -138,7 +136,6 @@ pub trait ExtensionPlanner { /// use datafusion::catalog::Session; /// use datafusion::error::Result; /// use datafusion_session::{ExtensionPlanner, PhysicalPlanner}; - /// use async_trait::async_trait; /// /// // Your custom table source type /// struct MyCustomTableSource { /* ... */ } @@ -148,51 +145,52 @@ pub trait ExtensionPlanner { /// /// struct MyExtensionPlanner; /// - /// #[async_trait] /// impl ExtensionPlanner for MyExtensionPlanner { - /// async fn plan_extension( - /// &self, - /// _planner: &dyn PhysicalPlanner, - /// _node: &dyn UserDefinedLogicalNode, - /// _logical_inputs: &[&LogicalPlan], - /// _physical_inputs: &[Arc], - /// _session: &dyn Session, - /// _planning_ctx: &PhysicalPlanningContext, - /// ) -> Result>> { - /// Ok(None) + /// fn plan_extension<'a>( + /// &'a self, + /// _planner: &'a dyn PhysicalPlanner, + /// _node: &'a dyn UserDefinedLogicalNode, + /// _logical_inputs: &'a [&'a LogicalPlan], + /// _physical_inputs: &'a [Arc], + /// _session: &'a dyn Session, + /// _planning_ctx: &'a PhysicalPlanningContext, + /// ) -> BoxFuture<'a, Result>>> { + /// Box::pin(async { Ok(None) }) /// } /// - /// async fn plan_table_scan( - /// &self, - /// _planner: &dyn PhysicalPlanner, - /// scan: &TableScan, - /// _session: &dyn Session, - /// _planning_ctx: &PhysicalPlanningContext, - /// ) -> Result>> { - /// // Check if this is your custom table source - /// if scan.source.is::() { - /// // Create a custom execution plan for your table source - /// let exec = MyCustomExec::new( - /// scan.table_name.clone(), - /// Arc::clone(scan.projected_schema.inner()), - /// ); - /// Ok(Some(Arc::new(exec))) - /// } else { - /// // Return None to let other extension planners handle it - /// Ok(None) - /// } + /// fn plan_table_scan<'a>( + /// &'a self, + /// _planner: &'a dyn PhysicalPlanner, + /// scan: &'a TableScan, + /// _session: &'a dyn Session, + /// _planning_ctx: &'a PhysicalPlanningContext, + /// ) -> BoxFuture<'a, Result>>> { + /// Box::pin(async move { + /// // Check if this is your custom table source + /// if scan.source.is::() { + /// // Create a custom execution plan for your table source + /// let exec = MyCustomExec::new( + /// scan.table_name.clone(), + /// Arc::clone(scan.projected_schema.inner()), + /// ); + /// Ok(Some(Arc::new(exec))) + /// } else { + /// // Return None to let other extension planners handle it + /// Ok(None) + /// } + /// }) /// } /// } /// ``` /// /// [`TableSource`]: datafusion_expr::TableSource - async fn plan_table_scan( - &self, - _planner: &dyn PhysicalPlanner, - _scan: &TableScan, - _session: &dyn Session, - _planning_ctx: &PhysicalPlanningContext, - ) -> Result>> { - Ok(None) + fn plan_table_scan<'a>( + &'a self, + _planner: &'a dyn PhysicalPlanner, + _scan: &'a TableScan, + _session: &'a dyn Session, + _planning_ctx: &'a PhysicalPlanningContext, + ) -> BoxFuture<'a, Result>>> { + Box::pin(async { Ok(None) }) } } diff --git a/datafusion/session/src/schema.rs b/datafusion/session/src/schema.rs index 7a66072bb4d8a..fc101da52a9aa 100644 --- a/datafusion/session/src/schema.rs +++ b/datafusion/session/src/schema.rs @@ -18,8 +18,8 @@ //! Describes the interface and built-in implementations of schemas, //! representing collections of named tables. -use async_trait::async_trait; use datafusion_common::{DataFusionError, exec_err}; +use futures::future::BoxFuture; use std::any::Any; use std::fmt::Debug; use std::sync::Arc; @@ -33,10 +33,9 @@ use datafusion_expr::TableType; /// Please see [`CatalogProvider`] for details of implementing a custom catalog. /// /// [`CatalogProvider`]: super::CatalogProvider -#[async_trait] pub trait SchemaProvider: Any + Debug + Sync + Send { /// Returns the owner of the Schema, default is None. This value is reported - /// as part of `information_schema.schemata`. + /// as part of `information_schema.schemata`.\ fn owner_name(&self) -> Option<&str> { None } @@ -46,17 +45,20 @@ pub trait SchemaProvider: Any + Debug + Sync + Send { /// Retrieves a specific table from the schema by name, if it exists, /// otherwise returns `None`. - async fn table( - &self, - name: &str, - ) -> Result>, DataFusionError>; + fn table<'a>( + &'a self, + name: &'a str, + ) -> BoxFuture<'a, Result>, DataFusionError>>; /// Retrieves the type of a specific table from the schema by name, if it exists, otherwise /// returns `None`. Implementations for which this operation is cheap but [Self::table] is /// expensive can override this to improve operations that only need the type, e.g. /// `SELECT * FROM information_schema.tables`. - async fn table_type(&self, name: &str) -> Result> { - self.table(name).await.map(|o| o.map(|t| t.table_type())) + fn table_type<'a>( + &'a self, + name: &'a str, + ) -> BoxFuture<'a, Result>> { + Box::pin(async move { self.table(name).await.map(|o| o.map(|t| t.table_type())) }) } /// If supported by the implementation, adds a new table named `name` to diff --git a/datafusion/session/src/session.rs b/datafusion/session/src/session.rs index f6143cc4a4d1d..0f1ee48b3e8c6 100644 --- a/datafusion/session/src/session.rs +++ b/datafusion/session/src/session.rs @@ -15,7 +15,6 @@ // specific language governing permissions and limitations // under the License. -use async_trait::async_trait; use datafusion_common::config::{ConfigOptions, TableOptions}; use datafusion_common::{DFSchema, Result}; use datafusion_execution::TaskContext; @@ -28,6 +27,7 @@ use datafusion_expr::{ }; use datafusion_physical_plan::operator_statistics::StatisticsRegistry; use datafusion_physical_plan::{ExecutionPlan, PhysicalExpr}; +use futures::future::BoxFuture; use crate::CatalogProviderList; use parking_lot::{Mutex, RwLock}; @@ -76,7 +76,6 @@ use crate::{PhysicalOptimizerRule, QueryPlanner, UnsupportedQueryPlanner}; /// /// [`SessionState`]: https://docs.rs/datafusion/latest/datafusion/execution/session_state/struct.SessionState.html /// [`TableProvider`]: https://docs.rs/datafusion/latest/datafusion/catalog/trait.TableProvider.html -#[async_trait] pub trait Session: Send + Sync { /// Return the session ID fn session_id(&self) -> &str; @@ -147,17 +146,18 @@ pub trait Session: Send + Sync { /// This function will error for [`LogicalPlan`]s such as catalog DDL like /// `CREATE TABLE`, which do not have corresponding physical plans and must /// be handled by another layer, typically the `SessionContext`. - async fn create_physical_plan( - &self, - logical_plan: &LogicalPlan, - ) -> Result>; + fn create_physical_plan<'a>( + &'a self, + logical_plan: &'a LogicalPlan, + ) -> BoxFuture<'a, Result>>; /// Create a [`PhysicalExpr`] from an [`Expr`] after applying type /// coercion, and function rewrites. /// - /// Note: The expression is not simplified or otherwise optimized: `a = 1 - /// + 2` will not be simplified to `a = 3` as this is a more involved process. - /// See the [expr_api] example for how to simplify expressions. + /// Note: The expression is not simplified or otherwise optimized: + /// `a = 1 + 2` will not be simplified to `a = 3` as this is a more + /// involved process. See the [expr_api] example for how to simplify + /// expressions. /// /// [expr_api]: https://github.com/apache/datafusion/blob/main/datafusion-examples/examples/query_planning/expr_api.rs fn create_physical_expr( diff --git a/datafusion/session/src/table.rs b/datafusion/session/src/table.rs index a6cddd10ce80a..b54d08f314d6b 100644 --- a/datafusion/session/src/table.rs +++ b/datafusion/session/src/table.rs @@ -18,12 +18,10 @@ use std::any::Any; use std::borrow::Cow; use std::fmt::Debug; -use std::future::ready; use std::sync::Arc; use crate::session::Session; use arrow_schema::SchemaRef; -use async_trait::async_trait; use datafusion_common::{Constraints, Statistics, not_impl_err}; use datafusion_common::{DFSchemaRef, Result, internal_err}; use datafusion_expr::Expr; @@ -50,7 +48,6 @@ use datafusion_physical_plan::ExecutionPlan; /// /// [`RecordBatch`]: https://docs.rs/arrow/latest/arrow/record_batch/struct.RecordBatch.html /// [`CatalogProvider`]: super::CatalogProvider -#[async_trait] pub trait TableProvider: Any + Debug + Sync + Send { /// Get a reference to the schema for this table fn schema(&self) -> SchemaRef; @@ -59,6 +56,7 @@ pub trait TableProvider: Any + Debug + Sync + Send { /// Returns: /// - `None` for tables that do not support constraints. /// - `Some(&Constraints)` for tables supporting constraints. + /// /// Therefore, a `Some(&Constraints::empty())` return value indicates that /// this table supports constraints, but there are no constraints. fn constraints(&self) -> Option<&Constraints> { @@ -184,13 +182,13 @@ pub trait TableProvider: Any + Debug + Sync + Send { /// /// As noted above, columns referenced only by pushed-down filters may be /// absent from `projection`. - async fn scan( - &self, - state: &dyn Session, - projection: Option<&[usize]>, - filters: &[Expr], + fn scan<'a>( + &'a self, + state: &'a dyn Session, + projection: Option<&'a [usize]>, + filters: &'a [Expr], limit: Option, - ) -> Result>; + ) -> BoxFuture<'a, Result>>; /// Create an [`ExecutionPlan`] for scanning the table using structured arguments. /// @@ -209,26 +207,21 @@ pub trait TableProvider: Any + Debug + Sync + Send { /// A [`ScanResult`] containing the [`ExecutionPlan`] for scanning the table /// /// See [`Self::scan`] for detailed documentation about projection, filters, and limits. - // Hand-written `#[async_trait]` expansion to reduce compile time. See - // . - fn scan_with_args<'a, 'life0, 'life1, 'async_trait>( - &'life0 self, - state: &'life1 dyn Session, + fn scan_with_args<'a>( + &'a self, + state: &'a dyn Session, args: ScanArgs<'a>, - ) -> BoxFuture<'async_trait, Result> - where - 'a: 'async_trait, - 'life0: 'async_trait, - 'life1: 'async_trait, - Self: 'async_trait, - { - let plan = self.scan( - state, - args.projection(), - args.filters().unwrap_or(&[]), - args.limit(), - ); - Box::pin(async move { Ok(plan.await?.into()) }) + ) -> BoxFuture<'a, Result> { + Box::pin(async move { + self.scan( + state, + args.projection(), + args.filters().unwrap_or(&[]), + args.limit(), + ) + .await + .map(Into::into) + }) } /// Specify if DataFusion should provide filter expressions to the @@ -249,7 +242,7 @@ pub trait TableProvider: Any + Debug + Sync + Send { /// /// Each element in the resulting `Vec` is one of the following: /// * [`Exact`] or [`Inexact`]: The TableProvider can apply the filter - /// during scan + /// during scan /// * [`Unsupported`]: The TableProvider cannot apply the filter during scan /// /// By default, this function returns [`Unsupported`] for all filters, @@ -264,7 +257,6 @@ pub trait TableProvider: Any + Debug + Sync + Send { /// # use std::any::Any; /// # use std::sync::Arc; /// # use arrow_schema::SchemaRef; - /// # use async_trait::async_trait; /// # use datafusion_session::{TableProvider, Session}; /// # use datafusion_common::Result; /// # use datafusion_expr::{Expr, TableProviderFilterPushDown, TableType}; @@ -273,7 +265,6 @@ pub trait TableProvider: Any + Debug + Sync + Send { /// #[derive(Debug)] /// struct TestDataSource {} /// - /// #[async_trait] /// impl TableProvider for TestDataSource { /// # fn schema(&self) -> SchemaRef { todo!() } /// # fn table_type(&self) -> TableType { todo!() } @@ -348,66 +339,55 @@ pub trait TableProvider: Any + Debug + Sync + Send { /// streams of `RecordBatch`es as files to an ObjectStore. /// /// [`DataSinkExec`]: https://docs.rs/datafusion-datasource/latest/datafusion_datasource/sink/struct.DataSinkExec.html - async fn insert_into( - &self, - _state: &dyn Session, + fn insert_into<'a>( + &'a self, + _state: &'a dyn Session, _input: Arc, _insert_op: InsertOp, - ) -> Result> { - not_impl_err!("Insert into not implemented for this table") + ) -> BoxFuture<'a, Result>> { + Box::pin(async { not_impl_err!("Insert into not implemented for this table") }) } /// Delete rows matching the filter predicates. /// /// Returns an [`ExecutionPlan`] producing a single row with `count` (UInt64). /// Empty `filters` deletes all rows. - // Hand-written `#[async_trait]` expansion to reduce compile time. See - // - fn delete_from<'life0, 'life1, 'async_trait>( - &'life0 self, - _state: &'life1 dyn Session, + fn delete_from<'a>( + &'a self, + _state: &'a dyn Session, _filters: Vec, - ) -> BoxFuture<'async_trait, Result>> - where - 'life0: 'async_trait, - 'life1: 'async_trait, - Self: 'async_trait, - { - Box::pin(ready(not_impl_err!( - "DELETE not supported for {} table", - self.table_type() - ))) + ) -> BoxFuture<'a, Result>> { + Box::pin(async move { + not_impl_err!("DELETE not supported for {} table", self.table_type()) + }) } /// Update rows matching the filter predicates. /// /// Returns an [`ExecutionPlan`] producing a single row with `count` (UInt64). /// Empty `filters` updates all rows. - // Hand-written `#[async_trait]` expansion to reduce compile time. See - // - fn update<'life0, 'life1, 'async_trait>( - &'life0 self, - _state: &'life1 dyn Session, + fn update<'a>( + &'a self, + _state: &'a dyn Session, _assignments: Vec<(String, Expr)>, _filters: Vec, - ) -> BoxFuture<'async_trait, Result>> - where - 'life0: 'async_trait, - 'life1: 'async_trait, - Self: 'async_trait, - { - Box::pin(ready(not_impl_err!( - "UPDATE not supported for {} table", - self.table_type() - ))) + ) -> BoxFuture<'a, Result>> { + Box::pin(async move { + not_impl_err!("UPDATE not supported for {} table", self.table_type()) + }) } /// Remove all rows from the table. /// /// Should return an [ExecutionPlan] producing a single row with count (UInt64), /// representing the number of rows removed. - async fn truncate(&self, _state: &dyn Session) -> Result> { - not_impl_err!("TRUNCATE not supported for {} table", self.table_type()) + fn truncate<'a>( + &'a self, + _state: &'a dyn Session, + ) -> BoxFuture<'a, Result>> { + Box::pin(async move { + not_impl_err!("TRUNCATE not supported for {} table", self.table_type()) + }) } /// Merge rows from a source into this table. @@ -421,25 +401,17 @@ pub trait TableProvider: Any + Debug + Sync + Send { /// The `clauses` describe the WHEN MATCHED / WHEN NOT MATCHED actions. /// /// Returns an [`ExecutionPlan`] producing a single row with `count` (UInt64). - // Hand-written `#[async_trait]` expansion to reduce compile time. See - // - fn merge_into<'life0, 'life1, 'async_trait>( - &'life0 self, - _state: &'life1 dyn Session, + fn merge_into<'a>( + &'a self, + _state: &'a dyn Session, _source: Arc, _merge_schema: DFSchemaRef, _on: Expr, _clauses: Vec, - ) -> BoxFuture<'async_trait, Result>> - where - 'life0: 'async_trait, - 'life1: 'async_trait, - Self: 'async_trait, - { - Box::pin(ready(not_impl_err!( - "MERGE INTO not supported for {} table", - self.table_type() - ))) + ) -> BoxFuture<'a, Result>> { + Box::pin(async move { + not_impl_err!("MERGE INTO not supported for {} table", self.table_type()) + }) } } @@ -602,14 +574,13 @@ impl From> for ScanResult { /// /// For example, this can be used to create a table "on the fly" /// from a directory of files only when that name is referenced. -#[async_trait] pub trait TableProviderFactory: Debug + Sync + Send { /// Create a TableProvider with the given url - async fn create( - &self, - state: &dyn Session, - cmd: &CreateExternalTable, - ) -> Result>; + fn create<'a>( + &'a self, + state: &'a dyn Session, + cmd: &'a CreateExternalTable, + ) -> BoxFuture<'a, Result>>; } /// Describes arguments provided to the table function call. diff --git a/datafusion/sqllogictest/Cargo.toml b/datafusion/sqllogictest/Cargo.toml index c43f78f846092..d4a06385bb41a 100644 --- a/datafusion/sqllogictest/Cargo.toml +++ b/datafusion/sqllogictest/Cargo.toml @@ -41,7 +41,6 @@ name = "datafusion_sqllogictest" [dependencies] arrow = { workspace = true } -async-trait = { workspace = true } bigdecimal = { workspace = true } bytes = { workspace = true, optional = true } chrono = { workspace = true, optional = true } diff --git a/datafusion/sqllogictest/src/engines/datafusion_engine/runner.rs b/datafusion/sqllogictest/src/engines/datafusion_engine/runner.rs index 470efa6f4f993..85016dcbcd7a4 100644 --- a/datafusion/sqllogictest/src/engines/datafusion_engine/runner.rs +++ b/datafusion/sqllogictest/src/engines/datafusion_engine/runner.rs @@ -17,6 +17,8 @@ use std::collections::HashMap; use std::fmt::Write as _; +use std::future::Future; +use std::pin::Pin; use std::sync::{Arc, Mutex}; use std::{path::PathBuf, time::Duration}; @@ -25,7 +27,6 @@ use crate::engines::currently_executed_sql::CurrentlyExecutingSqlTracker; use crate::engines::output::{DFColumnType, DFOutput}; use crate::is_spark_path; use arrow::record_batch::RecordBatch; -use async_trait::async_trait; use datafusion::physical_plan::common::collect; use datafusion::physical_plan::execute_stream; use datafusion::prelude::SessionContext; @@ -137,42 +138,52 @@ impl DataFusion { } } -#[async_trait] impl sqllogictest::AsyncDB for DataFusion { type Error = DFSqlLogicTestError; type ColumnType = DFColumnType; - async fn run(&mut self, sql: &str) -> Result { - if log_enabled!(Debug) { - debug!( - "[{}] Running query: \"{}\"", - self.relative_path.display(), - sql - ); - } + fn run<'life0, 'life1, 'async_trait>( + &'life0 mut self, + sql: &'life1 str, + ) -> Pin> + Send + 'async_trait>> + where + 'life0: 'async_trait, + 'life1: 'async_trait, + Self: 'async_trait, + { + Box::pin(async move { + if log_enabled!(Debug) { + debug!( + "[{}] Running query: \"{}\"", + self.relative_path.display(), + sql + ); + } - let tracked_sql = self.currently_executing_sql_tracker.set_sql(sql); + let tracked_sql = self.currently_executing_sql_tracker.set_sql(sql); - let start = Instant::now(); - let result = run_query(&self.ctx, is_spark_path(&self.relative_path), sql).await; - let duration = start.elapsed(); + let start = Instant::now(); + let result = + run_query(&self.ctx, is_spark_path(&self.relative_path), sql).await; + let duration = start.elapsed(); - self.currently_executing_sql_tracker.remove_sql(tracked_sql); + self.currently_executing_sql_tracker.remove_sql(tracked_sql); - if duration.gt(&Duration::from_millis(500)) { - self.update_slow_count(); - } + if duration.gt(&Duration::from_millis(500)) { + self.update_slow_count(); + } - self.pb.inc(1); + self.pb.inc(1); - if log_enabled!(Info) && duration.gt(&Duration::from_secs(2)) { - warn!( - "[{}] Running query took more than 2 sec ({duration:?}): \"{sql}\"", - self.relative_path.display() - ); - } + if log_enabled!(Info) && duration.gt(&Duration::from_secs(2)) { + warn!( + "[{}] Running query took more than 2 sec ({duration:?}): \"{sql}\"", + self.relative_path.display() + ); + } - result + result + }) } /// Engine name of current database. @@ -185,17 +196,32 @@ impl sqllogictest::AsyncDB for DataFusion { /// The default implementation is `std::thread::sleep`, which is universal to any async runtime /// but would block the current thread. If you are running in tokio runtime, you should override /// this by `tokio::time::sleep`. - async fn sleep(dur: Duration) { - tokio::time::sleep(dur).await; + fn sleep<'async_trait>( + dur: Duration, + ) -> Pin + Send + 'async_trait>> + where + Self: 'async_trait, + { + Box::pin(async move { + tokio::time::sleep(dur).await; + }) } /// Shutdown and check no DataFusion configuration has changed during test - async fn shutdown(&mut self) { - if let Some(config_change_errors) = self.config_change_errors.clone() - && let Err(error) = self.validate_config_unchanged() - { - config_change_errors.lock().unwrap().push(error.to_string()); - } + fn shutdown<'life0, 'async_trait>( + &'life0 mut self, + ) -> Pin + Send + 'async_trait>> + where + 'life0: 'async_trait, + Self: 'async_trait, + { + Box::pin(async move { + if let Some(config_change_errors) = self.config_change_errors.clone() + && let Err(error) = self.validate_config_unchanged() + { + config_change_errors.lock().unwrap().push(error.to_string()); + } + }) } } diff --git a/datafusion/sqllogictest/src/engines/datafusion_substrait_roundtrip_engine/runner.rs b/datafusion/sqllogictest/src/engines/datafusion_substrait_roundtrip_engine/runner.rs index 90ad6646dc906..0fd16539ea0b2 100644 --- a/datafusion/sqllogictest/src/engines/datafusion_substrait_roundtrip_engine/runner.rs +++ b/datafusion/sqllogictest/src/engines/datafusion_substrait_roundtrip_engine/runner.rs @@ -15,6 +15,8 @@ // specific language governing permissions and limitations // under the License. +use std::future::Future; +use std::pin::Pin; use std::sync::Arc; use std::{path::PathBuf, time::Duration}; @@ -23,7 +25,6 @@ use crate::engines::datafusion_engine::Result; use crate::engines::output::{DFColumnType, DFOutput}; use crate::{DFSqlLogicTestError, convert_batches, convert_schema_to_types}; use arrow::record_batch::RecordBatch; -use async_trait::async_trait; use datafusion::logical_expr::LogicalPlan; use datafusion::physical_plan::common::collect; use datafusion::physical_plan::execute_stream; @@ -83,42 +84,51 @@ impl DataFusionSubstraitRoundTrip { } } -#[async_trait] impl sqllogictest::AsyncDB for DataFusionSubstraitRoundTrip { type Error = DFSqlLogicTestError; type ColumnType = DFColumnType; - async fn run(&mut self, sql: &str) -> Result { - if log_enabled!(Debug) { - debug!( - "[{}] Running query: \"{}\"", - self.relative_path.display(), - sql - ); - } - - let tracked_sql = self.currently_executing_sql_tracker.set_sql(sql); - - let start = Instant::now(); - let result = run_query_substrait_round_trip(&self.ctx, sql).await; - let duration = start.elapsed(); - - self.currently_executing_sql_tracker.remove_sql(tracked_sql); - - if duration.gt(&Duration::from_millis(500)) { - self.update_slow_count(); - } - - self.pb.inc(1); - - if log_enabled!(Info) && duration.gt(&Duration::from_secs(2)) { - warn!( - "[{}] Running query took more than 2 sec ({duration:?}): \"{sql}\"", - self.relative_path.display() - ); - } - - result + fn run<'life0, 'life1, 'async_trait>( + &'life0 mut self, + sql: &'life1 str, + ) -> Pin> + Send + 'async_trait>> + where + 'life0: 'async_trait, + 'life1: 'async_trait, + Self: 'async_trait, + { + Box::pin(async move { + if log_enabled!(Debug) { + debug!( + "[{}] Running query: \"{}\"", + self.relative_path.display(), + sql + ); + } + + let tracked_sql = self.currently_executing_sql_tracker.set_sql(sql); + + let start = Instant::now(); + let result = run_query_substrait_round_trip(&self.ctx, sql).await; + let duration = start.elapsed(); + + self.currently_executing_sql_tracker.remove_sql(tracked_sql); + + if duration.gt(&Duration::from_millis(500)) { + self.update_slow_count(); + } + + self.pb.inc(1); + + if log_enabled!(Info) && duration.gt(&Duration::from_secs(2)) { + warn!( + "[{}] Running query took more than 2 sec ({duration:?}): \"{sql}\"", + self.relative_path.display() + ); + } + + result + }) } /// Engine name of current database. @@ -131,11 +141,26 @@ impl sqllogictest::AsyncDB for DataFusionSubstraitRoundTrip { /// The default implementation is `std::thread::sleep`, which is universal to any async runtime /// but would block the current thread. If you are running in tokio runtime, you should override /// this by `tokio::time::sleep`. - async fn sleep(dur: Duration) { - tokio::time::sleep(dur).await; + fn sleep<'async_trait>( + dur: Duration, + ) -> Pin + Send + 'async_trait>> + where + Self: 'async_trait, + { + Box::pin(async move { + tokio::time::sleep(dur).await; + }) } - async fn shutdown(&mut self) {} + fn shutdown<'life0, 'async_trait>( + &'life0 mut self, + ) -> Pin + Send + 'async_trait>> + where + 'life0: 'async_trait, + Self: 'async_trait, + { + Box::pin(async move {}) + } } async fn run_query_substrait_round_trip( diff --git a/datafusion/sqllogictest/src/engines/postgres_engine/mod.rs b/datafusion/sqllogictest/src/engines/postgres_engine/mod.rs index daa5acf165891..9c5d8fe59bf63 100644 --- a/datafusion/sqllogictest/src/engines/postgres_engine/mod.rs +++ b/datafusion/sqllogictest/src/engines/postgres_engine/mod.rs @@ -15,15 +15,16 @@ // specific language governing permissions and limitations // under the License. -use async_trait::async_trait; use bigdecimal::BigDecimal; use bytes::Bytes; use datafusion::common::runtime::SpawnedTask; use futures::{SinkExt, StreamExt}; use log::{debug, info}; use sqllogictest::DBOutput; +use std::future::Future; /// Postgres engine implementation for sqllogictest. use std::path::{Path, PathBuf}; +use std::pin::Pin; use std::str::FromStr; use std::time::Duration; @@ -250,98 +251,118 @@ fn schema_name(relative_path: &Path) -> String { .to_string() } -#[async_trait] impl sqllogictest::AsyncDB for Postgres { type Error = Error; type ColumnType = DFColumnType; - async fn run( - &mut self, - sql: &str, - ) -> Result, Self::Error> { - debug!( - "[{}] Running query: \"{}\"", - self.relative_path.display(), - sql - ); - - let tracked_sql = self.currently_executing_sql_tracker.set_sql(sql); - - let lower_sql = sql.trim_start().to_ascii_lowercase(); - - let is_query_sql = { - lower_sql.starts_with("select") - || lower_sql.starts_with("values") - || lower_sql.starts_with("show") - || lower_sql.starts_with("with") - || lower_sql.starts_with("describe") - || ((lower_sql.starts_with("insert") - || lower_sql.starts_with("update") - || lower_sql.starts_with("delete")) - && lower_sql.contains("returning")) - }; - - if lower_sql.starts_with("copy") { - self.pb.inc(1); - let result = self.run_copy_command(sql).await; - self.currently_executing_sql_tracker.remove_sql(tracked_sql); + fn run<'life0, 'life1, 'async_trait>( + &'life0 mut self, + sql: &'life1 str, + ) -> Pin< + Box< + dyn Future, Self::Error>> + + Send + + 'async_trait, + >, + > + where + 'life0: 'async_trait, + 'life1: 'async_trait, + Self: 'async_trait, + { + Box::pin(async move { + debug!( + "[{}] Running query: \"{}\"", + self.relative_path.display(), + sql + ); + + let tracked_sql = self.currently_executing_sql_tracker.set_sql(sql); + + let lower_sql = sql.trim_start().to_ascii_lowercase(); + + let is_query_sql = { + lower_sql.starts_with("select") + || lower_sql.starts_with("values") + || lower_sql.starts_with("show") + || lower_sql.starts_with("with") + || lower_sql.starts_with("describe") + || ((lower_sql.starts_with("insert") + || lower_sql.starts_with("update") + || lower_sql.starts_with("delete")) + && lower_sql.contains("returning")) + }; - return result; - } + if lower_sql.starts_with("copy") { + self.pb.inc(1); + let result = self.run_copy_command(sql).await; + self.currently_executing_sql_tracker.remove_sql(tracked_sql); - if !is_query_sql { - self.get_client().execute(sql, &[]).await?; - self.currently_executing_sql_tracker.remove_sql(tracked_sql); - self.pb.inc(1); - return Ok(DBOutput::StatementComplete(0)); - } - // Use a prepared statement to get the output column types - let statement = self.get_client().prepare(sql).await?; - let types: Vec = statement - .columns() - .iter() - .map(|c| c.type_().clone()) - .collect(); - - // Run the actual query using the "simple query" protocol that returns all - // rows as text. Doing this avoids having to convert values from the binary - // format to strings, which is somewhat tricky for numeric types. - // See https://github.com/apache/datafusion/pull/19666#discussion_r2668090587 - let start = Instant::now(); - let messages = self.get_client().simple_query(sql).await?; - let duration = start.elapsed(); - - if duration.gt(&Duration::from_millis(500)) { - self.update_slow_count(); - } + return result; + } - self.pb.inc(1); + if !is_query_sql { + self.get_client().execute(sql, &[]).await?; + self.currently_executing_sql_tracker.remove_sql(tracked_sql); + self.pb.inc(1); + return Ok(DBOutput::StatementComplete(0)); + } + // Use a prepared statement to get the output column types + let statement = self.get_client().prepare(sql).await?; + let types: Vec = statement + .columns() + .iter() + .map(|c| c.type_().clone()) + .collect(); + + // Run the actual query using the "simple query" protocol that returns all + // rows as text. Doing this avoids having to convert values from the binary + // format to strings, which is somewhat tricky for numeric types. + // See https://github.com/apache/datafusion/pull/19666#discussion_r2668090587 + let start = Instant::now(); + let messages = self.get_client().simple_query(sql).await?; + let duration = start.elapsed(); + + if duration.gt(&Duration::from_millis(500)) { + self.update_slow_count(); + } - self.currently_executing_sql_tracker.remove_sql(tracked_sql); + self.pb.inc(1); - let rows = convert_rows(&types, &messages); + self.currently_executing_sql_tracker.remove_sql(tracked_sql); - if rows.is_empty() && types.is_empty() { - Ok(DBOutput::StatementComplete(0)) - } else { - Ok(DBOutput::Rows { - types: convert_types(types), - rows, - }) - } + let rows = convert_rows(&types, &messages); + + if rows.is_empty() && types.is_empty() { + Ok(DBOutput::StatementComplete(0)) + } else { + Ok(DBOutput::Rows { + types: convert_types(types), + rows, + }) + } + }) } fn engine_name(&self) -> &str { "postgres" } - async fn shutdown(&mut self) { - if let Some(client) = self.client.take() { - drop(client); - } - if let Some(spawned_task) = self.spawned_task.take() { - spawned_task.join().await.ok(); - } + fn shutdown<'life0, 'async_trait>( + &'life0 mut self, + ) -> Pin + Send + 'async_trait>> + where + 'life0: 'async_trait, + Self: 'async_trait, + { + Box::pin(async move { + if let Some(client) = self.client.take() { + drop(client); + } + if let Some(spawned_task) = self.spawned_task.take() { + spawned_task.join().await.ok(); + } + }) } } diff --git a/datafusion/sqllogictest/src/test_context.rs b/datafusion/sqllogictest/src/test_context.rs index 307a8876abdd5..0927daf59d489 100644 --- a/datafusion/sqllogictest/src/test_context.rs +++ b/datafusion/sqllogictest/src/test_context.rs @@ -15,6 +15,7 @@ // specific language governing permissions and limitations // under the License. +use futures::future::BoxFuture; use std::collections::HashMap; use std::fs::File; use std::io::Write; @@ -58,7 +59,6 @@ use range_partitioning::{ register_range_partitioned_table, register_range_sorted_time_bin_table, }; -use async_trait::async_trait; use datafusion::common::cast::as_float64_array; use datafusion::execution::SessionStateBuilder; use datafusion::execution::runtime_env::RuntimeEnv; @@ -292,24 +292,24 @@ impl TestContext { struct StrictOrdersSchema { orders: Arc, } - -#[async_trait] impl SchemaProvider for StrictOrdersSchema { fn table_names(&self) -> Vec { vec!["orders".to_string()] } - async fn table( - &self, - name: &str, - ) -> Result>, DataFusionError> { - match name { - "orders" => Ok(Some(Arc::clone(&self.orders))), - other => panic!( - "unexpected table lookup: {other}. This maybe indicates a CTE reference was \ - incorrectly treated as a catalog table reference." - ), - } + fn table<'a>( + &'a self, + name: &'a str, + ) -> BoxFuture<'a, Result>, DataFusionError>> { + Box::pin(async move { + match name { + "orders" => Ok(Some(Arc::clone(&self.orders))), + other => panic!( + "unexpected table lookup: {other}. This maybe indicates a CTE reference was \ + incorrectly treated as a catalog table reference." + ), + } + }) } fn table_exist(&self, name: &str) -> bool { @@ -419,8 +419,6 @@ pub async fn register_partition_table(test_ctx: &mut TestContext) { pub fn register_temp_table(ctx: &SessionContext) { #[derive(Debug)] struct TestTable(TableType); - - #[async_trait] impl TableProvider for TestTable { fn schema(&self) -> SchemaRef { unimplemented!() @@ -430,14 +428,14 @@ pub fn register_temp_table(ctx: &SessionContext) { self.0 } - async fn scan( - &self, - _state: &dyn Session, - _: Option<&[usize]>, - _: &[Expr], + fn scan<'a>( + &'a self, + _state: &'a dyn Session, + _: Option<&'a [usize]>, + _: &'a [Expr], _: Option, - ) -> Result, DataFusionError> { - unimplemented!() + ) -> BoxFuture<'a, Result, DataFusionError>> { + Box::pin(async move { unimplemented!() }) } } @@ -801,13 +799,12 @@ fn register_async_abs_udf(ctx: &SessionContext) { not_impl_err!("{} can only be called from async contexts", self.name()) } } - #[async_trait] impl AsyncScalarUDFImpl for AsyncAbs { - async fn invoke_async_with_args( + fn invoke_async_with_args( &self, args: ScalarFunctionArgs, - ) -> Result { - return self.inner_abs.invoke_with_args(args); + ) -> BoxFuture<'_, Result> { + Box::pin(async move { self.inner_abs.invoke_with_args(args) }) } } let async_abs = AsyncAbs::new(); diff --git a/datafusion/substrait/Cargo.toml b/datafusion/substrait/Cargo.toml index a0f203cec8db6..500ad656e0ac5 100644 --- a/datafusion/substrait/Cargo.toml +++ b/datafusion/substrait/Cargo.toml @@ -35,8 +35,8 @@ workspace = true [dependencies] async-recursion = "1.0" -async-trait = { workspace = true } chrono = { workspace = true } +futures = { workspace = true } datafusion = { workspace = true, features = ["sql"] } half = { workspace = true } itertools = { workspace = true } diff --git a/datafusion/substrait/src/logical_plan/consumer/substrait_consumer.rs b/datafusion/substrait/src/logical_plan/consumer/substrait_consumer.rs index 36ef39561e7c3..8c85cd1006f9a 100644 --- a/datafusion/substrait/src/logical_plan/consumer/substrait_consumer.rs +++ b/datafusion/substrait/src/logical_plan/consumer/substrait_consumer.rs @@ -26,7 +26,6 @@ use crate::extensions::Extensions; use crate::logical_plan::consumer::{ field_from_substrait_type_without_names, from_lambda, }; -use async_trait::async_trait; use datafusion::arrow::datatypes::{DataType, FieldRef}; use datafusion::catalog::TableProvider; use datafusion::common::datatype::FieldExt; @@ -36,6 +35,7 @@ use datafusion::common::{ use datafusion::execution::{FunctionRegistry, SessionState}; use datafusion::logical_expr::expr::LambdaVariable; use datafusion::logical_expr::{Expr, Extension, LogicalPlan}; +use futures::future::BoxFuture; use std::collections::{HashMap, VecDeque}; use std::sync::{Arc, RwLock}; use substrait::proto::expression as substrait_expression; @@ -49,8 +49,6 @@ use substrait::proto::{ Expression, ExtensionLeafRel, ExtensionMultiRel, ExtensionSingleRel, FetchRel, FilterRel, JoinRel, ProjectRel, ReadRel, Rel, SetRel, SortRel, r#type, }; - -#[async_trait] /// This trait is used to consume Substrait plans, converting them into DataFusion Logical Plans. /// It can be implemented by users to allow for custom handling of relations, expressions, etc. /// @@ -60,13 +58,13 @@ use substrait::proto::{ /// # Example Usage /// /// ``` -/// # use async_trait::async_trait; -/// # use datafusion::catalog::TableProvider; +/// # /// # use datafusion::catalog::TableProvider; /// # use datafusion::common::{not_impl_err, substrait_err, DFSchema, ScalarValue, TableReference}; /// # use datafusion::error::Result; /// # use datafusion::execution::{FunctionRegistry, SessionState}; /// # use datafusion::logical_expr::{Expr, LogicalPlan, LogicalPlanBuilder}; /// # use std::sync::Arc; +/// # use futures::future::BoxFuture; /// # use substrait::proto; /// # use substrait::proto::{ExtensionLeafRel, FilterRel, ProjectRel, Type}; /// # use datafusion::arrow::datatypes::DataType; @@ -83,16 +81,17 @@ use substrait::proto::{ /// lambda_consumer: DefaultSubstraitLambdaConsumer, /// } /// -/// #[async_trait] /// impl SubstraitConsumer for CustomSubstraitConsumer { -/// async fn resolve_table_ref( -/// &self, -/// table_ref: &TableReference, -/// ) -> Result>> { -/// let table = table_ref.table().to_string(); -/// let schema = self.state.schema_for_ref(table_ref.clone())?; -/// let table_provider = schema.table(&table).await?; -/// Ok(table_provider) +/// fn resolve_table_ref<'a>( +/// &'a self, +/// table_ref: &'a TableReference, +/// ) -> BoxFuture<'a, Result>>> { +/// Box::pin(async move { +/// let table = table_ref.table().to_string(); +/// let schema = self.state.schema_for_ref(table_ref.clone())?; +/// let table_provider = schema.table(&table).await?; +/// Ok(table_provider) +/// }) /// } /// /// fn get_extensions(&self) -> &Extensions { @@ -128,37 +127,49 @@ use substrait::proto::{ /// } /// /// // You can reuse existing consumer code to assist in handling advanced extensions -/// async fn consume_project(&self, rel: &ProjectRel) -> Result { -/// let df_plan = from_project_rel(self, rel).await?; -/// if let Some(advanced_extension) = rel.advanced_extension.as_ref() { -/// not_impl_err!( -/// "decode and handle an advanced extension: {:?}", -/// advanced_extension -/// ) -/// } else { -/// Ok(df_plan) -/// } +/// fn consume_project<'a>( +/// &'a self, +/// rel: &'a ProjectRel, +/// ) -> BoxFuture<'a, Result> { +/// Box::pin(async move { +/// let df_plan = from_project_rel(self, rel).await?; +/// if let Some(advanced_extension) = rel.advanced_extension.as_ref() { +/// not_impl_err!( +/// "decode and handle an advanced extension: {:?}", +/// advanced_extension +/// ) +/// } else { +/// Ok(df_plan) +/// } +/// }) /// } /// /// // You can implement a fully custom consumer method if you need special handling -/// async fn consume_filter(&self, rel: &FilterRel) -> Result { -/// let input = self.consume_rel(rel.input.as_ref().unwrap()).await?; -/// let expression = -/// self.consume_expression(rel.condition.as_ref().unwrap(), input.schema()) -/// .await?; -/// // though this one is quite boring -/// LogicalPlanBuilder::from(input).filter(expression)?.build() +/// fn consume_filter<'a>( +/// &'a self, +/// rel: &'a FilterRel, +/// ) -> BoxFuture<'a, Result> { +/// Box::pin(async move { +/// let input = self.consume_rel(rel.input.as_ref().unwrap()).await?; +/// let expression = +/// self.consume_expression(rel.condition.as_ref().unwrap(), input.schema()) +/// .await?; +/// // though this one is quite boring +/// LogicalPlanBuilder::from(input).filter(expression)?.build() +/// }) /// } /// /// // You can add handlers for extension relations -/// async fn consume_extension_leaf( -/// &self, -/// rel: &ExtensionLeafRel, -/// ) -> Result { -/// not_impl_err!( -/// "handle protobuf Any {} as you need", -/// rel.detail.as_ref().unwrap().type_url -/// ) +/// fn consume_extension_leaf<'a>( +/// &'a self, +/// rel: &'a ExtensionLeafRel, +/// ) -> BoxFuture<'a, Result> { +/// Box::pin(async move { +/// not_impl_err!( +/// "handle protobuf Any {} as you need", +/// rel.detail.as_ref().unwrap().type_url +/// ) +/// }) /// } /// /// // and handlers for user-define types @@ -191,10 +202,10 @@ use substrait::proto::{ /// } /// ``` pub trait SubstraitConsumer: Send + Sync + Sized { - async fn resolve_table_ref( - &self, - table_ref: &TableReference, - ) -> datafusion::common::Result>>; + fn resolve_table_ref<'a>( + &'a self, + table_ref: &'a TableReference, + ) -> BoxFuture<'a, datafusion::common::Result>>>; // TODO: Remove these two methods // Ideally, the abstract consumer should not place any constraints on implementations. @@ -211,82 +222,90 @@ pub trait SubstraitConsumer: Send + Sync + Sized { /// All [Rel]s to be converted pass through this method. /// You can provide your own implementation if you wish to customize the conversion behaviour. - async fn consume_rel(&self, rel: &Rel) -> datafusion::common::Result { - from_substrait_rel(self, rel).await + fn consume_rel<'a>( + &'a self, + rel: &'a Rel, + ) -> BoxFuture<'a, datafusion::common::Result> { + Box::pin(from_substrait_rel(self, rel)) } - async fn consume_read( - &self, - rel: &ReadRel, - ) -> datafusion::common::Result { - from_read_rel(self, rel).await + fn consume_read<'a>( + &'a self, + rel: &'a ReadRel, + ) -> BoxFuture<'a, datafusion::common::Result> { + Box::pin(from_read_rel(self, rel)) } - async fn consume_filter( - &self, - rel: &FilterRel, - ) -> datafusion::common::Result { - from_filter_rel(self, rel).await + fn consume_filter<'a>( + &'a self, + rel: &'a FilterRel, + ) -> BoxFuture<'a, datafusion::common::Result> { + Box::pin(from_filter_rel(self, rel)) } - async fn consume_fetch( - &self, - rel: &FetchRel, - ) -> datafusion::common::Result { - from_fetch_rel(self, rel).await + fn consume_fetch<'a>( + &'a self, + rel: &'a FetchRel, + ) -> BoxFuture<'a, datafusion::common::Result> { + Box::pin(from_fetch_rel(self, rel)) } - async fn consume_aggregate( - &self, - rel: &AggregateRel, - ) -> datafusion::common::Result { - from_aggregate_rel(self, rel).await + fn consume_aggregate<'a>( + &'a self, + rel: &'a AggregateRel, + ) -> BoxFuture<'a, datafusion::common::Result> { + Box::pin(from_aggregate_rel(self, rel)) } - async fn consume_sort( - &self, - rel: &SortRel, - ) -> datafusion::common::Result { - from_sort_rel(self, rel).await + fn consume_sort<'a>( + &'a self, + rel: &'a SortRel, + ) -> BoxFuture<'a, datafusion::common::Result> { + Box::pin(from_sort_rel(self, rel)) } - async fn consume_join( - &self, - rel: &JoinRel, - ) -> datafusion::common::Result { - from_join_rel(self, rel).await + fn consume_join<'a>( + &'a self, + rel: &'a JoinRel, + ) -> BoxFuture<'a, datafusion::common::Result> { + Box::pin(from_join_rel(self, rel)) } - async fn consume_project( - &self, - rel: &ProjectRel, - ) -> datafusion::common::Result { - from_project_rel(self, rel).await + fn consume_project<'a>( + &'a self, + rel: &'a ProjectRel, + ) -> BoxFuture<'a, datafusion::common::Result> { + Box::pin(from_project_rel(self, rel)) } - async fn consume_set(&self, rel: &SetRel) -> datafusion::common::Result { - from_set_rel(self, rel).await + fn consume_set<'a>( + &'a self, + rel: &'a SetRel, + ) -> BoxFuture<'a, datafusion::common::Result> { + Box::pin(from_set_rel(self, rel)) } - async fn consume_cross( - &self, - rel: &CrossRel, - ) -> datafusion::common::Result { - from_cross_rel(self, rel).await + fn consume_cross<'a>( + &'a self, + rel: &'a CrossRel, + ) -> BoxFuture<'a, datafusion::common::Result> { + Box::pin(from_cross_rel(self, rel)) } - async fn consume_consistent_partition_window( - &self, - _rel: &ConsistentPartitionWindowRel, - ) -> datafusion::common::Result { - not_impl_err!("Consistent Partition Window Rel not supported") + fn consume_consistent_partition_window<'a>( + &'a self, + _rel: &'a ConsistentPartitionWindowRel, + ) -> BoxFuture<'a, datafusion::common::Result> { + Box::pin( + async move { not_impl_err!("Consistent Partition Window Rel not supported") }, + ) } - async fn consume_exchange( - &self, - rel: &ExchangeRel, - ) -> datafusion::common::Result { - from_exchange_rel(self, rel).await + fn consume_exchange<'a>( + &'a self, + rel: &'a ExchangeRel, + ) -> BoxFuture<'a, datafusion::common::Result> { + Box::pin(from_exchange_rel(self, rel)) } // Expression Methods @@ -296,132 +315,137 @@ pub trait SubstraitConsumer: Send + Sync + Sized { /// All [Expression]s to be converted pass through this method. /// You can provide your own implementation if you wish to customize the conversion behaviour. - async fn consume_expression( - &self, - expr: &Expression, - input_schema: &DFSchema, - ) -> datafusion::common::Result { - from_substrait_rex(self, expr, input_schema).await - } - - async fn consume_literal(&self, expr: &Literal) -> datafusion::common::Result { - from_literal(self, expr).await - } - - async fn consume_field_reference( - &self, - expr: &FieldReference, - input_schema: &DFSchema, - ) -> datafusion::common::Result { - from_field_reference(self, expr, input_schema).await - } - - async fn consume_scalar_function( - &self, - expr: &ScalarFunction, - input_schema: &DFSchema, - ) -> datafusion::common::Result { - from_scalar_function(self, expr, input_schema).await - } - - async fn consume_window_function( - &self, - expr: &WindowFunction, - input_schema: &DFSchema, - ) -> datafusion::common::Result { - from_window_function(self, expr, input_schema).await + fn consume_expression<'a>( + &'a self, + expr: &'a Expression, + input_schema: &'a DFSchema, + ) -> BoxFuture<'a, datafusion::common::Result> { + Box::pin(from_substrait_rex(self, expr, input_schema)) } - - async fn consume_if_then( - &self, - expr: &IfThen, - input_schema: &DFSchema, - ) -> datafusion::common::Result { - from_if_then(self, expr, input_schema).await - } - - async fn consume_switch( - &self, - _expr: &SwitchExpression, - _input_schema: &DFSchema, - ) -> datafusion::common::Result { - not_impl_err!("Switch expression not supported") - } - - async fn consume_singular_or_list( - &self, - expr: &SingularOrList, - input_schema: &DFSchema, - ) -> datafusion::common::Result { - from_singular_or_list(self, expr, input_schema).await - } - - async fn consume_multi_or_list( - &self, - _expr: &MultiOrList, - _input_schema: &DFSchema, - ) -> datafusion::common::Result { - not_impl_err!("Multi Or List expression not supported") - } - - async fn consume_cast( - &self, - expr: &substrait_expression::Cast, - input_schema: &DFSchema, - ) -> datafusion::common::Result { - from_cast(self, expr, input_schema).await - } - - async fn consume_subquery( - &self, - expr: &substrait_expression::Subquery, - input_schema: &DFSchema, - ) -> datafusion::common::Result { - from_subquery(self, expr, input_schema).await - } - - async fn consume_nested( - &self, - expr: &Nested, - input_schema: &DFSchema, - ) -> datafusion::common::Result { - from_nested(self, expr, input_schema).await - } - - async fn consume_enum( - &self, - _expr: &Enum, - _input_schema: &DFSchema, - ) -> datafusion::common::Result { - not_impl_err!("Enum expression not supported") - } - - async fn consume_dynamic_parameter( - &self, - expr: &DynamicParameter, - _input_schema: &DFSchema, - ) -> datafusion::common::Result { - let id = format!("${}", expr.parameter_reference + 1); - let field = expr - .r#type - .as_ref() - .map(|t| { - super::from_substrait_type_without_names(self, t).map(|dt| { - Arc::new(datafusion::arrow::datatypes::Field::new(&id, dt, true)) + + fn consume_literal<'a>( + &'a self, + expr: &'a Literal, + ) -> BoxFuture<'a, datafusion::common::Result> { + Box::pin(from_literal(self, expr)) + } + + fn consume_field_reference<'a>( + &'a self, + expr: &'a FieldReference, + input_schema: &'a DFSchema, + ) -> BoxFuture<'a, datafusion::common::Result> { + Box::pin(from_field_reference(self, expr, input_schema)) + } + + fn consume_scalar_function<'a>( + &'a self, + expr: &'a ScalarFunction, + input_schema: &'a DFSchema, + ) -> BoxFuture<'a, datafusion::common::Result> { + Box::pin(from_scalar_function(self, expr, input_schema)) + } + + fn consume_window_function<'a>( + &'a self, + expr: &'a WindowFunction, + input_schema: &'a DFSchema, + ) -> BoxFuture<'a, datafusion::common::Result> { + Box::pin(from_window_function(self, expr, input_schema)) + } + + fn consume_if_then<'a>( + &'a self, + expr: &'a IfThen, + input_schema: &'a DFSchema, + ) -> BoxFuture<'a, datafusion::common::Result> { + Box::pin(from_if_then(self, expr, input_schema)) + } + + fn consume_switch<'a>( + &'a self, + _expr: &'a SwitchExpression, + _input_schema: &'a DFSchema, + ) -> BoxFuture<'a, datafusion::common::Result> { + Box::pin(async move { not_impl_err!("Switch expression not supported") }) + } + + fn consume_singular_or_list<'a>( + &'a self, + expr: &'a SingularOrList, + input_schema: &'a DFSchema, + ) -> BoxFuture<'a, datafusion::common::Result> { + Box::pin(from_singular_or_list(self, expr, input_schema)) + } + + fn consume_multi_or_list<'a>( + &'a self, + _expr: &'a MultiOrList, + _input_schema: &'a DFSchema, + ) -> BoxFuture<'a, datafusion::common::Result> { + Box::pin(async move { not_impl_err!("Multi Or List expression not supported") }) + } + + fn consume_cast<'a>( + &'a self, + expr: &'a substrait_expression::Cast, + input_schema: &'a DFSchema, + ) -> BoxFuture<'a, datafusion::common::Result> { + Box::pin(from_cast(self, expr, input_schema)) + } + + fn consume_subquery<'a>( + &'a self, + expr: &'a substrait_expression::Subquery, + input_schema: &'a DFSchema, + ) -> BoxFuture<'a, datafusion::common::Result> { + Box::pin(from_subquery(self, expr, input_schema)) + } + + fn consume_nested<'a>( + &'a self, + expr: &'a Nested, + input_schema: &'a DFSchema, + ) -> BoxFuture<'a, datafusion::common::Result> { + Box::pin(from_nested(self, expr, input_schema)) + } + + fn consume_enum<'a>( + &'a self, + _expr: &'a Enum, + _input_schema: &'a DFSchema, + ) -> BoxFuture<'a, datafusion::common::Result> { + Box::pin(async move { not_impl_err!("Enum expression not supported") }) + } + + fn consume_dynamic_parameter<'a>( + &'a self, + expr: &'a DynamicParameter, + _input_schema: &'a DFSchema, + ) -> BoxFuture<'a, datafusion::common::Result> { + Box::pin(async move { + let id = format!("${}", expr.parameter_reference + 1); + let field = expr + .r#type + .as_ref() + .map(|t| { + super::from_substrait_type_without_names(self, t).map(|dt| { + Arc::new(datafusion::arrow::datatypes::Field::new(&id, dt, true)) + }) }) - }) - .transpose()?; - Ok(Expr::Placeholder( - datafusion::logical_expr::expr::Placeholder::new_with_field(id, field), - )) + .transpose()?; + Ok(Expr::Placeholder( + datafusion::logical_expr::expr::Placeholder::new_with_field(id, field), + )) + }) } - async fn consume_lambda( - &self, - expr: &proto::expression::Lambda, - input_schema: &DFSchema, - ) -> datafusion::common::Result { - from_lambda(self, expr, input_schema).await + fn consume_lambda<'a>( + &'a self, + expr: &'a proto::expression::Lambda, + input_schema: &'a DFSchema, + ) -> BoxFuture<'a, datafusion::common::Result> { + Box::pin(from_lambda(self, expr, input_schema)) } // Outer Schema Stack @@ -449,43 +473,49 @@ pub trait SubstraitConsumer: Send + Sync + Sized { // The details of extension relations, and how to handle them, are fully up to users to specify. // The following methods allow users to customize the consumer behaviour - async fn consume_extension_leaf( - &self, - rel: &ExtensionLeafRel, - ) -> datafusion::common::Result { - if let Some(detail) = rel.detail.as_ref() { - return substrait_err!( - "Missing handler for ExtensionLeafRel: {}", - detail.type_url - ); - } - substrait_err!("Missing handler for ExtensionLeafRel") - } - - async fn consume_extension_single( - &self, - rel: &ExtensionSingleRel, - ) -> datafusion::common::Result { - if let Some(detail) = rel.detail.as_ref() { - return substrait_err!( - "Missing handler for ExtensionSingleRel: {}", - detail.type_url - ); - } - substrait_err!("Missing handler for ExtensionSingleRel") - } - - async fn consume_extension_multi( - &self, - rel: &ExtensionMultiRel, - ) -> datafusion::common::Result { - if let Some(detail) = rel.detail.as_ref() { - return substrait_err!( - "Missing handler for ExtensionMultiRel: {}", - detail.type_url - ); - } - substrait_err!("Missing handler for ExtensionMultiRel") + fn consume_extension_leaf<'a>( + &'a self, + rel: &'a ExtensionLeafRel, + ) -> BoxFuture<'a, datafusion::common::Result> { + Box::pin(async move { + if let Some(detail) = rel.detail.as_ref() { + return substrait_err!( + "Missing handler for ExtensionLeafRel: {}", + detail.type_url + ); + } + substrait_err!("Missing handler for ExtensionLeafRel") + }) + } + + fn consume_extension_single<'a>( + &'a self, + rel: &'a ExtensionSingleRel, + ) -> BoxFuture<'a, datafusion::common::Result> { + Box::pin(async move { + if let Some(detail) = rel.detail.as_ref() { + return substrait_err!( + "Missing handler for ExtensionSingleRel: {}", + detail.type_url + ); + } + substrait_err!("Missing handler for ExtensionSingleRel") + }) + } + + fn consume_extension_multi<'a>( + &'a self, + rel: &'a ExtensionMultiRel, + ) -> BoxFuture<'a, datafusion::common::Result> { + Box::pin(async move { + if let Some(detail) = rel.detail.as_ref() { + return substrait_err!( + "Missing handler for ExtensionMultiRel: {}", + detail.type_url + ); + } + substrait_err!("Missing handler for ExtensionMultiRel") + }) } // Users can bring their own types to Substrait which require custom handling @@ -594,17 +624,17 @@ impl<'a> DefaultSubstraitConsumer<'a> { } } } - -#[async_trait] impl SubstraitConsumer for DefaultSubstraitConsumer<'_> { - async fn resolve_table_ref( - &self, - table_ref: &TableReference, - ) -> datafusion::common::Result>> { - let table = table_ref.table().to_string(); - let schema = self.state.schema_for_ref(table_ref.clone())?; - let table_provider = schema.table(&table).await?; - Ok(table_provider) + fn resolve_table_ref<'a>( + &'a self, + table_ref: &'a TableReference, + ) -> BoxFuture<'a, datafusion::common::Result>>> { + Box::pin(async move { + let table = table_ref.table().to_string(); + let schema = self.state.schema_for_ref(table_ref.clone())?; + let table_provider = schema.table(&table).await?; + Ok(table_provider) + }) } fn get_extensions(&self) -> &Extensions { @@ -633,59 +663,66 @@ impl SubstraitConsumer for DefaultSubstraitConsumer<'_> { .and_then(|idx| schemas.get(idx).cloned()) } - async fn consume_extension_leaf( - &self, - rel: &ExtensionLeafRel, - ) -> datafusion::common::Result { - let Some(ext_detail) = &rel.detail else { - return substrait_err!("Unexpected empty detail in ExtensionLeafRel"); - }; - let plan = self - .state - .serializer_registry() - .deserialize_logical_plan(&ext_detail.type_url, &ext_detail.value)?; - Ok(LogicalPlan::Extension(Extension { node: plan })) - } - - async fn consume_extension_single( - &self, - rel: &ExtensionSingleRel, - ) -> datafusion::common::Result { - let Some(ext_detail) = &rel.detail else { - return substrait_err!("Unexpected empty detail in ExtensionSingleRel"); - }; - let plan = self - .state - .serializer_registry() - .deserialize_logical_plan(&ext_detail.type_url, &ext_detail.value)?; - let Some(input_rel) = &rel.input else { - return substrait_err!( - "ExtensionSingleRel missing input rel, try using ExtensionLeafRel instead" - ); - }; - let input_plan = self.consume_rel(input_rel).await?; - let plan = plan.with_exprs_and_inputs(plan.expressions(), vec![input_plan])?; - Ok(LogicalPlan::Extension(Extension { node: plan })) - } - - async fn consume_extension_multi( - &self, - rel: &ExtensionMultiRel, - ) -> datafusion::common::Result { - let Some(ext_detail) = &rel.detail else { - return substrait_err!("Unexpected empty detail in ExtensionMultiRel"); - }; - let plan = self - .state - .serializer_registry() - .deserialize_logical_plan(&ext_detail.type_url, &ext_detail.value)?; - let mut inputs = Vec::with_capacity(rel.inputs.len()); - for input in &rel.inputs { - let input_plan = self.consume_rel(input).await?; - inputs.push(input_plan); - } - let plan = plan.with_exprs_and_inputs(plan.expressions(), inputs)?; - Ok(LogicalPlan::Extension(Extension { node: plan })) + fn consume_extension_leaf<'a>( + &'a self, + rel: &'a ExtensionLeafRel, + ) -> BoxFuture<'a, datafusion::common::Result> { + Box::pin(async move { + let Some(ext_detail) = &rel.detail else { + return substrait_err!("Unexpected empty detail in ExtensionLeafRel"); + }; + let plan = self + .state + .serializer_registry() + .deserialize_logical_plan(&ext_detail.type_url, &ext_detail.value)?; + Ok(LogicalPlan::Extension(Extension { node: plan })) + }) + } + + fn consume_extension_single<'a>( + &'a self, + rel: &'a ExtensionSingleRel, + ) -> BoxFuture<'a, datafusion::common::Result> { + Box::pin(async move { + let Some(ext_detail) = &rel.detail else { + return substrait_err!("Unexpected empty detail in ExtensionSingleRel"); + }; + let plan = self + .state + .serializer_registry() + .deserialize_logical_plan(&ext_detail.type_url, &ext_detail.value)?; + let Some(input_rel) = &rel.input else { + return substrait_err!( + "ExtensionSingleRel missing input rel, try using ExtensionLeafRel instead" + ); + }; + let input_plan = self.consume_rel(input_rel).await?; + let plan = + plan.with_exprs_and_inputs(plan.expressions(), vec![input_plan])?; + Ok(LogicalPlan::Extension(Extension { node: plan })) + }) + } + + fn consume_extension_multi<'a>( + &'a self, + rel: &'a ExtensionMultiRel, + ) -> BoxFuture<'a, datafusion::common::Result> { + Box::pin(async move { + let Some(ext_detail) = &rel.detail else { + return substrait_err!("Unexpected empty detail in ExtensionMultiRel"); + }; + let plan = self + .state + .serializer_registry() + .deserialize_logical_plan(&ext_detail.type_url, &ext_detail.value)?; + let mut inputs = Vec::with_capacity(rel.inputs.len()); + for input in &rel.inputs { + let input_plan = self.consume_rel(input).await?; + inputs.push(input_plan); + } + let plan = plan.with_exprs_and_inputs(plan.expressions(), inputs)?; + Ok(LogicalPlan::Extension(Extension { node: plan })) + }) } fn push_lambda_parameters( diff --git a/datafusion/substrait/src/logical_plan/consumer/types.rs b/datafusion/substrait/src/logical_plan/consumer/types.rs index e588734782741..d45db5e61334c 100644 --- a/datafusion/substrait/src/logical_plan/consumer/types.rs +++ b/datafusion/substrait/src/logical_plan/consumer/types.rs @@ -444,7 +444,6 @@ mod tests { use super::*; use crate::extensions::Extensions; use crate::logical_plan::consumer::DefaultSubstraitConsumer; - use async_trait::async_trait; use datafusion::catalog::TableProvider; use datafusion::common::TableReference; use datafusion::execution::{FunctionRegistry, SessionState, SessionStateBuilder}; @@ -460,14 +459,15 @@ mod tests { inner: DefaultSubstraitConsumer<'a>, metadata: Option>, } + use futures::future::BoxFuture; - #[async_trait] impl SubstraitConsumer for MetadataConsumer<'_> { - async fn resolve_table_ref( - &self, - table_ref: &TableReference, - ) -> datafusion::common::Result>> { - self.inner.resolve_table_ref(table_ref).await + fn resolve_table_ref<'a>( + &'a self, + table_ref: &'a TableReference, + ) -> BoxFuture<'a, datafusion::common::Result>>> + { + Box::pin(async move { self.inner.resolve_table_ref(table_ref).await }) } fn get_extensions(&self) -> &Extensions {