Consider changing the return type of compute to mutable arrays

[jorgecarleitao]

When an expression being evaluated is only known at runtime (e.g. a + exp(-b * 2) coming from SQL or equivalent coming from a Python expression), it would be optimal to evaluate it as follows:

• mul_scalar(&b, 2) -> MutablePrimitiveArray (= c)
• neg(&mut c)
• exp(&mut c)
• add(&a, &mut c)
• arc_buffers(c) -> PrimitiveArray

However, currently users have no way of supporting this behavior because all our kernels "arc" the buffers prior to returning them. Specifically, users must run the following:

• let c = mul_scalar(&b, 2);
• let c = neg(&c);
• let c = exp(&c);
• let c = add(&a, &c);

This uses an intermediary allocation per operation, which scales with O(N * E) where N is the length and E the number of operations.

Proposal

If we changed all kernels to return MutableArray (so, Box<dyn MutableArray> and Mutable*Array for dyn and typed versions respectively) and require users to "arc" them afterwards, we would allow an optimizer pass to convert a + exp(-b * 2) to the mutable version of the operations above.

So, the proposal here is: change all kernels to return MutablePrimitiveArray in our kernels, so that the user may continue to operate over it as shown above. The migration path is quite simple: sprinkle the code with Arc(.into()) to freeze the arrays once the user wants to make the array sharable.

Notes

• AFAIK this behavior is supported by numpy, whereby we can use the optional out to write the result to.

• such an optimizer does not require a compiler; all of the above is using the same binary and just changing behavior at runtime. There is a related optimization that requires a compiler: compile a + b * 2 to arrow2::compute::arity::binary(|a, b| a + b * 2) so that both operations could be performed on a single loop / vectorized together.

[jorgecarleitao]

thoughts @ritchie46 @Dandandan @nevi-me @houqp @sundy-li ?

[Dandandan]

Nice idea - seems like a missing piece indeed.

I think we should also revisit the kernels themselves, as we want to write the results to the input array instead of allocating a new array. This is different from what we have now - i.e. we won't change the input array.

I think you are suggesting to make the kernels always.write to an existing array? In that case the user also has to change his code to copy the input array if also the original, unchanged array is still of interest.

I am wondering if it doesn't make more sense / it's easier to have separate methods for the expressions on the mutable arrays instead? So we have two code paths - for immutable and mutable arrays. A lot of code could be shared between them - but the public signatures will be different in taking mutable inputs, and the semantics will be different in writing to an existing array.

[ritchie46]

This sounds really interesting! IMO we should rewrite all kernels to this:

// from 

fn some_comptute(in: &Array) -> ArrayRef { .. }

// to 

fn some_compute(int: &Array, out: &mut MutableArrray) { .. }

And maybe add wrapper kernels (for backwards compatibility) that call these inner kernels.

The wrappers can be really small.

fn original_compute(in: &Arrray) -> ArrayRef {
    let mut buffer = MutableBuffer::with_capacity(..);
    some_compute_new(in, buffer);
    Arc::new(buf.into())
}

[houqp]

I think this is a great idea. Promote mutable array into the core, then build immutable apis as thin wrappers unlocks a lot of optimization opportunities :+1:

[drivenow]

@jorgecarleitao @ritchie46 As this idea hasn't been realezed in Arrow2 crate, here is my case: How can I update MutablePrimitiveArray value after computing between a MutablePrimitiveArray and PrimitiveArray? My example is:

use arrow2::array::{Array, PrimitiveArray, MutablePrimitiveArray};
use arrow2::compute::arithmetics::basic;

 let a = PrimitiveArray::from([Some(1), Some(2), Some(3)]);
 let mut b = MutablePrimitiveArray::<i32>::from([Some(0), Some(0), Some(0)]);

//compute first, and convert MutablePrimitiveArray to PrimitiveArray
 let c = basic::add(&a, &PrimitiveArray::from(b));   //O(1) from MutablePrimitiveArray to PrimitiveArray

//then do some update, but failed. 
b.set(2, Some(30)); //Error, borrow of moved value: `b`, value b moved at `PrimitiveArray::from(b)`

• If I use b.clone(), it could solve the moved value problem, but it will allocate unnecessary memory .
• I also tried use RC<T> on variable b, but PrimitiveArray::from dosen't accept RC<T> type parameter, so it doesn't work.

Need your helps, Thanks!