From e155ff5b56b1f2f2d330378695b452acd1632840 Mon Sep 17 00:00:00 2001 From: R Tyler Croy Date: Wed, 9 Sep 2026 17:47:03 +0000 Subject: [PATCH 1/4] refactor: convert async_trait traits/impls to explicit BoxFuture signatures Convert trait definitions in session, datasource, execution, and expr crates from #[async_trait] to explicit BoxFuture<'a, ...> return types (the desugared form async_trait produces), and mechanically transform impl blocks across the workspace to match. Async blocks do not get the unsized coercion that async fn return positions did, so returned concrete Arc values gain explicit 'as Arc' casts where required. Ref: https://github.com/apache/datafusion/issues/17689 Co-Authored-By: Claude Fable 5 Signed-off-by: R Tyler Croy --- Cargo.lock | 4 + datafusion-cli/src/catalog.rs | 112 ++-- datafusion-cli/src/functions.rs | 106 ++-- datafusion-cli/tests/cli_integration.rs | 35 +- .../builtin_functions/function_factory.rs | 16 +- .../custom_data_source/custom_datasource.rs | 18 +- .../custom_data_source/custom_file_format.rs | 58 +- .../default_column_values.rs | 95 +-- .../examples/data_io/catalog.rs | 15 +- .../data_io/parquet_advanced_index.rs | 94 +-- .../data_io/parquet_embedded_index.rs | 118 ++-- .../data_io/parquet_encrypted_with_kms.rs | 94 +-- .../examples/data_io/parquet_index.rs | 89 +-- .../examples/data_io/remote_catalog.rs | 82 +-- .../examples/dataframe/cache_factory.rs | 127 ++-- .../examples/relation_planner/table_sample.rs | 75 +-- datafusion-examples/examples/udf/async_udf.rs | 82 +-- .../examples/udf/simple_udtf.rs | 64 +- datafusion/catalog-listing/src/table.rs | 55 +- datafusion/catalog/src/async.rs | 442 ++++++------- datafusion/catalog/src/cte_worktable.rs | 40 +- .../catalog/src/default_table_source.rs | 22 +- .../catalog/src/dynamic_file/catalog.rs | 31 +- datafusion/catalog/src/empty.rs | 30 +- datafusion/catalog/src/information_schema.rs | 67 +- datafusion/catalog/src/listing_schema.rs | 26 +- datafusion/catalog/src/memory/schema.rs | 19 +- datafusion/catalog/src/memory/table.rs | 69 +- datafusion/catalog/src/stream.rs | 76 +-- datafusion/catalog/src/streaming.rs | 24 +- datafusion/catalog/src/view.rs | 24 +- datafusion/core/src/dataframe/mod.rs | 24 +- .../core/src/datasource/dynamic_file.rs | 18 +- .../core/src/datasource/file_format/csv.rs | 1 - .../src/datasource/file_format/options.rs | 111 +--- .../src/datasource/file_format/parquet.rs | 1 - .../core/src/datasource/listing/table.rs | 40 +- .../src/datasource/listing_table_factory.rs | 32 +- datafusion/core/src/datasource/provider.rs | 21 +- datafusion/core/src/execution/context/mod.rs | 49 +- .../core/src/execution/session_state.rs | 38 +- datafusion/core/src/physical_planner.rs | 275 ++++---- datafusion/core/src/test_util/mod.rs | 39 +- .../custom_sources_cases/dml_planning.rs | 126 ++-- .../core/tests/custom_sources_cases/mod.rs | 25 +- .../provider_filter_pushdown.rs | 95 +-- .../tests/custom_sources_cases/statistics.rs | 73 ++- datafusion/core/tests/memory_limit/mod.rs | 34 +- datafusion/core/tests/parquet/encryption.rs | 74 +-- .../physical_optimizer/sanity_checker.rs | 42 +- .../tests/user_defined/insert_operation.rs | 32 +- .../tests/user_defined/statistics_requests.rs | 60 +- .../user_defined_async_scalar_functions.rs | 28 +- .../tests/user_defined/user_defined_plan.rs | 84 ++- .../user_defined_scalar_functions.rs | 33 +- .../user_defined_table_functions.rs | 67 +- .../datasource-arrow/src/file_format.rs | 345 +++++----- datafusion/datasource-avro/src/file_format.rs | 276 ++++---- datafusion/datasource-csv/src/file_format.rs | 286 ++++----- datafusion/datasource-json/src/file_format.rs | 256 ++++---- .../datasource-parquet/src/file_format.rs | 441 ++++++------- datafusion/datasource-parquet/src/sink.rs | 264 ++++---- datafusion/datasource/src/file_format.rs | 91 +-- datafusion/datasource/src/file_sink_config.rs | 46 +- datafusion/datasource/src/memory.rs | 50 +- datafusion/datasource/src/sink.rs | 11 +- datafusion/datasource/src/url.rs | 10 +- .../execution/src/parquet_encryption.rs | 25 +- datafusion/expr/Cargo.toml | 1 + datafusion/expr/src/async_udf.rs | 31 +- datafusion/ffi/src/physical_optimizer.rs | 3 - datafusion/ffi/src/query_planner.rs | 42 +- datafusion/ffi/src/schema_provider.rs | 26 +- datafusion/ffi/src/session/mod.rs | 44 +- datafusion/ffi/src/table_provider.rs | 320 +++++----- datafusion/ffi/src/table_provider_factory.rs | 100 +-- datafusion/ffi/src/tests/async_provider.rs | 59 +- datafusion/ffi/src/tests/catalog.rs | 11 +- datafusion/ffi/src/tests/mod.rs | 109 ++-- datafusion/ffi/src/tests/query_planner.rs | 221 +++---- .../ffi/src/tests/table_provider_factory.rs | 42 +- datafusion/functions-table/Cargo.toml | 1 + .../functions-table/src/generate_series.rs | 36 +- datafusion/optimizer/src/push_down_filter.rs | 2 - datafusion/physical-plan/src/streaming.rs | 3 - datafusion/physical-plan/src/unnest.rs | 3 - datafusion/proto/Cargo.toml | 1 + datafusion/proto/tests/cases/plans/sinks.rs | 16 +- datafusion/proto/tests/cases/plans/udfs.rs | 9 +- datafusion/session/src/planner.rs | 136 ++-- datafusion/session/src/schema.rs | 20 +- datafusion/session/src/session.rs | 11 +- datafusion/session/src/table.rs | 146 ++--- datafusion/sqllogictest/src/test_context.rs | 55 +- datafusion/substrait/Cargo.toml | 1 + .../consumer/substrait_consumer.rs | 589 +++++++++--------- .../src/logical_plan/consumer/types.rs | 15 +- 97 files changed, 3827 insertions(+), 3958 deletions(-) diff --git a/Cargo.lock b/Cargo.lock index 6d4c38d8d9d04..005868d57ab1b 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -2178,6 +2178,7 @@ dependencies = [ "datafusion-proto-common", "datafusion-proto-models", "env_logger", + "futures", "indexmap 2.14.2", "insta", "itertools 0.15.0", @@ -2339,6 +2340,7 @@ dependencies = [ "datafusion-expr", "datafusion-physical-expr", "datafusion-physical-plan", + "futures", "parking_lot", ] @@ -2555,6 +2557,7 @@ dependencies = [ "datafusion-proto-models", "doc-comment", "flate2", + "futures", "object_store", "pretty_assertions", "prost", @@ -2719,6 +2722,7 @@ dependencies = [ "chrono", "datafusion", "datafusion-functions-aggregate", + "futures", "half", "insta", "itertools 0.15.0", 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/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/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/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..eb89f19c85854 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}, @@ -563,49 +565,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..2d8d994c9e608 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; @@ -181,7 +181,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 +190,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( - &self, + fn invoke_async_with_args<'a>( + &'a 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<'a, 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/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/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..c1ade1cd07323 100644 --- a/datafusion/catalog/src/default_table_source.rs +++ b/datafusion/catalog/src/default_table_source.rs @@ -17,6 +17,7 @@ //! Default TableSource implementation used in DataFusion physical plans +use futures::future::BoxFuture; use std::borrow::Cow; use std::sync::Arc; @@ -99,13 +100,10 @@ pub fn source_as_provider( #[test] fn preserves_table_type() { - use async_trait::async_trait; use datafusion_common::DataFusionError; #[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/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..294add44eea8d 100644 --- a/datafusion/core/src/datasource/file_format/csv.rs +++ b/datafusion/core/src/datasource/file_format/csv.rs @@ -86,7 +86,6 @@ mod tests { write!(f, "VariableStream") } } - #[async_trait] impl ObjectStore for VariableStream { async fn put_opts( 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..4980f9e51e7a1 100644 --- a/datafusion/core/src/datasource/file_format/parquet.rs +++ b/datafusion/core/src/datasource/file_format/parquet.rs @@ -303,7 +303,6 @@ mod tests { self.clone() } } - #[async_trait] impl ObjectStore for RequestCountingObjectStore { async fn put_opts( diff --git a/datafusion/core/src/datasource/listing/table.rs b/datafusion/core/src/datasource/listing/table.rs index 982766dc88519..ca6895acacc69 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<'a>( self, - state: &dyn Session, - ) -> datafusion_common::Result; + state: &'a dyn Session, + ) -> BoxFuture<'a, datafusion_common::Result>; /// Convenience method to call both [`Self::infer_options`] and [`ListingTableConfig::infer_schema`] - async fn infer( + fn infer<'a>( self, - state: &dyn Session, - ) -> datafusion_common::Result; + state: &'a dyn Session, + ) -> BoxFuture<'a, 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<'a>( self, - state: &'life0 dyn Session, - ) -> BoxFuture<'async_trait, datafusion_common::Result> - where - 'life0: 'async_trait, - Self: 'async_trait, - { + state: &'a dyn Session, + ) -> BoxFuture<'a, 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<'a>( self, - state: &'life0 dyn Session, - ) -> BoxFuture<'async_trait, datafusion_common::Result> - where - 'life0: 'async_trait, - Self: 'async_trait, - { + state: &'a dyn Session, + ) -> BoxFuture<'a, 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_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/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/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..2e2537319cd36 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( - &self, + fn invoke_async_with_args<'a>( + &'a self, args: ScalarFunctionArgs, - ) -> Result { - Ok(args.args[0].clone()) + ) -> BoxFuture<'a, 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( - &self, + fn invoke_async_with_args<'a>( + &'a self, args: ScalarFunctionArgs, - ) -> Result { - let arg1 = &args.args[0]; - let results = call_external_service(arg1.clone()).await?; - Ok(results) + ) -> BoxFuture<'a, 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/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/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/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/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/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/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..162da2d3fb47a 100644 --- a/datafusion/datasource/src/url.rs +++ b/datafusion/datasource/src/url.rs @@ -1086,7 +1086,6 @@ mod tests { self.in_mem.fmt(f) } } - #[async_trait] impl ObjectStore for MockObjectStore { async fn put_opts( @@ -1182,7 +1181,6 @@ mod tests { } } - #[async_trait::async_trait] impl Session for MockSession { fn session_id(&self) -> &str { unimplemented!() @@ -1196,10 +1194,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/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..e4edb57e27ecb 100644 --- a/datafusion/expr/Cargo.toml +++ b/datafusion/expr/Cargo.toml @@ -53,6 +53,7 @@ sql = ["sqlparser"] 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..723c093f03ee7 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( - &self, + fn invoke_async_with_args<'a>( + &'a self, args: ScalarFunctionArgs, - ) -> Result; + ) -> BoxFuture<'a, 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( - &self, + fn invoke_async_with_args<'a>( + &'a self, _args: ScalarFunctionArgs, - ) -> Result { - todo!() + ) -> BoxFuture<'a, 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( - &self, + fn invoke_async_with_args<'a>( + &'a self, _args: ScalarFunctionArgs, - ) -> Result { - todo!() + ) -> BoxFuture<'a, Result> { + Box::pin(async move { todo!() }) } } 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..f6638c8958db7 100644 --- a/datafusion/functions-table/Cargo.toml +++ b/datafusion/functions-table/Cargo.toml @@ -45,6 +45,7 @@ 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/src/push_down_filter.rs b/datafusion/optimizer/src/push_down_filter.rs index 8a1dcc12ef874..e8a45afde94a0 100644 --- a/datafusion/optimizer/src/push_down_filter.rs +++ b/datafusion/optimizer/src/push_down_filter.rs @@ -3100,8 +3100,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/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..f8b68c2b4952e 100644 --- a/datafusion/proto/Cargo.toml +++ b/datafusion/proto/Cargo.toml @@ -81,6 +81,7 @@ 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..b0b8bf5dc6324 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( - &self, + fn invoke_async_with_args<'a>( + &'a self, args: ScalarFunctionArgs, - ) -> Result { - Ok(args.args[0].clone()) + ) -> futures::future::BoxFuture<'a, Result> { + Box::pin(async move { Ok(args.args[0].clone()) }) } } 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..6ca9d3d26f7d7 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,10 +146,10 @@ 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. diff --git a/datafusion/session/src/table.rs b/datafusion/session/src/table.rs index a6cddd10ce80a..419b4f8bc17bc 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; @@ -184,13 +181,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 +206,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 @@ -264,7 +256,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 +264,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 +338,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 +400,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 +573,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/src/test_context.rs b/datafusion/sqllogictest/src/test_context.rs index 307a8876abdd5..31b5f5972399f 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,14 @@ 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( - &self, + fn invoke_async_with_args<'a>( + &'a self, args: ScalarFunctionArgs, - ) -> Result { - return self.inner_abs.invoke_with_args(args); + ) -> BoxFuture<'a, Result> { + Box::pin(async move { + return 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..9efa3e701fb5a 100644 --- a/datafusion/substrait/Cargo.toml +++ b/datafusion/substrait/Cargo.toml @@ -37,6 +37,7 @@ workspace = true 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..848ce224aac6d 100644 --- a/datafusion/substrait/src/logical_plan/consumer/substrait_consumer.rs +++ b/datafusion/substrait/src/logical_plan/consumer/substrait_consumer.rs @@ -36,6 +36,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 +50,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. /// @@ -191,10 +190,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 +210,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 +303,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 + 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_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 - } - - 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 +461,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 +612,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 +651,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..f16c97b1b7a90 100644 --- a/datafusion/substrait/src/logical_plan/consumer/types.rs +++ b/datafusion/substrait/src/logical_plan/consumer/types.rs @@ -39,6 +39,7 @@ use datafusion::arrow::datatypes::{ use datafusion::common::{ DFSchema, not_impl_err, substrait_datafusion_err, substrait_err, }; +use futures::future::BoxFuture; use std::sync::Arc; use substrait::proto::{NamedStruct, Type, r#type}; @@ -444,7 +445,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 +460,13 @@ mod tests { inner: DefaultSubstraitConsumer<'a>, metadata: Option>, } - - #[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 { From d27c7bf91d73eb70728c0d14e59c8b3756382162 Mon Sep 17 00:00:00 2001 From: R Tyler Croy Date: Wed, 9 Sep 2026 21:46:34 +0000 Subject: [PATCH 2/4] refactor: remove async-trait dependency from Cargo.toml files Remove async-trait from 18 crates that no longer use it after the BoxFuture conversion. Move it to dev-dependencies in datafusion-datasource (only used in a cfg(test) ObjectStore impl). async-trait remains a dependency only where impls of external traits still require it: datafusion-cli and benchmarks (object_store's ObjectStore / CredentialProvider), datafusion-sqllogictest (sqllogictest's AsyncDB), and datafusion-core (ObjectStore test utilities). Also convert CliSessionContext (datafusion-cli) to BoxFuture signatures, update the SubstraitConsumer doc example, and drop now-unused async_trait imports and doc references. Ref: https://github.com/apache/datafusion/issues/17689 Co-Authored-By: Claude Fable 5 Signed-off-by: R Tyler Croy --- Cargo.lock | 21 ----- Cargo.toml | 1 - benchmarks/Cargo.toml | 1 - datafusion-cli/Cargo.toml | 1 - .../examples/cli-session-context.rs | 16 ++-- datafusion-cli/src/cli_context.rs | 13 ++- datafusion-examples/Cargo.toml | 1 - .../examples/relation_planner/table_sample.rs | 1 - datafusion-examples/examples/udf/async_udf.rs | 2 +- datafusion/catalog-listing/Cargo.toml | 1 - datafusion/catalog/Cargo.toml | 1 - .../catalog/src/default_table_source.rs | 2 +- datafusion/core/Cargo.toml | 2 - datafusion/datasource-arrow/Cargo.toml | 1 - datafusion/datasource-avro/Cargo.toml | 1 - datafusion/datasource-csv/Cargo.toml | 1 - datafusion/datasource-json/Cargo.toml | 1 - datafusion/datasource-parquet/Cargo.toml | 1 - datafusion/datasource/Cargo.toml | 1 - datafusion/execution/Cargo.toml | 1 - datafusion/expr/Cargo.toml | 1 - datafusion/expr/src/async_udf.rs | 18 ++-- datafusion/ffi/Cargo.toml | 1 - datafusion/functions-table/Cargo.toml | 1 - datafusion/optimizer/Cargo.toml | 1 - datafusion/optimizer/src/push_down_filter.rs | 1 - datafusion/physical-plan/Cargo.toml | 1 - datafusion/proto/Cargo.toml | 1 - datafusion/session/Cargo.toml | 1 - datafusion/sqllogictest/Cargo.toml | 1 - datafusion/substrait/Cargo.toml | 1 - .../consumer/substrait_consumer.rs | 86 +++++++++++-------- 32 files changed, 75 insertions(+), 109 deletions(-) diff --git a/Cargo.lock b/Cargo.lock index 005868d57ab1b..2e480d1d8f00d 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -1704,7 +1704,6 @@ version = "55.0.0" dependencies = [ "arrow", "arrow-schema", - "async-trait", "bytes", "bzip2", "chrono", @@ -1809,7 +1808,6 @@ name = "datafusion-catalog" version = "55.0.0" dependencies = [ "arrow", - "async-trait", "dashmap", "datafusion-common", "datafusion-common-runtime", @@ -1832,7 +1830,6 @@ name = "datafusion-catalog-listing" version = "55.0.0" dependencies = [ "arrow", - "async-trait", "chrono", "datafusion-catalog", "datafusion-common", @@ -1856,7 +1853,6 @@ name = "datafusion-cli" version = "55.0.0" dependencies = [ "arrow", - "async-trait", "aws-config", "aws-credential-types", "chrono", @@ -1928,7 +1924,6 @@ version = "55.0.0" dependencies = [ "arrow", "async-compression", - "async-trait", "bytes", "bzip2", "chrono", @@ -1967,7 +1962,6 @@ version = "55.0.0" dependencies = [ "arrow", "arrow-ipc", - "async-trait", "bytes", "chrono", "datafusion-common", @@ -1991,7 +1985,6 @@ version = "55.0.0" dependencies = [ "arrow", "arrow-avro", - "async-trait", "bytes", "datafusion-common", "datafusion-common-runtime", @@ -2013,7 +2006,6 @@ name = "datafusion-datasource-csv" version = "55.0.0" dependencies = [ "arrow", - "async-trait", "bytes", "datafusion-common", "datafusion-common-runtime", @@ -2035,7 +2027,6 @@ name = "datafusion-datasource-json" version = "55.0.0" dependencies = [ "arrow", - "async-trait", "bytes", "datafusion-common", "datafusion-common-runtime", @@ -2059,7 +2050,6 @@ version = "55.0.0" dependencies = [ "arrow", "arrow-schema", - "async-trait", "bytes", "chrono", "criterion", @@ -2099,7 +2089,6 @@ dependencies = [ "arrow", "arrow-flight", "arrow-schema", - "async-trait", "base64 0.23.1", "bytes", "dashmap", @@ -2139,7 +2128,6 @@ version = "55.0.0" dependencies = [ "arrow", "arrow-buffer", - "async-trait", "bytes", "chrono", "dashmap", @@ -2166,7 +2154,6 @@ version = "55.0.0" dependencies = [ "arrow", "arrow-schema", - "async-trait", "chrono", "ctor", "datafusion-common", @@ -2206,7 +2193,6 @@ dependencies = [ "arrow", "arrow-schema", "async-ffi", - "async-trait", "chrono", "datafusion", "datafusion-catalog", @@ -2334,7 +2320,6 @@ name = "datafusion-functions-table" version = "55.0.0" dependencies = [ "arrow", - "async-trait", "datafusion-catalog", "datafusion-common", "datafusion-expr", @@ -2382,7 +2367,6 @@ name = "datafusion-optimizer" version = "55.0.0" dependencies = [ "arrow", - "async-trait", "chrono", "criterion", "ctor", @@ -2494,7 +2478,6 @@ dependencies = [ "arrow-ipc", "arrow-ord", "arrow-schema", - "async-trait", "bytes", "criterion", "datafusion-common", @@ -2533,7 +2516,6 @@ name = "datafusion-proto" version = "55.0.0" dependencies = [ "arrow", - "async-trait", "datafusion", "datafusion-catalog", "datafusion-catalog-listing", @@ -2615,7 +2597,6 @@ name = "datafusion-session" version = "55.0.0" dependencies = [ "arrow-schema", - "async-trait", "datafusion-common", "datafusion-execution", "datafusion-expr", @@ -2685,7 +2666,6 @@ name = "datafusion-sqllogictest" version = "55.0.0" dependencies = [ "arrow", - "async-trait", "bigdecimal", "bytes", "chrono", @@ -2718,7 +2698,6 @@ name = "datafusion-substrait" version = "55.0.0" dependencies = [ "async-recursion", - "async-trait", "chrono", "datafusion", "datafusion-functions-aggregate", 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/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/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-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/relation_planner/table_sample.rs b/datafusion-examples/examples/relation_planner/table_sample.rs index eb89f19c85854..988dd03267633 100644 --- a/datafusion-examples/examples/relation_planner/table_sample.rs +++ b/datafusion-examples/examples/relation_planner/table_sample.rs @@ -100,7 +100,6 @@ use futures::{ stream::{Stream, StreamExt}, }; use rand::{Rng, SeedableRng, rngs::StdRng}; -use tonic::async_trait; use datafusion::{ catalog::Session, diff --git a/datafusion-examples/examples/udf/async_udf.rs b/datafusion-examples/examples/udf/async_udf.rs index 2d8d994c9e608..5678f8f92e3d8 100644 --- a/datafusion-examples/examples/udf/async_udf.rs +++ b/datafusion-examples/examples/udf/async_udf.rs @@ -158,7 +158,7 @@ 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" 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/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/default_table_source.rs b/datafusion/catalog/src/default_table_source.rs index c1ade1cd07323..5c9ee3907c8ba 100644 --- a/datafusion/catalog/src/default_table_source.rs +++ b/datafusion/catalog/src/default_table_source.rs @@ -17,7 +17,6 @@ //! Default TableSource implementation used in DataFusion physical plans -use futures::future::BoxFuture; use std::borrow::Cow; use std::sync::Arc; @@ -101,6 +100,7 @@ pub fn source_as_provider( #[test] fn preserves_table_type() { use datafusion_common::DataFusionError; + use futures::future::BoxFuture; #[derive(Debug)] struct TestTempTable; 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/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-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-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-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-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/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/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/expr/Cargo.toml b/datafusion/expr/Cargo.toml index e4edb57e27ecb..62118cb57a866 100644 --- a/datafusion/expr/Cargo.toml +++ b/datafusion/expr/Cargo.toml @@ -52,7 +52,6 @@ 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 } diff --git a/datafusion/expr/src/async_udf.rs b/datafusion/expr/src/async_udf.rs index 723c093f03ee7..78eb39957fe12 100644 --- a/datafusion/expr/src/async_udf.rs +++ b/datafusion/expr/src/async_udf.rs @@ -43,10 +43,10 @@ pub trait AsyncScalarUDFImpl: ScalarUDFImpl { } /// Invoke the function asynchronously with the async arguments - fn invoke_async_with_args<'a>( - &'a self, + fn invoke_async_with_args( + &self, args: ScalarFunctionArgs, - ) -> BoxFuture<'a, Result>; + ) -> BoxFuture<'_, Result>; } /// A scalar UDF that must be invoked using async methods @@ -168,10 +168,10 @@ mod tests { } } impl AsyncScalarUDFImpl for TestAsyncUDFImpl1 { - fn invoke_async_with_args<'a>( - &'a self, + fn invoke_async_with_args( + &self, _args: ScalarFunctionArgs, - ) -> BoxFuture<'a, Result> { + ) -> BoxFuture<'_, Result> { Box::pin(async move { todo!() }) } } @@ -199,10 +199,10 @@ mod tests { } } impl AsyncScalarUDFImpl for TestAsyncUDFImpl2 { - fn invoke_async_with_args<'a>( - &'a self, + fn invoke_async_with_args( + &self, _args: ScalarFunctionArgs, - ) -> BoxFuture<'a, Result> { + ) -> 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/functions-table/Cargo.toml b/datafusion/functions-table/Cargo.toml index f6638c8958db7..95b486c491f18 100644 --- a/datafusion/functions-table/Cargo.toml +++ b/datafusion/functions-table/Cargo.toml @@ -42,7 +42,6 @@ name = "datafusion_functions_table" [dependencies] arrow = { workspace = true } -async-trait = { workspace = true } datafusion-catalog = { workspace = true } datafusion-common = { workspace = true } futures = { workspace = true } 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 e8a45afde94a0..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; 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/proto/Cargo.toml b/datafusion/proto/Cargo.toml index f8b68c2b4952e..1619f7c7822ef 100644 --- a/datafusion/proto/Cargo.toml +++ b/datafusion/proto/Cargo.toml @@ -80,7 +80,6 @@ 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", 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/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/substrait/Cargo.toml b/datafusion/substrait/Cargo.toml index 9efa3e701fb5a..500ad656e0ac5 100644 --- a/datafusion/substrait/Cargo.toml +++ b/datafusion/substrait/Cargo.toml @@ -35,7 +35,6 @@ workspace = true [dependencies] async-recursion = "1.0" -async-trait = { workspace = true } chrono = { workspace = true } futures = { workspace = true } datafusion = { workspace = true, features = ["sql"] } diff --git a/datafusion/substrait/src/logical_plan/consumer/substrait_consumer.rs b/datafusion/substrait/src/logical_plan/consumer/substrait_consumer.rs index 848ce224aac6d..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; @@ -59,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; @@ -82,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 { @@ -127,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 From dddb685501162cc2e9044432fedf377095b6c86d Mon Sep 17 00:00:00 2001 From: R Tyler Croy Date: Wed, 9 Sep 2026 21:54:20 +0000 Subject: [PATCH 3/4] refactor: fix clippy lints from async-trait removal - Elide unnecessary explicit lifetimes where only &self is borrowed - Fix doc list indentation and markdown lazy continuations - Remove unused imports and needless return - Apply clippy --fix suggestions across the workspace cargo clippy --workspace --all-targets -- -D warnings now passes. Co-Authored-By: Claude Fable 5 Signed-off-by: R Tyler Croy --- datafusion-examples/examples/udf/async_udf.rs | 7 +++--- .../core/src/datasource/listing/table.rs | 24 +++++++++---------- .../user_defined_async_scalar_functions.rs | 12 +++++----- datafusion/proto/tests/cases/plans/udfs.rs | 6 ++--- datafusion/session/src/session.rs | 7 +++--- datafusion/session/src/table.rs | 3 ++- datafusion/sqllogictest/src/test_context.rs | 10 ++++---- .../src/logical_plan/consumer/types.rs | 3 ++- 8 files changed, 36 insertions(+), 36 deletions(-) diff --git a/datafusion-examples/examples/udf/async_udf.rs b/datafusion-examples/examples/udf/async_udf.rs index 5678f8f92e3d8..3500eecd3d1ac 100644 --- a/datafusion-examples/examples/udf/async_udf.rs +++ b/datafusion-examples/examples/udf/async_udf.rs @@ -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. - impl ScalarUDFImpl for AskLLM { fn name(&self) -> &str { "ask_llm" @@ -190,10 +189,10 @@ 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. - fn invoke_async_with_args<'a>( - &'a self, + fn invoke_async_with_args( + &self, args: ScalarFunctionArgs, - ) -> BoxFuture<'a, Result> { + ) -> 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 diff --git a/datafusion/core/src/datasource/listing/table.rs b/datafusion/core/src/datasource/listing/table.rs index ca6895acacc69..f1d2bbb873d8e 100644 --- a/datafusion/core/src/datasource/listing/table.rs +++ b/datafusion/core/src/datasource/listing/table.rs @@ -32,29 +32,29 @@ pub trait ListingTableConfigExt { /// Infer `ListingOptions` based on `table_path` and file suffix. /// /// The format is inferred based on the first `table_path`. - fn infer_options<'a>( + fn infer_options( self, - state: &'a dyn Session, - ) -> BoxFuture<'a, datafusion_common::Result>; + state: &dyn Session, + ) -> BoxFuture<'_, datafusion_common::Result>; /// Convenience method to call both [`Self::infer_options`] and [`ListingTableConfig::infer_schema`] - fn infer<'a>( + fn infer( self, - state: &'a dyn Session, - ) -> BoxFuture<'a, datafusion_common::Result>; + state: &dyn Session, + ) -> BoxFuture<'_, datafusion_common::Result>; } impl ListingTableConfigExt for ListingTableConfig { - fn infer_options<'a>( + fn infer_options( self, - state: &'a dyn Session, - ) -> BoxFuture<'a, datafusion_common::Result> { + state: &dyn Session, + ) -> BoxFuture<'_, datafusion_common::Result> { infer_options_boxed(self, state) } - fn infer<'a>( + fn infer( self, - state: &'a dyn Session, - ) -> BoxFuture<'a, datafusion_common::Result> { + state: &dyn Session, + ) -> BoxFuture<'_, datafusion_common::Result> { infer_boxed(self, state) } } 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 2e2537319cd36..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 @@ -159,10 +159,10 @@ async fn test_async_udf_preserves_result_field_metadata() -> Result<()> { } } impl AsyncScalarUDFImpl for AsyncExtensionUDF { - fn invoke_async_with_args<'a>( - &'a self, + fn invoke_async_with_args( + &self, args: ScalarFunctionArgs, - ) -> BoxFuture<'a, Result> { + ) -> BoxFuture<'_, Result> { Box::pin(async move { Ok(args.args[0].clone()) }) } } @@ -252,10 +252,10 @@ impl AsyncScalarUDFImpl for TestAsyncUDFImpl { fn ideal_batch_size(&self) -> Option { Some(self.batch_size) } - fn invoke_async_with_args<'a>( - &'a self, + fn invoke_async_with_args( + &self, args: ScalarFunctionArgs, - ) -> BoxFuture<'a, Result> { + ) -> BoxFuture<'_, Result> { Box::pin(async move { let arg1 = &args.args[0]; let results = call_external_service(arg1.clone()).await?; diff --git a/datafusion/proto/tests/cases/plans/udfs.rs b/datafusion/proto/tests/cases/plans/udfs.rs index b0b8bf5dc6324..db06eaa475717 100644 --- a/datafusion/proto/tests/cases/plans/udfs.rs +++ b/datafusion/proto/tests/cases/plans/udfs.rs @@ -548,10 +548,10 @@ async fn roundtrip_async_func_exec() -> Result<()> { } impl AsyncScalarUDFImpl for TestAsyncUDF { - fn invoke_async_with_args<'a>( - &'a self, + fn invoke_async_with_args( + &self, args: ScalarFunctionArgs, - ) -> futures::future::BoxFuture<'a, Result> { + ) -> futures::future::BoxFuture<'_, Result> { Box::pin(async move { Ok(args.args[0].clone()) }) } } diff --git a/datafusion/session/src/session.rs b/datafusion/session/src/session.rs index 6ca9d3d26f7d7..0f1ee48b3e8c6 100644 --- a/datafusion/session/src/session.rs +++ b/datafusion/session/src/session.rs @@ -154,9 +154,10 @@ pub trait Session: Send + Sync { /// 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 419b4f8bc17bc..b54d08f314d6b 100644 --- a/datafusion/session/src/table.rs +++ b/datafusion/session/src/table.rs @@ -56,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> { @@ -241,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, diff --git a/datafusion/sqllogictest/src/test_context.rs b/datafusion/sqllogictest/src/test_context.rs index 31b5f5972399f..0927daf59d489 100644 --- a/datafusion/sqllogictest/src/test_context.rs +++ b/datafusion/sqllogictest/src/test_context.rs @@ -800,13 +800,11 @@ fn register_async_abs_udf(ctx: &SessionContext) { } } impl AsyncScalarUDFImpl for AsyncAbs { - fn invoke_async_with_args<'a>( - &'a self, + fn invoke_async_with_args( + &self, args: ScalarFunctionArgs, - ) -> BoxFuture<'a, Result> { - Box::pin(async move { - 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/src/logical_plan/consumer/types.rs b/datafusion/substrait/src/logical_plan/consumer/types.rs index f16c97b1b7a90..d45db5e61334c 100644 --- a/datafusion/substrait/src/logical_plan/consumer/types.rs +++ b/datafusion/substrait/src/logical_plan/consumer/types.rs @@ -39,7 +39,6 @@ use datafusion::arrow::datatypes::{ use datafusion::common::{ DFSchema, not_impl_err, substrait_datafusion_err, substrait_err, }; -use futures::future::BoxFuture; use std::sync::Arc; use substrait::proto::{NamedStruct, Type, r#type}; @@ -460,6 +459,8 @@ mod tests { inner: DefaultSubstraitConsumer<'a>, metadata: Option>, } + use futures::future::BoxFuture; + impl SubstraitConsumer for MetadataConsumer<'_> { fn resolve_table_ref<'a>( &'a self, From ebcbbf6b14f3fcb3a1e3dcc36a3edfde8d8ba648 Mon Sep 17 00:00:00 2001 From: R Tyler Croy Date: Thu, 10 Sep 2026 02:06:24 +0000 Subject: [PATCH 4/4] refactor: hand-desugar ObjectStore impls to fully remove async-trait Signed-off-by: R Tyler Croy Co-Authored-By: Claude Fable 5 --- Cargo.lock | 1 - benchmarks/src/util/latency_object_store.rs | 109 ++++++--- datafusion-cli/src/object_storage.rs | 46 ++-- .../src/object_storage/instrumented.rs | 137 +++++++---- .../core/src/datasource/file_format/csv.rs | 180 +++++++++----- .../src/datasource/file_format/parquet.rs | 112 ++++++--- datafusion/core/src/test/object_store.rs | 140 +++++++---- .../tests/datasource/object_store_access.rs | 150 ++++++++---- datafusion/core/tests/sql/path_partition.rs | 226 +++++++++++------- .../tests/tracing/traceable_object_store.rs | 113 ++++++--- datafusion/datasource/src/url.rs | 146 +++++++---- .../src/engines/datafusion_engine/runner.rs | 94 +++++--- .../runner.rs | 97 +++++--- .../src/engines/postgres_engine/mod.rs | 175 ++++++++------ 14 files changed, 1136 insertions(+), 590 deletions(-) diff --git a/Cargo.lock b/Cargo.lock index 2e480d1d8f00d..f6a2942d45658 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -1777,7 +1777,6 @@ name = "datafusion-benchmarks" version = "55.0.0" dependencies = [ "arrow", - "async-trait", "bytes", "clap", "criterion", 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/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/core/src/datasource/file_format/csv.rs b/datafusion/core/src/datasource/file_format/csv.rs index 294add44eea8d..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,64 +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( @@ -153,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/parquet.rs b/datafusion/core/src/datasource/file_format/parquet.rs index 4980f9e51e7a1..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,32 +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( @@ -345,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/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/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/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/datasource/src/url.rs b/datafusion/datasource/src/url.rs index 162da2d3fb47a..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,46 +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( @@ -1142,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 }) } } 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(); + } + }) } }