datafusion-examples/examples/advanced_udf.rs - datafusion - Git at Google

 // 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 std::any::Any;
 use std::sync::Arc;

 use arrow::array::{
     new_null_array, Array, ArrayRef, AsArray, Float32Array, Float64Array,
 };
 use arrow::compute;
 use arrow::datatypes::{DataType, Float64Type};
 use arrow::record_batch::RecordBatch;
 use datafusion::error::Result;
 use datafusion::logical_expr::Volatility;
 use datafusion::prelude::*;
 use datafusion_common::{internal_err, ScalarValue};
 use datafusion_expr::sort_properties::{ExprProperties, SortProperties};
 use datafusion_expr::{ColumnarValue, ScalarUDF, ScalarUDFImpl, Signature};

 /// This example shows how to use the full ScalarUDFImpl API to implement a user
 /// defined function. As in the `simple_udf.rs` example, this struct implements
 /// a function that takes two arguments and returns the first argument raised to
 /// the power of the second argument `a^b`.
 ///
 /// To do so, we must implement the `ScalarUDFImpl` trait.
 #[derive(Debug, Clone)]
 struct PowUdf {
     signature: Signature,
     aliases: Vec<String>,
 }

 impl PowUdf {
     /// Create a new instance of the `PowUdf` struct
     fn new() -> Self {
         Self {
             signature: Signature::exact(
                 // this function will always take two arguments of type f64
                 vec![DataType::Float64, DataType::Float64],
                 // this function is deterministic and will always return the same
                 // result for the same input
                 Volatility::Immutable,
             ),
             // we will also add an alias of "my_pow"
             aliases: vec!["my_pow".to_string()],
         }
     }
 }

 impl ScalarUDFImpl for PowUdf {
     /// We implement as_any so that we can downcast the ScalarUDFImpl trait object
     fn as_any(&self) -> &dyn Any {
         self
     }

     /// Return the name of this function
     fn name(&self) -> &str {
         "pow"
     }

     /// Return the "signature" of this function -- namely what types of arguments it will take
     fn signature(&self) -> &Signature {
         &self.signature
     }

     /// What is the type of value that will be returned by this function? In
     /// this case it will always be a constant value, but it could also be a
     /// function of the input types.
     fn return_type(&self, _arg_types: &[DataType]) -> Result<DataType> {
         Ok(DataType::Float64)
     }

     /// This is the function that actually calculates the results.
     ///
     /// This is the same way that functions built into DataFusion are invoked,
     /// which permits important special cases when one or both of the arguments
     /// are single values (constants). For example `pow(a, 2)`
     ///
     /// However, it also means the implementation is more complex than when
     /// using `create_udf`.
     fn invoke(&self, args: &[ColumnarValue]) -> Result<ColumnarValue> {
         // DataFusion has arranged for the correct inputs to be passed to this
         // function, but we check again to make sure
         assert_eq!(args.len(), 2);
         let (base, exp) = (&args[0], &args[1]);
         assert_eq!(base.data_type(), DataType::Float64);
         assert_eq!(exp.data_type(), DataType::Float64);

         match (base, exp) {
             // For demonstration purposes we also implement the scalar / scalar
             // case here, but it is not typically required for high performance.
             //
             // For performance it is most important to optimize cases where at
             // least one argument is an array. If all arguments are constants,
             // the DataFusion expression simplification logic will often invoke
             // this path once during planning, and simply use the result during
             // execution.
             (
                 ColumnarValue::Scalar(ScalarValue::Float64(base)),
                 ColumnarValue::Scalar(ScalarValue::Float64(exp)),
             ) => {
                 // compute the output. Note DataFusion treats `None` as NULL.
                 let res = match (base, exp) {
                     (Some(base), Some(exp)) => Some(base.powf(*exp)),
                     // one or both arguments were NULL
                     _ => None,
                 };
                 Ok(ColumnarValue::Scalar(ScalarValue::from(res)))
             }
             // special case if the exponent is a constant
             (
                 ColumnarValue::Array(base_array),
                 ColumnarValue::Scalar(ScalarValue::Float64(exp)),
             ) => {
                 let result_array = match exp {
                     // a ^ null = null
                     None => new_null_array(base_array.data_type(), base_array.len()),
                     // a ^ exp
                     Some(exp) => {
                         // DataFusion has ensured both arguments are Float64:
                         let base_array = base_array.as_primitive::<Float64Type>();
                         // calculate the result for every row. The `unary`
                         // kernel creates very fast "vectorized" code and
                         // handles things like null values for us.
                         let res: Float64Array =
                             compute::unary(base_array, |base| base.powf(*exp));
                         Arc::new(res)
                     }
                 };
                 Ok(ColumnarValue::Array(result_array))
             }

             // special case if the base is a constant (note this code is quite
             // similar to the previous case, so we omit comments)
             (
                 ColumnarValue::Scalar(ScalarValue::Float64(base)),
                 ColumnarValue::Array(exp_array),
             ) => {
                 let res = match base {
                     None => new_null_array(exp_array.data_type(), exp_array.len()),
                     Some(base) => {
                         let exp_array = exp_array.as_primitive::<Float64Type>();
                         let res: Float64Array =
                             compute::unary(exp_array, |exp| base.powf(exp));
                         Arc::new(res)
                     }
                 };
                 Ok(ColumnarValue::Array(res))
             }
             // Both arguments are arrays so we have to perform the calculation for every row
             (ColumnarValue::Array(base_array), ColumnarValue::Array(exp_array)) => {
                 let res: Float64Array = compute::binary(
                     base_array.as_primitive::<Float64Type>(),
                     exp_array.as_primitive::<Float64Type>(),
                     |base, exp| base.powf(exp),
                 )?;
                 Ok(ColumnarValue::Array(Arc::new(res)))
             }
             // if the types were not float, it is a bug in DataFusion
             _ => {
                 internal_err!("Invalid argument types to pow function")
             }
         }
     }

     /// We will also add an alias of "my_pow"
     fn aliases(&self) -> &[String] {
         &self.aliases
     }

     fn output_ordering(&self, input: &[ExprProperties]) -> Result<SortProperties> {
         // The POW function preserves the order of its argument.
         Ok(input[0].sort_properties)
     }
 }

 /// In this example we register `PowUdf` as a user defined function
 /// and invoke it via the DataFrame API and SQL
 #[tokio::main]
 async fn main() -> Result<()> {
     let ctx = create_context()?;

     // create the UDF
     let pow = ScalarUDF::from(PowUdf::new());

     // register the UDF with the context so it can be invoked by name and from SQL
     ctx.register_udf(pow.clone());

     // get a DataFrame from the context for scanning the "t" table
     let df = ctx.table("t").await?;

     // Call pow(a, 10) using the DataFrame API
     let df = df.select(vec![pow.call(vec![col("a"), lit(10i32)])])?;

     // note that the second argument is passed as an i32, not f64. DataFusion
     // automatically coerces the types to match the UDF's defined signature.

     // print the results
     df.show().await?;

     // You can also invoke both pow(2, 10)  and its alias my_pow(a, b) using SQL
     let sql_df = ctx.sql("SELECT pow(2, 10), my_pow(a, b) FROM t").await?;
     sql_df.show().await?;

     Ok(())
 }

 /// create local execution context with an in-memory table:
 ///
 /// ```text
 /// +-----+-----+
 /// | a   | b   |
 /// +-----+-----+
 /// | 2.1 | 1.0 |
 /// | 3.1 | 2.0 |
 /// | 4.1 | 3.0 |
 /// | 5.1 | 4.0 |
 /// +-----+-----+
 /// ```
 fn create_context() -> Result<SessionContext> {
     // define data.
     let a: ArrayRef = Arc::new(Float32Array::from(vec![2.1, 3.1, 4.1, 5.1]));
     let b: ArrayRef = Arc::new(Float64Array::from(vec![1.0, 2.0, 3.0, 4.0]));
     let batch = RecordBatch::try_from_iter(vec![("a", a), ("b", b)])?;

     // declare a new context. In Spark API, this corresponds to a new SparkSession
     let ctx = SessionContext::new();

     // declare a table in memory. In Spark API, this corresponds to createDataFrame(...).
     ctx.register_batch("t", batch)?;
     Ok(ctx)
 }
	// 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 std::any::Any;
	use std::sync::Arc;

	use arrow::array::{
	new_null_array, Array, ArrayRef, AsArray, Float32Array, Float64Array,
	};
	use arrow::compute;
	use arrow::datatypes::{DataType, Float64Type};
	use arrow::record_batch::RecordBatch;
	use datafusion::error::Result;
	use datafusion::logical_expr::Volatility;
	use datafusion::prelude::*;
	use datafusion_common::{internal_err, ScalarValue};
	use datafusion_expr::sort_properties::{ExprProperties, SortProperties};
	use datafusion_expr::{ColumnarValue, ScalarUDF, ScalarUDFImpl, Signature};

	/// This example shows how to use the full ScalarUDFImpl API to implement a user
	/// defined function. As in the `simple_udf.rs` example, this struct implements
	/// a function that takes two arguments and returns the first argument raised to
	/// the power of the second argument `a^b`.
	///
	/// To do so, we must implement the `ScalarUDFImpl` trait.
	#[derive(Debug, Clone)]
	struct PowUdf {
	signature: Signature,
	aliases: Vec<String>,
	}

	impl PowUdf {
	/// Create a new instance of the `PowUdf` struct
	fn new() -> Self {
	Self {
	signature: Signature::exact(
	// this function will always take two arguments of type f64
	vec![DataType::Float64, DataType::Float64],
	// this function is deterministic and will always return the same
	// result for the same input
	Volatility::Immutable,
	),
	// we will also add an alias of "my_pow"
	aliases: vec!["my_pow".to_string()],
	}
	}
	}

	impl ScalarUDFImpl for PowUdf {
	/// We implement as_any so that we can downcast the ScalarUDFImpl trait object
	fn as_any(&self) -> &dyn Any {
	self
	}

	/// Return the name of this function
	fn name(&self) -> &str {
	"pow"
	}

	/// Return the "signature" of this function -- namely what types of arguments it will take
	fn signature(&self) -> &Signature {
	&self.signature
	}

	/// What is the type of value that will be returned by this function? In
	/// this case it will always be a constant value, but it could also be a
	/// function of the input types.
	fn return_type(&self, _arg_types: &[DataType]) -> Result<DataType> {
	Ok(DataType::Float64)
	}

	/// This is the function that actually calculates the results.
	///
	/// This is the same way that functions built into DataFusion are invoked,
	/// which permits important special cases when one or both of the arguments
	/// are single values (constants). For example `pow(a, 2)`
	///
	/// However, it also means the implementation is more complex than when
	/// using `create_udf`.
	fn invoke(&self, args: &[ColumnarValue]) -> Result<ColumnarValue> {
	// DataFusion has arranged for the correct inputs to be passed to this
	// function, but we check again to make sure
	assert_eq!(args.len(), 2);
	let (base, exp) = (&args[0], &args[1]);
	assert_eq!(base.data_type(), DataType::Float64);
	assert_eq!(exp.data_type(), DataType::Float64);

	match (base, exp) {
	// For demonstration purposes we also implement the scalar / scalar
	// case here, but it is not typically required for high performance.
	//
	// For performance it is most important to optimize cases where at
	// least one argument is an array. If all arguments are constants,
	// the DataFusion expression simplification logic will often invoke
	// this path once during planning, and simply use the result during
	// execution.
	(
	ColumnarValue::Scalar(ScalarValue::Float64(base)),
	ColumnarValue::Scalar(ScalarValue::Float64(exp)),
	) => {
	// compute the output. Note DataFusion treats `None` as NULL.
	let res = match (base, exp) {
	(Some(base), Some(exp)) => Some(base.powf(*exp)),
	// one or both arguments were NULL
	_ => None,
	};
	Ok(ColumnarValue::Scalar(ScalarValue::from(res)))
	}
	// special case if the exponent is a constant
	(
	ColumnarValue::Array(base_array),
	ColumnarValue::Scalar(ScalarValue::Float64(exp)),
	) => {
	let result_array = match exp {
	// a ^ null = null
	None => new_null_array(base_array.data_type(), base_array.len()),
	// a ^ exp
	Some(exp) => {
	// DataFusion has ensured both arguments are Float64:
	let base_array = base_array.as_primitive::<Float64Type>();
	// calculate the result for every row. The `unary`
	// kernel creates very fast "vectorized" code and
	// handles things like null values for us.
	let res: Float64Array =
	compute::unary(base_array, \|base\| base.powf(*exp));
	Arc::new(res)
	}
	};
	Ok(ColumnarValue::Array(result_array))
	}

	// special case if the base is a constant (note this code is quite
	// similar to the previous case, so we omit comments)
	(
	ColumnarValue::Scalar(ScalarValue::Float64(base)),
	ColumnarValue::Array(exp_array),
	) => {
	let res = match base {
	None => new_null_array(exp_array.data_type(), exp_array.len()),
	Some(base) => {
	let exp_array = exp_array.as_primitive::<Float64Type>();
	let res: Float64Array =
	compute::unary(exp_array, \|exp\| base.powf(exp));
	Arc::new(res)
	}
	};
	Ok(ColumnarValue::Array(res))
	}
	// Both arguments are arrays so we have to perform the calculation for every row
	(ColumnarValue::Array(base_array), ColumnarValue::Array(exp_array)) => {
	let res: Float64Array = compute::binary(
	base_array.as_primitive::<Float64Type>(),
	exp_array.as_primitive::<Float64Type>(),
	\|base, exp\| base.powf(exp),
	)?;
	Ok(ColumnarValue::Array(Arc::new(res)))
	}
	// if the types were not float, it is a bug in DataFusion
	_ => {
	internal_err!("Invalid argument types to pow function")
	}
	}
	}

	/// We will also add an alias of "my_pow"
	fn aliases(&self) -> &[String] {
	&self.aliases
	}

	fn output_ordering(&self, input: &[ExprProperties]) -> Result<SortProperties> {
	// The POW function preserves the order of its argument.
	Ok(input[0].sort_properties)
	}
	}

	/// In this example we register `PowUdf` as a user defined function
	/// and invoke it via the DataFrame API and SQL
	#[tokio::main]
	async fn main() -> Result<()> {
	let ctx = create_context()?;

	// create the UDF
	let pow = ScalarUDF::from(PowUdf::new());

	// register the UDF with the context so it can be invoked by name and from SQL
	ctx.register_udf(pow.clone());

	// get a DataFrame from the context for scanning the "t" table
	let df = ctx.table("t").await?;

	// Call pow(a, 10) using the DataFrame API
	let df = df.select(vec![pow.call(vec![col("a"), lit(10i32)])])?;

	// note that the second argument is passed as an i32, not f64. DataFusion
	// automatically coerces the types to match the UDF's defined signature.

	// print the results
	df.show().await?;

	// You can also invoke both pow(2, 10) and its alias my_pow(a, b) using SQL
	let sql_df = ctx.sql("SELECT pow(2, 10), my_pow(a, b) FROM t").await?;
	sql_df.show().await?;

	Ok(())
	}

	/// create local execution context with an in-memory table:
	///
	/// ```text
	/// +-----+-----+
	/// \| a \| b \|
	/// +-----+-----+
	/// \| 2.1 \| 1.0 \|
	/// \| 3.1 \| 2.0 \|
	/// \| 4.1 \| 3.0 \|
	/// \| 5.1 \| 4.0 \|
	/// +-----+-----+
	/// ```
	fn create_context() -> Result<SessionContext> {
	// define data.
	let a: ArrayRef = Arc::new(Float32Array::from(vec![2.1, 3.1, 4.1, 5.1]));
	let b: ArrayRef = Arc::new(Float64Array::from(vec![1.0, 2.0, 3.0, 4.0]));
	let batch = RecordBatch::try_from_iter(vec![("a", a), ("b", b)])?;

	// declare a new context. In Spark API, this corresponds to a new SparkSession
	let ctx = SessionContext::new();

	// declare a table in memory. In Spark API, this corresponds to createDataFrame(...).
	ctx.register_batch("t", batch)?;
	Ok(ctx)
	}