ENH: optimize join/merge on integer keys, close #682

Author: wesm

RELEASE.rst

@@ -57,6 +57,7 @@ pandas 0.8.0
   - Improve the speed of "square" reindexing of homogeneous DataFrame objects
     by significant margin (#836)
   - Handle more dtypes when passed MaskedArrays in DataFrame constructor (#406)
+  - Improved performance of join operations on integer keys (#682)
 
 **API Changes**
 

pandas/src/hashtable.pyx

@@ -823,9 +823,10 @@ cdef class Int64Factorizer:
     def get_count(self):
         return self.count
 
-    def factorize(self, ndarray[int64_t] values, sort=False):
+    def factorize(self, ndarray[int64_t] values, sort=False,
+                  na_sentinel=-1):
         labels, counts = self.table.get_labels(values, self.uniques,
-                                               self.count, -1)
+                                               self.count, na_sentinel)
 
         # sort on
         if sort:

pandas/tools/merge.py

@@ -304,7 +304,7 @@ def _get_group_keys(self):
         group_sizes = []
 
         for lk, rk in zip(left_keys, right_keys):
-            llab, rlab, count = _factorize_objects(lk, rk, sort=self.sort)
+            llab, rlab, count = _factorize_keys(lk, rk, sort=self.sort)
 
             left_labels.append(llab)
             right_labels.append(rlab)
@@ -321,24 +321,24 @@ def _get_group_keys(self):
             raise Exception('Combinatorial explosion! (boom)')
 
         left_group_key, right_group_key, max_groups = \
-            _factorize_int64(left_group_key, right_group_key,
+            _factorize_keys(left_group_key, right_group_key,
                              sort=self.sort)
         return left_group_key, right_group_key, max_groups
 
 def _get_multiindex_indexer(join_keys, index, sort=False):
     shape = []
     labels = []
     for level, key in zip(index.levels, join_keys):
-        llab, rlab, count = _factorize_objects(level, key, sort=False)
+        llab, rlab, count = _factorize_keys(level, key, sort=False)
         labels.append(rlab)
         shape.append(count)
 
     left_group_key = get_group_index(labels, shape)
     right_group_key = get_group_index(index.labels, shape)
 
     left_group_key, right_group_key, max_groups = \
-        _factorize_int64(left_group_key, right_group_key,
-                         sort=False)
+        _factorize_keys(left_group_key, right_group_key,
+                        sort=False)
 
     left_indexer, right_indexer = \
         lib.left_outer_join(com._ensure_int64(left_group_key),
@@ -348,7 +348,7 @@ def _get_multiindex_indexer(join_keys, index, sort=False):
     return left_indexer, right_indexer
 
 def _get_single_indexer(join_key, index, sort=False):
-    left_key, right_key, count = _factorize_objects(join_key, index, sort=sort)
+    left_key, right_key, count = _factorize_keys(join_key, index, sort=sort)
 
     left_indexer, right_indexer = \
         lib.left_outer_join(com._ensure_int64(left_key),
@@ -394,26 +394,21 @@ def _left_join_on_index(left_ax, right_ax, join_keys, sort=False):
     'outer' : lib.full_outer_join,
 }
 
-def _factorize_int64(left_index, right_index, sort=True):
-    rizer = lib.Int64Factorizer(max(len(left_index), len(right_index)))
 
-    # 32-bit compatibility
-    left_index = com._ensure_int64(left_index)
-    right_index = com._ensure_int64(right_index)
-
-    llab, _ = rizer.factorize(left_index)
-    rlab, _ = rizer.factorize(right_index)
-
-    if sort:
-        llab, rlab = _sort_labels(np.array(rizer.uniques), llab, rlab)
-
-    return llab, rlab, rizer.get_count()
+def _factorize_keys(lk, rk, sort=True):
+    if com.is_integer_dtype(lk) and com.is_integer_dtype(rk):
+        klass = lib.Int64Factorizer
+        lk = com._ensure_int64(lk)
+        rk = com._ensure_int64(rk)
+    else:
+        klass = lib.Factorizer
+        lk = com._ensure_object(lk)
+        rk = com._ensure_object(rk)
 
-def _factorize_objects(left_index, right_index, sort=True):
-    rizer = lib.Factorizer(max(len(left_index), len(right_index)))
+    rizer = klass(max(len(lk), len(rk)))
 
-    llab, _ = rizer.factorize(left_index.astype('O'))
-    rlab, _ = rizer.factorize(right_index.astype('O'))
+    llab, _ = rizer.factorize(lk)
+    rlab, _ = rizer.factorize(rk)
 
     count = rizer.get_count()
 

vb_suite/join_merge.py

@@ -66,6 +66,12 @@
               name='join_dataframe_index_multi',
               start_date=datetime(2011, 10, 20))
 
+#----------------------------------------------------------------------
+# Joins on integer keys
+
+join_dataframe_integer_key = Benchmark("merge(df, df2, on='key')", setup,
+                                       start_date=datetime(2011, 10, 20))
+
 #----------------------------------------------------------------------
 # DataFrame joins on index