WIP: high memory in MI

xref pandas-dev#13904

Author: jreback

pandas/hashtable.pxd

@@ -1,3 +1,4 @@
+from numpy cimport ndarray
 from khash cimport (kh_int64_t, kh_uint64_t, kh_float64_t, kh_pymap_t,
                     kh_str_t, uint64_t, int64_t, float64_t)
 
@@ -30,6 +31,15 @@ cdef class PyObjectHashTable(HashTable):
     cpdef get_item(self, object val)
     cpdef set_item(self, object key, Py_ssize_t val)
 
+cdef class MultiIndexHashTable(HashTable):
+    cdef:
+        kh_int64_t *table
+        object values
+        int64_t nlevels
+
+    cpdef get_item(self, object val)
+    cpdef set_item(self, object key, Py_ssize_t val)
+
 cdef class StringHashTable(HashTable):
     cdef kh_str_t *table
 

pandas/index.pyx

@@ -1,4 +1,4 @@
-# cython: profile=False
+# cython: profile=True
 
 from numpy cimport ndarray
 
@@ -275,7 +275,6 @@ cdef class IndexEngine:
         if not self.is_mapping_populated:
 
             values = self._get_index_values()
-
             self.mapping = self._make_hash_table(len(values))
             self.mapping.map_locations(values)
 
@@ -545,5 +544,60 @@ cdef inline bint _is_utc(object tz):
     return tz is UTC or isinstance(tz, _du_utc)
 
 
+cdef class MultiIndexEngine(IndexEngine):
+
+    def _call_monotonic(self, values):
+        return algos.is_monotonic_object(values, timelike=False)
+
+    def get_backfill_indexer(self, other, limit=None):
+        return algos.backfill_object(self._get_index_values(),
+                                        other, limit=limit)
+
+    def get_pad_indexer(self, other, limit=None):
+        return algos.pad_object(self._get_index_values(),
+                                   other, limit=limit)
+
+    cdef _make_hash_table(self, n):
+        return _hash.MultiIndexHashTable(n)
+
+    cdef _check_type(self, object val):
+        hash(val)
+        if util.is_bool_object(val):
+            raise KeyError(val)
+        elif util.is_float_object(val):
+            raise KeyError(val)
+
+    cdef _maybe_get_bool_indexer(self, object val):
+        cdef:
+            ndarray[uint8_t, cast=True] indexer
+            ndarray[object] values
+            int count = 0
+            Py_ssize_t i, n
+            int last_true
+
+        if not util.is_integer_object(val):
+            raise KeyError(val)
+
+        values = self._get_index_values()
+        n = len(values)
+
+        result = np.empty(n, dtype=bool)
+        indexer = result.view(np.uint8)
+
+        for i in range(n):
+            if values[i] == val:
+                count += 1
+                indexer[i] = 1
+                last_true = i
+            else:
+                indexer[i] = 0
+
+        if count == 0:
+            raise KeyError(val)
+        if count == 1:
+            return last_true
+
+        return result
+
 # Generated from template.
 include "index_class_helper.pxi"

pandas/indexes/multi.py

@@ -73,6 +73,7 @@ class MultiIndex(Index):
     _levels = FrozenList()
     _labels = FrozenList()
     _comparables = ['names']
+    _engine_type = _index.MultiIndexEngine
     rename = Index.set_names
 
     def __new__(cls, levels=None, labels=None, sortorder=None, names=None,
@@ -114,7 +115,6 @@ def __new__(cls, levels=None, labels=None, sortorder=None, names=None,
             result._verify_integrity()
         if _set_identity:
             result._reset_identity()
-
         return result
 
     def _verify_integrity(self, labels=None, levels=None):
@@ -613,8 +613,13 @@ def _get_level_number(self, level):
 
     _tuples = None
 
+    @cache_readonly
+    def _engine(self):
+        return self._engine_type(lambda: self, len(self))
+
     @property
     def values(self):
+        print("values called")
         if self._tuples is not None:
             return self._tuples
 

pandas/src/hashtable_class_helper.pxi.in

@@ -843,3 +843,153 @@ cdef class PyObjectHashTable(HashTable):
                 count += 1
 
         return np.asarray(labels)
+
+
+cdef class MultiIndexHashTable(HashTable):
+
+    def __init__(self, size_hint=1):
+        self.table = kh_init_int64()
+        self.values = None
+        self.nlevels = 0
+        kh_resize_int64(self.table, size_hint)
+
+    def __dealloc__(self):
+        if self.table is not NULL:
+            kh_destroy_int64(self.table)
+            self.table = NULL
+            self.values = None
+            self.nlevels = 0
+
+    def __len__(self):
+        return self.table.size
+
+    def __contains__(self, object key):
+        cdef khiter_t k
+        hash(key)
+        if key != key or key is None:
+            key = na_sentinel
+        k = kh_get_int64(self.table, hash(key))
+        return k != self.table.n_buckets
+
+    cpdef get_item(self, object val):
+        cdef:
+            khiter_t k
+            Py_ssize_t level
+            int64_t i
+
+        if val != val or val is None:
+            val = na_sentinel
+        k = kh_get_int64(self.table, hash(val))
+        if k != self.table.n_buckets:
+            i = self.table.vals[k]
+            return tuple([self.values[level][i] for level in range(self.nlevels)])
+        else:
+            raise KeyError(val)
+
+    def get_iter_test(self, object key, Py_ssize_t iterations):
+        cdef Py_ssize_t i, val
+        if key != key or key is None:
+            key = na_sentinel
+        for i in range(iterations):
+            k = kh_get_int64(self.table, hash(key))
+            if k != self.table.n_buckets:
+                val = self.table.vals[k]
+
+    cpdef set_item(self, object key, Py_ssize_t val):
+        raise NotImplementedError
+
+    def map_locations(self, object mi):
+        cdef:
+            Py_ssize_t i, n, level
+            int64_t val
+            int ret = 0
+            khiter_t k
+
+        # the values are a rec-array
+        self.nlevels = len(mi.levels)
+        self.values = [mi.get_level_values(level).values
+                       for level in range(self.nlevels)]
+
+        n = len(self.values[0])
+        for i in range(n):
+            val = hash(tuple([self.values[level][i]
+                              for level in range(self.nlevels)]))
+            k = kh_put_int64(self.table, val, &ret)
+            self.table.vals[k] = i
+
+    def lookup(self, ndarray[object] values):
+        cdef:
+            Py_ssize_t i, n = len(values)
+            int ret = 0
+            int64_t val
+            khiter_t k
+            int64_t[:] locs = np.empty(n, dtype=np.int64)
+
+        for i in range(n):
+            val = hash(values[i])
+            if val != val or val is None:
+                val = na_sentinel
+
+            k = kh_get_int64(self.table, val)
+            if k != self.table.n_buckets:
+                locs[i] = self.table.vals[k]
+            else:
+                locs[i] = -1
+
+        return np.asarray(locs)
+
+    def unique(self, ndarray[object] values):
+        cdef:
+            Py_ssize_t i, n = len(values)
+            int ret = 0
+            int64_t val
+            khiter_t k
+            ObjectVector uniques = ObjectVector()
+            bint seen_na = 0
+
+        for i in range(n):
+            val = hash(values[i])
+
+            if not _checknan(val):
+                k = kh_get_int64(self.table, val)
+                if k == self.table.n_buckets:
+                    kh_put_int64(self.table, val, &ret)
+                    uniques.append(val)
+            elif not seen_na:
+                seen_na = 1
+                uniques.append(nan)
+
+        return uniques.to_array()
+
+    def get_labels(self, ndarray[object] values, ObjectVector uniques,
+                   Py_ssize_t count_prior, int64_t na_sentinel,
+                   bint check_null=True):
+        cdef:
+            Py_ssize_t i, n = len(values)
+            int64_t[:] labels
+            Py_ssize_t idx, count = count_prior
+            int ret = 0
+            int64_t val
+            khiter_t k
+
+        labels = np.empty(n, dtype=np.int64)
+
+        for i in range(n):
+            val = hash(values[i])
+
+            if check_null and val != val or val is None:
+                labels[i] = na_sentinel
+                continue
+
+            k = kh_get_int64(self.table, val)
+            if k != self.table.n_buckets:
+                idx = self.table.vals[k]
+                labels[i] = idx
+            else:
+                k = kh_put_int64(self.table, val, &ret)
+                self.table.vals[k] = count
+                uniques.append(val)
+                labels[i] = count
+                count += 1
+
+        return np.asarray(labels)