Repository navigation
Specialized / Pre-compiled / Prepared ScalarUDFs #8051
Description
Activity
@thinkharderdev suggests
it would be useful to be able to serialize constant parameters into the user-defined scalar function themselves rather than pass them in as expressions. So for instance if you had to create a UDF to do something with a regex that you have as a static constant. Currently the way to do that is pass it as a literal expression. But then you have to compile the regex again for every batch you process through the UDF. Ideally you could have something like:
struct MyRegexUdf { regex: Regex } impl ScalarFunction for MyRegexUdf { // use regex on each value somehow }The regex would only need to be compiled once during deserialization (or construction) instead of once for each batch.
- changed the title
[-]Pre-compiled / Prepared ScalarUDFs[/-][+]Specialized / Pre-compiled / Prepared ScalarUDFs[/+]on Nov 4, 2023 One way to achieve this might be a PhysicalOptimizerRule that replaces relevant instances of ScalarFunctionExpr.
However this is likely somewhat awkward to write and is not clear that these expressions would serialize well (as serialization matches a name to an expr). It is probably possible to do this during derialization with the PhysicalExtensionCodec
@2010YOUY01 noted that in order to use the same APIs for built in functions and ScalarUDFs we will need to have some way to handle:
Figure out the interface to let core pass information to function definitions (e.g. now() requires to be passed query start time from core)
This sounds very similar to this method
Maybe we could add a
prepare(ctx: &TaskContext)type method on ScalarUDF 🤔 that can potentially return a new version of the function to invokeI think the trick here will be to ensure we can still serialize such "precompiled" functions
What I was thinking was maybe we can make a new
PhysicalExprthat is like/// A function that is "precompiled" in some way /// for example, for a regular expression that has a constant argument /// constant can be pre-compiled into a Regexpr match instance once per query /// rather than once per batch /// /// Somtimes precompiling make it hard/impossible to serialize the function again (e.g. the prepared regular expressions) /// so this structure contains the original PhysicalExpr that can be used to serialize the function struct PrecompiledExpr { precompiled: Arc<dyn PhysicalExpr>, original: Arc<dyn PhysicalExpr> }
That seems a bit heavy handed? Why not simply augment the scalar function expr node with an interior mutable "cache" cell, that is invisible to serialization? The cell can be automatically populated as needed by the function implementation upon node instantiation/deserialization.
That seems a bit heavy handed? Why not simply augment the scalar function expr node with an interior mutable "cache" cell, that is invisible to serialization? The cell can be automatically populated as needed by the function implementation upon node instantiation/deserialization.
@sadboy that is a (really) good idea
That seems a bit heavy handed? Why not simply augment the scalar function expr node with an interior mutable "cache" cell, that is invisible to serialization? The cell can be automatically populated as needed by the function implementation upon node instantiation/deserialization.
@sadboy that is a (really) good idea
BTW I thought more about this and one challenge is that currently the scalar function implementations that are passed around may be shared -- so using interior mutability would get confusing as each function invocation would be from the same instance of the function
the scalar function implementations that are passed around may be shared -- so using interior mutability would get confusing as each function invocation would be from the same instance of the function
Hmm, not sure I'm following? I was thinking of the
ScalarFunctionExprphysical node (https://docs.rs/datafusion-physical-expr/34.0.0/src/datafusion_physical_expr/scalar_function.rs.html#51), e.g. adding something likepub struct ScalarFunctionExpr { [...] args: Vec<Arc<dyn PhysicalExpr>>, prepared_args: OnceCell<Vec<Box<dyn PreparedArgs>>>, [...] }where
prepared_argsis populated by the return value of callingScalarFunction::prepareon the set of physicalargs. The function implementation objects themselves do not need to mutate -- they only need to provide an optional implementation ofScalarFunction::preparemethod if they wish to pre-process the arguments.prepared_argsis purely an optimization construct, it should have no effect on the semantics ofScalarFunctionExpr-- the node should compute the exact same result whether or notprepared_argsis populated. This way, any part of the system (e.g.Serialize,Clone, etc.) is always free to simply drop it without affecting correctness.Reacted by Andrew LambBTW now that @jayzhan211 and I have implemented
ScalarUDF::simplifyin #9298 and we have ported the regular_expression functions to useScalarUDF, I think we could actually use that API to implement precompiled functionsNote sure if that would meet your requirements @thinkharderdev
For example, to implement "precompiled regexp functions" we could do something like this (would be sweet if someone wanted to prototype this):
/// A new UDF that has a precompiled pattern impl PrecompiledRegexpReplace { precompiled_match: Arc<Pattern> } impl ScalarUDFImpl for PrecompiledRegexpReplace { // invoke function uses `self.precompiled_match` directly ... } // Update the existing RegexpReplace function to implement `simplify` impl ScalarUDFImpl for RegexpReplace { /// if the pattern argument is a scalar, rewrite the function to a new scalar UDF that /// contains a pre-compiled regular-expression fn simplify(&self) .. { match (args[1], args[2]) { (ScalarValue::Utf8(pattern), ScalarValue::Utf8(flags)) => { let pattern = // create regexp match SImplified::Rewritten(ScalarUdf::new(PrecompiledRegexpMatch { precompiled } ))) .call(args) }, _ => Simplified::Original(args) } }
We could then run some gnarly regular expression case, such as what is found on #8492 and see if it helps or not.
If it doesn't help performance, then the extra complexity isn't worth it for regexp_replace
Reacted by Marko MilenkovićReacted by Jay ZhanNote sure if that would meet your requirements @thinkharderdev
Yeah, although I was mainly using regexes as an example originally. Our particular issues are with some UDFs we have implemented that require some hilarious hacks to pass state as expressions.
But the idea does seem good to me.
Reacted by Andrew LambIt looks somebody needs to prototype what @alamb suggested in
I believe there is small change, some args will be used to create pattern so they should be removed from passing them down to result expression
match (args[1], args[2]) { (ScalarValue::Utf8(pattern), ScalarValue::Utf8(flags)) => { let pattern = // create regexp match, probably from args[1], and arg[2] let new_args = vec![args[0]] SImplified::Rewritten(ScalarUdf::new(PrecompiledRegexpMatch { precompiled } ))) .call(new_args) },
The only downside I see is that the "state" is not going to be serialised if it has to be distributed in systems like ballista. It would make sense to have generic "serialise to physical plan" at some point, it would help with distributing something like python udfs but, thats probably different discussion.
also note: #12270
The only downside I see is that the "state" is not going to be serialised if it has to be distributed in systems like ballista.
That's because we don't use expressions.
What about using simplify all the way down? ThePrecompiledRegexpMatchcould get dumped as a bytes buffer (varbinary) into an expression and then cast back toPrecompiledRegexpMatch. This will work as long as it's a flat structure. It won't work when it's something that has pointers internally and requires actual serialization.Alternatively, we can avoid all this complexity -- at the cost of different complexity, but conceptually simpler.
Let's imagine ScalarUDF invoke gets an option to create a thread local scratch space that it can reuse on all invocations. That would make reusing compiled pattern easier without having to serialize it in the plan.
The downside would be that the implementation would need to explicitly check whether the pattern is the same on every invocation (equality check once per batch).Sorry but I'm not sure I understand why do you need thread local storage.
What @alamb proposed will work with datafusion execution, no need for any other state handling, state is in the struct's property, which will be re-used for each and every batch this function participate,
PrecompiledRegexpMatchwill be "specialised" for a given logical plan.Or to put it in a different way, if there is no need to serialise your plan to be distributed, like with ballista, I believe proposed solution totally make sense.
Do you have use-case where you need to distribute function expression to different executors, like ballista ?
I am not using ballista currently. I realized the plan serialization concern is easy to address if we separate simplify into phases: the Expr-constraint simplify (e.g. pruning args known to be null, etc) that would be run during plan optimization phase. And then, local execution simplify which allows a function to "compile itself" into most optimal form, without any needs for serialization anymore. Ballista would need to serialize and distribute the plans in between these phases.
BTW we focused so far on compiling regular expressions, but we didn't think about memory needs for their execution.
Internallyregex::Regex::is_matchuses a synchronized pool of "caches" (regex execution scratch space) underneath. I don't know if this is a perf problem (probably not!), but let me use this as an example. It would probably be good if at runtime a scalar could have its own thread local "scratch space" / "local buffer". And without having to use thread locals which aren't great if DF is embedded and doesn't control thread creation.Why am I mentioning this? I thought that maybe if we had "scratch space" / "local buffer" support, we wouldn't have need to "compile functions" during planning.
Perhaps the
ScalarFunctionArgsadded in- Add
ScalarUDFImpl::invoke_with_argsto support passing the return type created for the udf instance #13290 by @joseph-isaacs would be a natural way / mechanism to pass pre-compiled state to the ScalarUDF 🤔
- Add
Copying a comment from #13290
Yes, I was thinking something like
trait ScalarUDFImpl { /// prepares to run the function, returning any state (such as /// a precompiled regex). Called once per instance of function in the query fn prepare(&self, args: &ScalarFunctionArgs) -> Option<Box<dyn Any>> { None } /// `prepared` field in ScalarFunctonArgs has the result of calling `prepare` fn invoke_with_args(&self, args: &ScalarFunctionArgs) -> ...
pub struct ScalarFunctionArgs<'a> { ... /// The result from a call to `prepare` for this function instance prepared: Option<Box<dyn Any>>, }
Is your feature request related to a problem or challenge?
Currently, scalar UDF functions can not be "specialized"
What happens is that the regexp string
'[a-z].*'gets passed as a literal expression to each invocation of the function, then you have to compile the regex again and again for every batch you process through the UDF.Since it is expensive to compile these RegExps, it would be nice if there was something that could compile the RegExp once per plan rather than once per batch
Describe the solution you'd like
No response
Describe alternatives you've considered
See comments below
Additional context
Suggested by @thinkharderdev on #8045 (comment)