lincc-frameworks / lincc-frameworks/nested-pandas

Improve performance of series.ext_array.replace_with_mask()

Open
#52 3 comments 0 reactions 0 assignees View on GitHub

Nobody has claimed this yet.

bug performance
Dominant language
Python
Stars
26
Forks
8
Avg merge
2d 2h
Merged PRs (30d)
9

Description

Currently, arrow misses the support of `pyarrow.compute.replace_with_mask` for struct arrays:
https://github.com/apache/arrow/issues/29558

That's why we have our own implementation used by `NestedExtenstionArray.__setitem__()`. The implementation has an overhead of creating a `len(self)`-sized struct array to perform the replacement. This approach would work well when we are going to replace many elements, but when we replacing just few, it would produce a large memory foot-print and probably take a while.

An alternative approach would be copying the original array to `np.ndarray[pa.StructScalar]`, replace the elements in-place, and convert it back:
```python
def replace_with_mask(array: pa.ChunkedArray, mask: pa.BooleanArray, value: pa.Array) -> pa.ChunkedArray:
"""Replace the elements of the array with the value where the mask is True"""
np_array = np.fromiter(array, dtype=object)
np_array[mask] = value
new_pa_array = pa.array(np_array)
return pa.chunked_array([new_pa_array])
```

We should create a benchmark and see what works faster and have smaller memory foot-print.

Contributor guide

No contributing guide indexed for this repository

First steps

  1. Read the whole issue, then the project's contributing guide.
  2. Comment on the issue to say you are picking it up — it saves two people doing the same work.
  3. Fork the repository and make your change on a branch.
  4. Open a pull request that references the issue number.

Research direction

Start at series.ext_array.replace_with_mask() and its use by NestedExtenstionArray.__setitem__(). Benchmark the current struct-array implementation against the proposed NumPy-based approach for replacing many and few elements, measuring runtime and memory footprint. Done means the benchmark identifies which approach performs better for each case and supports an implementation decision.

Written by the indexing model from the issue text.

Assessment

Tech stack
python
Domain
performance
Issue type
Refactor
Difficulty
4/5
Estimated time
3-5 days
Activity status
Stale
Clarity
Mostly clear
Newbie friendliness
38/100

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.