apache / datafusion

Apache DataFusion SQL Query Engine
https://datafusion.apache.org/
Apache License 2.0
6.33k stars 1.2k forks source link

Perf: Allow User defined functions to potentially reuse their argument arrays (to avoid new allocations) #13516

Open alamb opened 15 hours ago

alamb commented 15 hours ago

Is your feature request related to a problem or challenge?

Arrow Arrays are designed to be immutable and use shared references extensively, but it is possible to reuse the underlying buffer in some cases when there are no other references (see the arrow unary_mut kernel for example)

At the time of writing, DataFusion scalar functions (ScalarFunctionImpl must always allocate a new array when generating output. They can not reuse the existing underlying memory, even if the source array will never be used again

This is because the invoke signature gets the arguments as reference (slice of ColumnarValue) rather than by ownership

fn invoke_batch(
    &self,
    args: &[ColumnarValue],
    number_rows: usize,
) -> Result<ColumnarValue, DataFusionError>

For example, an expression like (a + b) + c will be evaluated like

Resulting in two new allocations

Describe the solution you'd like

It would be really nice if it were possible to evaluate (a + b) + c like this (with no new allocations)

And the result would be a new array that re-used the original allocation of the a array

Describe alternatives you've considered

Now that this is merged

I think we can make it possible in the future to reuse allocations by changing what is passed into ScalarFunctionArgs

Since we haven't yet released a version with ScalarFunctionArgs we can change its signature without breaking APIs until DataFusion 44 is released

Additional context

I have a draft of the basic idea here: