From 7ed1db7f348bbed894295a42fb1f7efaf4c0b4dc Mon Sep 17 00:00:00 2001 From: JasonLi-cn Date: Thu, 26 Sep 2024 10:46:08 +0800 Subject: [PATCH] modify register table functions args --- datafusion/core/src/execution/context/avro.rs | 10 +++++----- datafusion/core/src/execution/context/csv.rs | 10 +++++----- datafusion/core/src/execution/context/json.rs | 10 +++++----- datafusion/core/src/execution/context/mod.rs | 7 ++----- datafusion/core/src/execution/context/parquet.rs | 7 ++++--- 5 files changed, 21 insertions(+), 23 deletions(-) diff --git a/datafusion/core/src/execution/context/avro.rs b/datafusion/core/src/execution/context/avro.rs index e829f6123eab..a31f2af642d0 100644 --- a/datafusion/core/src/execution/context/avro.rs +++ b/datafusion/core/src/execution/context/avro.rs @@ -15,10 +15,10 @@ // specific language governing permissions and limitations // under the License. -use std::sync::Arc; - use super::super::options::{AvroReadOptions, ReadOptions}; use super::{DataFilePaths, DataFrame, Result, SessionContext}; +use datafusion_common::TableReference; +use std::sync::Arc; impl SessionContext { /// Creates a [`DataFrame`] for reading an Avro data source. @@ -39,15 +39,15 @@ impl SessionContext { /// SQL statements executed against this context. pub async fn register_avro( &self, - name: &str, - table_path: &str, + table_ref: impl Into, + table_path: impl AsRef, options: AvroReadOptions<'_>, ) -> Result<()> { let listing_options = options .to_listing_options(&self.copied_config(), self.copied_table_options()); self.register_listing_table( - name, + table_ref, table_path, listing_options, options.schema.map(|s| Arc::new(s.to_owned())), diff --git a/datafusion/core/src/execution/context/csv.rs b/datafusion/core/src/execution/context/csv.rs index 08e93cb61305..e97c70ef9812 100644 --- a/datafusion/core/src/execution/context/csv.rs +++ b/datafusion/core/src/execution/context/csv.rs @@ -15,9 +15,9 @@ // specific language governing permissions and limitations // under the License. -use std::sync::Arc; - use crate::datasource::physical_plan::plan_to_csv; +use datafusion_common::TableReference; +use std::sync::Arc; use super::super::options::{CsvReadOptions, ReadOptions}; use super::{DataFilePaths, DataFrame, ExecutionPlan, Result, SessionContext}; @@ -55,15 +55,15 @@ impl SessionContext { /// statements executed against this context. pub async fn register_csv( &self, - name: &str, - table_path: &str, + table_ref: impl Into, + table_path: impl AsRef, options: CsvReadOptions<'_>, ) -> Result<()> { let listing_options = options .to_listing_options(&self.copied_config(), self.copied_table_options()); self.register_listing_table( - name, + table_ref, table_path, listing_options, options.schema.map(|s| Arc::new(s.to_owned())), diff --git a/datafusion/core/src/execution/context/json.rs b/datafusion/core/src/execution/context/json.rs index c21e32cfdefb..c9a9492f9162 100644 --- a/datafusion/core/src/execution/context/json.rs +++ b/datafusion/core/src/execution/context/json.rs @@ -15,9 +15,9 @@ // specific language governing permissions and limitations // under the License. -use std::sync::Arc; - use crate::datasource::physical_plan::plan_to_json; +use datafusion_common::TableReference; +use std::sync::Arc; use super::super::options::{NdJsonReadOptions, ReadOptions}; use super::{DataFilePaths, DataFrame, ExecutionPlan, Result, SessionContext}; @@ -41,15 +41,15 @@ impl SessionContext { /// from SQL statements executed against this context. pub async fn register_json( &self, - name: &str, - table_path: &str, + table_ref: impl Into, + table_path: impl AsRef, options: NdJsonReadOptions<'_>, ) -> Result<()> { let listing_options = options .to_listing_options(&self.copied_config(), self.copied_table_options()); self.register_listing_table( - name, + table_ref, table_path, listing_options, options.schema.map(|s| Arc::new(s.to_owned())), diff --git a/datafusion/core/src/execution/context/mod.rs b/datafusion/core/src/execution/context/mod.rs index 53eb7c431b47..12ce71768f91 100644 --- a/datafusion/core/src/execution/context/mod.rs +++ b/datafusion/core/src/execution/context/mod.rs @@ -1264,7 +1264,7 @@ impl SessionContext { /// [`ObjectStore`]: object_store::ObjectStore pub async fn register_listing_table( &self, - name: &str, + table_ref: impl Into, table_path: impl AsRef, options: ListingOptions, provided_schema: Option, @@ -1279,10 +1279,7 @@ impl SessionContext { .with_listing_options(options) .with_schema(resolved_schema); let table = ListingTable::try_new(config)?.with_definition(sql_definition); - self.register_table( - TableReference::Bare { table: name.into() }, - Arc::new(table), - )?; + self.register_table(table_ref, Arc::new(table))?; Ok(()) } diff --git a/datafusion/core/src/execution/context/parquet.rs b/datafusion/core/src/execution/context/parquet.rs index 1d83c968c1a8..3f23c150be83 100644 --- a/datafusion/core/src/execution/context/parquet.rs +++ b/datafusion/core/src/execution/context/parquet.rs @@ -21,6 +21,7 @@ use super::super::options::{ParquetReadOptions, ReadOptions}; use super::{DataFilePaths, DataFrame, ExecutionPlan, Result, SessionContext}; use crate::datasource::physical_plan::parquet::plan_to_parquet; +use datafusion_common::TableReference; use parquet::file::properties::WriterProperties; impl SessionContext { @@ -42,15 +43,15 @@ impl SessionContext { /// statements executed against this context. pub async fn register_parquet( &self, - name: &str, - table_path: &str, + table_ref: impl Into, + table_path: impl AsRef, options: ParquetReadOptions<'_>, ) -> Result<()> { let listing_options = options .to_listing_options(&self.copied_config(), self.copied_table_options()); self.register_listing_table( - name, + table_ref, table_path, listing_options, options.schema.map(|s| Arc::new(s.to_owned())),