Skip to content

Commit 7b5842b

Browse files
author
Brent Gardner
authored
Allow listing tables to be created via TableFactories (#4112)
* Allow listing tables to be created via TableFactories * Add a test
1 parent c1fc732 commit 7b5842b

11 files changed

Lines changed: 192 additions & 43 deletions

File tree

‎datafusion-cli/src/main.rs‎

Lines changed: 28 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -16,6 +16,9 @@
1616
// under the License.
1717

1818
use clap::Parser;
19+
use datafusion::datasource::datasource::TableProviderFactory;
20+
use datafusion::datasource::file_format::file_type::FileType;
21+
use datafusion::datasource::listing_table_factory::ListingTableFactory;
1922
use datafusion::datasource::object_store::ObjectStoreRegistry;
2023
use datafusion::error::{DataFusionError, Result};
2124
use datafusion::execution::context::SessionConfig;
@@ -26,6 +29,7 @@ use datafusion_cli::{
2629
exec, print_format::PrintFormat, print_options::PrintOptions, DATAFUSION_CLI_VERSION,
2730
};
2831
use mimalloc::MiMalloc;
32+
use std::collections::HashMap;
2933
use std::env;
3034
use std::path::Path;
3135
use std::sync::Arc;
@@ -93,7 +97,7 @@ pub async fn main() -> Result<()> {
9397

9498
if let Some(ref path) = args.data_path {
9599
let p = Path::new(path);
96-
env::set_current_dir(&p).unwrap();
100+
env::set_current_dir(p).unwrap();
97101
};
98102

99103
let mut session_config = SessionConfig::from_env().with_information_schema(true);
@@ -105,6 +109,7 @@ pub async fn main() -> Result<()> {
105109
let runtime_env = create_runtime_env()?;
106110
let mut ctx =
107111
SessionContext::with_config_rt(session_config.clone(), Arc::new(runtime_env));
112+
ctx.refresh_catalogs().await?;
108113

109114
let mut print_options = PrintOptions {
110115
format: args.format,
@@ -142,11 +147,31 @@ pub async fn main() -> Result<()> {
142147
}
143148

144149
fn create_runtime_env() -> Result<RuntimeEnv> {
150+
let mut table_factories: HashMap<String, Arc<dyn TableProviderFactory>> =
151+
HashMap::new();
152+
table_factories.insert(
153+
"csv".to_string(),
154+
Arc::new(ListingTableFactory::new(FileType::CSV)),
155+
);
156+
table_factories.insert(
157+
"parquet".to_string(),
158+
Arc::new(ListingTableFactory::new(FileType::PARQUET)),
159+
);
160+
table_factories.insert(
161+
"avro".to_string(),
162+
Arc::new(ListingTableFactory::new(FileType::AVRO)),
163+
);
164+
table_factories.insert(
165+
"json".to_string(),
166+
Arc::new(ListingTableFactory::new(FileType::JSON)),
167+
);
168+
145169
let object_store_provider = DatafusionCliObjectStoreProvider {};
146170
let object_store_registry =
147171
ObjectStoreRegistry::new_with_provider(Some(Arc::new(object_store_provider)));
148-
let rn_config =
149-
RuntimeConfig::new().with_object_store_registry(Arc::new(object_store_registry));
172+
let rn_config = RuntimeConfig::new()
173+
.with_object_store_registry(Arc::new(object_store_registry))
174+
.with_table_factories(table_factories);
150175
RuntimeEnv::new(rn_config)
151176
}
152177

‎datafusion-cli/src/object_storage.rs‎

Lines changed: 7 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -141,7 +141,13 @@ mod tests {
141141
assert!(err.to_string().contains("Generic S3 error: Missing region"));
142142

143143
env::set_var("AWS_REGION", "us-east-1");
144-
assert!(provider.get_by_url(&Url::from_str(s3).unwrap()).is_ok());
144+
let url = Url::from_str(s3).expect("Unable to parse s3 url");
145+
let res = provider.get_by_url(&url);
146+
let msg = match res {
147+
Err(e) => format!("{}", e),
148+
Ok(_) => "".to_string()
149+
};
150+
assert_eq!("".to_string(), msg); // Fail with error message
145151
env::remove_var("AWS_REGION");
146152
}
147153
}

‎datafusion-cli/src/print_format.rs‎

Lines changed: 6 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -107,9 +107,9 @@ mod tests {
107107
let batch = RecordBatch::try_new(
108108
schema,
109109
vec![
110-
Arc::new(Int32Array::from_slice(&[1, 2, 3])),
111-
Arc::new(Int32Array::from_slice(&[4, 5, 6])),
112-
Arc::new(Int32Array::from_slice(&[7, 8, 9])),
110+
Arc::new(Int32Array::from_slice([1, 2, 3])),
111+
Arc::new(Int32Array::from_slice([4, 5, 6])),
112+
Arc::new(Int32Array::from_slice([7, 8, 9])),
113113
],
114114
)
115115
.unwrap();
@@ -137,9 +137,9 @@ mod tests {
137137
let batch = RecordBatch::try_new(
138138
schema,
139139
vec![
140-
Arc::new(Int32Array::from_slice(&[1, 2, 3])),
141-
Arc::new(Int32Array::from_slice(&[4, 5, 6])),
142-
Arc::new(Int32Array::from_slice(&[7, 8, 9])),
140+
Arc::new(Int32Array::from_slice([1, 2, 3])),
141+
Arc::new(Int32Array::from_slice([4, 5, 6])),
142+
Arc::new(Int32Array::from_slice([7, 8, 9])),
143143
],
144144
)
145145
.unwrap();

‎datafusion/core/src/catalog/listing_schema.rs‎

Lines changed: 16 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -19,8 +19,10 @@
1919
use crate::catalog::schema::SchemaProvider;
2020
use crate::datasource::datasource::TableProviderFactory;
2121
use crate::datasource::TableProvider;
22-
use datafusion_common::DataFusionError;
22+
use crate::execution::context::SessionState;
23+
use datafusion_common::{context, DataFusionError};
2324
use futures::TryStreamExt;
25+
use itertools::Itertools;
2426
use object_store::ObjectStore;
2527
use std::any::Any;
2628
use std::collections::{HashMap, HashSet};
@@ -72,7 +74,7 @@ impl ListingSchemaProvider {
7274
}
7375

7476
/// Reload table information from ObjectStore
75-
pub async fn refresh(&self) -> datafusion_common::Result<()> {
77+
pub async fn refresh(&self, state: &SessionState) -> datafusion_common::Result<()> {
7678
let entries: Vec<_> = self
7779
.store
7880
.list(Some(&self.path))
@@ -100,13 +102,20 @@ impl ListingSchemaProvider {
100102
.ok_or_else(|| {
101103
DataFusionError::Internal("Cannot parse file name!".to_string())
102104
})?;
103-
let table_name = table.to_str().ok_or_else(|| {
105+
let table_name = file_name.split('.').collect_vec()[0];
106+
let table_path = table.to_str().ok_or_else(|| {
104107
DataFusionError::Internal("Cannot parse file name!".to_string())
105108
})?;
106-
if !self.table_exist(file_name) {
107-
let table_name = format!("{}/{}", self.authority, table_name);
108-
let provider = self.factory.create(table_name.as_str()).await?;
109-
let _ = self.register_table(file_name.to_string(), provider.clone())?;
109+
if !self.table_exist(table_name) {
110+
let table_url = format!("{}/{}", self.authority, table_path);
111+
let provider = self
112+
.factory
113+
.create(state, table_url.as_str())
114+
.await
115+
.map_err(|e| {
116+
context!(format!("Could not create table for {}", table_url), e)
117+
})?;
118+
let _ = self.register_table(table_name.to_string(), provider.clone())?;
110119
}
111120
}
112121
Ok(())

‎datafusion/core/src/datasource/datasource.rs‎

Lines changed: 5 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -92,5 +92,9 @@ pub trait TableProvider: Sync + Send {
9292
#[async_trait]
9393
pub trait TableProviderFactory: Sync + Send {
9494
/// Create a TableProvider with the given url
95-
async fn create(&self, url: &str) -> Result<Arc<dyn TableProvider>>;
95+
async fn create(
96+
&self,
97+
ctx: &SessionState,
98+
url: &str,
99+
) -> Result<Arc<dyn TableProvider>>;
96100
}

‎datafusion/core/src/datasource/listing_table_factory.rs‎

Lines changed: 79 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,79 @@
1+
// Licensed to the Apache Software Foundation (ASF) under one
2+
// or more contributor license agreements. See the NOTICE file
3+
// distributed with this work for additional information
4+
// regarding copyright ownership. The ASF licenses this file
5+
// to you under the Apache License, Version 2.0 (the
6+
// "License"); you may not use this file except in compliance
7+
// with the License. You may obtain a copy of the License at
8+
//
9+
// http://www.apache.org/licenses/LICENSE-2.0
10+
//
11+
// Unless required by applicable law or agreed to in writing,
12+
// software distributed under the License is distributed on an
13+
// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
14+
// KIND, either express or implied. See the License for the
15+
// specific language governing permissions and limitations
16+
// under the License.
17+
18+
//! Factory for creating ListingTables with default options
19+
20+
use crate::datasource::datasource::TableProviderFactory;
21+
use crate::datasource::file_format::avro::AvroFormat;
22+
use crate::datasource::file_format::csv::CsvFormat;
23+
use crate::datasource::file_format::file_type::{FileType, GetExt};
24+
use crate::datasource::file_format::json::JsonFormat;
25+
use crate::datasource::file_format::parquet::ParquetFormat;
26+
use crate::datasource::file_format::FileFormat;
27+
use crate::datasource::listing::{
28+
ListingOptions, ListingTable, ListingTableConfig, ListingTableUrl,
29+
};
30+
use crate::datasource::TableProvider;
31+
use crate::execution::context::SessionState;
32+
use async_trait::async_trait;
33+
use std::sync::Arc;
34+
35+
/// A `TableProviderFactory` capable of creating new `ListingTable`s
36+
pub struct ListingTableFactory {
37+
file_type: FileType,
38+
}
39+
40+
impl ListingTableFactory {
41+
/// Creates a new `ListingTableFactory`
42+
pub fn new(file_type: FileType) -> Self {
43+
Self { file_type }
44+
}
45+
}
46+
47+
#[async_trait]
48+
impl TableProviderFactory for ListingTableFactory {
49+
async fn create(
50+
&self,
51+
state: &SessionState,
52+
url: &str,
53+
) -> datafusion_common::Result<Arc<dyn TableProvider>> {
54+
let file_extension = self.file_type.get_ext();
55+
56+
let file_format: Arc<dyn FileFormat> = match self.file_type {
57+
FileType::CSV => Arc::new(CsvFormat::default()),
58+
FileType::PARQUET => Arc::new(ParquetFormat::default()),
59+
FileType::AVRO => Arc::new(AvroFormat::default()),
60+
FileType::JSON => Arc::new(JsonFormat::default()),
61+
};
62+
63+
let options = ListingOptions {
64+
format: file_format,
65+
collect_stat: true,
66+
file_extension: file_extension.to_owned(),
67+
target_partitions: 1,
68+
table_partition_cols: vec![],
69+
};
70+
71+
let table_path = ListingTableUrl::parse(url)?;
72+
let resolved_schema = options.infer_schema(state, &table_path).await?;
73+
let config = ListingTableConfig::new(table_path)
74+
.with_listing_options(options)
75+
.with_schema(resolved_schema);
76+
let table = ListingTable::try_new(config)?;
77+
Ok(Arc::new(table))
78+
}
79+
}

‎datafusion/core/src/datasource/mod.rs‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -23,6 +23,7 @@ pub mod default_table_source;
2323
pub mod empty;
2424
pub mod file_format;
2525
pub mod listing;
26+
pub mod listing_table_factory;
2627
pub mod memory;
2728
pub mod object_store;
2829
pub mod view;

‎datafusion/core/src/execution/context.rs‎

Lines changed: 42 additions & 18 deletions
Original file line numberDiff line numberDiff line change
@@ -175,6 +175,26 @@ impl SessionContext {
175175
Self::with_config(SessionConfig::new())
176176
}
177177

178+
/// Finds any ListSchemaProviders and instructs them to reload tables from "disk"
179+
pub async fn refresh_catalogs(&self) -> Result<()> {
180+
let cat_names = self.catalog_names().clone();
181+
for cat_name in cat_names.iter() {
182+
let cat = self.catalog(cat_name.as_str()).ok_or_else(|| {
183+
DataFusionError::Internal("Catalog not found!".to_string())
184+
})?;
185+
for schema_name in cat.schema_names() {
186+
let schema = cat.schema(schema_name.as_str()).ok_or_else(|| {
187+
DataFusionError::Internal("Schema not found!".to_string())
188+
})?;
189+
let lister = schema.as_any().downcast_ref::<ListingSchemaProvider>();
190+
if let Some(lister) = lister {
191+
lister.refresh(&self.state()).await?;
192+
}
193+
}
194+
}
195+
Ok(())
196+
}
197+
178198
/// Creates a new session context using the provided session configuration.
179199
pub fn with_config(config: SessionConfig) -> Self {
180200
let runtime = Arc::new(RuntimeEnv::default());
@@ -486,7 +506,7 @@ impl SessionContext {
486506
cmd.file_type
487507
))
488508
})?;
489-
let table = (*factory).create(cmd.location.as_str()).await?;
509+
let table = (*factory).create(&state, cmd.location.as_str()).await?;
490510
self.register_table(cmd.name.as_str(), table)?;
491511
let plan = LogicalPlanBuilder::empty(false).build()?;
492512
Ok(Arc::new(DataFrame::new(self.state.clone(), &plan)))
@@ -1764,7 +1784,7 @@ impl ContextProvider for SessionState {
17641784
Ok(schema) => {
17651785
let provider = schema.table(resolved_ref.table).ok_or_else(|| {
17661786
DataFusionError::Plan(format!(
1767-
"'{}.{}.{}' not found",
1787+
"table '{}.{}.{}' not found",
17681788
resolved_ref.catalog, resolved_ref.schema, resolved_ref.table
17691789
))
17701790
})?;
@@ -2005,11 +2025,12 @@ mod tests {
20052025
use super::*;
20062026
use crate::assert_batches_eq;
20072027
use crate::datasource::datasource::TableProviderFactory;
2028+
use crate::datasource::listing_table_factory::ListingTableFactory;
20082029
use crate::execution::context::QueryPlanner;
20092030
use crate::execution::runtime_env::RuntimeConfig;
20102031
use crate::physical_plan::expressions::AvgAccumulator;
20112032
use crate::test;
2012-
use crate::test_util::{parquet_test_data, TestTableFactory};
2033+
use crate::test_util::parquet_test_data;
20132034
use crate::variable::VarType;
20142035
use arrow::array::ArrayRef;
20152036
use arrow::datatypes::*;
@@ -2267,30 +2288,33 @@ mod tests {
22672288

22682289
let mut table_factories: HashMap<String, Arc<dyn TableProviderFactory>> =
22692290
HashMap::new();
2270-
table_factories.insert("test".to_string(), Arc::new(TestTableFactory {}));
2291+
let factory = Arc::new(ListingTableFactory::new(FileType::CSV));
2292+
table_factories.insert("test".to_string(), factory);
22712293
let rt_cfg = RuntimeConfig::new().with_table_factories(table_factories);
22722294
let runtime = Arc::new(RuntimeEnv::new(rt_cfg).unwrap());
22732295
let cfg = SessionConfig::new()
22742296
.set_str("datafusion.catalog.location", url.as_str())
22752297
.set_str("datafusion.catalog.type", "test");
22762298
let session_state = SessionState::with_config_rt(cfg, runtime);
22772299
let ctx = SessionContext::with_state(session_state);
2300+
ctx.refresh_catalogs().await?;
22782301

2279-
let mut table_count = 0;
2280-
for cat_name in ctx.catalog_names().iter() {
2281-
let cat = ctx.catalog(cat_name).unwrap();
2282-
for s_name in cat.schema_names().iter() {
2283-
let schema = cat.schema(s_name).unwrap();
2284-
if let Some(listing) =
2285-
schema.as_any().downcast_ref::<ListingSchemaProvider>()
2286-
{
2287-
listing.refresh().await.unwrap();
2288-
table_count = schema.table_names().len();
2289-
}
2290-
}
2291-
}
2302+
let result =
2303+
plan_and_collect(&ctx, "select c_name from default.customer limit 3;")
2304+
.await?;
2305+
2306+
let actual = arrow::util::pretty::pretty_format_batches(&result)
2307+
.unwrap()
2308+
.to_string();
2309+
let expected = r#"+--------------------+
2310+
| c_name |
2311+
+--------------------+
2312+
| Customer#000000002 |
2313+
| Customer#000000003 |
2314+
| Customer#000000004 |
2315+
+--------------------+"#;
2316+
assert_eq!(actual, expected);
22922317

2293-
assert_eq!(table_count, 8);
22942318
Ok(())
22952319
}
22962320

‎datafusion/core/src/test_util.rs‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -284,6 +284,7 @@ pub struct TestTableFactory {}
284284
impl TableProviderFactory for TestTableFactory {
285285
async fn create(
286286
&self,
287+
_state: &SessionState,
287288
url: &str,
288289
) -> datafusion_common::Result<Arc<dyn TableProvider>> {
289290
Ok(Arc::new(TestTableProvider {

‎datafusion/core/tests/sql/errors.rs‎

Lines changed: 3 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -80,7 +80,7 @@ async fn query_cte_incorrect() -> Result<()> {
8080
assert!(plan.is_err());
8181
assert_eq!(
8282
format!("{}", plan.unwrap_err()),
83-
"Error during planning: 'datafusion.public.t' not found"
83+
"Error during planning: table 'datafusion.public.t' not found"
8484
);
8585

8686
// forward referencing
@@ -89,7 +89,7 @@ async fn query_cte_incorrect() -> Result<()> {
8989
assert!(plan.is_err());
9090
assert_eq!(
9191
format!("{}", plan.unwrap_err()),
92-
"Error during planning: 'datafusion.public.u' not found"
92+
"Error during planning: table 'datafusion.public.u' not found"
9393
);
9494

9595
// wrapping should hide u
@@ -98,7 +98,7 @@ async fn query_cte_incorrect() -> Result<()> {
9898
assert!(plan.is_err());
9999
assert_eq!(
100100
format!("{}", plan.unwrap_err()),
101-
"Error during planning: 'datafusion.public.u' not found"
101+
"Error during planning: table 'datafusion.public.u' not found"
102102
);
103103

104104
Ok(())

0 commit comments

Comments
 (0)