ENH: merge/join functions, compress group index with possible number of groups is insanely large, speed enhancements, refactoring

Author: wesm

TODO.rst

@@ -1,4 +1,4 @@
-DONE
+meDONE
 ----
 - SparseSeries name integration + tests
 - Refactor Series.repr

bench/bench_merge.R

@@ -0,0 +1,11 @@
+N <- 10000
+indices = rep(NA, N)
+for (i in 1:N)
+  indices[i] <- paste(sample(letters, 10), collapse="")
+
+left <- data.frame(key=rep(indices, 10),
+                   key2=sample(rep(indices, 10)),
+                   value=rnorm(100000))
+right <- data.frame(key=indices,
+                    key2=sample(indices),
+                    value2=rnorm(10000))

bench/bench_merge.py

@@ -1,4 +1,5 @@
 from pandas import *
+from pandas.util.testing import rands
 import random
 
 N = 10000
@@ -30,3 +31,19 @@ def get_test_data(ngroups=100, n=N):
 reload(merge)
 
 result = merge.merge(df, df2, on='key2')
+
+from pandas.util.testing import rands
+N = 10000
+indices = np.array([rands(10) for _ in xrange(N)], dtype='O')
+
+key = np.tile(indices, 10)
+key2 = key.copy()
+random.shuffle(key2)
+indices2 = indices.copy()
+random.shuffle(indices2)
+
+
+left = DataFrame({'key' : key, 'key2':key2,
+                  'value' : np.random.randn(100000)})
+right = DataFrame({'key': indices, 'key2':indices2,
+                   'value2' : np.random.randn(10000)})

pandas/core/generic.py

@@ -432,23 +432,6 @@ def cumprod(self, axis=None, skipna=True):
             result = y.cumprod(axis)
         return self._wrap_array(result, self.axes, copy=False)
 
-    def _values_aggregate(self, func, axis, fill_value, skipna=True):
-        axis = self._get_axis_number(axis)
-
-        values = self.values
-        mask = np.isfinite(values)
-
-        if skipna and fill_value is not None:
-            values = values.copy()
-            values[-mask] = fill_value
-
-        result = func(values, axis=axis)
-        count = mask.sum(axis=axis)
-
-        result[count == 0] = np.NaN
-
-        return result
-
     def copy(self, deep=True):
         """
         Make a copy of this object

pandas/src/hashtable.pyx

@@ -239,6 +239,113 @@ cdef class StringHashTable:
         # return None
         return reverse, labels, counts[:count].copy()
 
+cdef class Int32HashTable:
+
+    cdef:
+        kh_int32_t *table
+
+    def __init__(self, size_hint=1):
+        if size_hint is not None:
+            kh_resize_int32(self.table, size_hint)
+
+    def __cinit__(self):
+        self.table = kh_init_int32()
+
+    def __dealloc__(self):
+        kh_destroy_int32(self.table)
+
+    cdef inline int check_type(self, object val):
+        return PyString_Check(val)
+
+    cpdef get_item(self, int32_t val):
+        cdef khiter_t k
+        k = kh_get_int32(self.table, val)
+        if k != self.table.n_buckets:
+            return self.table.vals[k]
+        else:
+            raise KeyError(val)
+
+    def get_iter_test(self, int32_t key, Py_ssize_t iterations):
+        cdef Py_ssize_t i, val
+        for i in range(iterations):
+            k = kh_get_int32(self.table, val)
+            if k != self.table.n_buckets:
+                val = self.table.vals[k]
+
+    cpdef set_item(self, int32_t key, Py_ssize_t val):
+        cdef:
+            khiter_t k
+            int ret
+
+        k = kh_put_int32(self.table, key, &ret)
+        self.table.keys[k] = key
+        if kh_exist_int32(self.table, k):
+            self.table.vals[k] = val
+        else:
+            raise KeyError(key)
+
+    def map_locations(self, ndarray[int32_t] values):
+        cdef:
+            Py_ssize_t i, n = len(values)
+            int ret
+            int32_t val
+            khiter_t k
+
+        for i in range(n):
+            val = values[i]
+            k = kh_put_int32(self.table, val, &ret)
+            # print 'putting %s, %s' % (val, count)
+            self.table.vals[k] = i
+
+    def lookup_locations(self, ndarray[int32_t] values):
+        cdef:
+            Py_ssize_t i, n = len(values)
+            int ret
+            int32_t val
+            khiter_t k
+            ndarray[int32_t] locs = np.empty(n, dtype='i4')
+
+        for i in range(n):
+            val = values[i]
+            k = kh_get_int32(self.table, val)
+            if k != self.table.n_buckets:
+                locs[i] = self.table.vals[k]
+            else:
+                locs[i] = -1
+
+        return locs
+
+    def factorize(self, ndarray[int32_t] values):
+        cdef:
+            Py_ssize_t i, n = len(values)
+            ndarray[int32_t] labels = np.empty(n, dtype=np.int32)
+            ndarray[int32_t] counts = np.empty(n, dtype=np.int32)
+            dict reverse = {}
+            Py_ssize_t idx, count = 0
+            int ret
+            int32_t val
+            khiter_t k
+
+        for i in range(n):
+            val = values[i]
+            k = kh_get_int32(self.table, val)
+            if k != self.table.n_buckets:
+                idx = self.table.vals[k]
+                labels[i] = idx
+                counts[idx] = counts[idx] + 1
+            else:
+                k = kh_put_int32(self.table, val, &ret)
+                if not ret:
+                    kh_del_int32(self.table, k)
+                self.table.vals[k] = count
+                reverse[count] = val
+                labels[i] = count
+                counts[count] = 1
+                count += 1
+
+        # return None
+        return reverse, labels, counts[:count].copy()
+
 cdef class Int64HashTable:
 
     cdef:
@@ -315,17 +422,25 @@ cdef class Int64HashTable:
 
         return locs
 
-    def factorize(self, ndarray[int64_t] values):
+    def factorize(self, ndarray[object] values):
+        reverse = {}
+        labels, counts = self.get_labels(values, reverse, 0)
+        return reverse, labels, counts
+
+    def get_labels(self, ndarray[int64_t] values, list uniques,
+                   Py_ssize_t count_prior):
         cdef:
             Py_ssize_t i, n = len(values)
-            ndarray[int32_t] labels = np.empty(n, dtype=np.int32)
-            ndarray[int32_t] counts = np.empty(n, dtype=np.int32)
-            dict reverse = {}
-            Py_ssize_t idx, count = 0
+            ndarray[int32_t] labels
+            ndarray[int32_t] counts
+            Py_ssize_t idx, count = count_prior
             int ret
             int64_t val
             khiter_t k
 
+        labels = np.empty(n, dtype=np.int32)
+        counts = np.empty(count_prior + n, dtype=np.int32)
+
         for i in range(n):
             val = values[i]
             k = kh_get_int64(self.table, val)
@@ -335,16 +450,13 @@ cdef class Int64HashTable:
                 counts[idx] = counts[idx] + 1
             else:
                 k = kh_put_int64(self.table, val, &ret)
-                if not ret:
-                    kh_del_int64(self.table, k)
                 self.table.vals[k] = count
-                reverse[count] = val
+                uniques.append(val)
                 labels[i] = count
                 counts[count] = 1
                 count += 1
 
-        # return None
-        return reverse, labels, counts[:count].copy()
+        return labels, counts[:count].copy()
 
 cdef class PyObjectHashTable:
 
@@ -526,6 +638,37 @@ cdef class Factorizer:
         self.count = len(counts)
         return labels, counts
 
+cdef class Int64Factorizer:
+
+    cdef public:
+        Int64HashTable table
+        list uniques
+        Py_ssize_t count
+
+    def __init__(self, size_hint):
+        self.table = Int64HashTable(size_hint)
+        self.uniques = []
+        self.count = 0
+
+    def get_count(self):
+        return self.count
+
+    def factorize(self, ndarray[int64_t] values, sort=False):
+        labels, counts = self.table.get_labels(values, self.uniques,
+                                               self.count)
+
+        # sort on
+        if sort:
+            sorter = list_to_object_array(self.uniques).argsort()
+            reverse_indexer = np.empty(len(sorter), dtype=np.int32)
+            reverse_indexer.put(sorter, np.arange(len(sorter)))
+
+            labels = reverse_indexer.take(labels)
+            counts = counts.take(sorter)
+
+        self.count = len(counts)
+        return labels, counts
+
 def lookup_locations2(ndarray[object] values):
     cdef:
         Py_ssize_t i, n = len(values)

pandas/src/iterator.pyx

@@ -1,3 +1,4 @@
+import time
 
 def inner_join(ndarray[int32_t] left, ndarray[int32_t] right,
                Py_ssize_t max_groups):
@@ -170,3 +171,34 @@ def _get_result_indexer(sorter, indexer):
     res = sorter.take(indexer)
     np.putmask(res, indexer == -1, -1)
     return res
+
+
+@cython.boundscheck(False)
+@cython.wraparound(False)
+def join_sorter(ndarray[int32_t] index, Py_ssize_t ngroups):
+    cdef:
+        Py_ssize_t i, loc, label, n
+        ndarray[int32_t] counts, where, result
+
+    # count group sizes, location 0 for NA
+    counts = np.zeros(ngroups + 1, dtype='i4')
+    n = len(index)
+    for i from 0 <= i < n:
+        counts[index[i] + 1] += 1
+
+    # mark the start of each contiguous group of like-indexed data
+    where = np.zeros(ngroups + 1, dtype='i4')
+    for i from 1 <= i < ngroups + 1:
+        where[i] = where[i - 1] + counts[i - 1]
+
+    # this is our indexer
+    result = np.zeros(n, dtype='i4')
+    for i from 0 <= i < n:
+        label = index[i] + 1
+        result[where[label]] = i
+        where[label] += 1
+
+    return result, counts
+
+def _big_join_sorter(index):
+    pass

pandas/src/join.pyx

@@ -585,9 +585,11 @@ KHASH_SET_INIT_PYOBJECT(pyset)
 #define kh_exist_pyset(h, k) (kh_exist(h, k))
 #define kh_exist_str(h, k) (kh_exist(h, k))
 #define kh_exist_int64(h, k) (kh_exist(h, k))
+#define kh_exist_int32(h, k) (kh_exist(h, k))
 
 KHASH_MAP_INIT_STR(str, Py_ssize_t)
 
+KHASH_MAP_INIT_INT(int32, Py_ssize_t)
 KHASH_MAP_INIT_INT64(int64, Py_ssize_t)
 
 #endif /* __AC_KHASH_H */

pandas/src/khash.h

@@ -1,5 +1,5 @@
 from cpython cimport PyObject
-from numpy cimport int64_t, uint32_t
+from numpy cimport int64_t, int32_t, uint32_t
 
 cdef extern from "khash.h":
     ctypedef uint32_t khint_t
@@ -71,3 +71,19 @@ cdef extern from "khash.h":
 
     bint kh_exist_int64(kh_int64_t*, khiter_t)
 
+    ctypedef struct kh_int32_t:
+        khint_t n_buckets, size, n_occupied, upper_bound
+        uint32_t *flags
+        int32_t *keys
+        Py_ssize_t *vals
+
+    inline kh_int32_t* kh_init_int32()
+    inline void kh_destroy_int32(kh_int32_t*)
+    inline void kh_clear_int32(kh_int32_t*)
+    inline khint_t kh_get_int32(kh_int32_t*, int32_t)
+    inline void kh_resize_int32(kh_int32_t*, khint_t)
+    inline khint_t kh_put_int32(kh_int32_t*, int32_t, int*)
+    inline void kh_del_int32(kh_int32_t*, khint_t)
+
+    bint kh_exist_int32(kh_int32_t*, khiter_t)
+

pandas/src/khash.pxd

@@ -75,6 +75,7 @@ cdef class IndexableSkiplist:
                 i -= node.width[level]
                 node = node.next[level]
 
+
         return node.value
 
     cpdef insert(self, double value):