optimized iteration speed over LArray, Group and Axis

Author: gdementen

larray/core/array.py

@@ -64,7 +64,7 @@
 from larray.core.axis import Axis, AxisReference, AxisCollection, X, _make_axis
 from larray.util.misc import (table2str, size2str, basestring, izip, rproduct, ReprString, duplicates,
                               float_error_handler_factory, _isnoneslice, light_product, unique_list, common_type,
-                              renamed_to, deprecate_kwarg, LHDFStore, lazy_attribute)
+                              renamed_to, deprecate_kwarg, LHDFStore, lazy_attribute, PY2)
 from larray.util.options import _OPTIONS, DISPLAY_MAXLINES, DISPLAY_EDGEITEMS, DISPLAY_WIDTH, DISPLAY_PRECISION
 
 
@@ -312,20 +312,21 @@ def concat(arrays, axis=0, dtype=None):
 
 class LArrayIterator(object):
     def __init__(self, array):
-        self.array = array
-        self.index = 0
+        data_iter = iter(array.data)
+        self.nextfunc = data_iter.next if PY2 else data_iter.__next__
+        self.axes = array.axes[1:]
 
     def __iter__(self):
         return self
 
     def __next__(self):
-        array = self.array
-        if self.index == len(self.array):
-            raise StopIteration
-        # result = array.i[array.axes[0].i[self.index]]
-        result = array.i[self.index]
-        self.index += 1
-        return result
+        data = self.nextfunc()
+        axes = self.axes
+        if len(axes):
+            return LArray(data, axes)
+        else:
+            return data
+
     # Python 2
     next = __next__
 
@@ -2327,7 +2328,11 @@ def __str__(self):
     __repr__ = __str__
 
     def __iter__(self):
-        return LArrayIterator(self)
+        # fast path for 1D arrays where we return elements
+        if self.ndim <= 1:
+            return iter(self.data)
+        else:
+            return LArrayIterator(self)
 
     def __contains__(self, key):
         return any(key in axis for axis in self.axes)

larray/core/axis.py

@@ -721,7 +721,7 @@ def __len__(self):
         return self._length
 
     def __iter__(self):
-        return iter([self.i[i] for i in range(self._length)])
+        return iter([IGroup(i, None, self) for i in range(self._length)])
 
     def __getitem__(self, key):
         """

larray/core/group.py

@@ -869,7 +869,8 @@ def __len__(self):
     def __iter__(self):
         # XXX: use translate/IGroup instead, so that it works even in the presence of duplicate labels
         #      possibly, only if axis is set?
-        return iter([LGroup(v, axis=self.axis) for v in self.eval()])
+        axis = self.axis
+        return iter([LGroup(v, axis=axis) for v in self.eval()])
 
     def named(self, name):
         """Returns group with a different name.