apache / apache/arrow-rs

Hard to use `PrimitiveArray::unary_mut`, `PrimitiveArray:try_unary_mut`, etc

Open
#8,808 0 comments 0 reactions 0 assignees View on GitHub
enhancement
Dominant language
Rust
Stars
3.6k
Forks
1.3k
Avg merge
2d 18h
Merged PRs (30d)
169

Description

**Is your feature request related to a problem or challenge? Please describe what you are trying to do.**
This came up in the context of this PR in DataFusion:
- https://github.com/apache/datafusion/pull/18360

In that case we are applying some operations to a `PrimitiveArray` and would like to reuse the allocation if possible

However, the current API of [PrimitiveArray::unary_mut](https://docs.rs/arrow/latest/arrow/array/struct.PrimitiveArray.html#method.unary_mut) and similar functions makes this awkward to do as the caller must handle the case where the allocation can not be reused

```rust
// want to apply an operation to arr, reusing allocation if possible
let arr: PrimitiveArray = ...
// to do so we call try_unary but also must handle when the allocation is shared
let new_arr = match arr.unary_mut(|a| a+ 1) {
Ok(arr) => arr,
Err(old_arr) => old_arr.unary(|a| a+1)
}
```

This can be done, but it is hard to use.

I proposed the following function in DataFusion
```rust
/// Applies the unary operation in place if possible, or cloning the array if not
fn try_unary_mut_or_clone(
array: PrimitiveArray,
op: F,
) -> Result>
where
F: Fn(i64) -> Result,
{
match array.try_unary_mut(&op) {
Ok(result) => result,
// on error, make a new array
Err(array) => array.try_unary(op),
}
}
```

but quoting @findepi on https://github.com/apache/datafusion/pull/18360/files#r2475557450:
> can this be made more flexible with a more generous use of generics?
> perhaps it could even be in arrow-rs. it makes try_unary_mut significantly more approachable

**Describe the solution you'd like**

I would like it to be easier to apply unary and binary operations on PrimitiveArrays and reuse the allocation if possble

**Describe alternatives you've considered**

One alternative would be to follow the API of [`Arc::unwrap_or_clone`](https://doc.rust-lang.org/std/sync/struct.Arc.html#method.unwrap_or_clone)

So that would mean functions something like
* `PrimitiveArray::unary_mut_or_clone`
* `PrimitiveArray::try_unary_mut_or_clone`
* `PrimitiveArray::binary_mut_or_clone`
* `PrimitiveArray::try_binary_mut_or_clone`

Which would be implemented like the function above

I think this would make it much easier to use these APIs

Contributor guide

Open the contributing guide

Research direction

Start with PrimitiveArray::unary_mut, try_unary_mut, and the corresponding binary APIs, then review the linked DataFusion PR and the Arc::unwrap_or_clone alternative. The desired result is a consistent set of unary and binary operations that reuse the allocation when possible and clone when necessary, including fallible variants.

Written by the indexing model from the issue text.

Assessment

Tech stack
rust
Domain
data
Issue type
Feature
Difficulty
4/5
Estimated time
3-5 days
Activity status
Stale
Clarity
Mostly clear
Newbie friendliness
45/100

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.