Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
32 changes: 31 additions & 1 deletion datafusion/core/src/dataframe/mod.rs
Original file line number Diff line number Diff line change
Expand Up @@ -70,6 +70,7 @@ use datafusion_functions_aggregate::expr_fn::{
use async_trait::async_trait;
use datafusion_catalog::Session;
use datafusion_expr::extension_types::DFArrayFormatterFactory;
use futures::future::BoxFuture;

/// Contains options that control how data is
/// written out from a DataFrame
Expand Down Expand Up @@ -2716,7 +2717,36 @@ impl TableProvider for DataFrameTableProvider {
self.table_type
}

async fn scan(
fn scan<'life0, 'life1, 'life2, 'life3, 'async_trait>(
&'life0 self,
state: &'life1 dyn Session,
projection: Option<&'life2 Vec<usize>>,
filters: &'life3 [Expr],
limit: Option<usize>,
) -> BoxFuture<'async_trait, Result<Arc<dyn ExecutionPlan>>>
where
'life0: 'async_trait,
'life1: 'async_trait,
'life2: 'async_trait,
'life3: 'async_trait,
Self: 'async_trait,
{
self.scan_boxed(state, projection, filters, limit)
}
}

impl DataFrameTableProvider {
fn scan_boxed<'a>(
&'a self,
state: &'a dyn Session,
projection: Option<&'a Vec<usize>>,
filters: &'a [Expr],
limit: Option<usize>,
) -> BoxFuture<'a, Result<Arc<dyn ExecutionPlan>>> {
Box::pin(self.scan_inner(state, projection, filters, limit))
}

async fn scan_inner(
&self,
state: &dyn Session,
projection: Option<&Vec<usize>>,
Expand Down
24 changes: 23 additions & 1 deletion datafusion/core/src/datasource/dynamic_file.rs
Original file line number Diff line number Diff line change
Expand Up @@ -31,6 +31,7 @@ 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.
#[derive(Default, Debug)]
Expand All @@ -53,7 +54,28 @@ impl DynamicListTableFactory {

#[async_trait]
impl UrlTableFactory for DynamicListTableFactory {
async fn try_new(&self, url: &str) -> Result<Option<Arc<dyn TableProvider>>> {
fn try_new<'life0, 'life1, 'async_trait>(
&'life0 self,
url: &'life1 str,
) -> BoxFuture<'async_trait, Result<Option<Arc<dyn TableProvider>>>>
where
'life0: 'async_trait,
'life1: 'async_trait,
Self: 'async_trait,
{
self.try_new_boxed(url)
}
}

impl DynamicListTableFactory {
fn try_new_boxed<'a>(
&'a self,
url: &'a str,
) -> BoxFuture<'a, Result<Option<Arc<dyn TableProvider>>>> {
Box::pin(self.try_new_inner(url))
}

async fn try_new_inner(&self, url: &str) -> Result<Option<Arc<dyn TableProvider>>> {
let Ok(table_url) = ListingTableUrl::parse(url) else {
return Ok(None);
};
Expand Down
97 changes: 65 additions & 32 deletions datafusion/core/src/datasource/file_format/options.rs
Original file line number Diff line number Diff line change
Expand Up @@ -17,6 +17,7 @@

//! User facing options for the file formats readers

use std::future::ready;
use std::sync::Arc;

#[cfg(feature = "avro")]
Expand Down Expand Up @@ -44,6 +45,7 @@ use datafusion_common::{
use async_trait::async_trait;
use datafusion_datasource_json::file_format::JsonFormat;
use datafusion_expr::SortExpr;
use futures::future::BoxFuture;

/// Options that control the reading of CSV files.
///
Expand Down Expand Up @@ -602,26 +604,37 @@ 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.
async fn _get_resolved_schema(
fn _get_resolved_schema<'life0, 'async_trait>(
&'a self,
config: &SessionConfig,
config: &'life0 SessionConfig,
state: SessionState,
table_path: ListingTableUrl,
schema: Option<&'a Schema>,
) -> Result<SchemaRef>
) -> BoxFuture<'async_trait, Result<SchemaRef>>
where
'a: 'async_trait,
'life0: 'async_trait,
Self: 'async_trait,
{
if let Some(s) = schema {
return Ok(Arc::new(s.to_owned()));
return Box::pin(ready(Ok(Arc::new(s.to_owned()))));
}

self.to_listing_options(config, state.default_table_options())
.infer_schema(&state, &table_path)
.await
let listing_options =
self.to_listing_options(config, state.default_table_options());
infer_schema_boxed(listing_options, state, table_path)
}
}

/// Infers a schema from `table_path`, boxed for [`ReadOptions::_get_resolved_schema`].
fn infer_schema_boxed(
listing_options: ListingOptions,
state: SessionState,
table_path: ListingTableUrl,
) -> BoxFuture<'static, Result<SchemaRef>> {
Box::pin(async move { listing_options.infer_schema(&state, &table_path).await })
}

#[async_trait]
impl ReadOptions<'_> for CsvReadOptions<'_> {
fn to_listing_options(
Expand Down Expand Up @@ -649,14 +662,18 @@ impl ReadOptions<'_> for CsvReadOptions<'_> {
.with_file_sort_order(self.file_sort_order.clone())
}

async fn get_resolved_schema(
&self,
config: &SessionConfig,
fn get_resolved_schema<'life0, 'life1, 'async_trait>(
&'life0 self,
config: &'life1 SessionConfig,
state: SessionState,
table_path: ListingTableUrl,
) -> Result<SchemaRef> {
) -> BoxFuture<'async_trait, Result<SchemaRef>>
where
'life0: 'async_trait,
'life1: 'async_trait,
Self: 'async_trait,
{
self._get_resolved_schema(config, state, table_path, self.schema)
.await
}

fn schema_source(&self) -> SchemaSource {
Expand Down Expand Up @@ -696,14 +713,18 @@ impl ReadOptions<'_> for ParquetReadOptions<'_> {
.with_file_sort_order(self.file_sort_order.clone())
}

async fn get_resolved_schema(
&self,
config: &SessionConfig,
fn get_resolved_schema<'life0, 'life1, 'async_trait>(
&'life0 self,
config: &'life1 SessionConfig,
state: SessionState,
table_path: ListingTableUrl,
) -> Result<SchemaRef> {
) -> BoxFuture<'async_trait, Result<SchemaRef>>
where
'life0: 'async_trait,
'life1: 'async_trait,
Self: 'async_trait,
{
self._get_resolved_schema(config, state, table_path, self.schema)
.await
}

fn schema_source(&self) -> SchemaSource {
Expand All @@ -730,14 +751,18 @@ impl ReadOptions<'_> for JsonReadOptions<'_> {
.with_file_sort_order(self.file_sort_order.clone())
}

async fn get_resolved_schema(
&self,
config: &SessionConfig,
fn get_resolved_schema<'life0, 'life1, 'async_trait>(
&'life0 self,
config: &'life1 SessionConfig,
state: SessionState,
table_path: ListingTableUrl,
) -> Result<SchemaRef> {
) -> BoxFuture<'async_trait, Result<SchemaRef>>
where
'life0: 'async_trait,
'life1: 'async_trait,
Self: 'async_trait,
{
self._get_resolved_schema(config, state, table_path, self.schema)
.await
}

fn schema_source(&self) -> SchemaSource {
Expand All @@ -760,14 +785,18 @@ impl ReadOptions<'_> for AvroReadOptions<'_> {
.with_table_partition_cols(self.table_partition_cols.clone())
}

async fn get_resolved_schema(
&self,
config: &SessionConfig,
fn get_resolved_schema<'life0, 'life1, 'async_trait>(
&'life0 self,
config: &'life1 SessionConfig,
state: SessionState,
table_path: ListingTableUrl,
) -> Result<SchemaRef> {
) -> BoxFuture<'async_trait, Result<SchemaRef>>
where
'life0: 'async_trait,
'life1: 'async_trait,
Self: 'async_trait,
{
self._get_resolved_schema(config, state, table_path, self.schema)
.await
}

fn schema_source(&self) -> SchemaSource {
Expand All @@ -789,14 +818,18 @@ impl ReadOptions<'_> for ArrowReadOptions<'_> {
.with_table_partition_cols(self.table_partition_cols.clone())
}

async fn get_resolved_schema(
&self,
config: &SessionConfig,
fn get_resolved_schema<'life0, 'life1, 'async_trait>(
&'life0 self,
config: &'life1 SessionConfig,
state: SessionState,
table_path: ListingTableUrl,
) -> Result<SchemaRef> {
) -> BoxFuture<'async_trait, Result<SchemaRef>>
where
'life0: 'async_trait,
'life1: 'async_trait,
Self: 'async_trait,
{
self._get_resolved_schema(config, state, table_path, self.schema)
.await
}

fn schema_source(&self) -> SchemaSource {
Expand Down
54 changes: 42 additions & 12 deletions datafusion/core/src/datasource/listing/table.rs
Original file line number Diff line number Diff line change
Expand Up @@ -21,6 +21,7 @@ use datafusion_catalog_listing::{ListingOptions, ListingTableConfig};
use datafusion_common::{config_datafusion_err, internal_datafusion_err};
use datafusion_session::Session;
use futures::StreamExt;
use futures::future::BoxFuture;
use std::collections::HashMap;

/// Extension trait for [`ListingTableConfig`] that supports inferring schemas
Expand All @@ -47,17 +48,50 @@ pub trait ListingTableConfigExt {

#[async_trait]
impl ListingTableConfigExt for ListingTableConfig {
async fn infer_options(
fn infer_options<'life0, 'async_trait>(
self,
state: &dyn Session,
) -> datafusion_common::Result<ListingTableConfig> {
let store = if let Some(url) = self.table_paths.first() {
state: &'life0 dyn Session,
) -> BoxFuture<'async_trait, datafusion_common::Result<ListingTableConfig>>
where
'life0: 'async_trait,
Self: 'async_trait,
{
infer_options_boxed(self, state)
}

fn infer<'life0, 'async_trait>(
self,
state: &'life0 dyn Session,
) -> BoxFuture<'async_trait, datafusion_common::Result<Self>>
where
'life0: 'async_trait,
Self: 'async_trait,
{
infer_boxed(self, state)
}
}

/// Body of [`ListingTableConfigExt::infer`].
fn infer_boxed<'a>(
config: ListingTableConfig,
state: &'a dyn Session,
) -> BoxFuture<'a, datafusion_common::Result<ListingTableConfig>> {
Box::pin(async move { config.infer_options(state).await?.infer_schema(state).await })
}

/// Body of [`ListingTableConfigExt::infer_options`].
fn infer_options_boxed<'a>(
config: ListingTableConfig,
state: &'a dyn Session,
) -> BoxFuture<'a, datafusion_common::Result<ListingTableConfig>> {
Box::pin(async move {
let store = if let Some(url) = config.table_paths.first() {
state.runtime_env().object_store(url)?
} else {
return Ok(self);
return Ok(config);
};

let file = self
let file = config
.table_paths
.first()
.unwrap()
Expand Down Expand Up @@ -95,12 +129,8 @@ impl ListingTableConfigExt for ListingTableConfig {
let listing_options =
ListingOptions::new(file_format).with_file_extension(listing_file_extension);

Ok(self.with_listing_options(listing_options))
}

async fn infer(self, state: &dyn Session) -> datafusion_common::Result<Self> {
self.infer_options(state).await?.infer_schema(state).await
}
Ok(config.with_listing_options(listing_options))
})
}

#[cfg(test)]
Expand Down
27 changes: 26 additions & 1 deletion datafusion/core/src/datasource/listing_table_factory.rs
Original file line number Diff line number Diff line change
Expand Up @@ -36,6 +36,7 @@ use datafusion_expr::CreateExternalTable;

use async_trait::async_trait;
use datafusion_catalog::Session;
use futures::future::BoxFuture;

/// A `TableProviderFactory` capable of creating new `ListingTable`s
#[derive(Debug, Default)]
Expand All @@ -50,7 +51,31 @@ impl ListingTableFactory {

#[async_trait]
impl TableProviderFactory for ListingTableFactory {
async fn create(
fn create<'life0, 'life1, 'life2, 'async_trait>(
&'life0 self,
state: &'life1 dyn Session,
cmd: &'life2 CreateExternalTable,
) -> BoxFuture<'async_trait, Result<Arc<dyn TableProvider>>>
where
'life0: 'async_trait,
'life1: 'async_trait,
'life2: 'async_trait,
Self: 'async_trait,
{
self.create_boxed(state, cmd)
}
}

impl ListingTableFactory {
fn create_boxed<'a>(
&'a self,
state: &'a dyn Session,
cmd: &'a CreateExternalTable,
) -> BoxFuture<'a, Result<Arc<dyn TableProvider>>> {
Box::pin(self.create_inner(state, cmd))
}

async fn create_inner(
&self,
state: &dyn Session,
cmd: &CreateExternalTable,
Expand Down
Loading
Loading