[Rust] [DataFusion] Add support for scalar UDFs

[asfimport]

As a user, I would like to be able to define my own functions and then use them in SQL statements.

 

Reporter: Andy Grove / @andygrove
Assignee: Andy Grove / @andygrove

PRs and other links:

• GitHub Pull Request #6749

_{Note: This issue was originally created as ARROW-6947. Please see the migration documentation for further details.}

[asfimport]

Kyle McCarthy / @kyle-mccarthy:
I would like to provide some assistance with implementing UDFs. Is the idea to implement them similarly to gandiva and generate LLVM IR or WebAssembly? Or to take a different approach and provide a trait for users to implement for their UDF and then register?

[asfimport]

Jorge Leitão / @jorgecarleitao:
@andygrove, can I take a shoot at this?

My strategy to tackle this:

1. add support for unary udfs

2. support binary udfs

3. add macro to support arbitrary argument udfs (I do not see an easy way to implement variadic generics)

   From what I gathered, we need to

4. Implement a new kernel (ala arrow::compute::kernels::arithmetic) to execute a closure

5. Implement a generic expression that implements PhysicalExpr

6. Implement a FunctionProvider and add it to ExecutionContext (like the attribute datasources)

7. bits and pieces for integration

   The design I was planning for the expression:

    

   pub struct UnaryFunctionExpr<T, R>
   where
    T: ArrowPrimitiveType, // argument type
    R: ArrowPrimitiveType, // return type
   {
    name: String,
   /// name of the function
    func:
    Arc<dyn Fn(T::Native) -> arrow::error::Result<R::Native> + Sync + Send + 'static>,
    arg: Arc<dyn PhysicalExpr>,
    arg_type: DataType,
    return_type: DataType,
   }

    

   for the kernel:

   pub fn unary_op<T, R>(
    op: &dyn Fn(T::Native) -> Result<R::Native>,
    arg: &PrimitiveArray<T>,
   ) -> Result<PrimitiveArray<R>>

    

[asfimport]

Jorge Leitão / @jorgecarleitao:
I finished a first draft of this that I am minimally happy with, see https://github.com/jorgecarleitao/arrow/commit/0c77ff531cee191712d08633080de9caed786b62 . Not all function signatures are implemented on purpose, as I see we need to take a design decision with respect to them.

The API is used as shown in this test:

// declare a function
let f = FunctionSignature::UInt64UInt64(Arc::new(Box::new(|a: &u64| Ok(a * a))));

// register function
ctx.register_function("pow2", f);

let results = collect(&mut ctx, "SELECT c2, pow2(c2) FROM test")?;

The important aspect here is that FunctionSignature::UInt64UInt64 is an enum variant of the enum FunctionSignature (this is the design choice that we need to discuss).

I took this design decision based on the following hypothesis:

1. there are limited use-cases for UDFs with may arguments

2. there are a limited number of different types in arrow (25ish)

   under these hypothesis, we can enumerate all variations on the relevant match. I am not entirely convinced that this is the correct approach, but it is the one that is the most aligned with enumerating all types in match currently in the code base and does not resort to dynamic typing.

   An alternative design is to use std::any::Any. I can also play around with this and see how it feels. Assuming that we can do it, I suspect that the trade-off is between run-time and compilation time+binary size.

[asfimport]

Andy Grove / @andygrove:
Sorry for taking so long to get to this. I think it will be important for this to be more dynamic, where the function receives a batch of data or a vector of columns to operate on so we can have a function signature along the lines of:

trait ScalarFunction {
  fn evaluate(args: Vec<Array>) -> Result<Array>;
} 

[asfimport]

Andy Grove / @andygrove:
It's important to remember that these UDFs will operate on columns, not rows, so although there is overhead in checking the data types of the input arguments, this happens once per batch and not once per row.

[asfimport]

Andy Grove / @andygrove:
Issue resolved by pull request 6749
#6749