ENH: start building RangeIndex to avoid creating so many np.arange(N) indexes

Author: wesm

pandas/core/common.py

@@ -525,8 +525,8 @@ def _is_bool_indexer(key):
     return False
 
 def _default_index(n):
-    from pandas.core.index import Index
-    return Index(np.arange(n))
+    from pandas.core.index import RangeIndex
+    return RangeIndex(n)
 
 def ensure_float(arr):
     if issubclass(arr.dtype.type, np.integer):

pandas/core/frame.py

@@ -401,7 +401,7 @@ def __init__(self, data=None, index=None, columns=None, dtype=None,
                     conv_data, columns = _to_sdict(data, columns)
                     if isinstance(conv_data, dict):
                         if len(conv_data) == 0 and index is None:
-                            index = np.arange(len(data))
+                            index = _default_index(len(data))
                         mgr = self._init_dict(conv_data, index, columns,
                                               dtype=dtype)
                     else:
@@ -876,9 +876,9 @@ def from_records(cls, data, index=None, exclude=None, columns=None,
                     result_index = index
         elif isinstance(data, dict) and len(data) > 0:
             # utilize first element of sdict to get length
-            result_index = np.arange(len(data.values()[0]))
+            result_index = _default_index(len(data.values()[0]))
         else:
-            result_index = np.arange(len(data))
+            result_index = _default_index(len(data))
 
         return cls(sdict, index=result_index, columns=columns)
 
@@ -2505,7 +2505,7 @@ def _maybe_cast(values):
                 values = lib.maybe_convert_objects(values)
             return values
 
-        new_index = np.arange(len(new_obj))
+        new_index = _default_index(len(new_obj))
         if isinstance(self.index, MultiIndex):
             if level is not None:
                 if not isinstance(level, (tuple, list)):
@@ -4910,7 +4910,7 @@ def extract_index(data):
                           % (lengths[0], len(index)))
                     raise ValueError(msg)
             else:
-                index = Index(np.arange(lengths[0]))
+                index = _default_index(lengths[0])
 
     return _ensure_index(index)
 

pandas/core/index.py

@@ -1110,7 +1110,7 @@ def delete(self, loc):
         new_index : Index
         """
         arr = np.delete(np.asarray(self), loc)
-        return Index(arr)
+        return Index(arr, name=self.name)
 
     def insert(self, loc, item):
         """
@@ -1151,7 +1151,11 @@ def drop(self, labels):
         return self.delete(indexer)
 
 
-class Int64Index(Index):
+class IntIndex(object):
+    pass
+
+
+class Int64Index(Index, IntIndex):
 
     _groupby = _algos.groupby_int64
     _arrmap = _algos.arrmap_int64
@@ -1219,7 +1223,10 @@ def equals(self, other):
         #     return False
 
         try:
-            return np.array_equal(self, other)
+            if isinstance(other, Index):
+                other = other.values
+
+            return np.array_equal(self.values, other)
         except TypeError:
             # e.g. fails in numpy 1.6 with DatetimeIndex #1681
             return False
@@ -1229,6 +1236,131 @@ def _wrap_joined_index(self, joined, other):
         return Int64Index(joined, name=name)
 
 
+class RangeIndex(Int64Index):
+    """
+    Compactly represents regular integer range (instead of generating a full
+    array of integers)
+    """
+
+    def __new__(cls, start=None, stop=None, step=None, name=None):
+        if stop is None and step is None:
+            #
+            stop = start
+            start = 0
+
+        if step is None:
+            step = 1
+
+        result = np.array([], dtype=np.int_).view(RangeIndex)
+        result.start = start
+        result.stop = stop
+        result.step = step
+        result.name = name
+
+        return result
+
+    def __array_finalize__(self, obj):
+        if not isinstance(obj, type(self)):
+            # Only relevant if array being created from an Index instance
+            return
+
+        self.start = getattr(obj, 'start', None)
+        self.stop = getattr(obj, 'stop', None)
+        self.step = getattr(obj, 'step', None)
+        self.name = getattr(obj, 'name', None)
+
+    @property
+    def values(self):
+        return np.arange(self.start, self.stop, self.step)
+
+    def __len__(self):
+        return (self.stop - self.start - 1) // self.step + 1
+
+    def __repr__(self):
+        if self.step != 1:
+            return 'RangeIndex(%d, %d, step=%d)' % (self.start, self.stop,
+                                                    self.step)
+        else:
+            return 'RangeIndex(%d, %d)' % (self.start, self.stop)
+
+    def __getslice__(self, i, j):
+        return self.__getitem__(slice(i, j))
+
+    def __getitem__(self, key):
+        if com.is_integer(key):
+            if key >= 0:
+                return self.start + self.step * key
+            else:
+                return self.stop + self.step * key
+        elif isinstance(key, slice):
+            kstart = key.start or 0
+            kstop = self.stop if key.stop is None else key.stop
+            kstep = key.step or 1
+
+            if kstop >= 0:
+                new_stop = min(self.start + kstop * self.step, self.stop)
+            else:
+                new_stop = self.stop + kstop * self.step
+
+            if kstart is None:
+                new_start = self.start
+            else:
+                if kstart >= 0:
+                    new_start = self.start + kstart * self.step
+                else:
+                    new_start = self.step + kstart * self.step
+
+            new_step = self.step * kstep
+            if kstep < 0:
+                return RangeIndex(new_stop - self.step, new_start - self.step,
+                                  new_step, name=self.name)
+            else:
+                return RangeIndex(new_start, new_stop, new_step, name=self.name)
+
+        elif isinstance(key, (list, np.ndarray)):
+            key = np.asarray(key)
+
+            if key.dtype == np.bool_:
+                return Int64Index(self.values[key], name=self.name)
+
+            return np.where(key >= 0, self.start + self.step * key,
+                            self.stop + self.step * key)
+        else: # pragma: no cover
+            raise TypeError('Unhandled getitem type: %s' % type(key))
+
+    def insert(self, loc, item):
+        return Int64Index(self.values, name=self.name).insert(loc, item)
+
+    def delete(self, loc):
+        return Int64Index(self.values, name=self.name).delete(loc)
+
+    def take(self, key):
+        # XXX
+        return self[key]
+
+    def equals(self, other):
+        """
+        Returns True if the indexes are equivalent
+
+        Parameters
+        ----------
+        other : Index
+
+        Returns
+        -------
+        is_equal : True or False
+        """
+        if not isinstance(other, RangeIndex):
+            if isinstance(other, Index):
+                other = other.values
+            return np.array_equal(self.values, other)
+
+        return (self.start == other.start and self.stop == other.stop
+                and self.step == other.step)
+
+    @property
+    def _constructor(self):
+        return RangeIndex
 
 
 class MultiIndex(Index):

pandas/tests/test_index.py

@@ -10,8 +10,7 @@
 import numpy as np
 from numpy.testing import assert_array_equal
 
-from pandas.core.categorical import Factor
-from pandas.core.index import Index, Int64Index, MultiIndex
+from pandas.core.index import Index, Int64Index, MultiIndex, RangeIndex
 from pandas.util.testing import assert_almost_equal
 from pandas.util import py3compat
 import pandas.core.common as com
@@ -848,6 +847,156 @@ def test_int_name_format(self):
         repr(s)
         repr(df)
 
+
+class TestRangeIndex(unittest.TestCase):
+
+    def test_constructor(self):
+        pass
+
+    def test_repr(self):
+        index = RangeIndex(0, 20, 1)
+        result = repr(index)
+        exp = 'RangeIndex(0, 20)'
+        self.assertEqual(result, exp)
+
+        index = RangeIndex(0, 20, 3)
+        result = repr(index)
+        exp = 'RangeIndex(0, 20, step=3)'
+        self.assertEqual(result, exp)
+
+    def test_len(self):
+        for i in range(20):
+            index = RangeIndex(0, i, 3)
+            self.assertEqual(len(index), len(list(range(0, i, 3))))
+
+    def test_values(self):
+        index = RangeIndex(1000)
+
+        values = index.values
+        expected = np.arange(1000)
+        assert_almost_equal(values, expected)
+
+    def test_getitem(self):
+        index = RangeIndex(0, 20, 2)
+
+        self.assertEqual(index[4], 8)
+        self.assertEqual(index[-1], 18)
+
+    def test_fancy_index(self):
+        index = RangeIndex(0, 20, 2)
+        result = index[[4, 1, 3, -1]]
+        expected = np.array([8, 2, 6, 18])
+        assert_almost_equal(result, expected)
+
+    def test_boolean_index(self):
+        index = RangeIndex(0, 20, 2)
+
+        mask = np.ones(len(index), dtype=bool)
+        mask[::2] = False
+
+        result = index[mask]
+        exp = index.values[mask]
+        assert_almost_equal(result, exp)
+
+    def test_take(self):
+        index = RangeIndex(0, 20, 2)
+        result = index.take([4, 1, 3, -1])
+        expected = np.array([8, 2, 6, 18])
+        assert_almost_equal(result, expected)
+
+    def test_slice(self):
+        index = RangeIndex(5, 20, name='foo')
+
+        result = index[2:]
+        self.assertTrue(result.equals(RangeIndex(7, 20)))
+        self.assertTrue(result.name == 'foo')
+
+        result = index[:7]
+        self.assertTrue(result.equals(RangeIndex(5, 12)))
+
+        result = index[3:10]
+        self.assertTrue(result.equals(RangeIndex(8, 15)))
+
+        result = index[-5:]
+        self.assertTrue(result.equals(RangeIndex(15, 20)))
+
+        result = index[-5:-2]
+        self.assertTrue(result.equals(RangeIndex(15, 18)))
+
+        result = index[::-1]
+        expected = RangeIndex(19, 4, -1)
+        self.assertTrue(result.equals(expected))
+
+        result = result[::-1]
+        self.assertTrue(result.equals(index))
+
+        result = index[2::2]
+        exp = RangeIndex(7, 20, 2)
+        self.assertTrue(result.equals(exp))
+
+        # with a step
+        index = RangeIndex(5, 41, 2)
+        result = index[2:]
+        self.assertTrue(result.equals(RangeIndex(9, 41, 2)))
+
+        result = index[:7]
+        self.assertTrue(result.equals(RangeIndex(5, 19, 2)))
+
+        result = index[3:10]
+        self.assertTrue(result.equals(RangeIndex(11, 25, 2)))
+
+        result = index[-5:]
+        self.assertTrue(result.equals(RangeIndex(31, 41, 2)))
+
+        result = index[-5:-2]
+        self.assertTrue(result.equals(RangeIndex(31, 37, 2)))
+
+        result = index[::-1]
+        expected = RangeIndex(39, 3, -2)
+        self.assertTrue(result.equals(expected))
+
+        result = result[::-1]
+        self.assertTrue(result.equals(index))
+
+        # corner case
+        index = RangeIndex(0, 1)
+        self.assertTrue(index[:0].equals(RangeIndex(0, 0)))
+
+    def test_insert(self):
+        index = RangeIndex(5, name='bar')
+
+        result = index.insert(2, 'foo')
+        exp = Index([0, 1, 'foo', 2, 3, 4])
+        self.assertTrue(result.equals(exp))
+        self.assertEqual(result.name, 'bar')
+
+    def test_delete(self):
+        index = RangeIndex(5, name='bar')
+        result = index.delete(2)
+        exp = Index([0, 1, 3, 4])
+        self.assertTrue(result.equals(exp))
+        self.assertEqual(result.name, 'bar')
+
+    def test_union(self):
+        pass
+
+    def test_intersection(self):
+        pass
+
+    def test_pickle(self):
+        pass
+
+    def test_equals(self):
+        left = RangeIndex(0, 20)
+
+        self.assertTrue(left.equals(RangeIndex(0, 20)))
+        self.assertFalse(left.equals(RangeIndex(0, 10)))
+        self.assertFalse(left.equals(RangeIndex(0, 20, 2)))
+
+        left = RangeIndex(0, 20, 2)
+        self.assertTrue(left.equals(RangeIndex(0, 20, 2)))
+
+
 class TestMultiIndex(unittest.TestCase):
 
     def setUp(self):