ARROW-9950: [Rust] [DataFusion] Made UDFs usable without a registry

rust/datafusion/examples/simple_udf.rs

@@ -111,20 +111,28 @@ fn main() -> Result<()> {
  pow,
  );
 
- // finally, register the UDF
- ctx.register_udf(pow);
+ // at this point, we can use it or register it, depending on the use-case:
+ // * if the UDF is expected to be used throughout the program in different contexts,
+ // we can register it, and call it later:
+ ctx.register_udf(pow.clone()); // clone is only required in this example because we show both usages
 
- // at this point, we can use it. Note that the code below can be in a
-// scope on which we do not have access to `pow`.
+ // * if the UDF is expected to be used directly in the scope, `.call` it directly:
+let expr = pow.call(vec![col("a"), col("b")]);
 
  // get a DataFrame from the context
  let df = ctx.table("t")?;
 
- // get the udf registry.
- let f = df.registry();
-
- // equivalent to `'SELECT pow(a, b) FROM t'`
- let df = df.select(vec![f.udf("pow", vec![col("a"), col("b")])?])?;
+ // if we do not have `pow` in the scope and we registered it, we can get it from the registry
+ let pow = df.registry().udf("pow")?;
+ // equivalent to expr
+ let expr1 = pow.call(vec![col("a"), col("b")]);
+
+ // equivalent to `'SELECT pow(a, b), pow(a, b) AS pow1 FROM t'`
+ let df = df.select(vec![
+ expr,
+ // alias so that they have different column names
+ expr1.alias("pow1"),
+ ])?;
 
  // note that "b" is f32, not f64. DataFusion coerces the types to match the UDF's signature.
 

rust/datafusion/src/execution/context.rs

@@ -36,7 +36,7 @@ use crate::datasource::parquet::ParquetTable;
 use crate::datasource::TableProvider;
 use crate::error::{ExecutionError, Result};
 use crate::execution::dataframe_impl::DataFrameImpl;
-use crate::logical_plan::{Expr, FunctionRegistry, LogicalPlan, LogicalPlanBuilder};
+use crate::logical_plan::{FunctionRegistry, LogicalPlan, LogicalPlanBuilder};
 use crate::optimizer::filter_push_down::FilterPushDown;
 use crate::optimizer::optimizer::OptimizerRule;
 use crate::optimizer::projection_push_down::ProjectionPushDown;
@@ -481,17 +481,14 @@ impl FunctionRegistry for ExecutionContextState {
  self.scalar_functions.keys().cloned().collect()
  }
 
- fn udf(&self, name: &str, args: Vec<Expr>) -> Result<Expr> {
+ fn udf(&self, name: &str) -> Result<&ScalarUDF> {
  let result = self.scalar_functions.get(name);
  if result.is_none() {
  Err(ExecutionError::General(
  format!("There is no UDF named \"{}\" in the registry", name).to_string(),
  ))
  } else {
- Ok(Expr::ScalarUDF {
- fun: result.unwrap().clone(),
- args,
- })
+ Ok(result.unwrap())
  }
  }
 }
@@ -1051,7 +1048,7 @@ mod tests {
  .project(vec![
  col("a"),
  col("b"),
- ctx.registry().udf("my_add", vec![col("a"), col("b")])?,
+ ctx.registry().udf("my_add")?.call(vec![col("a"), col("b")]),
  ])?
  .build()?;
 

rust/datafusion/src/execution/dataframe_impl.rs

@@ -259,7 +259,7 @@ mod tests {
 
  let f = df.registry();
 
- let df = df.select(vec![f.udf("my_fn", vec![col("c12")])?])?;
+ let df = df.select(vec![f.udf("my_fn")?.call(vec![col("c12")])])?;
  let plan = df.to_logical_plan();
 
  // build query using SQL

rust/datafusion/src/logical_plan/mod.rs

@@ -1133,8 +1133,8 @@ pub trait FunctionRegistry {
  /// Set of all available udfs.
  fn udfs(&self) -> HashSet<String>;
 
- /// Constructs a logical expression with a call to the udf.
- fn udf(&self, name: &str, args: Vec<Expr>) -> Result<Expr>;
+ /// Returns a reference to the udf named `name`.
+ fn udf(&self, name: &str) -> Result<&ScalarUDF>;
 }
 
 /// Builder for logical plans

rust/datafusion/src/physical_plan/udf.rs

@@ -23,7 +23,7 @@ use std::fmt;
 use arrow::datatypes::Schema;
 
 use crate::error::Result;
-use crate::physical_plan::PhysicalExpr;
+use crate::{logical_plan::Expr, physical_plan::PhysicalExpr};
 
 use super::{
  functions::{
@@ -71,6 +71,15 @@ impl ScalarUDF {
  fun: fun.clone(),
  }
  }
+
+ /// creates a logical expression with a call of the UDF
+ /// This utility allows using the UDF without requiring access to the registry.
+ pub fn call(&self, args: Vec<Expr>) -> Expr {
+ Expr::ScalarUDF {
+ fun: Arc::new(self.clone()),
+ args,
+ }
+ }
 }
 
 /// Create a physical expression of the UDF.

rust/datafusion/examples/simple_udf.rs

@@ -111,20 +111,28 @@ fn main() -> Result<()> {
  pow,
  );
 
- // finally, register the UDF
- ctx.register_udf(pow);
+ // at this point, we can use it or register it, depending on the use-case:
+ // * if the UDF is expected to be used throughout the program in different contexts,
+ // we can register it, and call it later:
+ ctx.register_udf(pow.clone()); // clone is only required in this example because we show both usages
 
- // at this point, we can use it. Note that the code below can be in a
-// scope on which we do not have access to `pow`.
+ // * if the UDF is expected to be used directly in the scope, `.call` it directly:
+let expr = pow.call(vec![col("a"), col("b")]);
 
  // get a DataFrame from the context
  let df = ctx.table("t")?;
 
- // get the udf registry.
- let f = df.registry();
-
- // equivalent to `'SELECT pow(a, b) FROM t'`
- let df = df.select(vec![f.udf("pow", vec![col("a"), col("b")])?])?;
+ // if we do not have `pow` in the scope and we registered it, we can get it from the registry
+ let pow = df.registry().udf("pow")?;
+ // equivalent to expr
+ let expr1 = pow.call(vec![col("a"), col("b")]);
+
+ // equivalent to `'SELECT pow(a, b), pow(a, b) AS pow1 FROM t'`
+ let df = df.select(vec![
+ expr,
+ // alias so that they have different column names
+ expr1.alias("pow1"),
+ ])?;
 
  // note that "b" is f32, not f64. DataFusion coerces the types to match the UDF's signature.
 

rust/datafusion/src/execution/context.rs

@@ -36,7 +36,7 @@ use crate::datasource::parquet::ParquetTable;
 use crate::datasource::TableProvider;
 use crate::error::{ExecutionError, Result};
 use crate::execution::dataframe_impl::DataFrameImpl;
-use crate::logical_plan::{Expr, FunctionRegistry, LogicalPlan, LogicalPlanBuilder};
+use crate::logical_plan::{FunctionRegistry, LogicalPlan, LogicalPlanBuilder};
 use crate::optimizer::filter_push_down::FilterPushDown;
 use crate::optimizer::optimizer::OptimizerRule;
 use crate::optimizer::projection_push_down::ProjectionPushDown;
@@ -481,17 +481,14 @@ impl FunctionRegistry for ExecutionContextState {
  self.scalar_functions.keys().cloned().collect()
  }
 
- fn udf(&self, name: &str, args: Vec<Expr>) -> Result<Expr> {
+ fn udf(&self, name: &str) -> Result<&ScalarUDF> {
  let result = self.scalar_functions.get(name);
  if result.is_none() {
  Err(ExecutionError::General(
  format!("There is no UDF named \"{}\" in the registry", name).to_string(),
  ))
  } else {
- Ok(Expr::ScalarUDF {
- fun: result.unwrap().clone(),
- args,
- })
+ Ok(result.unwrap())
  }
  }
 }
@@ -1051,7 +1048,7 @@ mod tests {
  .project(vec![
  col("a"),
  col("b"),
- ctx.registry().udf("my_add", vec![col("a"), col("b")])?,
+ ctx.registry().udf("my_add")?.call(vec![col("a"), col("b")]),
  ])?
  .build()?;
 

rust/datafusion/src/execution/dataframe_impl.rs

@@ -259,7 +259,7 @@ mod tests {
 
  let f = df.registry();
 
- let df = df.select(vec![f.udf("my_fn", vec![col("c12")])?])?;
+ let df = df.select(vec![f.udf("my_fn")?.call(vec![col("c12")])])?;
  let plan = df.to_logical_plan();
 
  // build query using SQL

rust/datafusion/src/logical_plan/mod.rs

@@ -1133,8 +1133,8 @@ pub trait FunctionRegistry {
  /// Set of all available udfs.
  fn udfs(&self) -> HashSet<String>;
 
- /// Constructs a logical expression with a call to the udf.
- fn udf(&self, name: &str, args: Vec<Expr>) -> Result<Expr>;
+ /// Returns a reference to the udf named `name`.
+ fn udf(&self, name: &str) -> Result<&ScalarUDF>;
 }
 
 /// Builder for logical plans

rust/datafusion/src/physical_plan/udf.rs

@@ -23,7 +23,7 @@ use std::fmt;
 use arrow::datatypes::Schema;
 
 use crate::error::Result;
-use crate::physical_plan::PhysicalExpr;
+use crate::{logical_plan::Expr, physical_plan::PhysicalExpr};
 
 use super::{
  functions::{
@@ -71,6 +71,15 @@ impl ScalarUDF {
  fun: fun.clone(),
  }
  }
+
+ /// creates a logical expression with a call of the UDF
+ /// This utility allows using the UDF without requiring access to the registry.
+ pub fn call(&self, args: Vec<Expr>) -> Expr {
+ Expr::ScalarUDF {
+ fun: Arc::new(self.clone()),
+ args,
+ }
+ }
 }
 
 /// Create a physical expression of the UDF.

rust/datafusion/examples/simple_udf.rs

@@ -111,20 +111,28 @@ fn main() -> Result<()> {
  pow,
  );
 
- // finally, register the UDF
- ctx.register_udf(pow);
+ // at this point, we can use it or register it, depending on the use-case:
+ // * if the UDF is expected to be used throughout the program in different contexts,
+ // we can register it, and call it later:
+ ctx.register_udf(pow.clone()); // clone is only required in this example because we show both usages
 
- // at this point, we can use it. Note that the code below can be in a
-// scope on which we do not have access to `pow`.
+ // * if the UDF is expected to be used directly in the scope, `.call` it directly:
+let expr = pow.call(vec![col("a"), col("b")]);
 
  // get a DataFrame from the context
  let df = ctx.table("t")?;
 
- // get the udf registry.
- let f = df.registry();
-
- // equivalent to `'SELECT pow(a, b) FROM t'`
- let df = df.select(vec![f.udf("pow", vec![col("a"), col("b")])?])?;
+ // if we do not have `pow` in the scope and we registered it, we can get it from the registry
+ let pow = df.registry().udf("pow")?;
+ // equivalent to expr
+ let expr1 = pow.call(vec![col("a"), col("b")]);
+
+ // equivalent to `'SELECT pow(a, b), pow(a, b) AS pow1 FROM t'`
+ let df = df.select(vec![
+ expr,
+ // alias so that they have different column names
+ expr1.alias("pow1"),
+ ])?;
 
  // note that "b" is f32, not f64. DataFusion coerces the types to match the UDF's signature.
 

rust/datafusion/src/execution/context.rs

@@ -36,7 +36,7 @@ use crate::datasource::parquet::ParquetTable;
 use crate::datasource::TableProvider;
 use crate::error::{ExecutionError, Result};
 use crate::execution::dataframe_impl::DataFrameImpl;
-use crate::logical_plan::{Expr, FunctionRegistry, LogicalPlan, LogicalPlanBuilder};
+use crate::logical_plan::{FunctionRegistry, LogicalPlan, LogicalPlanBuilder};
 use crate::optimizer::filter_push_down::FilterPushDown;
 use crate::optimizer::optimizer::OptimizerRule;
 use crate::optimizer::projection_push_down::ProjectionPushDown;
@@ -481,17 +481,14 @@ impl FunctionRegistry for ExecutionContextState {
  self.scalar_functions.keys().cloned().collect()
  }
 
- fn udf(&self, name: &str, args: Vec<Expr>) -> Result<Expr> {
+ fn udf(&self, name: &str) -> Result<&ScalarUDF> {
  let result = self.scalar_functions.get(name);
  if result.is_none() {
  Err(ExecutionError::General(
  format!("There is no UDF named \"{}\" in the registry", name).to_string(),
  ))
  } else {
- Ok(Expr::ScalarUDF {
- fun: result.unwrap().clone(),
- args,
- })
+ Ok(result.unwrap())
  }
  }
 }
@@ -1051,7 +1048,7 @@ mod tests {
  .project(vec![
  col("a"),
  col("b"),
- ctx.registry().udf("my_add", vec![col("a"), col("b")])?,
+ ctx.registry().udf("my_add")?.call(vec![col("a"), col("b")]),
  ])?
  .build()?;
 

rust/datafusion/src/execution/dataframe_impl.rs

@@ -259,7 +259,7 @@ mod tests {
 
  let f = df.registry();
 
- let df = df.select(vec![f.udf("my_fn", vec![col("c12")])?])?;
+ let df = df.select(vec![f.udf("my_fn")?.call(vec![col("c12")])])?;
  let plan = df.to_logical_plan();
 
  // build query using SQL

rust/datafusion/src/logical_plan/mod.rs

@@ -1133,8 +1133,8 @@ pub trait FunctionRegistry {
  /// Set of all available udfs.
  fn udfs(&self) -> HashSet<String>;
 
- /// Constructs a logical expression with a call to the udf.
- fn udf(&self, name: &str, args: Vec<Expr>) -> Result<Expr>;
+ /// Returns a reference to the udf named `name`.
+ fn udf(&self, name: &str) -> Result<&ScalarUDF>;
 }
 
 /// Builder for logical plans

rust/datafusion/src/physical_plan/udf.rs

@@ -23,7 +23,7 @@ use std::fmt;
 use arrow::datatypes::Schema;
 
 use crate::error::Result;
-use crate::physical_plan::PhysicalExpr;
+use crate::{logical_plan::Expr, physical_plan::PhysicalExpr};
 
 use super::{
  functions::{
@@ -71,6 +71,15 @@ impl ScalarUDF {
  fun: fun.clone(),
  }
  }
+
+ /// creates a logical expression with a call of the UDF
+ /// This utility allows using the UDF without requiring access to the registry.
+ pub fn call(&self, args: Vec<Expr>) -> Expr {
+ Expr::ScalarUDF {
+ fun: Arc::new(self.clone()),
+ args,
+ }
+ }
 }
 
 /// Create a physical expression of the UDF.

rust/datafusion/examples/simple_udf.rs

@@ -111,20 +111,28 @@ fn main() -> Result<()> {
  pow,
  );
 
- // finally, register the UDF
- ctx.register_udf(pow);
+ // at this point, we can use it or register it, depending on the use-case:
+ // * if the UDF is expected to be used throughout the program in different contexts,
+ // we can register it, and call it later:
+ ctx.register_udf(pow.clone()); // clone is only required in this example because we show both usages
 
- // at this point, we can use it. Note that the code below can be in a
-// scope on which we do not have access to `pow`.
+ // * if the UDF is expected to be used directly in the scope, `.call` it directly:
+let expr = pow.call(vec![col("a"), col("b")]);
 
  // get a DataFrame from the context
  let df = ctx.table("t")?;
 
- // get the udf registry.
- let f = df.registry();
-
- // equivalent to `'SELECT pow(a, b) FROM t'`
- let df = df.select(vec![f.udf("pow", vec![col("a"), col("b")])?])?;
+ // if we do not have `pow` in the scope and we registered it, we can get it from the registry
+ let pow = df.registry().udf("pow")?;
+ // equivalent to expr
+ let expr1 = pow.call(vec![col("a"), col("b")]);
+
+ // equivalent to `'SELECT pow(a, b), pow(a, b) AS pow1 FROM t'`
+ let df = df.select(vec![
+ expr,
+ // alias so that they have different column names
+ expr1.alias("pow1"),
+ ])?;
 
  // note that "b" is f32, not f64. DataFusion coerces the types to match the UDF's signature.
 

rust/datafusion/src/execution/context.rs

@@ -36,7 +36,7 @@ use crate::datasource::parquet::ParquetTable;
 use crate::datasource::TableProvider;
 use crate::error::{ExecutionError, Result};
 use crate::execution::dataframe_impl::DataFrameImpl;
-use crate::logical_plan::{Expr, FunctionRegistry, LogicalPlan, LogicalPlanBuilder};
+use crate::logical_plan::{FunctionRegistry, LogicalPlan, LogicalPlanBuilder};
 use crate::optimizer::filter_push_down::FilterPushDown;
 use crate::optimizer::optimizer::OptimizerRule;
 use crate::optimizer::projection_push_down::ProjectionPushDown;
@@ -481,17 +481,14 @@ impl FunctionRegistry for ExecutionContextState {
  self.scalar_functions.keys().cloned().collect()
  }
 
- fn udf(&self, name: &str, args: Vec<Expr>) -> Result<Expr> {
+ fn udf(&self, name: &str) -> Result<&ScalarUDF> {
  let result = self.scalar_functions.get(name);
  if result.is_none() {
  Err(ExecutionError::General(
  format!("There is no UDF named \"{}\" in the registry", name).to_string(),
  ))
  } else {
- Ok(Expr::ScalarUDF {
- fun: result.unwrap().clone(),
- args,
- })
+ Ok(result.unwrap())
  }
  }
 }
@@ -1051,7 +1048,7 @@ mod tests {
  .project(vec![
  col("a"),
  col("b"),
- ctx.registry().udf("my_add", vec![col("a"), col("b")])?,
+ ctx.registry().udf("my_add")?.call(vec![col("a"), col("b")]),
  ])?
  .build()?;
 

rust/datafusion/src/execution/dataframe_impl.rs

@@ -259,7 +259,7 @@ mod tests {
 
  let f = df.registry();
 
- let df = df.select(vec![f.udf("my_fn", vec![col("c12")])?])?;
+ let df = df.select(vec![f.udf("my_fn")?.call(vec![col("c12")])])?;
  let plan = df.to_logical_plan();
 
  // build query using SQL

rust/datafusion/src/logical_plan/mod.rs

@@ -1133,8 +1133,8 @@ pub trait FunctionRegistry {
  /// Set of all available udfs.
  fn udfs(&self) -> HashSet<String>;
 
- /// Constructs a logical expression with a call to the udf.
- fn udf(&self, name: &str, args: Vec<Expr>) -> Result<Expr>;
+ /// Returns a reference to the udf named `name`.
+ fn udf(&self, name: &str) -> Result<&ScalarUDF>;
 }
 
 /// Builder for logical plans

rust/datafusion/src/physical_plan/udf.rs

@@ -23,7 +23,7 @@ use std::fmt;
 use arrow::datatypes::Schema;
 
 use crate::error::Result;
-use crate::physical_plan::PhysicalExpr;
+use crate::{logical_plan::Expr, physical_plan::PhysicalExpr};
 
 use super::{
  functions::{
@@ -71,6 +71,15 @@ impl ScalarUDF {
  fun: fun.clone(),
  }
  }
+
+ /// creates a logical expression with a call of the UDF
+ /// This utility allows using the UDF without requiring access to the registry.
+ pub fn call(&self, args: Vec<Expr>) -> Expr {
+ Expr::ScalarUDF {
+ fun: Arc::new(self.clone()),
+ args,
+ }
+ }
 }
 
 /// Create a physical expression of the UDF.

rust/datafusion/examples/simple_udf.rs

@@ -111,20 +111,28 @@ fn main() -> Result<()> {
  pow,
  );
 
- // finally, register the UDF
- ctx.register_udf(pow);
+ // at this point, we can use it or register it, depending on the use-case:
+ // * if the UDF is expected to be used throughout the program in different contexts,
+ // we can register it, and call it later:
+ ctx.register_udf(pow.clone()); // clone is only required in this example because we show both usages
 
- // at this point, we can use it. Note that the code below can be in a
-// scope on which we do not have access to `pow`.
+ // * if the UDF is expected to be used directly in the scope, `.call` it directly:
+let expr = pow.call(vec![col("a"), col("b")]);
 
  // get a DataFrame from the context
  let df = ctx.table("t")?;
 
- // get the udf registry.
- let f = df.registry();
-
- // equivalent to `'SELECT pow(a, b) FROM t'`
- let df = df.select(vec![f.udf("pow", vec![col("a"), col("b")])?])?;
+ // if we do not have `pow` in the scope and we registered it, we can get it from the registry
+ let pow = df.registry().udf("pow")?;
+ // equivalent to expr
+ let expr1 = pow.call(vec![col("a"), col("b")]);
+
+ // equivalent to `'SELECT pow(a, b), pow(a, b) AS pow1 FROM t'`
+ let df = df.select(vec![
+ expr,
+ // alias so that they have different column names
+ expr1.alias("pow1"),
+ ])?;
 
  // note that "b" is f32, not f64. DataFusion coerces the types to match the UDF's signature.
 

rust/datafusion/src/execution/context.rs

@@ -36,7 +36,7 @@ use crate::datasource::parquet::ParquetTable;
 use crate::datasource::TableProvider;
 use crate::error::{ExecutionError, Result};
 use crate::execution::dataframe_impl::DataFrameImpl;
-use crate::logical_plan::{Expr, FunctionRegistry, LogicalPlan, LogicalPlanBuilder};
+use crate::logical_plan::{FunctionRegistry, LogicalPlan, LogicalPlanBuilder};
 use crate::optimizer::filter_push_down::FilterPushDown;
 use crate::optimizer::optimizer::OptimizerRule;
 use crate::optimizer::projection_push_down::ProjectionPushDown;
@@ -481,17 +481,14 @@ impl FunctionRegistry for ExecutionContextState {
  self.scalar_functions.keys().cloned().collect()
  }
 
- fn udf(&self, name: &str, args: Vec<Expr>) -> Result<Expr> {
+ fn udf(&self, name: &str) -> Result<&ScalarUDF> {
  let result = self.scalar_functions.get(name);
  if result.is_none() {
  Err(ExecutionError::General(
  format!("There is no UDF named \"{}\" in the registry", name).to_string(),
  ))
  } else {
- Ok(Expr::ScalarUDF {
- fun: result.unwrap().clone(),
- args,
- })
+ Ok(result.unwrap())
  }
  }
 }
@@ -1051,7 +1048,7 @@ mod tests {
  .project(vec![
  col("a"),
  col("b"),
- ctx.registry().udf("my_add", vec![col("a"), col("b")])?,
+ ctx.registry().udf("my_add")?.call(vec![col("a"), col("b")]),
  ])?
  .build()?;
 

rust/datafusion/src/execution/dataframe_impl.rs

@@ -259,7 +259,7 @@ mod tests {
 
  let f = df.registry();
 
- let df = df.select(vec![f.udf("my_fn", vec![col("c12")])?])?;
+ let df = df.select(vec![f.udf("my_fn")?.call(vec![col("c12")])])?;
  let plan = df.to_logical_plan();
 
  // build query using SQL

rust/datafusion/src/logical_plan/mod.rs

@@ -1133,8 +1133,8 @@ pub trait FunctionRegistry {
  /// Set of all available udfs.
  fn udfs(&self) -> HashSet<String>;
 
- /// Constructs a logical expression with a call to the udf.
- fn udf(&self, name: &str, args: Vec<Expr>) -> Result<Expr>;
+ /// Returns a reference to the udf named `name`.
+ fn udf(&self, name: &str) -> Result<&ScalarUDF>;
 }
 
 /// Builder for logical plans

rust/datafusion/src/physical_plan/udf.rs

@@ -23,7 +23,7 @@ use std::fmt;
 use arrow::datatypes::Schema;
 
 use crate::error::Result;
-use crate::physical_plan::PhysicalExpr;
+use crate::{logical_plan::Expr, physical_plan::PhysicalExpr};
 
 use super::{
  functions::{
@@ -71,6 +71,15 @@ impl ScalarUDF {
  fun: fun.clone(),
  }
  }
+
+ /// creates a logical expression with a call of the UDF
+ /// This utility allows using the UDF without requiring access to the registry.
+ pub fn call(&self, args: Vec<Expr>) -> Expr {
+ Expr::ScalarUDF {
+ fun: Arc::new(self.clone()),
+ args,
+ }
+ }
 }
 
 /// Create a physical expression of the UDF.

rust/datafusion/examples/simple_udf.rs

@@ -111,20 +111,28 @@ fn main() -> Result<()> {
  pow,
  );
 
- // finally, register the UDF
- ctx.register_udf(pow);
+ // at this point, we can use it or register it, depending on the use-case:
+ // * if the UDF is expected to be used throughout the program in different contexts,
+ // we can register it, and call it later:
+ ctx.register_udf(pow.clone()); // clone is only required in this example because we show both usages
 
- // at this point, we can use it. Note that the code below can be in a
-// scope on which we do not have access to `pow`.
+ // * if the UDF is expected to be used directly in the scope, `.call` it directly:
+let expr = pow.call(vec![col("a"), col("b")]);
 
  // get a DataFrame from the context
  let df = ctx.table("t")?;
 
- // get the udf registry.
- let f = df.registry();
-
- // equivalent to `'SELECT pow(a, b) FROM t'`
- let df = df.select(vec![f.udf("pow", vec![col("a"), col("b")])?])?;
+ // if we do not have `pow` in the scope and we registered it, we can get it from the registry
+ let pow = df.registry().udf("pow")?;
+ // equivalent to expr
+ let expr1 = pow.call(vec![col("a"), col("b")]);
+
+ // equivalent to `'SELECT pow(a, b), pow(a, b) AS pow1 FROM t'`
+ let df = df.select(vec![
+ expr,
+ // alias so that they have different column names
+ expr1.alias("pow1"),
+ ])?;
 
  // note that "b" is f32, not f64. DataFusion coerces the types to match the UDF's signature.
 

rust/datafusion/src/execution/context.rs

@@ -36,7 +36,7 @@ use crate::datasource::parquet::ParquetTable;
 use crate::datasource::TableProvider;
 use crate::error::{ExecutionError, Result};
 use crate::execution::dataframe_impl::DataFrameImpl;
-use crate::logical_plan::{Expr, FunctionRegistry, LogicalPlan, LogicalPlanBuilder};
+use crate::logical_plan::{FunctionRegistry, LogicalPlan, LogicalPlanBuilder};
 use crate::optimizer::filter_push_down::FilterPushDown;
 use crate::optimizer::optimizer::OptimizerRule;
 use crate::optimizer::projection_push_down::ProjectionPushDown;
@@ -481,17 +481,14 @@ impl FunctionRegistry for ExecutionContextState {
  self.scalar_functions.keys().cloned().collect()
  }
 
- fn udf(&self, name: &str, args: Vec<Expr>) -> Result<Expr> {
+ fn udf(&self, name: &str) -> Result<&ScalarUDF> {
  let result = self.scalar_functions.get(name);
  if result.is_none() {
  Err(ExecutionError::General(
  format!("There is no UDF named \"{}\" in the registry", name).to_string(),
  ))
  } else {
- Ok(Expr::ScalarUDF {
- fun: result.unwrap().clone(),
- args,
- })
+ Ok(result.unwrap())
  }
  }
 }
@@ -1051,7 +1048,7 @@ mod tests {
  .project(vec![
  col("a"),
  col("b"),
- ctx.registry().udf("my_add", vec![col("a"), col("b")])?,
+ ctx.registry().udf("my_add")?.call(vec![col("a"), col("b")]),
  ])?
  .build()?;
 

rust/datafusion/src/execution/dataframe_impl.rs

@@ -259,7 +259,7 @@ mod tests {
 
  let f = df.registry();
 
- let df = df.select(vec![f.udf("my_fn", vec![col("c12")])?])?;
+ let df = df.select(vec![f.udf("my_fn")?.call(vec![col("c12")])])?;
  let plan = df.to_logical_plan();
 
  // build query using SQL

rust/datafusion/src/logical_plan/mod.rs

@@ -1133,8 +1133,8 @@ pub trait FunctionRegistry {
  /// Set of all available udfs.
  fn udfs(&self) -> HashSet<String>;
 
- /// Constructs a logical expression with a call to the udf.
- fn udf(&self, name: &str, args: Vec<Expr>) -> Result<Expr>;
+ /// Returns a reference to the udf named `name`.
+ fn udf(&self, name: &str) -> Result<&ScalarUDF>;
 }
 
 /// Builder for logical plans

rust/datafusion/src/physical_plan/udf.rs

@@ -23,7 +23,7 @@ use std::fmt;
 use arrow::datatypes::Schema;
 
 use crate::error::Result;
-use crate::physical_plan::PhysicalExpr;
+use crate::{logical_plan::Expr, physical_plan::PhysicalExpr};
 
 use super::{
  functions::{
@@ -71,6 +71,15 @@ impl ScalarUDF {
  fun: fun.clone(),
  }
  }
+
+ /// creates a logical expression with a call of the UDF
+ /// This utility allows using the UDF without requiring access to the registry.
+ pub fn call(&self, args: Vec<Expr>) -> Expr {
+ Expr::ScalarUDF {
+ fun: Arc::new(self.clone()),
+ args,
+ }
+ }
 }
 
 /// Create a physical expression of the UDF.

rust/datafusion/examples/simple_udf.rs

@@ -111,20 +111,28 @@ fn main() -> Result<()> {
  pow,
  );
 
- // finally, register the UDF
- ctx.register_udf(pow);
+ // at this point, we can use it or register it, depending on the use-case:
+ // * if the UDF is expected to be used throughout the program in different contexts,
+ // we can register it, and call it later:
+ ctx.register_udf(pow.clone()); // clone is only required in this example because we show both usages
 
- // at this point, we can use it. Note that the code below can be in a
-// scope on which we do not have access to `pow`.
+ // * if the UDF is expected to be used directly in the scope, `.call` it directly:
+let expr = pow.call(vec![col("a"), col("b")]);
 
  // get a DataFrame from the context
  let df = ctx.table("t")?;
 
- // get the udf registry.
- let f = df.registry();
-
- // equivalent to `'SELECT pow(a, b) FROM t'`
- let df = df.select(vec![f.udf("pow", vec![col("a"), col("b")])?])?;
+ // if we do not have `pow` in the scope and we registered it, we can get it from the registry
+ let pow = df.registry().udf("pow")?;
+ // equivalent to expr
+ let expr1 = pow.call(vec![col("a"), col("b")]);
+
+ // equivalent to `'SELECT pow(a, b), pow(a, b) AS pow1 FROM t'`
+ let df = df.select(vec![
+ expr,
+ // alias so that they have different column names
+ expr1.alias("pow1"),
+ ])?;
 
  // note that "b" is f32, not f64. DataFusion coerces the types to match the UDF's signature.
 

rust/datafusion/src/execution/context.rs

@@ -36,7 +36,7 @@ use crate::datasource::parquet::ParquetTable;
 use crate::datasource::TableProvider;
 use crate::error::{ExecutionError, Result};
 use crate::execution::dataframe_impl::DataFrameImpl;
-use crate::logical_plan::{Expr, FunctionRegistry, LogicalPlan, LogicalPlanBuilder};
+use crate::logical_plan::{FunctionRegistry, LogicalPlan, LogicalPlanBuilder};
 use crate::optimizer::filter_push_down::FilterPushDown;
 use crate::optimizer::optimizer::OptimizerRule;
 use crate::optimizer::projection_push_down::ProjectionPushDown;
@@ -481,17 +481,14 @@ impl FunctionRegistry for ExecutionContextState {
  self.scalar_functions.keys().cloned().collect()
  }
 
- fn udf(&self, name: &str, args: Vec<Expr>) -> Result<Expr> {
+ fn udf(&self, name: &str) -> Result<&ScalarUDF> {
  let result = self.scalar_functions.get(name);
  if result.is_none() {
  Err(ExecutionError::General(
  format!("There is no UDF named \"{}\" in the registry", name).to_string(),
  ))
  } else {
- Ok(Expr::ScalarUDF {
- fun: result.unwrap().clone(),
- args,
- })
+ Ok(result.unwrap())
  }
  }
 }
@@ -1051,7 +1048,7 @@ mod tests {
  .project(vec![
  col("a"),
  col("b"),
- ctx.registry().udf("my_add", vec![col("a"), col("b")])?,
+ ctx.registry().udf("my_add")?.call(vec![col("a"), col("b")]),
  ])?
  .build()?;
 

rust/datafusion/src/execution/dataframe_impl.rs

@@ -259,7 +259,7 @@ mod tests {
 
  let f = df.registry();
 
- let df = df.select(vec![f.udf("my_fn", vec![col("c12")])?])?;
+ let df = df.select(vec![f.udf("my_fn")?.call(vec![col("c12")])])?;
  let plan = df.to_logical_plan();
 
  // build query using SQL

rust/datafusion/src/logical_plan/mod.rs

@@ -1133,8 +1133,8 @@ pub trait FunctionRegistry {
  /// Set of all available udfs.
  fn udfs(&self) -> HashSet<String>;
 
- /// Constructs a logical expression with a call to the udf.
- fn udf(&self, name: &str, args: Vec<Expr>) -> Result<Expr>;
+ /// Returns a reference to the udf named `name`.
+ fn udf(&self, name: &str) -> Result<&ScalarUDF>;
 }
 
 /// Builder for logical plans

rust/datafusion/src/physical_plan/udf.rs

@@ -23,7 +23,7 @@ use std::fmt;
 use arrow::datatypes::Schema;
 
 use crate::error::Result;
-use crate::physical_plan::PhysicalExpr;
+use crate::{logical_plan::Expr, physical_plan::PhysicalExpr};
 
 use super::{
  functions::{
@@ -71,6 +71,15 @@ impl ScalarUDF {
  fun: fun.clone(),
  }
  }
+
+ /// creates a logical expression with a call of the UDF
+ /// This utility allows using the UDF without requiring access to the registry.
+ pub fn call(&self, args: Vec<Expr>) -> Expr {
+ Expr::ScalarUDF {
+ fun: Arc::new(self.clone()),
+ args,
+ }
+ }
 }
 
 /// Create a physical expression of the UDF.

rust/datafusion/examples/simple_udf.rs

@@ -111,20 +111,28 @@ fn main() -> Result<()> {
  pow,
  );
 
- // finally, register the UDF
- ctx.register_udf(pow);
+ // at this point, we can use it or register it, depending on the use-case:
+ // * if the UDF is expected to be used throughout the program in different contexts,
+ // we can register it, and call it later:
+ ctx.register_udf(pow.clone()); // clone is only required in this example because we show both usages
 
- // at this point, we can use it. Note that the code below can be in a
-// scope on which we do not have access to `pow`.
+ // * if the UDF is expected to be used directly in the scope, `.call` it directly:
+let expr = pow.call(vec![col("a"), col("b")]);
 
  // get a DataFrame from the context
  let df = ctx.table("t")?;
 
- // get the udf registry.
- let f = df.registry();
-
- // equivalent to `'SELECT pow(a, b) FROM t'`
- let df = df.select(vec![f.udf("pow", vec![col("a"), col("b")])?])?;
+ // if we do not have `pow` in the scope and we registered it, we can get it from the registry
+ let pow = df.registry().udf("pow")?;
+ // equivalent to expr
+ let expr1 = pow.call(vec![col("a"), col("b")]);
+
+ // equivalent to `'SELECT pow(a, b), pow(a, b) AS pow1 FROM t'`
+ let df = df.select(vec![
+ expr,
+ // alias so that they have different column names
+ expr1.alias("pow1"),
+ ])?;
 
  // note that "b" is f32, not f64. DataFusion coerces the types to match the UDF's signature.
 

rust/datafusion/src/execution/context.rs

@@ -36,7 +36,7 @@ use crate::datasource::parquet::ParquetTable;
 use crate::datasource::TableProvider;
 use crate::error::{ExecutionError, Result};
 use crate::execution::dataframe_impl::DataFrameImpl;
-use crate::logical_plan::{Expr, FunctionRegistry, LogicalPlan, LogicalPlanBuilder};
+use crate::logical_plan::{FunctionRegistry, LogicalPlan, LogicalPlanBuilder};
 use crate::optimizer::filter_push_down::FilterPushDown;
 use crate::optimizer::optimizer::OptimizerRule;
 use crate::optimizer::projection_push_down::ProjectionPushDown;
@@ -481,17 +481,14 @@ impl FunctionRegistry for ExecutionContextState {
  self.scalar_functions.keys().cloned().collect()
  }
 
- fn udf(&self, name: &str, args: Vec<Expr>) -> Result<Expr> {
+ fn udf(&self, name: &str) -> Result<&ScalarUDF> {
  let result = self.scalar_functions.get(name);
  if result.is_none() {
  Err(ExecutionError::General(
  format!("There is no UDF named \"{}\" in the registry", name).to_string(),
  ))
  } else {
- Ok(Expr::ScalarUDF {
- fun: result.unwrap().clone(),
- args,
- })
+ Ok(result.unwrap())
  }
  }
 }
@@ -1051,7 +1048,7 @@ mod tests {
  .project(vec![
  col("a"),
  col("b"),
- ctx.registry().udf("my_add", vec![col("a"), col("b")])?,
+ ctx.registry().udf("my_add")?.call(vec![col("a"), col("b")]),
  ])?
  .build()?;
 

rust/datafusion/src/execution/dataframe_impl.rs

@@ -259,7 +259,7 @@ mod tests {
 
  let f = df.registry();
 
- let df = df.select(vec![f.udf("my_fn", vec![col("c12")])?])?;
+ let df = df.select(vec![f.udf("my_fn")?.call(vec![col("c12")])])?;
  let plan = df.to_logical_plan();
 
  // build query using SQL

rust/datafusion/src/logical_plan/mod.rs

@@ -1133,8 +1133,8 @@ pub trait FunctionRegistry {
  /// Set of all available udfs.
  fn udfs(&self) -> HashSet<String>;
 
- /// Constructs a logical expression with a call to the udf.
- fn udf(&self, name: &str, args: Vec<Expr>) -> Result<Expr>;
+ /// Returns a reference to the udf named `name`.
+ fn udf(&self, name: &str) -> Result<&ScalarUDF>;
 }
 
 /// Builder for logical plans

rust/datafusion/src/physical_plan/udf.rs

@@ -23,7 +23,7 @@ use std::fmt;
 use arrow::datatypes::Schema;
 
 use crate::error::Result;
-use crate::physical_plan::PhysicalExpr;
+use crate::{logical_plan::Expr, physical_plan::PhysicalExpr};
 
 use super::{
  functions::{
@@ -71,6 +71,15 @@ impl ScalarUDF {
  fun: fun.clone(),
  }
  }
+
+ /// creates a logical expression with a call of the UDF
+ /// This utility allows using the UDF without requiring access to the registry.
+ pub fn call(&self, args: Vec<Expr>) -> Expr {
+ Expr::ScalarUDF {
+ fun: Arc::new(self.clone()),
+ args,
+ }
+ }
 }
 
 /// Create a physical expression of the UDF.