Simple Functions

[findepi]

Design doc

• https://docs.google.com/document/d/1JyfbYfpmjvd3N3LEVAN7tX1D9ze6N98yGx3g4po_WoQ

Is your feature request related to a problem or challenge?

Verbosity

Currently implementing a scalar function is a pretty involved process. For example a simple function calculating greatest common divisor looks like this (a lot of detail omitted for brevity):

impl ScalarUDFImpl for GcdFunc {  
    fn name(&self) -> &str {
        "gcd"
    }

    fn return_type(&self, _arg_types: &[DataType]) -> Result<DataType> {
        Ok(Int64)
    }

    fn invoke(&self, args: &[ColumnarValue]) -> Result<ColumnarValue> {
        make_scalar_function(gcd, vec![])(args)
    }
}

/// Gcd SQL function
fn gcd(args: &[ArrayRef]) -> Result<ArrayRef> {
    match args[0].data_type() {
        Int64 => {
            let arg1 = downcast_arg!(&args[0], "x", Int64Array);
            let arg2 = downcast_arg!(&args[1], "y", Int64Array);

            Ok(arg1
                .iter()
                .zip(arg2.iter())
                .map(|(a1, a2)| match (a1, a2) {
                    (Some(a1), Some(a2)) => Ok(Some(compute_gcd(a1, a2)?)),
                    _ => Ok(None),
                })
                .collect::<Result<Int64Array>>()
                .map(Arc::new)? as ArrayRef)
        }
        other => exec_err!("Unsupported data type {other:?} for function gcd"),
    }
}

/// Computes greatest common divisor using Binary GCD algorithm.
pub fn compute_gcd(x: i64, y: i64) -> Result<i64> {
   ...
}

This function is still "relatively simple" because:

• it handles only int64 type (has no overloads)
  • a more complicated function would calculate return type based in input types in return_type and again in invoke
• doesn't not support constant x or constant y parameters
  • the make_scalar_function helper expands scalar constant argument into an array, thus gcd(x, 5) is less efficient than gcd(x, y) because it first needs to allocate a temporary buffer for 5s repeated times the batch length
• it doesn't handle dictionaries or run-end (REE / RLE) encoded values
  • i doubt people frequently calculate gcd on partitioning keys, but this is just an example
• it doesn't handle all-non-null case, so probably misses optimization/vectorization opportunity

It should be possible to express functions using simple row-level operations, because in query engines most function logic is structured like that anyway (compute_gcd in the example).

Local data structures

Currently DataFusion functions are singletons plugged into the execution engine. They have no way to store and reuse buffers or compiled regular expressions, etc.
Thread local storage is not an option because DataFusion, when embedded, does not control creation of application threads.

Describe the solution you'd like

Simple

It should be possible to separate function logic from the mechanics.
Exemplary GCD function needs to provide fn compute_gcd(x: i64, y: i64) -> Result<i64> and the rest of boilerplate should not be hand-written for every function separately.

It should be possible to implement a function that accepts string values, without having to deal with the 5 different runtime representations that can carry string values: StringArray, LargeStringArray, StringViewArray, DictionaryArray<Key>, RunArray<Key> (maybe more than 5 because they are recursive in theory: can DictionaryArray contain a DictionaryArray? can it contain RunArray?)

Focused

Because SQL is statically typed, it is necessary to select function overload during planning time, so that return type is also known (the return_type function). The invoke function needs to implement same logic again. This process should be less error-prone: once the function is bound to the query call site, its signature is known and the implementation should not need to do type checks.

Performant / Efficient

It should be the compiler's / framework's job to provide vectorization, without hand-writing same logic in every function separately.

It should be the compiler's / framework's to do null checks, providing efficient tight loops for the all-non-null case without having to write such logic in every function separately.

It should be possible to write efficient functions that need thread local data structures, for example for regular expressions, without having to use thread locals and/or shared pools which introduce contention.

Describe alternatives you've considered

• Prototype implementing DataFusion functions / operators using arrow-udf liibrary #11413

However, direct use of the library is not possible because

• it operates on Arrow arrays directly, doesn't fit DataFusion execution model
• it operates on Arrow data types, but DataFusion needs its own type system: [Proposal] Decouple logical from physical types #11513
• The problem space for a query engine is much more complicated

The library could still be part of the solution, but doesn't have to be and it's a non-goal to use a particular library.

Additional context

• [Proposal] Decouple logical from physical types #11513
• Change DataFusion to use LogicalTypes and hide "encodings" (like Dictionary and REE) lower down #7421
• How function definitions look like in Trino, for example a function operating on Int64-equivalent https://github.com/trinodb/trino/blob/e6aabfec55fcf5e4dd3c283b9a29da97c1301f6c/core/trino-main/src/main/java/io/trino/operator/scalar/MathFunctions.java#L128-L135
• Velox "Simple Functions" https://vldb.org/pvldb/vol15/p3372-pedreira.pdf

[findepi]

cc @alamb, @andygrove, @jayzhan211, @ozankabak, @notfilippo, @comphead

[findepi]

FYI: I am doing some experiments how this could look like

[alamb]

I think the idea of making it easier to write functions that include specialized implementations for different types is a great idea. This would likely both make our code faster (more specializations, more uniformly handle constants) as well as reduce the replication (there are several different patterns for dispatch / specialization now -- consolidating would be really nice)

[alamb]

Currently DataFusion functions are singletons plugged into the execution engine. They have no way to store and reuse buffers or compiled regular expressions, etc.

here is some prior work on the precompiled regular expressions: #11146 (note compiling the regex hasn't ever shown up in any performance trace I did, probably because the cost of actually matching is so much bigger than the cost of compiling and the cost of compiling is spread over 8192 rows

[comphead]

Thanks @findepi I think this process go through iterations, and easier than was before but still far from perfect.

The ScalarUDFImpl common trait is already a huge help, but for example dynamic typing, like MAX/MIN/GREATEST when return type depends on input is gonna be challenging, in Trino it seems also implemented in runtime rather than static.

Rust can benefit from macros and generate functions with all possible output types combinations, but this is just 1 problem.

I'm totally down if we can make it easier

[alamb]

Just to be clear, what I am imagining comes out of this work is:

1. No changes to ScalarUDFImpl
2. Some new way (generic functions, macros, etc) that would make creating the specialized function implementations (perhaps following the pattern of declarative macros on Prototype implementing DataFusion functions / operators using arrow-udf liibrary #11413)

[findepi]

FYI i touched upon the topic of types on DataFusion meetup in Belgrade yesterday.
The slides are here if anyone is interested: https://docs.google.com/presentation/d/1VW_JCGbN22lrGUOMRvUXGpAmlJopbG02hn_SDYJouiY . It was an attempt to summarize why we need both: simpler types (#11513), more types (#12644), and simple function "SDK" (#12635).
The document has comments disabled to avoid diverging the discussion from the github issue.

[alamb]

. It was an attempt to summarize why we need both: simpler types (#11513), more types (#12644), and simple function "SDK" (#12635).

It was also a very nice and understandable talk

[findepi]

I experimented with this on the way from DataFusion meetup in Belgrade.

i came up with something like this

function author would write this

row-oriented function logic

// hand-written
pub struct MyGcd {}
impl MyGcd {
    pub fn call(x: i64, y: i64) -> Result<i64> {
        datafusion::functions::math::gcd::compute_gcd(x, y)
    }
}

a macro would generate

generated (syntactically derived)

impl datafusion::logical_expr::ScalarUDFImpl for MyGcd {
    fn as_any(&self) -> &dyn Any {
        self
    }

    fn name(&self) -> &str {
        "MyGcd"
    }

    fn signature(&self) -> &Signature {
        todo!() // TODO here we should use logical typyes
    }

    fn return_type(&self, _arg_types: &[DataType]) -> Result<DataType> {
        todo!() // TODO here we should use logical typyes
    }

    fn invoke(&self, args: &[ColumnarValue]) -> Result<ColumnarValue> {
        lifted_invoke::<MyGcd>(self, args)
    }
}

impl Invocable for MyGcd {
    const ARGUMENT_COUNT: u8 = 2;
    type ArgumentTypeList = (i64, (i64, ()));
    type OutArgType = ();
    type FormalReturnType = Result<i64>;

    fn invoke(&self, regular_args: Self::ArgumentTypeList, _out_arg: &mut Self::OutArgType) -> Self::FormalReturnType {
        let (a, (b, ())) = regular_args;
        MyGcd::call(a, b)
    }
}

All the heavy lifting happens inside lifted_invoke (source https://gist.github.com/findepi/f1b0b84fc5f2002cb3d31b19a5bb983e#file-library-rs)

Notes on API design:

• type ArgumentTypeList = (i64, (i64, ())); -- this form of list presentation yields to template programming
• OutArgType + FormalReturnType -- we need both to support plain returning functions (numerics) and out args (for returning allocatable stuff like strings or structs)
• FormalReturnType = Result<i64>; -- lifted_invoke supports functions fallible and infallible functions, as well as nullable return values (via Option<T>)

Note:

• for continuing work on this would be best to have logical types defined ([EPIC] Decouple logical from physical types #12622). @notfilippo could we perhaps have a logical types stub in main at some point?

[notfilippo]

@notfilippo could we perhaps have a logical types stub in main at some point?

Planning on opening the PR on main soon. Then once it gets merged I can merge main in logical-types and continue my work.

[comphead]

I experimented with this on the way from DataFusion meetup in Belgrade.

i came up with something like this

function author would write this

row-oriented function logic

// hand-written
pub struct MyGcd {}
impl MyGcd {
    pub fn call(x: i64, y: i64) -> Result<i64> {
        datafusion::functions::math::gcd::compute_gcd(x, y)
    }
}

a macro would generate

generated (syntactically derived)

impl datafusion::logical_expr::ScalarUDFImpl for MyGcd {
    fn as_any(&self) -> &dyn Any {
        self
    }

    fn name(&self) -> &str {
        "MyGcd"
    }

    fn signature(&self) -> &Signature {
        todo!() // TODO here we should use logical typyes
    }

    fn return_type(&self, _arg_types: &[DataType]) -> Result<DataType> {
        todo!() // TODO here we should use logical typyes
    }

    fn invoke(&self, args: &[ColumnarValue]) -> Result<ColumnarValue> {
        lifted_invoke::<MyGcd>(self, args)
    }
}

impl Invocable for MyGcd {
    const ARGUMENT_COUNT: u8 = 2;
    type ArgumentTypeList = (i64, (i64, ()));
    type OutArgType = ();
    type FormalReturnType = Result<i64>;

    fn invoke(&self, regular_args: Self::ArgumentTypeList, _out_arg: &mut Self::OutArgType) -> Self::FormalReturnType {
        let (a, (b, ())) = regular_args;
        MyGcd::call(a, b)
    }
}

All the heavy lifting happens inside lifted_invoke (source https://gist.github.com/findepi/f1b0b84fc5f2002cb3d31b19a5bb983e#file-library-rs)

Notes on API design:

• type ArgumentTypeList = (i64, (i64, ())); -- this form of list presentation yields to template programming
• OutArgType + FormalReturnType -- we need both to support plain returning functions (numerics) and out args (for returning allocatable stuff like strings or structs)
• FormalReturnType = Result<i64>; -- lifted_invoke supports functions fallible and infallible functions, as well as nullable return values (via Option<T>)

Note:

• for continuing work on this would be best to have logical types defined ([EPIC] Decouple logical from physical types #12622). @notfilippo could we perhaps have a logical types stub in main at some point?

Thanks @findepi I was thinking about proc macros as well, having the the original function I hope we can extract necessary information and generate the DF code and the documentation, I do some testing for documentation in #12822 on top of what @Omega359 has already built.

If we think about simplication on this phase as to generate boilerplate code leaving the developer to implement the logic it sounds achievable

[Omega359]

    pub fn call(x: i64, y: i64) -> Result<i64> {
        datafusion::functions::math::gcd::compute_gcd(x, y)
    }

Wouldn't this incur a significant amount of overhead calling a function for every single row?

[findepi]

@Omega359 this is how the function logic is structured anyway --

datafusion/datafusion/functions/src/math/gcd.rs

Line 85 in a08f923

(Some(a1), Some(a2)) => Ok(Some(compute_gcd(a1, a2)?)),

[Omega359]

@Omega359 this is how the function logic is structured anyway

🤔

[adriangb]

I'll let @davidhewitt chime in but we've experience a lot of generic bloat from having to implement functions that operate on scalars, arrays, dictionary arrays and take multiple parameters: you end up with an exponential explosion of generics which is very painful at compile time.

[Omega359]

https://github.com/apache/datafusion/blob/main/datafusion/functions-nested/src/string.rs is a good example of what can result with supporting many types and args

[davidhewitt]

Thanks, looks similar to what we've done in datafusions-functions-json but we tried to handle dictionaries without casting them away. I guess we should just cast them and live with it 👍

[davidhewitt]

It seems like datafusion-contrib/datafusion-functions-json#66 might provide a possible avenue; by widening input dictionary keys to i64 size we can avoid having the key type as a generic parameter.

[findepi]

Going back to this topic. Not much happened on the Simple Functions front, but a lot happened in the world. New type system changes (Logical types can be found in type signatures #13372; preliminary steps towards extension types #14247). We internally implemented quite a few functions and it allowed me to experience first-hand how important it is to make function authoring easier.

Here is a design doc for what I would love to implement

• https://docs.google.com/document/d/1JyfbYfpmjvd3N3LEVAN7tX1D9ze6N98yGx3g4po_WoQ

It's open to comment by anyone.
Your comments are more than welcome.

Teaser

Wouldn't it be cool to be able to implement a function like this?
The engine should be able to automatically handle the common constant pattern case, as well as other situations.

// SQL regexp_like(text, pattern, flags) -> bool
#[vectorize]
fn regexp_like(
    #[arg(2)] pattern: &str,
    #[arg(3)] flags: &str,
) -> Result<impl FnMut(#[arg(1)] &str) -> bool> {
    let regex = Regex::new(format!("(?{flags})^({pattern})$").as_str())?;
    Ok(move |text: &str| regex.is_match(text))
}

All the regexp and many date/time, and many other functions would greatly benefit if we can pull this up.

[Omega359]

I love the work that you've put into this! You use the regexp_like as an example of what could be handled by if you look at that function it actually doesn't operate on individual values but delegates the whole arrays off to arrow-string.

My primary concern with this is not the direction at all - I love it. I am concerned primarily how to handle functions mixing string arrays and scalars. The whole casting logic for the various string types can get quite laboured and until Rust has Enum values as types I don't see a way around it unless we just force everything to a common type (GenericStringArray ?) which would remove the benefit of using something like StringViewArray.

I'm going to review the doc again later today and think about it a bit more.

[comphead]

Thanks @findepi and everyone, this work is epic, literally. in DataFusion it was always needed to unify the builtin functions as they implemented by different developers in different times.

Please also consider that functions can be declared in separate or even external crates and be plugged dynamically like now we discussing #5600

Thanks for the google doc, would you mind attach here the final result how procedural macros are unwrapped into the function code supporting all the attributes like #vectorized #arg(n), #constant, etc.

Another thing coming to my mind is to optimize builtin function if we know that the function parameter set by user is actually a literal which provides further optimizations