C++ refactoring: num()

src/awkward/_v2/contents/bitmaskedarray.py

@@ -324,6 +324,9 @@ def simplify_optiontype(self):
         else:
             return self
 
+    def num(self, axis, depth=0):
+        return self.toByteMaskedArray.num(axis, depth)
+
     def mergeable(self, other, mergebool):
         if not _parameters_equal(self._parameters, other._parameters):
             return False

src/awkward/_v2/contents/bytemaskedarray.py

@@ -409,6 +409,24 @@ def simplify_optiontype(self):
         else:
             return self
 
+    def num(self, axis, depth=0):
+        posaxis = self.axis_wrap_if_negative(axis)
+        if posaxis == depth:
+            out = ak._v2.index.Index64.empty(1, self.nplike)
+            out[0] = len(self)
+            return ak._v2.contents.numpyarray.NumpyArray(out)[0]
+        else:
+            numnull = ak._v2.index.Index64.empty(1, self.nplike, dtype=np.int64)
+            nextcarry, outindex = self._nextcarry_outindex(numnull)
+
+            next = self._content._carry(nextcarry, False, NestedIndexError)
+            out = next.num(posaxis, depth)
+
+            out2 = ak._v2.contents.indexedoptionarray.IndexedOptionArray(
+                outindex, out, parameters=self.parameters
+            )
+            return out2.simplify_optiontype()
+
     def mergeable(self, other, mergebool):
         if not _parameters_equal(self._parameters, other._parameters):
             return False

src/awkward/_v2/contents/emptyarray.py

@@ -127,6 +127,17 @@ def _getitem_next(self, head, tail, advanced):
         else:
             raise AssertionError(repr(head))
 
+    def num(self, axis, depth=0):
+        posaxis = self.axis_wrap_if_negative(axis)
+
+        if posaxis == depth:
+            out = ak._v2.index.Index64.empty(1, self.nplike)
+            out[0] = len(self)
+            return ak._v2.contents.numpyarray.NumpyArray(out)[0]
+        else:
+            out = ak._v2.index.Index64.empty(0, self.nplike)
+            return ak._v2.contents.numpyarray.NumpyArray(out)
+
     def mergeable(self, other, mergebool):
         if not _parameters_equal(self._parameters, other._parameters):
             return False

src/awkward/_v2/contents/indexedarray.py

@@ -359,6 +359,15 @@ def simplify_optiontype(self):
         else:
             return self
 
+    def num(self, axis, depth=0):
+        posaxis = self.axis_wrap_if_negative(axis)
+        if posaxis == depth:
+            out = ak._v2.index.Index64.empty(1, self.nplike)
+            out[0] = len(self)
+            return ak._v2.contents.numpyarray.NumpyArray(out)[0]
+        else:
+            return self.project().num(posaxis, depth)
+
     def mergeable(self, other, mergebool):
         if not _parameters_equal(self._parameters, other._parameters):
             return False

src/awkward/_v2/contents/indexedoptionarray.py

@@ -405,6 +405,20 @@ def simplify_optiontype(self):
         else:
             return self
 
+    def num(self, axis, depth=0):
+        posaxis = self.axis_wrap_if_negative(axis)
+        if posaxis == depth:
+            out = ak._v2.index.Index64.empty(1, self.nplike)
+            out[0] = len(self)
+            return ak._v2.contents.numpyarray.NumpyArray(out)[0]
+        _, nextcarry, outindex = self._nextcarry_outindex(self.nplike)
+        next = self._content._carry(nextcarry, False, NestedIndexError)
+        out = next.num(posaxis, depth)
+        out2 = ak._v2.contents.indexedoptionarray.IndexedOptionArray(
+            outindex, out, parameters=self.parameters
+        )
+        return out2.simplify_optiontype()
+
     def mergeable(self, other, mergebool):
         if not _parameters_equal(self._parameters, other._parameters):
             return False

src/awkward/_v2/contents/listarray.py

@@ -671,6 +671,31 @@ def _getitem_next(self, head, tail, advanced):
         else:
             raise AssertionError(repr(head))
 
+    def num(self, axis, depth=0):
+        posaxis = self.axis_wrap_if_negative(axis)
+        if posaxis == depth:
+            out = ak._v2.index.Index64.empty(1, self.nplike)
+            out[0] = len(self)
+            return ak._v2.contents.numpyarray.NumpyArray(out)[0]
+        elif posaxis == depth + 1:
+            tonum = ak._v2.index.Index64.empty(len(self), self.nplike)
+            self._handle_error(
+                self.nplike[
+                    "awkward_ListArray_num",
+                    tonum.dtype.type,
+                    self._starts.dtype.type,
+                    self._stops.dtype.type,
+                ](
+                    tonum.to(self.nplike),
+                    self._starts.to(self.nplike),
+                    self._stops.to(self.nplike),
+                    len(self),
+                )
+            )
+            return ak._v2.contents.numpyarray.NumpyArray(tonum)
+        else:
+            return self.toListOffsetArray64(True).num(posaxis, depth)
+
     def mergeable(self, other, mergebool):
         if not _parameters_equal(self._parameters, other._parameters):
             return False

src/awkward/_v2/contents/listoffsetarray.py

@@ -508,6 +508,35 @@ def _getitem_next(self, head, tail, advanced):
         else:
             raise AssertionError(repr(head))
 
+    def num(self, axis, depth=0):
+        posaxis = self.axis_wrap_if_negative(axis)
+        if posaxis == depth:
+            out = ak._v2.index.Index64.empty(1, self.nplike)
+            out[0] = len(self)
+            return ak._v2.contents.numpyarray.NumpyArray(out)[0]
+        elif posaxis == depth + 1:
+            tonum = ak._v2.index.Index64.empty(len(self), self.nplike)
+            self._handle_error(
+                self.nplike[
+                    "awkward_ListArray_num",
+                    tonum.dtype.type,
+                    self.starts.dtype.type,
+                    self.stops.dtype.type,
+                ](
+                    tonum.to(self.nplike),
+                    self.starts.to(self.nplike),
+                    self.stops.to(self.nplike),
+                    len(self),
+                )
+            )
+            return ak._v2.contents.numpyarray.NumpyArray(tonum)
+        else:
+            next = self._content.num(posaxis, depth + 1)
+            offsets = self._compact_offsets64(True)
+            return ak._v2.contents.listoffsetarray.ListOffsetArray(
+                offsets, next, None, self.parameters
+            )
+
     def mergeable(self, other, mergebool):
         if not _parameters_equal(self._parameters, other._parameters):
             return False

src/awkward/_v2/contents/numpyarray.py

@@ -297,6 +297,35 @@ def _getitem_next(self, head, tail, advanced):
         else:
             raise AssertionError(repr(head))
 
+    def num(self, axis, depth=0):
+        posaxis = self.axis_wrap_if_negative(axis)
+        if posaxis == depth:
+            out = ak._v2.index.Index64.empty(1, self.nplike)
+            out[0] = len(self)
+            return ak._v2.contents.numpyarray.NumpyArray(out)[0]
+        shape = []
+        reps = 1
+        size = len(self)
+        i = 0
+        while i < self._data.ndim - 1 and depth < posaxis:
+            shape.append(self.shape[i])
+            reps *= self.shape[i]
+            size = self.shape[i + 1]
+            i += 1
+            depth += 1
+        if posaxis > depth:
+            raise ValueError("'axis' out of range for 'num'")
+
+        tonum = ak._v2.index.Index64.empty(reps, self.nplike)
+        self._handle_error(
+            self.nplike["awkward_RegularArray_num", tonum.dtype.type](
+                tonum.to(self.nplike), size, reps
+            )
+        )
+        return ak._v2.contents.numpyarray.NumpyArray(
+            tonum.data.reshape(shape), None, self.parameters
+        )
+
     def mergeable(self, other, mergebool):
         if not _parameters_equal(self._parameters, other._parameters):
             return False

src/awkward/_v2/contents/recordarray.py

@@ -371,6 +371,26 @@ def _getitem_next(self, head, tail, advanced):
             )
             return next._getitem_next(nexthead, nexttail, advanced)
 
+    def num(self, axis, depth=0):
+        posaxis = self.axis_wrap_if_negative(axis)
+        if posaxis == depth:
+            single = ak._v2.index.Index64.empty(1, self.nplike)
+            single[0] = len(self)
+            singleton = ak._v2.contents.numpyarray.NumpyArray(single)
+            contents = [singleton] * len(self._contents)
+
+            record = ak._v2.contents.recordarray.RecordArray(
+                contents, self._fields, 1, None, self._parameters
+            )
+            return record[0]
+        else:
+            contents = []
+            for content in self._contents:
+                contents.append(content.num(posaxis, depth))
+            return ak._v2.contents.recordarray.RecordArray(
+                contents, self._fields, self._length, None, self._parameters
+            )
+
     def mergeable(self, other, mergebool=True):
         if not _parameters_equal(self._parameters, other._parameters):
             return False

src/awkward/_v2/contents/regulararray.py

@@ -513,10 +513,29 @@ def _getitem_next(self, head, tail, advanced):
         else:
             raise AssertionError(repr(head))
 
+    def num(self, axis, depth=0):
+        posaxis = self.axis_wrap_if_negative(axis)
+        if posaxis == depth:
+            out = ak._v2.index.Index64.empty(1, self.nplike)
+            out[0] = len(self)
+            return ak._v2.contents.numpyarray.NumpyArray(out)[0]
+        elif posaxis == depth + 1:
+            tonum = ak._v2.index.Index64.empty(len(self), self.nplike)
+            self._handle_error(
+                self.nplike["awkward_RegularArray_num", tonum.dtype.type](
+                    tonum.to(self.nplike), self._size, len(self)
+                )
+            )
+            return ak._v2.contents.numpyarray.NumpyArray(tonum)
+        else:
+            next = self._content.num(posaxis, depth + 1)
+            return ak._v2.contents.regulararray.RegularArray(
+                next, self._size, len(self), None, self._parameters
+            )
+
     def mergeable(self, other, mergebool):
         if not _parameters_equal(self._parameters, other._parameters):
             return False
-
         if isinstance(
             other,
             (

src/awkward/_v2/contents/unionarray.py

@@ -470,6 +470,19 @@ def simplify_uniontype(self, merge=True, mergebool=False):
         else:
             return UnionArray(tags, index, contents, self.identifier, self.parameters)
 
+    def num(self, axis, depth=0):
+        posaxis = self.axis_wrap_if_negative(axis)
+        if posaxis == depth:
+            out = ak._v2.index.Index64.empty(1, self.nplike)
+            out[0] = len(self)
+            return ak._v2.contents.numpyarray.NumpyArray(out)[0]
+        else:
+            contents = []
+            for content in self._contents:
+                contents.append(content.num(posaxis, depth))
+            out = UnionArray(self._tags, self._index, contents, None, self._parameters)
+            return out.simplify_uniontype(True, False)
+
     def mergeable(self, other, mergebool):
         if not _parameters_equal(self._parameters, other._parameters):
             return False

src/awkward/_v2/contents/unmaskedarray.py

@@ -199,6 +199,17 @@ def simplify_optiontype(self):
         else:
             return self
 
+    def num(self, axis, depth=0):
+        posaxis = self.axis_wrap_if_negative(axis)
+        if posaxis == depth:
+            out = ak._v2.index.Index64.empty(1, self.nplike)
+            out[0] = len(self)
+            return ak._v2.contents.numpyarray.NumpyArray(out)[0]
+        else:
+            return ak._v2.contents.unmaskedarray.UnmaskedArray(
+                self._content.num(posaxis, depth), None, self._parameters
+            )
+
     def mergeable(self, other, mergebool):
         if not _parameters_equal(self._parameters, other._parameters):
             return False