Repository: scikit-hep/awkward

Issue:
Buggy behavior in ak.num() on TTrees
Not sure whether this should be an issue in uproot or awkward. `ak.num()` doesn't work with `uproot.lazy()`:

```python
>>> import uproot
>>> import awkward as ak
>>> uproot.__version__
'4.0.4'
>>> ak.__version__
'1.1.2'
>>> uproot.lazy('file.root')
<Array [{int_branch: 0, ... bool_branch: True}] type='2 * {"int_branch": int32, ...'>
>>> ak.num(uproot.lazy('file.root'), axis=0)
Traceback (most recent call last):
  File "<stdin>", line 1, in <module>
  File "/home/user/miniconda3/envs/iris-hep/lib/python3.9/site-packages/awkward/operations/structure.py", line 256, in num
    out = layout.num(axis=axis)
  File "/home/user/miniconda3/envs/iris-hep/lib/python3.9/site-packages/awkward/partition.py", line 481, in num
    return sum(x.num(axis) for x in self.partitions)
TypeError: unsupported operand type(s) for +: 'int' and 'awkward._ext.Record'
```

Note that you can apply `ak.num()` to a TTree in other ways, albeit with inconsistent results:

```python
>>> ak.num(uproot.open('file.root')['tree'], axis=0)
5
```
(there are five branches)

```python
>>> ak.num(uproot.open('file.root')['tree'].arrays(), axis=0)
<Record ... double_branch: 2, bool_branch: 2} type='{"int_branch": int64, "long_...'>
```
(there are two entries in each branch)

Intuitively I would have expected all of the above options to give me just `2`.


Discussion:

Resolving patch:
diff --git a/src/awkward/partition.py b/src/awkward/partition.py
--- a/src/awkward/partition.py
+++ b/src/awkward/partition.py
@@ -478,7 +478,24 @@ def fillna(self, what):
 
     def num(self, axis):
         if first(self).axis_wrap_if_negative(axis) == 0:
-            return sum(x.num(axis) for x in self.partitions)
+            prepared = [x.num(axis) for x in self.partitions]
+            if any(isinstance(x, ak.layout.Record) for x in prepared):
+                names = None
+                counts = None
+                for x in prepared:
+                    assert isinstance(x, ak.layout.Record)
+                    n = ak.operations.describe.fields(x)
+                    c = ak.operations.structure.unzip(x)
+                    if names is None:
+                        names, counts = n, list(c)
+                    else:
+                        assert n == names
+                        for i, ci in enumerate(c):
+                            counts[i] += ci
+                arraycounts = [ak.layout.NumpyArray(numpy.array([x])) for x in counts]
+                return ak.layout.Record(ak.layout.RecordArray(arraycounts, names), 0)
+            else:
+                return sum(prepared)
         else:
             return self.replace_partitions([x.num(axis) for x in self.partitions])
 


Test patch:
diff --git a/tests/test_0756-num-of-partitioned-recordarrays.py b/tests/test_0756-num-of-partitioned-recordarrays.py
new file mode 100644
--- /dev/null
+++ b/tests/test_0756-num-of-partitioned-recordarrays.py
@@ -0,0 +1,12 @@
+# BSD 3-Clause License; see https://github.com/scikit-hep/awkward-1.0/blob/main/LICENSE
+
+from __future__ import absolute_import
+
+import pytest  # noqa: F401
+import numpy as np  # noqa: F401
+import awkward as ak  # noqa: F401
+
+
+def test():
+    array = ak.repartition([{"x": x, "y": x * 10} for x in range(10)], 2)
+    assert ak.to_list(ak.num(array.layout, axis=0)) == {"x": 10, "y": 10}