Simple Functions

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

@alamb

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

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

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)

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

@findepi

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

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

No changes to ScalarUDFImpl
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)

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.

. 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

@notfilippo

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.

@notfilippo

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

    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?

@Omega359

@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

🤔

@davidhewitt

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.

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

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 👍

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.

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.

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.

@findepi

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

@findepi

@findepi Thank you for thoroughly exploring this topic and creating such a detailed design document!

@linhr and I have reviewed it and have some initial thoughts:

This approach is promising!

The design would significantly reduce boilerplate code when authoring row-wise functions. However, implementing it through proc_macro could introduce significant complexity into the codebase.

We'd like to better understand how this design would help with function signature customization. Could you clarify if there are specific goals or non-goals around this aspect? For example, we're wondering about scenarios where downstream crates might need to override the default signature to express constraints or define allowed coercible behavior.

Additionally, it would be helpful to understand how the design handles cases where we might want to customize return types based on specific requirements (e.g., choosing between Int32 and Int64 outputs).

downstream crates might need to override the default signature

they would define own function

either via implementation or via "simplify" (not covered by the current design)

define allowed coercible behavior.

it's function role to declare what it accepts, it's engine role to define how arguments are fitted to a function call. this is called coercions. in some systems (eg Trino) coercions are pretty consistent and in some (eg Snowflake) are contextual: different for union, for equality for function calls. This complexity isn't going away, but my hope is that functions don't need to have that complexity baked into them. At least, so is the current plan.

Sharing WIP status here

Simple Functions Preview #14668

Navigation Menu

Repository navigation

Simple Functions #12635

Design doc

Is your feature request related to a problem or challenge?

Verbosity

Local data structures

Describe the solution you'd like

Simple

Focused

Performant / Efficient

Describe alternatives you've considered

Additional context

function author would write this

a macro would generate

function author would write this

a macro would generate

Teaser

Metadata

Assignees

Labels

Type

Projects

Milestone

Relationships

Development

Simple Functions #12635

Description

Design doc

Is your feature request related to a problem or challenge?

Verbosity

Local data structures

Describe the solution you'd like

Simple

Focused

Performant / Efficient

Describe alternatives you've considered

Additional context

Activity

findepi commented on Sep 26, 2024

findepi commented on Sep 26, 2024

alamb commented on Sep 26, 2024

alamb commented on Sep 26, 2024

comphead commented on Sep 26, 2024

alamb commented on Sep 27, 2024

findepi commented on Sep 28, 2024

alamb commented on Sep 28, 2024

findepi commented on Oct 9, 2024

function author would write this

a macro would generate

notfilippo commented on Oct 10, 2024

comphead commented on Oct 10, 2024

function author would write this

a macro would generate

Omega359 commented on Oct 10, 2024

findepi commented on Oct 11, 2024

Omega359 commented on Oct 11, 2024

adriangb commented on Jan 29, 2025

Omega359 commented on Jan 29, 2025

davidhewitt commented on Jan 31, 2025

davidhewitt commented on Jan 31, 2025

findepi commented on Feb 11, 2025

Teaser

Omega359 commented on Feb 11, 2025

comphead commented on Feb 12, 2025

shehabgamin commented on Feb 13, 2025

findepi commented on Feb 13, 2025

findepi commented on Feb 14, 2025

Metadata

Metadata

Assignees

Labels

Type

Projects

Milestone

Relationships

Development

Issue actions