Optimising regex functions (and literal function arguments in general)

[pepijnve]

I'm very new to both Rust and the DataFusion codebase, so apologies up front if I'm misinterpreting the code.

If I understand it correctly regexes are currently compiled at least once per batch and sometimes cached per batch (e.g. in the datafusion_functions::regex::regexpcount::compile_and_cache_regex). The cache lookup seems to be allocating a new String per hash table lookup. My intuition says this is probably not going to be very fast, and the simple experiments I ran today seem to confirm that.

If the regex argument is a string literal, it seems like it should in theory be feasible to optimize this by compiling the regex up front and reusing the compiled version for every invocation of the function. I spent some time today digging around the codebase trying to figure out how to implement this idea, but it wasn't immediately obvious to me if it's actually possible within the constraints of the current APIs.

The best option I came up with was to instantiate a curried version of the scalar function in the simplify implementation for the function in question. The part I couldn't quite figure out is what the appropriate place would be to store the compiled regex. Would a custom PhysicalExpr be the way to go?

Taking a step back, the more general question I had is if there is already some facility present in the codebase that supports performing some initial precomputation work per individual call site of a function as a way to improve performance.
The other instance of this I was looking at were string functions like start_with and ends_with when invoked with a literal. Currently the code expands the scalar to an array per batch. Perhaps that array could be created once and then sliced to the appropriate length per invocation instead.

Does any of this make sense or am I looking at this problem the wrong way?

[findepi]

The best option I came up with was to instantiate a curried version of the scalar function in the simplify implementation for the function in question.

That's my understanding too.

what the appropriate place would be to store the compiled regex.

would it work, if this is stored on a field of a ScalarUDFImpl instance?
like here

datafusion/datafusion/core/tests/user_defined/user_defined_scalar_functions.rs

Line 1214 in 749cd6f

regex: Regex,

The other instance of this I was looking at were string functions like start_with and ends_with when invoked with a literal. Currently the code expands the scalar to an array per batch.

The patter to use for these is to recognize ColumnarValue::Scalar and use that. The caveat is this makes code much more complex.
The mid ground is to use make_scalar_function as starts_with does, but pass the AcceptsSingular hint. As a result, the "expansion" is O(1) size

[pepijnve]

would it work, if this is stored on a field of a ScalarUDFImpl instance? like here

Yep that will work. I had assumed I needed mutable state and that would not work easily in ScalarUDFImpl since invoke takes a non mutable self. But Regex instances can be shared freely. Something like the test implementation is exactly what I had in mind. I'll see if I can put together a PR for the currying idea. I found #13789 in the meantime which I'l use to measure the performance difference.

make_scalar_function as starts_with does, but pass the AcceptsSingular hint

I don't think that will work. As far as I can tell the arrow kernel that's used for these operations does not support an array/scalar combination. I'll see if I can apply the currying idea here as well if I can get the regex function working.

[findepi]

would not work easily in ScalarUDFImpl since invoke takes a non mutable self.

yes, instances may be shared between threads

Regex instances can be shared freely.

but see #8051 (comment)

I don't think that will work. As far as I can tell the arrow kernel that's used for these operations does not support an array/scalar combination.

you're right. Whether this works or doesn't depends on how the main logic of the function is implemented.
In particular, if the function is implemented without delegating to a library (eg without calling an arrow kernel), the approach of "expanding" to 1-element array works pretty neatly.

[pepijnve]

Thanks @findepi. I spent some time following hyperlinks and found the fn prepare(&self, args: &ScalarFunctionArgs) work that was merged not so long ago. That was exactly the 'per call site' hook I had in mind.

[alamb]

This is a good question. Others have asked about this too @pepijnve

See - #11146

@zhuliquan seems to be actively working on it:

• feat: scalar regex match physical expr #12270

FWIW it has never been a particular priority for me personally as I have never seen a performance profile where compiling the regex showed up (if a query has regex the vast majority of the time is in the evaluation of the regex, not the compialtion)

[pepijnve]

@lamb thanks for the pointer, seems like I'm still missing the right vocabulary to find these in the big pile or issues/PRs. That's an interesting alternative approach. I was looking for a solution that I could implement without making changes in the datafusion crate itself. I might still try that as a learning exercise

Fair point on regex compilation not being the bottleneck. I hadn't gotten round to profiling yet. I'll be sure to check before/after results to see if there's any measurable difference.

[pepijnve]

I did some further testing and it turns out that you're absolutely correct and performance is dominated by IO in the case of

select count(1) from '<input_file>' where <text_column> ~ '<suffix>$'

Instruments is showing ~75% unmarshalling / 25% regex testing

My first experiment that got me thinking about optimising was

select count(1) from '<input file>' where regexp_count(<text_column>, '<suffix>$') > 0

I hadn't realised the binary operator ~ was available. The query is not entirely the same, but you would think it's close enough. The performance profile is radically different though. The first one takes about 18s while the second one takes about 97s and spends only ~10% on IO and everything else in the regex automata.

[alamb]

Yes, the string matching operations in DataFusion are quite highliy optimized by now -- so ~ (or doing LIKE / ILIKE) are about as fast as we know how to make them :)

[pepijnve]

Any idea what could explains the large gap between like and regexp_count? Is there an intrinsic performance penalty in calling scalar UDFs or is it just that that function hasn't received as much attention as like and friends?

[alamb]

Any idea what could explains the large gap between like and regexp_count? Is there an intrinsic performance penalty in calling scalar UDFs or is it just that that function hasn't received as much attention as like and friends?

I think it is the latter (not as much optimization)

Specifically there are a bunch of special cases where certain patterns are optimized in like -- for example prefix or substring matching that avoid the need to use a general purpose regexp

Perhaps similar optimizations could be applied to regexp_count perhaps