Repository: scikit-hep/awkward

Issue:
Remove pybind11 bottleneck in `ak.from_iter` (v1 and v2)
After a few hours of reading the docs, I came up with a few simple scenarios where I struggle to achieve performance improvements with `awkward` `1.8.0` with ragged data, and for the most part see slow downs compared to simple iteration and application of NumPy ufuncs, even when I increase the number of records and degree of "jaggedness" in the data.

I'll place a reproducer below, and I think my question is:
- are there prospects for improvement here?
- if not, maybe this is something to document re: appropriate use cases?

```python
from time import perf_counter

import numpy as np
from numpy.testing import assert_allclose
import awkward as ak


def setup_nested_list(num_records):
    # as you increase num_records, you get both
    # more records and more jagged data
    nested_list = []
    for size in range(3, num_records):
        nested_list.append(np.ones(shape=(size, 3)))
    return nested_list


def avg_no_awk(nested_list):
    result_arr = np.empty((len(nested_list), 3),
                          dtype=np.float64)
    for position, subarray in enumerate(nested_list):
        result_arr[position] = np.average(subarray, axis=0)
    return result_arr


def avg_with_awk(nested_list):
    awk_arr = ak.Array(nested_list)
    result_arr = ak.mean(awk_arr, axis=1)
    return ak.to_numpy(result_arr)


if __name__ == "__main__":
    for num_records in [500,
                        5000,
                        7000]:
        print("num jagged records:", num_records)
        nested_list = setup_nested_list(num_records)
        start = perf_counter()
        result_no_awk = avg_no_awk(nested_list=nested_list)
        end = perf_counter()
        without_awkard = end - start
        print("without awkward (s):", without_awkard)
        start = perf_counter()
        result_with_awk = avg_with_awk(nested_list=nested_list)
        end = perf_counter()
        with_awkward = end - start
        print("with awkward (s):", with_awkward)
        if with_awkward < without_awkard:
            print("speedup factor with awkward:", without_awkard / with_awkward)
        else:
            print("slowdown factor with awkward:", with_awkward / without_awkard)
        assert_allclose(result_with_awk, result_no_awk)

```

```
num jagged records: 500
without awkward (s): 0.00580483900012041
with awkward (s): 0.3663976220013865
slowdown factor with awkward: 63.119342671482585
num jagged records: 5000
without awkward (s): 0.2445533330010221
with awkward (s): 36.80084666699986
slowdown factor with awkward: 150.48188554782874
num jagged records: 7000
without awkward (s): 0.45862196600137395
with awkward (s): 72.65730339400034
slowdown factor with awkward: 158.42525823061573

```

This is a simplified version of an attempt to achieve performance improvements downstream (https://github.com/MDAnalysis/mdanalysis/issues/3776). Those cases were smaller, and `ak.Array` creation dominated, but increasing the record counts here doesn't seem too helpful in switching the balance--in fact, it is the opposite?


Discussion:
The primary performance goal for Awkward is to scale linearly in the data, rather than exponentially. Despite that goal, and the fact that we want to be faster than pure-Python, there are some places where it might be noticeably slower than e.g. NumPy. As long as the performance scales ~linearly with size, then it's behaving as expected. That said, sometimes the constant of proportionality may be larger than necessary - i.e. with performance optimisations we could lower it.

As you're noticing, array creation time can dominate for large Python objects. Awkward is not doing anything particularly magic here - it's walking through the list and growing a series of buffers to build a layout object. Really, you don't want to call `from_iter` unless you don't know what the data look like. The *fastest* that we could convert these data using the public API would be if we had a flat buffer of floats that we could unflatten with arrays of lengths.[^array] I tested this vs `from_iter`, and they both follow the same scaling relationship, although the direct approach is ~8x faster.  
```python3
content = [x for l2 in l1 for l3 in l2 for x in l3]
s1 = [len(l2) for l2 in l1]
s2 = [len(l3) for l2 in l1 for l3 in l2];
array = ak.unflatten(ak.unflatten(np.array(content), np.array(s2)), np.array(s1))
```
If you don't call `ak.unflatten`, this excerpt of code still takes the same amount of time - the cost of allocating and looping in Python dominate here. It would be even faster if we used NumPy's flatten on the inner lists, but that's not the point of my example :)


So, the first take-away is avoid `from_iter` array construction, at least, if its runtime is comparable to the operations you need to perform. If you already have a flat NumPy array, use `unflatten`, if you can restore it from (fast) disk, do that, etc. Looking at your linked issue (https://github.com/MDAnalysis/mdanalysis/issues/3776) I *think* you already have a nice NumPy array and an array of split positions, so you can actually use `ak.unflatten`, e.g.

```python3
splits = np.cumsum(np.unique(atoms.resindices, return_counts=True)[1])[:-1] # Assuming this is correct
positions = atoms.residues.atoms.positions

# Compute lengths of sublists
index = np.empty(len(splits)+2, dtype=np.int64)
index[0] = 0
index[-1] = len(positions)
index[1:-1] = splits
count = np.diff(index)

array = ak.unflatten(positions, count)
```

As to the actual reduction, `ak.mean` doesn't actually call a dedicated mean kernel; it manually computes a weighted mean. This involves calls to at least three kernels (to compute `sumw`, `sumwx`, and `sumwx/sumw`). So, there is a slight performance penalty relative to an optimised implementation, and then you can start to account for SIMD and other optimisations that we don't currently do in our CPU kernels to explain away the rest of the performance difference.

[^array]: NB: This would also work with `ak.from_iter` / `ak.Array` instead of `np.array`, except that for flat arrays we have a performance bug (it could be faster by some constant if we moved the loop into C++). @jpivarski incidentally, do you know why we have [a loop in Python for `ak.from_iter`](https://github.com/scikit-hep/awkward/blob/385f2926cf50c881b4c7851329708f0149cdbce6/src/awkward/operations/convert.py#L890-L891)? If it's just a legacy artefact, I could submit a PR.
Converting Python objects into an Awkward Array necessarily requires iteration over those Python objects. In this case,

```python
def avg_with_awk(nested_list):
    awk_arr = ak.Array(nested_list)
    result_arr = ak.mean(awk_arr, axis=1)
    return ak.to_numpy(result_arr)
```

the `ak.Array(nested_list)` step has to involve a Python for loop somewhere. (This version of the constructor calls [ak.from_iter](https://awkward-array.readthedocs.io/en/latest/_auto/ak.from_iter.html).) In v1, it's in the lines @agoose77 pointed out:

https://github.com/scikit-hep/awkward/blob/385f2926cf50c881b4c7851329708f0149cdbce6/src/awkward/operations/convert.py#L889-L891
and in v2, it's here:

https://github.com/scikit-hep/awkward/blob/d86fa74e6724539f323ac65395a5c6c82bf11156/src/awkward/_v2/operations/ak_from_iter.py#L85-L87

The alternative would be to put the `for x in iterable` inside of the pybind11-compiled code, which _in principle_ shouldn't be any better than leaving it outside like this. Writing `for x in iterable` using the Python C API is not a performance boost: it still has to evaluate each step on the Python virtual machine and deal with runtime types and all that. However, I've heard recently (at SciPy) that _in practice_ jumping through the pybind11 interface "adds quite a bit of overhead."

Okay, so I'll test it now:

```python
>>> iterable = np.random.normal(0, 1, 100000000).tolist()
>>> builder = ak.layout.ArrayBuilder()
>>> starttime = time.time()
>>> for x in iterable:
...     builder.fromiter(x)
... 
>>> time.time() - starttime
33.889479637145996
```

versus

```python
>>> builder = ak.layout.ArrayBuilder()
>>> starttime = time.time()
>>> builder.fromiter(iterable)
>>> time.time() - starttime
2.1149463653564453
```

Yeah, that's considerable. Both cases iterate over a large Python list (`iterable`), one in Python and the other in the Python C API, but probably the thing that matters is that it crosses the pybind11 boundary only once in the second case. Doing that with an ArrayBuilder creates a different data structure—in the second case, it's inside an extra layer of ListOffsetArray—but that unnecessary structure can just be discarded.

Okay, one more: what if the `iterable` is not a flat list of numbers but a list of lists of numbers?

```python
>>> iterable = np.random.normal(0, 1, (10000000, 10)).tolist()
>>> builder = ak.layout.ArrayBuilder()
>>> starttime = time.time()
>>> for x in iterable:
...     builder.fromiter(x)
... 
>>> time.time() - starttime
19.925483465194702
```

versus

```python
>>> builder = ak.layout.ArrayBuilder()
>>> starttime = time.time()
>>> builder.fromiter(iterable)
>>> time.time() - starttime
14.69497561454773
```

In this case, it's crossing the pybind11 boundary 10× less often, so that becomes less of a burden, and it has to unpack Python lists of lists, so that becomes more of a burden, and so this test that focuses on the effect of the pybind11 boundary sees less of a difference. At least, removing the pybind11 boundary does no harm, even when it's not the bottleneck.

So to answer your question, @agoose77: it went through `for x in iterable` due to ignorance on my part. Fixing that would sometimes be a performance boost for `ak.from_iter`. We've never worried much about optimizing this function because it will always have a Python for loop in it somewhere, but there's no reason to hamper it further with many passes through the pybind11 interface.

I'll change the title of this issue to say it's about removing the pybind11 bottleneck.

----------------------------

Can [ak.mean](https://awkward-array.readthedocs.io/en/latest/_auto/ak.mean.html) be optimized? Without a weight, it computes one [ak.count](https://awkward-array.readthedocs.io/en/latest/_auto/ak.count.html) and one [ak.sum](https://awkward-array.readthedocs.io/en/latest/_auto/ak.sum.html), so I don't see how. There is an open issue (#1323) to add a fast-path for reduction over `axis=-1`, but that's not @tylerjereddy's use-case here: `axis=1` is equivalent to `axis=-2`, which needs to create the intermediate structures that #1323 hopes to avoid.

----------------------------

That said, if your actual problem resembles the test you've made here, in which you know the structure of the nested lists, you'd be better off doing what @agoose77 suggested, instead of making many Python objects and then iterating over them to construct the Awkward array.
> Writing for x in iterable using the Python C API is not a performance boost: it still has to evaluate each step on the Python virtual machine and deal with runtime types and all that. 

Yes, I heard something similar and hence took a look at this line. I found the same outcomes as you - only really significant for large arrays.

> Can [ak.mean](https://awkward-array.readthedocs.io/en/latest/_auto/ak.mean.html) be optimized? Without a weight, it computes one [ak.count](https://awkward-array.readthedocs.io/en/latest/_auto/ak.count.html) and one [ak.sum](https://awkward-array.readthedocs.io/en/latest/_auto/ak.sum.html), so I don't see how.

I thought about this again, and I guess not actually. We could compute a running average, and trade memory for compute, but it might not be worth it. Clearly a sizeable chunk of the perf here is simply from having a non-regular list length. 


Resolving patch:
diff --git a/src/awkward/_v2/operations/ak_from_iter.py b/src/awkward/_v2/operations/ak_from_iter.py
--- a/src/awkward/_v2/operations/ak_from_iter.py
+++ b/src/awkward/_v2/operations/ak_from_iter.py
@@ -83,12 +83,11 @@ def _impl(iterable, highlevel, behavior, allow_record, initial, resize):
             )
 
     builder = ak.layout.ArrayBuilder(initial=initial, resize=resize)
-    for x in iterable:
-        builder.fromiter(x)
+    builder.fromiter(iterable)
 
     formstr, length, buffers = builder.to_buffers()
     form = ak._v2.forms.from_json(formstr)
 
     return ak._v2.operations.from_buffers(
         form, length, buffers, highlevel=highlevel, behavior=behavior
-    )
+    )[0]
diff --git a/src/awkward/operations/convert.py b/src/awkward/operations/convert.py
--- a/src/awkward/operations/convert.py
+++ b/src/awkward/operations/convert.py
@@ -887,10 +887,8 @@ def from_iter(
             )
 
     out = ak.layout.ArrayBuilder(initial=initial, resize=resize)
-    for x in iterable:
-        out.fromiter(x)
-    layout = out.snapshot()
-    return ak._util.maybe_wrap(layout, behavior, highlevel)
+    out.fromiter(iterable)
+    return ak._util.maybe_wrap(out[0], behavior, highlevel)
 
 
 def to_list(array):


Test patch: