Support User Defined Table Function

Signed-off-by: veeupup <[email protected]>

Author: Veeupup

datafusion-examples/examples/simple_udtf.rs

@@ -0,0 +1,156 @@
+// 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::streaming::StreamingTable;
+use datafusion::datasource::TableProvider;
+use datafusion::error::Result;
+use datafusion::execution::context::SessionState;
+use datafusion::physical_plan::memory::MemoryExec;
+use datafusion::physical_plan::stream::RecordBatchStreamAdapter;
+use datafusion::physical_plan::streaming::PartitionStream;
+use datafusion::physical_plan::ExecutionPlan;
+use datafusion::prelude::SessionContext;
+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. Define a function that returns a Result<Arc<dyn TableProvider>>
+//      maybe you can just implement your own TableProvider
+//      and you can even implement your own PartitionStream and return it with inner StreamTable
+// 2. 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(read_csv));
+    ctx.register_udtf("read_csv_stream", Arc::new(read_csv_stream));
+
+    let testdata = datafusion::test_util::arrow_test_data();
+    let csv_file = format!("{testdata}/csv/aggregate_test_100.csv");
+
+    let df = ctx
+        .sql(format!("SELECT * FROM read_csv('{csv_file}');").as_str())
+        .await?;
+    df.show().await?;
+
+    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)
+//      Define your own TableProvider and make a function return Arc<dyn Provider>
+struct LocalCsvTable {
+    schema: SchemaRef,
+    batches: Vec<RecordBatch>,
+}
+
+#[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<usize>>,
+        _filters: &[Expr],
+        _limit: Option<usize>,
+    ) -> Result<Arc<dyn ExecutionPlan>> {
+        Ok(Arc::new(MemoryExec::try_new(
+            &[self.batches.clone()],
+            TableProvider::schema(self),
+            projection.cloned(),
+        )?))
+    }
+}
+
+fn read_csv(args: &[String]) -> Result<Arc<dyn TableProvider>> {
+    let (schema, batches) = read_csv_batches(args[0].clone())?;
+    let table = LocalCsvTable { schema, batches };
+    Ok(Arc::new(table))
+}
+
+// Option2: (use StreamingTable to make it simpler)
+//      Implement PartitionStream and make a function return it with StreamTable
+impl PartitionStream for LocalCsvTable {
+    fn schema(&self) -> &SchemaRef {
+        &self.schema
+    }
+
+    fn execute(
+        &self,
+        _ctx: Arc<datafusion::execution::TaskContext>,
+    ) -> 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<Item = ArrowResult<RecordBatch>> + Send + Sync + 'static
+            futures::stream::iter(self.batches.clone().into_iter().map(Ok)),
+        ))
+    }
+}
+
+fn read_csv_stream(args: &[String]) -> Result<Arc<dyn TableProvider>> {
+    let (schema, batches) = read_csv_batches(args[0].clone())?;
+    let stream = LocalCsvTable {
+        schema: schema.clone(),
+        batches,
+    };
+    let table = StreamingTable::try_new(schema, vec![Arc::new(stream)])?;
+    Ok(Arc::new(table))
+}
+
+fn read_csv_batches(csv_path: impl AsRef<Path>) -> Result<(SchemaRef, Vec<RecordBatch>)> {
+    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))
+}

datafusion/core/src/datasource/function.rs

@@ -0,0 +1,55 @@
+// 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 std::sync::Arc;
+
+/// Table Function implementation
+pub type TableFunctionImplementation =
+    Arc<dyn Fn(&[String]) -> Result<Arc<dyn TableProvider>> + Send + Sync>;
+
+/// A table that uses a function to generate data
+pub struct TableFunction {
+    /// Name of the table function
+    name: String,
+    /// Function implementation
+    fun: TableFunctionImplementation,
+}
+
+impl TableFunction {
+    /// Create a new table function
+    pub fn new(name: String, fun: TableFunctionImplementation) -> 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: &[String],
+    ) -> Result<Arc<dyn TableProvider>> {
+        let table = (self.fun)(args)?;
+        Ok(table)
+    }
+}

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;

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

@@ -26,6 +26,7 @@ mod parquet;
 use crate::{
     catalog::{CatalogList, MemoryCatalogList},
     datasource::{
+        function::{TableFunction, TableFunctionImplementation},
         listing::{ListingOptions, ListingTable},
         provider::TableProviderFactory,
     },
@@ -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: TableFunctionImplementation) {
+        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<dyn QueryPlanner + Send + Sync>,
     /// Collection of catalogs containing schemas and ultimately TableProviders
     catalog_list: Arc<dyn CatalogList>,
+    /// Table Functions
+    table_functions: HashMap<String, Arc<TableFunction>>,
     /// Scalar functions that are registered with the context
     scalar_functions: HashMap<String, Arc<ScalarUDF>>,
     /// 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<String>,
+    ) -> Result<Arc<dyn TableSource>> {
+        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<Arc<ScalarUDF>> {
         self.state.scalar_functions().get(name).cloned()
     }

datafusion/sql/src/planner.rs

@@ -51,6 +51,15 @@ pub trait ContextProvider {
     }
     /// Getter for a datasource
     fn get_table_source(&self, name: TableReference) -> Result<Arc<dyn TableSource>>;
+    /// Getter for a table function
+    fn get_table_function_source(
+        &self,
+        _name: &str,
+        _args: Vec<String>,
+    ) -> Result<Arc<dyn TableSource>> {
+        unimplemented!()
+    }
+
     /// Getter for a UDF description
     fn get_function_meta(&self, name: &str) -> Option<Arc<ScalarUDF>>;
     /// Getter for a UDAF description

datafusion/sql/src/relation/mod.rs

@@ -16,9 +16,9 @@
 // under the License.
 
 use crate::planner::{ContextProvider, PlannerContext, SqlToRel};
-use datafusion_common::{not_impl_err, DataFusionError, Result};
+use datafusion_common::{not_impl_err, DataFusionError, Result, TableReference};
 use datafusion_expr::{LogicalPlan, LogicalPlanBuilder};
-use sqlparser::ast::TableFactor;
+use sqlparser::ast::{Expr, FunctionArgExpr, TableFactor, Value};
 
 mod join;
 
@@ -30,24 +30,56 @@ impl<'a, S: ContextProvider> SqlToRel<'a, S> {
         planner_context: &mut PlannerContext,
     ) -> Result<LogicalPlan> {
         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, ..
+            } => {
+                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 {
+                        if let sqlparser::ast::FunctionArg::Unnamed(
+                            FunctionArgExpr::Expr(Expr::Value(
+                                Value::SingleQuotedString(val),
+                            )),
+                        ) = arg
+                        {
+                            args.push(val);
+                        } else {
+                            unimplemented!("Unsupported function argument type")
                         }
-                        (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, ..