From 32ecce42dad45db9d8a6b3a6fa2a1e3a231a9ead Mon Sep 17 00:00:00 2001 From: veeupup Date: Sat, 25 Nov 2023 15:47:19 +0800 Subject: [PATCH 1/8] Support User Defined Table Function Signed-off-by: veeupup --- datafusion-examples/examples/simple_udtf.rs | 224 +++++++++++++++++++ datafusion/core/src/datasource/function.rs | 56 +++++ datafusion/core/src/datasource/mod.rs | 1 + datafusion/core/src/execution/context/mod.rs | 30 ++- datafusion/sql/src/planner.rs | 9 + datafusion/sql/src/relation/mod.rs | 81 +++++-- 6 files changed, 381 insertions(+), 20 deletions(-) create mode 100644 datafusion-examples/examples/simple_udtf.rs create mode 100644 datafusion/core/src/datasource/function.rs diff --git a/datafusion-examples/examples/simple_udtf.rs b/datafusion-examples/examples/simple_udtf.rs new file mode 100644 index 000000000000..1fa48611b4cd --- /dev/null +++ b/datafusion-examples/examples/simple_udtf.rs @@ -0,0 +1,224 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +use arrow::csv::reader::Format; +use arrow::csv::ReaderBuilder; +use async_trait::async_trait; +use datafusion::arrow::datatypes::SchemaRef; +use datafusion::arrow::record_batch::RecordBatch; +use datafusion::datasource::function::TableFunctionImpl; +use datafusion::datasource::streaming::StreamingTable; +use datafusion::datasource::TableProvider; +use datafusion::error::Result; +use datafusion::execution::context::SessionState; +use datafusion::execution::TaskContext; +use datafusion::physical_plan::memory::MemoryExec; +use datafusion::physical_plan::stream::RecordBatchStreamAdapter; +use datafusion::physical_plan::streaming::PartitionStream; +use datafusion::physical_plan::{collect, ExecutionPlan}; +use datafusion::prelude::SessionContext; +use datafusion_common::{DFSchema, ScalarValue}; +use datafusion_expr::{EmptyRelation, Expr, LogicalPlan, Projection, TableType}; +use std::fs::File; +use std::io::Seek; +use std::path::Path; +use std::sync::Arc; + +// To define your own table function, you only need to do the following 3 things: +// 1. Implement your own TableProvider +// 2. Implement your own TableFunctionImpl and return your TableProvider +// 3. Register the function using ctx.register_udtf + +/// This example demonstrates how to register a TableFunction +#[tokio::main] +async fn main() -> Result<()> { + // create local execution context + let ctx = SessionContext::new(); + + ctx.register_udtf("read_csv", Arc::new(LocalCsvTableFunc {})); + ctx.register_udtf("read_csv_stream", Arc::new(LocalStreamCsvTable {})); + + let testdata = datafusion::test_util::arrow_test_data(); + let csv_file = format!("{testdata}/csv/aggregate_test_100.csv"); + + // run it with println now() + let df = ctx + .sql(format!("SELECT * FROM read_csv('{csv_file}', now());").as_str()) + .await?; + df.show().await?; + + // just run + let df = ctx + .sql(format!("SELECT * FROM read_csv('{csv_file}');").as_str()) + .await?; + df.show().await?; + + // stream csv table + let df2 = ctx + .sql(format!("SELECT * FROM read_csv_stream('{csv_file}');").as_str()) + .await?; + df2.show().await?; + + Ok(()) +} + +// Option1: (full implmentation of a TableProvider) +struct LocalCsvTable { + schema: SchemaRef, + exprs: Vec, + batches: Vec, +} + +#[async_trait] +impl TableProvider for LocalCsvTable { + fn as_any(&self) -> &dyn std::any::Any { + self + } + + fn schema(&self) -> SchemaRef { + self.schema.clone() + } + + fn table_type(&self) -> TableType { + TableType::Base + } + + async fn scan( + &self, + state: &SessionState, + projection: Option<&Vec>, + _filters: &[Expr], + _limit: Option, + ) -> Result> { + if !self.exprs.is_empty() { + self.interpreter_expr(state).await?; + } + Ok(Arc::new(MemoryExec::try_new( + &[self.batches.clone()], + TableProvider::schema(self), + projection.cloned(), + )?)) + } +} + +impl LocalCsvTable { + // TODO(veeupup): maybe we can make interpreter Expr this more simpler for users + // TODO(veeupup): maybe we can support more type of exprs + async fn interpreter_expr(&self, state: &SessionState) -> Result<()> { + use datafusion::logical_expr::expr_rewriter::normalize_col; + use datafusion::logical_expr::utils::columnize_expr; + let plan = LogicalPlan::EmptyRelation(EmptyRelation { + produce_one_row: true, + schema: Arc::new(DFSchema::empty()), + }); + let logical_plan = Projection::try_new( + vec![columnize_expr( + normalize_col(self.exprs[0].clone(), &plan)?, + plan.schema(), + )], + Arc::new(plan), + ) + .map(LogicalPlan::Projection)?; + let rbs = collect( + state.create_physical_plan(&logical_plan).await?, + Arc::new(TaskContext::from(state)), + ) + .await?; + println!("time now: {:?}", rbs[0].column(0)); + Ok(()) + } +} + +struct LocalCsvTableFunc {} + +impl TableFunctionImpl for LocalCsvTableFunc { + fn call(&self, exprs: &[Expr]) -> Result> { + let mut new_exprs = vec![]; + let mut filepath = String::new(); + for expr in exprs { + match expr { + Expr::Literal(ScalarValue::Utf8(Some(ref path))) => { + filepath = path.clone() + } + expr => new_exprs.push(expr.clone()), + } + } + let (schema, batches) = read_csv_batches(filepath)?; + let table = LocalCsvTable { + schema, + exprs: new_exprs.clone(), + batches, + }; + Ok(Arc::new(table)) + } +} + +// Option2: (use StreamingTable to make it simpler) +// Implement PartitionStream and Use StreamTable to return streaming table +impl PartitionStream for LocalCsvTable { + fn schema(&self) -> &SchemaRef { + &self.schema + } + + fn execute( + &self, + _ctx: Arc, + ) -> datafusion::physical_plan::SendableRecordBatchStream { + Box::pin(RecordBatchStreamAdapter::new( + self.schema.clone(), + // You can even read data from network or else anywhere, using async is also ok + // In Fact, you can even implement your own SendableRecordBatchStream + // by implementing Stream> + Send + Sync + 'static + futures::stream::iter(self.batches.clone().into_iter().map(Ok)), + )) + } +} + +struct LocalStreamCsvTable {} + +impl TableFunctionImpl for LocalStreamCsvTable { + fn call(&self, args: &[Expr]) -> Result> { + let filepath = match args[0] { + Expr::Literal(ScalarValue::Utf8(Some(ref path))) => path.clone(), + _ => unimplemented!(), + }; + let (schema, batches) = read_csv_batches(filepath)?; + let stream = LocalCsvTable { + schema: schema.clone(), + batches, + exprs: vec![], + }; + let table = StreamingTable::try_new(schema, vec![Arc::new(stream)])?; + Ok(Arc::new(table)) + } +} + +fn read_csv_batches(csv_path: impl AsRef) -> Result<(SchemaRef, Vec)> { + let mut file = File::open(csv_path)?; + let (schema, _) = Format::default().infer_schema(&mut file, None)?; + file.rewind()?; + + let reader = ReaderBuilder::new(Arc::new(schema.clone())) + .with_header(true) + .build(file)?; + let mut batches = vec![]; + for bacth in reader { + batches.push(bacth?); + } + let schema = Arc::new(schema); + Ok((schema, batches)) +} diff --git a/datafusion/core/src/datasource/function.rs b/datafusion/core/src/datasource/function.rs new file mode 100644 index 000000000000..d87c482911c0 --- /dev/null +++ b/datafusion/core/src/datasource/function.rs @@ -0,0 +1,56 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +//! A table that uses a function to generate data + +use super::TableProvider; + +use datafusion_common::Result; +use datafusion_expr::Expr; + +use std::sync::Arc; + +/// A trait for table function implementations +pub trait TableFunctionImpl: Sync + Send { + /// Create a table provider + fn call(&self, args: &[Expr]) -> Result>; +} + +/// A table that uses a function to generate data +pub struct TableFunction { + /// Name of the table function + name: String, + /// Function implementation + fun: Arc, +} + +impl TableFunction { + /// Create a new table function + pub fn new(name: String, fun: Arc) -> Self { + Self { name, fun } + } + + /// Get the name of the table function + pub fn name(&self) -> String { + self.name.clone() + } + + /// Get the function implementation and generate a table + pub fn create_table_provider(&self, args: &[Expr]) -> Result> { + self.fun.call(args) + } +} diff --git a/datafusion/core/src/datasource/mod.rs b/datafusion/core/src/datasource/mod.rs index 45f9bee6a58b..2e516cc36a01 100644 --- a/datafusion/core/src/datasource/mod.rs +++ b/datafusion/core/src/datasource/mod.rs @@ -23,6 +23,7 @@ pub mod avro_to_arrow; pub mod default_table_source; pub mod empty; pub mod file_format; +pub mod function; pub mod listing; pub mod listing_table_factory; pub mod memory; diff --git a/datafusion/core/src/execution/context/mod.rs b/datafusion/core/src/execution/context/mod.rs index b8e111d361b1..70b1643a722f 100644 --- a/datafusion/core/src/execution/context/mod.rs +++ b/datafusion/core/src/execution/context/mod.rs @@ -26,6 +26,7 @@ mod parquet; use crate::{ catalog::{CatalogList, MemoryCatalogList}, datasource::{ + function::{TableFunction, TableFunctionImpl}, listing::{ListingOptions, ListingTable}, provider::TableProviderFactory, }, @@ -42,7 +43,7 @@ use datafusion_common::{ use datafusion_execution::registry::SerializerRegistry; use datafusion_expr::{ logical_plan::{DdlStatement, Statement}, - StringifiedPlan, UserDefinedLogicalNode, WindowUDF, + Expr, StringifiedPlan, UserDefinedLogicalNode, WindowUDF, }; pub use datafusion_physical_expr::execution_props::ExecutionProps; use datafusion_physical_expr::var_provider::is_system_variables; @@ -795,6 +796,14 @@ impl SessionContext { .add_var_provider(variable_type, provider); } + /// Register a table UDF with this context + pub fn register_udtf(&self, name: &str, fun: Arc) { + self.state.write().table_functions.insert( + name.to_owned(), + Arc::new(TableFunction::new(name.to_owned(), fun)), + ); + } + /// Registers a scalar UDF within this context. /// /// Note in SQL queries, function names are looked up using @@ -1224,6 +1233,8 @@ pub struct SessionState { query_planner: Arc, /// Collection of catalogs containing schemas and ultimately TableProviders catalog_list: Arc, + /// Table Functions + table_functions: HashMap>, /// Scalar functions that are registered with the context scalar_functions: HashMap>, /// Aggregate functions registered in the context @@ -1322,6 +1333,7 @@ impl SessionState { physical_optimizers: PhysicalOptimizer::new(), query_planner: Arc::new(DefaultQueryPlanner {}), catalog_list, + table_functions: HashMap::new(), scalar_functions: HashMap::new(), aggregate_functions: HashMap::new(), window_functions: HashMap::new(), @@ -1860,6 +1872,22 @@ impl<'a> ContextProvider for SessionContextProvider<'a> { .ok_or_else(|| plan_datafusion_err!("table '{name}' not found")) } + fn get_table_function_source( + &self, + name: &str, + args: Vec, + ) -> Result> { + let tbl_func = self + .state + .table_functions + .get(name) + .cloned() + .ok_or_else(|| plan_datafusion_err!("table function '{name}' not found"))?; + let provider = tbl_func.create_table_provider(&args)?; + + Ok(provider_as_source(provider)) + } + fn get_function_meta(&self, name: &str) -> Option> { self.state.scalar_functions().get(name).cloned() } diff --git a/datafusion/sql/src/planner.rs b/datafusion/sql/src/planner.rs index ca5e260aee05..fbaea92c9f30 100644 --- a/datafusion/sql/src/planner.rs +++ b/datafusion/sql/src/planner.rs @@ -51,6 +51,15 @@ pub trait ContextProvider { } /// Getter for a datasource fn get_table_source(&self, name: TableReference) -> Result>; + /// Getter for a table function + fn get_table_function_source( + &self, + _name: &str, + _args: Vec, + ) -> Result> { + unimplemented!() + } + /// Getter for a UDF description fn get_function_meta(&self, name: &str) -> Option>; /// Getter for a UDAF description diff --git a/datafusion/sql/src/relation/mod.rs b/datafusion/sql/src/relation/mod.rs index 180743d19b7b..84818c8de4b1 100644 --- a/datafusion/sql/src/relation/mod.rs +++ b/datafusion/sql/src/relation/mod.rs @@ -16,9 +16,11 @@ // under the License. use crate::planner::{ContextProvider, PlannerContext, SqlToRel}; -use datafusion_common::{not_impl_err, DataFusionError, Result}; +use datafusion_common::{ + not_impl_err, DFSchema, DataFusionError, Result, TableReference, +}; use datafusion_expr::{LogicalPlan, LogicalPlanBuilder}; -use sqlparser::ast::TableFactor; +use sqlparser::ast::{FunctionArg, FunctionArgExpr, TableFactor}; mod join; @@ -30,24 +32,65 @@ impl<'a, S: ContextProvider> SqlToRel<'a, S> { planner_context: &mut PlannerContext, ) -> Result { let (plan, alias) = match relation { - TableFactor::Table { name, alias, .. } => { - // normalize name and alias - let table_ref = self.object_name_to_table_reference(name)?; - let table_name = table_ref.to_string(); - let cte = planner_context.get_cte(&table_name); - ( - match ( - cte, - self.context_provider.get_table_source(table_ref.clone()), - ) { - (Some(cte_plan), _) => Ok(cte_plan.clone()), - (_, Ok(provider)) => { - LogicalPlanBuilder::scan(table_ref, provider, None)?.build() + TableFactor::Table { + name, alias, args, .. + } => { + // this maybe a little diffcult to resolve others tables' schema, so we only supprt value and scalar functions now + if let Some(func_args) = args { + let tbl_func_name = name.0.get(0).unwrap().value.to_string(); + let mut args = vec![]; + for arg in func_args { + match arg { + FunctionArg::Unnamed(FunctionArgExpr::Expr(expr)) => { + let expr = self.sql_expr_to_logical_expr( + expr, + // TODO(veeupup): for now, maybe it's little diffcult to resolve tables' schema before create provider + // maybe we can put all relations schema in + &DFSchema::empty(), + planner_context, + )?; + args.push(expr); + } + arg => { + unimplemented!( + "Unsupported function argument type: {:?}", + arg + ) + } } - (None, Err(e)) => Err(e), - }?, - alias, - ) + } + let provider = self + .context_provider + .get_table_function_source(&tbl_func_name, args)?; + let plan = LogicalPlanBuilder::scan( + TableReference::Bare { + table: std::borrow::Cow::Borrowed("tmp_table"), + }, + provider, + None, + )? + .build()?; + (plan, alias) + } else { + // normalize name and alias + let table_ref = self.object_name_to_table_reference(name)?; + let table_name = table_ref.to_string(); + let cte = planner_context.get_cte(&table_name); + ( + match ( + cte, + self.context_provider.get_table_source(table_ref.clone()), + ) { + (Some(cte_plan), _) => Ok(cte_plan.clone()), + (_, Ok(provider)) => { + LogicalPlanBuilder::scan(table_ref, provider, None)? + .build() + } + (None, Err(e)) => Err(e), + }?, + alias, + ) + } } TableFactor::Derived { subquery, alias, .. From 2f4e0893ddc702595766f7f0dd4fa18bfa316dba Mon Sep 17 00:00:00 2001 From: veeupup Date: Tue, 28 Nov 2023 21:46:47 +0800 Subject: [PATCH 2/8] fix comments Signed-off-by: veeupup --- datafusion-examples/examples/simple_udtf.rs | 96 +++++++-------------- datafusion/core/src/datasource/function.rs | 4 +- datafusion/sql/src/planner.rs | 2 +- datafusion/sql/src/relation/mod.rs | 29 +++---- 4 files changed, 47 insertions(+), 84 deletions(-) diff --git a/datafusion-examples/examples/simple_udtf.rs b/datafusion-examples/examples/simple_udtf.rs index 1fa48611b4cd..260c6cbee710 100644 --- a/datafusion-examples/examples/simple_udtf.rs +++ b/datafusion-examples/examples/simple_udtf.rs @@ -15,20 +15,18 @@ // specific language governing permissions and limitations // under the License. +use arrow::array::Int64Array; use arrow::csv::reader::Format; use arrow::csv::ReaderBuilder; use async_trait::async_trait; use datafusion::arrow::datatypes::SchemaRef; use datafusion::arrow::record_batch::RecordBatch; use datafusion::datasource::function::TableFunctionImpl; -use datafusion::datasource::streaming::StreamingTable; use datafusion::datasource::TableProvider; use datafusion::error::Result; use datafusion::execution::context::SessionState; use datafusion::execution::TaskContext; use datafusion::physical_plan::memory::MemoryExec; -use datafusion::physical_plan::stream::RecordBatchStreamAdapter; -use datafusion::physical_plan::streaming::PartitionStream; use datafusion::physical_plan::{collect, ExecutionPlan}; use datafusion::prelude::SessionContext; use datafusion_common::{DFSchema, ScalarValue}; @@ -50,33 +48,25 @@ async fn main() -> Result<()> { let ctx = SessionContext::new(); ctx.register_udtf("read_csv", Arc::new(LocalCsvTableFunc {})); - ctx.register_udtf("read_csv_stream", Arc::new(LocalStreamCsvTable {})); let testdata = datafusion::test_util::arrow_test_data(); let csv_file = format!("{testdata}/csv/aggregate_test_100.csv"); - // run it with println now() + // read csv with at most 2 rows let df = ctx - .sql(format!("SELECT * FROM read_csv('{csv_file}', now());").as_str()) + .sql(format!("SELECT * FROM read_csv('{csv_file}', 2);").as_str()) .await?; df.show().await?; - // just run + // just run, return all rows let df = ctx .sql(format!("SELECT * FROM read_csv('{csv_file}');").as_str()) .await?; df.show().await?; - // stream csv table - let df2 = ctx - .sql(format!("SELECT * FROM read_csv_stream('{csv_file}');").as_str()) - .await?; - df2.show().await?; - Ok(()) } -// Option1: (full implmentation of a TableProvider) struct LocalCsvTable { schema: SchemaRef, exprs: Vec, @@ -104,11 +94,28 @@ impl TableProvider for LocalCsvTable { _filters: &[Expr], _limit: Option, ) -> Result> { - if !self.exprs.is_empty() { - self.interpreter_expr(state).await?; - } + 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(Arc::new(MemoryExec::try_new( - &[self.batches.clone()], + &[batches], TableProvider::schema(self), projection.cloned(), )?)) @@ -116,9 +123,7 @@ impl TableProvider for LocalCsvTable { } impl LocalCsvTable { - // TODO(veeupup): maybe we can make interpreter Expr this more simpler for users - // TODO(veeupup): maybe we can support more type of exprs - async fn interpreter_expr(&self, state: &SessionState) -> Result<()> { + async fn interpreter_expr(&self, state: &SessionState) -> Result { use datafusion::logical_expr::expr_rewriter::normalize_col; use datafusion::logical_expr::utils::columnize_expr; let plan = LogicalPlan::EmptyRelation(EmptyRelation { @@ -138,8 +143,13 @@ impl LocalCsvTable { Arc::new(TaskContext::from(state)), ) .await?; - println!("time now: {:?}", rbs[0].column(0)); - Ok(()) + let limit = rbs[0] + .column(0) + .as_any() + .downcast_ref::() + .unwrap() + .value(0); + Ok(limit) } } @@ -167,46 +177,6 @@ impl TableFunctionImpl for LocalCsvTableFunc { } } -// Option2: (use StreamingTable to make it simpler) -// Implement PartitionStream and Use StreamTable to return streaming table -impl PartitionStream for LocalCsvTable { - fn schema(&self) -> &SchemaRef { - &self.schema - } - - fn execute( - &self, - _ctx: Arc, - ) -> datafusion::physical_plan::SendableRecordBatchStream { - Box::pin(RecordBatchStreamAdapter::new( - self.schema.clone(), - // You can even read data from network or else anywhere, using async is also ok - // In Fact, you can even implement your own SendableRecordBatchStream - // by implementing Stream> + Send + Sync + 'static - futures::stream::iter(self.batches.clone().into_iter().map(Ok)), - )) - } -} - -struct LocalStreamCsvTable {} - -impl TableFunctionImpl for LocalStreamCsvTable { - fn call(&self, args: &[Expr]) -> Result> { - let filepath = match args[0] { - Expr::Literal(ScalarValue::Utf8(Some(ref path))) => path.clone(), - _ => unimplemented!(), - }; - let (schema, batches) = read_csv_batches(filepath)?; - let stream = LocalCsvTable { - schema: schema.clone(), - batches, - exprs: vec![], - }; - let table = StreamingTable::try_new(schema, vec![Arc::new(stream)])?; - Ok(Arc::new(table)) - } -} - fn read_csv_batches(csv_path: impl AsRef) -> Result<(SchemaRef, Vec)> { let mut file = File::open(csv_path)?; let (schema, _) = Format::default().infer_schema(&mut file, None)?; diff --git a/datafusion/core/src/datasource/function.rs b/datafusion/core/src/datasource/function.rs index d87c482911c0..2fd352ee4eb3 100644 --- a/datafusion/core/src/datasource/function.rs +++ b/datafusion/core/src/datasource/function.rs @@ -45,8 +45,8 @@ impl TableFunction { } /// Get the name of the table function - pub fn name(&self) -> String { - self.name.clone() + pub fn name(&self) -> &str { + &self.name } /// Get the function implementation and generate a table diff --git a/datafusion/sql/src/planner.rs b/datafusion/sql/src/planner.rs index fbaea92c9f30..f43c908dc0c1 100644 --- a/datafusion/sql/src/planner.rs +++ b/datafusion/sql/src/planner.rs @@ -57,7 +57,7 @@ pub trait ContextProvider { _name: &str, _args: Vec, ) -> Result> { - unimplemented!() + not_impl_err!("Table Functions are not supported") } /// Getter for a UDF description diff --git a/datafusion/sql/src/relation/mod.rs b/datafusion/sql/src/relation/mod.rs index 84818c8de4b1..6fc7e9601243 100644 --- a/datafusion/sql/src/relation/mod.rs +++ b/datafusion/sql/src/relation/mod.rs @@ -17,7 +17,7 @@ use crate::planner::{ContextProvider, PlannerContext, SqlToRel}; use datafusion_common::{ - not_impl_err, DFSchema, DataFusionError, Result, TableReference, + not_impl_err, plan_err, DFSchema, DataFusionError, Result, TableReference, }; use datafusion_expr::{LogicalPlan, LogicalPlanBuilder}; use sqlparser::ast::{FunctionArg, FunctionArgExpr, TableFactor}; @@ -35,30 +35,23 @@ impl<'a, S: ContextProvider> SqlToRel<'a, S> { TableFactor::Table { name, alias, args, .. } => { - // this maybe a little diffcult to resolve others tables' schema, so we only supprt value and scalar functions now if let Some(func_args) = args { let tbl_func_name = name.0.get(0).unwrap().value.to_string(); - let mut args = vec![]; - for arg in func_args { - match arg { - FunctionArg::Unnamed(FunctionArgExpr::Expr(expr)) => { - let expr = self.sql_expr_to_logical_expr( + let args = func_args + .into_iter() + .flat_map(|arg| { + if let FunctionArg::Unnamed(FunctionArgExpr::Expr(expr)) = arg + { + self.sql_expr_to_logical_expr( expr, - // TODO(veeupup): for now, maybe it's little diffcult to resolve tables' schema before create provider - // maybe we can put all relations schema in &DFSchema::empty(), planner_context, - )?; - args.push(expr); - } - arg => { - unimplemented!( - "Unsupported function argument type: {:?}", - arg ) + } else { + plan_err!("Unsupported function argument type: {:?}", arg) } - } - } + }) + .collect::>(); let provider = self .context_provider .get_table_function_source(&tbl_func_name, args)?; From 7a85a434a5072a4024cd8a47276e8d8d1d556a66 Mon Sep 17 00:00:00 2001 From: veeupup Date: Tue, 28 Nov 2023 22:13:30 +0800 Subject: [PATCH 3/8] add udtf test Signed-off-by: veeupup --- datafusion-examples/examples/simple_udtf.rs | 4 +- datafusion/core/tests/user_defined/mod.rs | 3 + .../user_defined_table_functions.rs | 202 ++++++++++++++++++ 3 files changed, 208 insertions(+), 1 deletion(-) create mode 100644 datafusion/core/tests/user_defined/user_defined_table_functions.rs diff --git a/datafusion-examples/examples/simple_udtf.rs b/datafusion-examples/examples/simple_udtf.rs index 260c6cbee710..c2bb86a0f760 100644 --- a/datafusion-examples/examples/simple_udtf.rs +++ b/datafusion-examples/examples/simple_udtf.rs @@ -179,7 +179,9 @@ impl TableFunctionImpl for LocalCsvTableFunc { fn read_csv_batches(csv_path: impl AsRef) -> Result<(SchemaRef, Vec)> { let mut file = File::open(csv_path)?; - let (schema, _) = Format::default().infer_schema(&mut file, None)?; + let (schema, _) = Format::default() + .with_header(true) + .infer_schema(&mut file, None)?; file.rewind()?; let reader = ReaderBuilder::new(Arc::new(schema.clone())) diff --git a/datafusion/core/tests/user_defined/mod.rs b/datafusion/core/tests/user_defined/mod.rs index 09c7c3d3266b..6c6d966cc3aa 100644 --- a/datafusion/core/tests/user_defined/mod.rs +++ b/datafusion/core/tests/user_defined/mod.rs @@ -26,3 +26,6 @@ mod user_defined_plan; /// Tests for User Defined Window Functions mod user_defined_window_functions; + +/// Tests for User Defined Table Functions +mod user_defined_table_functions; diff --git a/datafusion/core/tests/user_defined/user_defined_table_functions.rs b/datafusion/core/tests/user_defined/user_defined_table_functions.rs new file mode 100644 index 000000000000..4ac3fefa9191 --- /dev/null +++ b/datafusion/core/tests/user_defined/user_defined_table_functions.rs @@ -0,0 +1,202 @@ +use arrow::array::Int64Array; +use arrow::csv::reader::Format; +use arrow::csv::ReaderBuilder; +use async_trait::async_trait; +use datafusion::arrow::datatypes::SchemaRef; +use datafusion::arrow::record_batch::RecordBatch; +use datafusion::datasource::function::TableFunctionImpl; +use datafusion::datasource::TableProvider; +use datafusion::error::Result; +use datafusion::execution::context::SessionState; +use datafusion::execution::TaskContext; +use datafusion::physical_plan::memory::MemoryExec; +use datafusion::physical_plan::{collect, ExecutionPlan}; +use datafusion::prelude::SessionContext; +use datafusion_common::{assert_batches_eq, DFSchema, ScalarValue}; +use datafusion_expr::{EmptyRelation, Expr, LogicalPlan, Projection, TableType}; +use std::fs::File; +use std::io::Seek; +use std::path::Path; +use std::sync::Arc; + +/// test simple udtf with define read_csv with parameters +#[tokio::test] +async fn test_simple_read_csv_udtf() -> Result<()> { + let ctx = SessionContext::new(); + + ctx.register_udtf("read_csv", Arc::new(SimpleCsvTableFunc {})); + + let csv_file = "tests/tpch-csv/nation.csv"; + // read csv with at most 2 rows + let rbs = ctx + .sql(format!("SELECT * FROM read_csv('{csv_file}', 5);").as_str()) + .await? + .collect() + .await?; + + let excepted = [ + "+-------------+-----------+-------------+-------------------------------------------------------------------------------------------------------------+", + "| n_nationkey | n_name | n_regionkey | n_comment |", + "+-------------+-----------+-------------+-------------------------------------------------------------------------------------------------------------+", + "| 1 | ARGENTINA | 1 | al foxes promise slyly according to the regular accounts. bold requests alon |", + "| 2 | BRAZIL | 1 | y alongside of the pending deposits. carefully special packages are about the ironic forges. slyly special |", + "| 3 | CANADA | 1 | eas hang ironic, silent packages. slyly regular packages are furiously over the tithes. fluffily bold |", + "| 4 | EGYPT | 4 | y above the carefully unusual theodolites. final dugouts are quickly across the furiously regular d |", + "| 5 | ETHIOPIA | 0 | ven packages wake quickly. regu |", + "+-------------+-----------+-------------+-------------------------------------------------------------------------------------------------------------+", ]; + assert_batches_eq!(excepted, &rbs); + + // just run, return all rows + let rbs = ctx + .sql(format!("SELECT * FROM read_csv('{csv_file}');").as_str()) + .await? + .collect() + .await?; + let excepted = [ + "+-------------+-----------+-------------+--------------------------------------------------------------------------------------------------------------------+", + "| n_nationkey | n_name | n_regionkey | n_comment |", + "+-------------+-----------+-------------+--------------------------------------------------------------------------------------------------------------------+", + "| 1 | ARGENTINA | 1 | al foxes promise slyly according to the regular accounts. bold requests alon |", + "| 2 | BRAZIL | 1 | y alongside of the pending deposits. carefully special packages are about the ironic forges. slyly special |", + "| 3 | CANADA | 1 | eas hang ironic, silent packages. slyly regular packages are furiously over the tithes. fluffily bold |", + "| 4 | EGYPT | 4 | y above the carefully unusual theodolites. final dugouts are quickly across the furiously regular d |", + "| 5 | ETHIOPIA | 0 | ven packages wake quickly. regu |", + "| 6 | FRANCE | 3 | refully final requests. regular, ironi |", + "| 7 | GERMANY | 3 | l platelets. regular accounts x-ray: unusual, regular acco |", + "| 8 | INDIA | 2 | ss excuses cajole slyly across the packages. deposits print aroun |", + "| 9 | INDONESIA | 2 | slyly express asymptotes. regular deposits haggle slyly. carefully ironic hockey players sleep blithely. carefull |", + "| 10 | IRAN | 4 | efully alongside of the slyly final dependencies. |", + "+-------------+-----------+-------------+--------------------------------------------------------------------------------------------------------------------+" + ]; + assert_batches_eq!(excepted, &rbs); + + Ok(()) +} + +struct SimpleCsvTable { + schema: SchemaRef, + exprs: Vec, + batches: Vec, +} + +#[async_trait] +impl TableProvider for SimpleCsvTable { + fn as_any(&self) -> &dyn std::any::Any { + self + } + + fn schema(&self) -> SchemaRef { + self.schema.clone() + } + + fn table_type(&self) -> TableType { + TableType::Base + } + + async fn scan( + &self, + state: &SessionState, + projection: Option<&Vec>, + _filters: &[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; + } + } + batches + } else { + self.batches.clone() + }; + Ok(Arc::new(MemoryExec::try_new( + &[batches], + TableProvider::schema(self), + projection.cloned(), + )?)) + } +} + +impl SimpleCsvTable { + async fn interpreter_expr(&self, state: &SessionState) -> Result { + use datafusion::logical_expr::expr_rewriter::normalize_col; + use datafusion::logical_expr::utils::columnize_expr; + let plan = LogicalPlan::EmptyRelation(EmptyRelation { + produce_one_row: true, + schema: Arc::new(DFSchema::empty()), + }); + let logical_plan = Projection::try_new( + vec![columnize_expr( + normalize_col(self.exprs[0].clone(), &plan)?, + plan.schema(), + )], + Arc::new(plan), + ) + .map(LogicalPlan::Projection)?; + let rbs = collect( + state.create_physical_plan(&logical_plan).await?, + Arc::new(TaskContext::from(state)), + ) + .await?; + let limit = rbs[0] + .column(0) + .as_any() + .downcast_ref::() + .unwrap() + .value(0); + Ok(limit) + } +} + +struct SimpleCsvTableFunc {} + +impl TableFunctionImpl for SimpleCsvTableFunc { + fn call(&self, exprs: &[Expr]) -> Result> { + let mut new_exprs = vec![]; + let mut filepath = String::new(); + for expr in exprs { + match expr { + Expr::Literal(ScalarValue::Utf8(Some(ref path))) => { + filepath = path.clone() + } + expr => new_exprs.push(expr.clone()), + } + } + let (schema, batches) = read_csv_batches(filepath)?; + let table = SimpleCsvTable { + schema, + exprs: new_exprs.clone(), + batches, + }; + Ok(Arc::new(table)) + } +} + +fn read_csv_batches(csv_path: impl AsRef) -> Result<(SchemaRef, Vec)> { + let mut file = File::open(csv_path)?; + let (schema, _) = Format::default() + .with_header(true) + .infer_schema(&mut file, None)?; + file.rewind()?; + + let reader = ReaderBuilder::new(Arc::new(schema.clone())) + .with_header(true) + .build(file)?; + let mut batches = vec![]; + for bacth in reader { + batches.push(bacth?); + } + let schema = Arc::new(schema); + Ok((schema, batches)) +} From 1fe0c8bac9a2b56ae5511b984d36aa2f8c59c5b4 Mon Sep 17 00:00:00 2001 From: veeupup Date: Tue, 28 Nov 2023 22:32:31 +0800 Subject: [PATCH 4/8] add file header --- datafusion-examples/examples/simple_udtf.rs | 4 +--- .../user_defined_table_functions.rs | 17 +++++++++++++++++ 2 files changed, 18 insertions(+), 3 deletions(-) diff --git a/datafusion-examples/examples/simple_udtf.rs b/datafusion-examples/examples/simple_udtf.rs index c2bb86a0f760..260c6cbee710 100644 --- a/datafusion-examples/examples/simple_udtf.rs +++ b/datafusion-examples/examples/simple_udtf.rs @@ -179,9 +179,7 @@ impl TableFunctionImpl for LocalCsvTableFunc { fn read_csv_batches(csv_path: impl AsRef) -> Result<(SchemaRef, Vec)> { let mut file = File::open(csv_path)?; - let (schema, _) = Format::default() - .with_header(true) - .infer_schema(&mut file, None)?; + let (schema, _) = Format::default().infer_schema(&mut file, None)?; file.rewind()?; let reader = ReaderBuilder::new(Arc::new(schema.clone())) 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 4ac3fefa9191..dc1200aee607 100644 --- a/datafusion/core/tests/user_defined/user_defined_table_functions.rs +++ b/datafusion/core/tests/user_defined/user_defined_table_functions.rs @@ -1,3 +1,20 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + use arrow::array::Int64Array; use arrow::csv::reader::Format; use arrow::csv::ReaderBuilder; From 169b13e99b3e14b08c6e144982796f00163a5dcc Mon Sep 17 00:00:00 2001 From: Andrew Lamb Date: Wed, 29 Nov 2023 16:09:04 -0500 Subject: [PATCH 5/8] Simply table function example, add some comments --- datafusion-examples/examples/simple_udtf.rs | 89 ++++++++------------- 1 file changed, 33 insertions(+), 56 deletions(-) diff --git a/datafusion-examples/examples/simple_udtf.rs b/datafusion-examples/examples/simple_udtf.rs index 260c6cbee710..4911d4c28076 100644 --- a/datafusion-examples/examples/simple_udtf.rs +++ b/datafusion-examples/examples/simple_udtf.rs @@ -15,7 +15,6 @@ // specific language governing permissions and limitations // under the License. -use arrow::array::Int64Array; use arrow::csv::reader::Format; use arrow::csv::ReaderBuilder; use async_trait::async_trait; @@ -25,21 +24,20 @@ use datafusion::datasource::function::TableFunctionImpl; use datafusion::datasource::TableProvider; use datafusion::error::Result; use datafusion::execution::context::SessionState; -use datafusion::execution::TaskContext; use datafusion::physical_plan::memory::MemoryExec; -use datafusion::physical_plan::{collect, ExecutionPlan}; +use datafusion::physical_plan::{ExecutionPlan}; use datafusion::prelude::SessionContext; -use datafusion_common::{DFSchema, ScalarValue}; -use datafusion_expr::{EmptyRelation, Expr, LogicalPlan, Projection, TableType}; +use datafusion_common::{plan_err, DataFusionError, ScalarValue}; +use datafusion_expr::{Expr, TableType}; use std::fs::File; use std::io::Seek; use std::path::Path; use std::sync::Arc; // To define your own table function, you only need to do the following 3 things: -// 1. Implement your own TableProvider -// 2. Implement your own TableFunctionImpl and return your TableProvider -// 3. Register the function using ctx.register_udtf +// 1. Implement your own [`TableProvider`] +// 2. Implement your own [`TableFunctionImpl`] and return your [`TableProvider`] +// 3. Register the function using [`SessionContext::register_udtf`] /// This example demonstrates how to register a TableFunction #[tokio::main] @@ -47,12 +45,13 @@ async fn main() -> Result<()> { // create local execution context let ctx = SessionContext::new(); + // register the table function that will be called in SQL statements by `read_csv` ctx.register_udtf("read_csv", Arc::new(LocalCsvTableFunc {})); let testdata = datafusion::test_util::arrow_test_data(); let csv_file = format!("{testdata}/csv/aggregate_test_100.csv"); - // read csv with at most 2 rows + // Pass 2 arguments, read csv with at most 2 rows let df = ctx .sql(format!("SELECT * FROM read_csv('{csv_file}', 2);").as_str()) .await?; @@ -67,9 +66,14 @@ async fn main() -> Result<()> { Ok(()) } +/// Table Function that mimics the [`read_csv`] function in DuckDB. +/// +/// Usage: `read_csv(filename, [limit])` +/// +/// [`read_csv`]: https://duckdb.org/docs/data/csv/overview.html struct LocalCsvTable { schema: SchemaRef, - exprs: Vec, + limit: Option, batches: Vec, } @@ -89,13 +93,12 @@ impl TableProvider for LocalCsvTable { async fn scan( &self, - state: &SessionState, + _state: &SessionState, projection: Option<&Vec>, _filters: &[Expr], _limit: Option, ) -> Result> { - let batches = if !self.exprs.is_empty() { - let max_return_lines = self.interpreter_expr(state).await?; + 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; @@ -121,56 +124,30 @@ impl TableProvider for LocalCsvTable { )?)) } } - -impl LocalCsvTable { - async fn interpreter_expr(&self, state: &SessionState) -> Result { - use datafusion::logical_expr::expr_rewriter::normalize_col; - use datafusion::logical_expr::utils::columnize_expr; - let plan = LogicalPlan::EmptyRelation(EmptyRelation { - produce_one_row: true, - schema: Arc::new(DFSchema::empty()), - }); - let logical_plan = Projection::try_new( - vec![columnize_expr( - normalize_col(self.exprs[0].clone(), &plan)?, - plan.schema(), - )], - Arc::new(plan), - ) - .map(LogicalPlan::Projection)?; - let rbs = collect( - state.create_physical_plan(&logical_plan).await?, - Arc::new(TaskContext::from(state)), - ) - .await?; - let limit = rbs[0] - .column(0) - .as_any() - .downcast_ref::() - .unwrap() - .value(0); - Ok(limit) - } -} - struct LocalCsvTableFunc {} impl TableFunctionImpl for LocalCsvTableFunc { fn call(&self, exprs: &[Expr]) -> Result> { - let mut new_exprs = vec![]; - let mut filepath = String::new(); - for expr in exprs { - match expr { - Expr::Literal(ScalarValue::Utf8(Some(ref path))) => { - filepath = path.clone() + let Some(Expr::Literal(ScalarValue::Utf8(Some(ref path)))) = exprs.get(0) else { + return plan_err!("read_csv requires at least one string argument"); + }; + + let limit = exprs + .get(1) + .map(|expr| { + if let Expr::Literal(ScalarValue::Int64(Some(limit))) = expr { + Ok(*limit as usize) + } else { + plan_err!("Limit must be an integer") } - expr => new_exprs.push(expr.clone()), - } - } - let (schema, batches) = read_csv_batches(filepath)?; + }) + .transpose()?; + + let (schema, batches) = read_csv_batches(path)?; + let table = LocalCsvTable { schema, - exprs: new_exprs.clone(), + limit, batches, }; Ok(Arc::new(table)) From 2df01dbae67fe60746901d07f1729e1d2d8e6bd6 Mon Sep 17 00:00:00 2001 From: Andrew Lamb Date: Wed, 29 Nov 2023 16:19:30 -0500 Subject: [PATCH 6/8] Simplfy exprs --- datafusion-examples/examples/simple_udtf.rs | 16 +++++++++++----- 1 file changed, 11 insertions(+), 5 deletions(-) diff --git a/datafusion-examples/examples/simple_udtf.rs b/datafusion-examples/examples/simple_udtf.rs index 4911d4c28076..05ad231949a8 100644 --- a/datafusion-examples/examples/simple_udtf.rs +++ b/datafusion-examples/examples/simple_udtf.rs @@ -23,12 +23,13 @@ use datafusion::arrow::record_batch::RecordBatch; use datafusion::datasource::function::TableFunctionImpl; use datafusion::datasource::TableProvider; use datafusion::error::Result; -use datafusion::execution::context::SessionState; +use datafusion::execution::context::{ExecutionProps, SessionState}; use datafusion::physical_plan::memory::MemoryExec; -use datafusion::physical_plan::{ExecutionPlan}; +use datafusion::physical_plan::ExecutionPlan; use datafusion::prelude::SessionContext; use datafusion_common::{plan_err, DataFusionError, ScalarValue}; use datafusion_expr::{Expr, TableType}; +use datafusion_optimizer::simplify_expressions::{ExprSimplifier, SimplifyContext}; use std::fs::File; use std::io::Seek; use std::path::Path; @@ -51,9 +52,9 @@ async fn main() -> Result<()> { let testdata = datafusion::test_util::arrow_test_data(); let csv_file = format!("{testdata}/csv/aggregate_test_100.csv"); - // Pass 2 arguments, read csv with at most 2 rows + // Pass 2 arguments, read csv with at most 2 rows (simplify logic makes 1+1 --> 2) let df = ctx - .sql(format!("SELECT * FROM read_csv('{csv_file}', 2);").as_str()) + .sql(format!("SELECT * FROM read_csv('{csv_file}', 1 + 1);").as_str()) .await?; df.show().await?; @@ -135,8 +136,13 @@ impl TableFunctionImpl for LocalCsvTableFunc { let limit = exprs .get(1) .map(|expr| { + // try to simpify the expression, so 1+2 becomes 3, for example + let execution_props = ExecutionProps::new(); + let info = SimplifyContext::new(&execution_props); + let expr = ExprSimplifier::new(info).simplify(expr.clone())?; + if let Expr::Literal(ScalarValue::Int64(Some(limit))) = expr { - Ok(*limit as usize) + Ok(limit as usize) } else { plan_err!("Limit must be an integer") } From c2e7cb03a06db0784fa40a4b2955cc4a4924445e Mon Sep 17 00:00:00 2001 From: veeupup Date: Thu, 30 Nov 2023 09:13:11 +0800 Subject: [PATCH 7/8] make clippy happy --- datafusion-examples/examples/simple_udtf.rs | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/datafusion-examples/examples/simple_udtf.rs b/datafusion-examples/examples/simple_udtf.rs index 05ad231949a8..bce633765281 100644 --- a/datafusion-examples/examples/simple_udtf.rs +++ b/datafusion-examples/examples/simple_udtf.rs @@ -105,8 +105,8 @@ impl TableProvider for LocalCsvTable { 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; + if lines + batch_lines > max_return_lines { + let batch_lines = max_return_lines - lines; batches.push(batch.slice(0, batch_lines)); break; } else { From 0e98c0b4220b1e30d8f7c48e5cb030ed946dab3a Mon Sep 17 00:00:00 2001 From: Andrew Lamb Date: Thu, 30 Nov 2023 17:14:06 -0500 Subject: [PATCH 8/8] Update datafusion/core/tests/user_defined/user_defined_table_functions.rs --- .../core/tests/user_defined/user_defined_table_functions.rs | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) 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 dc1200aee607..b5d10b1c5b9b 100644 --- a/datafusion/core/tests/user_defined/user_defined_table_functions.rs +++ b/datafusion/core/tests/user_defined/user_defined_table_functions.rs @@ -44,7 +44,7 @@ async fn test_simple_read_csv_udtf() -> Result<()> { ctx.register_udtf("read_csv", Arc::new(SimpleCsvTableFunc {})); let csv_file = "tests/tpch-csv/nation.csv"; - // read csv with at most 2 rows + // read csv with at most 5 rows let rbs = ctx .sql(format!("SELECT * FROM read_csv('{csv_file}', 5);").as_str()) .await?