udtf: provide session state ref to the call

This patch adds session state arg for the [`TableFunctionImpl::call`] method.
It is useful to implement table functions depending on other tables from
the state. For example, a table functions that return current list of
all views in the state.

Author: askalt

datafusion-cli/src/functions.rs

@@ -317,7 +317,11 @@ fn fixed_len_byte_array_to_string(val: Option<&FixedLenByteArray>) -> Option<Str
 pub struct ParquetMetadataFunc {}
 
 impl TableFunctionImpl for ParquetMetadataFunc {
-    fn call(&self, exprs: &[Expr]) -> Result<Arc<dyn TableProvider>> {
+    fn call(
+        &self,
+        _state: &SessionState,
+        exprs: &[Expr],
+    ) -> Result<Arc<dyn TableProvider>> {
         let filename = match exprs.first() {
             Some(Expr::Literal(ScalarValue::Utf8(Some(s)))) => s, // single quote: parquet_metadata('x.parquet')
             Some(Expr::Column(Column { name, .. })) => name, // double quote: parquet_metadata("x.parquet")

datafusion-examples/examples/simple_udtf.rs

@@ -25,6 +25,7 @@ use datafusion::datasource::function::TableFunctionImpl;
 use datafusion::datasource::TableProvider;
 use datafusion::error::Result;
 use datafusion::execution::context::ExecutionProps;
+use datafusion::execution::SessionState;
 use datafusion::physical_plan::memory::MemoryExec;
 use datafusion::physical_plan::ExecutionPlan;
 use datafusion::prelude::SessionContext;
@@ -130,7 +131,11 @@ impl TableProvider for LocalCsvTable {
 struct LocalCsvTableFunc {}
 
 impl TableFunctionImpl for LocalCsvTableFunc {
-    fn call(&self, exprs: &[Expr]) -> Result<Arc<dyn TableProvider>> {
+    fn call(
+        &self,
+        _state: &SessionState,
+        exprs: &[Expr],
+    ) -> Result<Arc<dyn TableProvider>> {
         let Some(Expr::Literal(ScalarValue::Utf8(Some(ref path)))) = exprs.first() else {
             return plan_err!("read_csv requires at least one string argument");
         };

datafusion/core/src/datasource/function.rs

@@ -17,6 +17,8 @@
 
 //! A table that uses a function to generate data
 
+use crate::execution::SessionState;
+
 use super::TableProvider;
 
 use datafusion_common::Result;
@@ -27,7 +29,8 @@ 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<Arc<dyn TableProvider>>;
+    fn call(&self, state: &SessionState, args: &[Expr])
+        -> Result<Arc<dyn TableProvider>>;
 }
 
 /// A table that uses a function to generate data
@@ -55,7 +58,11 @@ impl TableFunction {
     }
 
     /// Get the function implementation and generate a table
-    pub fn create_table_provider(&self, args: &[Expr]) -> Result<Arc<dyn TableProvider>> {
-        self.fun.call(args)
+    pub fn create_table_provider(
+        &self,
+        state: &SessionState,
+        args: &[Expr],
+    ) -> Result<Arc<dyn TableProvider>> {
+        self.fun.call(state, args)
     }
 }

datafusion/core/src/execution/session_state.rs

@@ -1543,7 +1543,7 @@ impl ContextProvider for SessionContextProvider<'_> {
             .get(name)
             .cloned()
             .ok_or_else(|| plan_datafusion_err!("table function '{name}' not found"))?;
-        let provider = tbl_func.create_table_provider(&args)?;
+        let provider = tbl_func.create_table_provider(self.state, &args)?;
 
         Ok(provider_as_source(provider))
     }

datafusion/core/tests/user_defined/user_defined_table_functions.rs

@@ -24,7 +24,7 @@ use datafusion::arrow::record_batch::RecordBatch;
 use datafusion::datasource::function::TableFunctionImpl;
 use datafusion::datasource::TableProvider;
 use datafusion::error::Result;
-use datafusion::execution::TaskContext;
+use datafusion::execution::{SessionState, TaskContext};
 use datafusion::physical_plan::memory::MemoryExec;
 use datafusion::physical_plan::{collect, ExecutionPlan};
 use datafusion::prelude::SessionContext;
@@ -194,7 +194,11 @@ impl SimpleCsvTable {
 struct SimpleCsvTableFunc {}
 
 impl TableFunctionImpl for SimpleCsvTableFunc {
-    fn call(&self, exprs: &[Expr]) -> Result<Arc<dyn TableProvider>> {
+    fn call(
+        &self,
+        _state: &SessionState,
+        exprs: &[Expr],
+    ) -> Result<Arc<dyn TableProvider>> {
         let mut new_exprs = vec![];
         let mut filepath = String::new();
         for expr in exprs {

datafusion/sql/src/query.rs

@@ -86,15 +86,21 @@ impl<S: ContextProvider> SqlToRel<'_, S> {
         }
 
         let skip = match skip {
-            Some(skip_expr) => self.get_constant_usize_result(skip_expr.value, input.schema(), "OFFSET"),
+            Some(skip_expr) => {
+                self.get_constant_usize_result(skip_expr.value, input.schema(), "OFFSET")
+            }
             _ => Ok(0),
         }?;
 
         let fetch = match fetch {
             Some(limit_expr)
                 if limit_expr != sqlparser::ast::Expr::Value(Value::Null) =>
             {
-                Some(self.get_constant_usize_result(limit_expr, input.schema(), "LIMIT")?)
+                Some(self.get_constant_usize_result(
+                    limit_expr,
+                    input.schema(),
+                    "LIMIT",
+                )?)
             }
             _ => None,
         };
@@ -159,7 +165,12 @@ impl<S: ContextProvider> SqlToRel<'_, S> {
     ///
     /// * `Result<usize>` - An `Ok` variant containing the constant result if evaluation is successful,
     ///   or an `Err` variant containing an error message if evaluation fails.
-    pub(super) fn get_constant_usize_result(&self, expr: SQLExpr, schema: &datafusion_common::DFSchema, arg_name: &str) -> Result<usize> {
+    pub(super) fn get_constant_usize_result(
+        &self,
+        expr: SQLExpr,
+        schema: &datafusion_common::DFSchema,
+        arg_name: &str,
+    ) -> Result<usize> {
         let expr = self.sql_to_expr(expr, schema, &mut PlannerContext::new())?;
         let value = get_constant_result(&expr, arg_name)?;
         convert_usize_with_check(value, arg_name)

datafusion/sql/src/statement.rs

@@ -1219,7 +1219,7 @@ impl<S: ContextProvider> SqlToRel<'_, S> {
         &self,
         table_name: ObjectName,
         predicate_expr: Option<SQLExpr>,
-        limit: Option<SQLExpr>
+        limit: Option<SQLExpr>,
     ) -> Result<LogicalPlan> {
         // Do a table lookup to verify the table exists
         let table_ref = self.object_name_to_table_reference(table_name.clone())?;
@@ -1251,9 +1251,11 @@ impl<S: ContextProvider> SqlToRel<'_, S> {
             }
         };
 
-        if let Some(limit_expr) = limit {        
+        if let Some(limit_expr) = limit {
             let limit = (limit_expr != SQLExpr::Value(Value::Null))
-                .then(|| self.get_constant_usize_result(limit_expr, source.schema(), "LIMIT"))
+                .then(|| {
+                    self.get_constant_usize_result(limit_expr, source.schema(), "LIMIT")
+                })
                 .transpose()?;
 
             source = LogicalPlanBuilder::from(source).limit(0, limit)?.build()?

docs/source/library-user-guide/adding-udfs.md

@@ -562,14 +562,18 @@ In the `call` method, you parse the input `Expr`s and return a `TableProvider`.
 ```rust
 use datafusion::common::plan_err;
 use datafusion::datasource::function::TableFunctionImpl;
+use datafusion::execution::SessionState;
 // Other imports here
 
 /// A table function that returns a table provider with the value as a single column
 #[derive(Default)]
 pub struct EchoFunction {}
 
 impl TableFunctionImpl for EchoFunction {
-    fn call(&self, exprs: &[Expr]) -> Result<Arc<dyn TableProvider>> {
+    fn call(&self, 
+        _state: &SessionState,
+        exprs: &[Expr],
+    ) -> Result<Arc<dyn TableProvider>> {
         let Some(Expr::Literal(ScalarValue::Int64(Some(value)))) = exprs.get(0) else {
             return plan_err!("First argument must be an integer");
         };