ENH: inner_join_indexer/outer_join_indexer functions

Author: wesm

pandas/src/reindex.pyx

@@ -409,36 +409,131 @@ def ordered_left_join_put(ndarray[int64_t] left, ndarray[int64_t] right,
             for k from kleft <= k < kleft + kright:
                 out[i, k] = NaN
 
-
-def ordered_outer_join(ndarray[int64] left, ndarray[int64] right):
+def inner_join_indexer(ndarray[int64_t] left, ndarray[int64_t] right):
+    '''
+    Two-pass algorithm?
+    '''
     cdef:
-        Py_ssize_t i, j, k, nright, nleft, kright, kleft
+        Py_ssize_t i, j, k, nright, nleft, count
         int64_t val
+        ndarray[int32_t] lindexer, rindexer
+        ndarray[int64_t] result
 
-    nleft, kleft = (<object> lvalues).shape
-    nright, kright = (<object> rvalues).shape
+    nleft = len(left)
+    nright = len(right)
 
+    lindexer = np.empty(min(nleft, nright), dtype=np.int32)
+    rindexer = np.empty(min(nleft, nright), dtype=np.int32)
+    result = np.empty(min(nleft, nright), dtype=np.int64)
+
+    i = 0
     j = 0
-    for i from 0 <= i < nleft:
-        for k from 0 <= k < kleft:
-            out[i, k] = lvalues[i, k]
+    count = 0
+    while True:
+        if i == nleft:
+            break
 
         val = left[i]
 
         while j < nright and right[j] < val:
             j += 1
 
         if j == nright:
-            for k from kleft <= k < kleft + kright:
-                out[i, k] = NaN
-            continue
+            break
 
         if val == right[j]:
-            for k from kleft <= k < kleft + kright:
-                out[i, k] = rvalues[j, k - kleft]
+            lindexer[count] = i
+            rindexer[count] = j
+            result[count] = val
+            count += 1
+            i += 1
+            j += 1
         else:
-            for k from kleft <= k < kleft + kright:
-                out[i, k] = NaN
+            while left[i] < right[j]:
+                i += 1
+
+    return (result[:count].copy(),
+            lindexer[:count].copy(),
+            rindexer[:count].copy())
+
+def _inner_join_count(ndarray[int64_t] left, ndarray[int64_t] right):
+    pass
+
+def outer_join_indexer(ndarray[int64_t] left, ndarray[int64_t] right):
+    cdef:
+        Py_ssize_t i, j, nright, nleft, tot, count
+        int64_t val
+        ndarray[int32_t] lindexer, rindexer
+        ndarray[int64_t] result
+
+    nleft = len(left)
+    nright = len(right)
+    tot = nleft + nright
+
+    lindexer = np.empty(tot, dtype=np.int32)
+    rindexer = np.empty(tot, dtype=np.int32)
+    result = np.empty(tot, dtype=np.int64)
+
+    i = 0
+    j = 0
+    count = 0
+    while True:
+        if i == nleft:
+            if j == nright:
+                # we are done
+                break
+            else:
+                while j < nright:
+                    lindexer[count] = -1
+                    rindexer[count] = j
+                    result[count] = right[j]
+                    j += 1
+                    count += 1
+        elif j == nright:
+            while i < nleft:
+                lindexer[count] = i
+                rindexer[count] = -1
+                result[count] = left[j]
+                i += 1
+                count += 1
+            break
+        else:
+            lval = left[i]
+            rval = right[j]
+            if lval == rval:
+                lindexer[count] = i
+                rindexer[count] = j
+                result[count] = lval
+                i += 1
+                j += 1
+            elif left[i] < right[j]:
+                lindexer[count] = i
+                rindexer[count] = -1
+                result[count] = lval
+                i += 1
+            else:
+                lindexer[count] = -1
+                rindexer[count] = j
+                result[count] = rval
+                j += 1
+
+            count += 1
+
+    return (result[:count].copy(),
+            lindexer[:count].copy(),
+            rindexer[:count].copy())
+
+def ordered_put_indexer(ndarray[int64_t] left, ndarray[int64_t] right,
+                        ndarray[float64_t, ndim=2] lvalues,
+                        ndarray[float64_t, ndim=2] rvalues,
+                        ndarray[float64_t, ndim=2] out):
+    pass
+
+def ordered_outer_join(ndarray[int64_t] left, ndarray[int64_t] right):
+    cdef:
+        Py_ssize_t i, j, k, nright, nleft, kright, kleft
+        int64_t val
+    pass
 
 
 def ordered_inner_join(ndarray[object] left, ndarray[object] right):

pandas/tests/test_tseries.py

@@ -2,8 +2,9 @@
 
 import numpy as np
 from pandas import Index
+from pandas.util.testing import assert_almost_equal
 import pandas.util.testing as common
-import pandas._tseries as tseries
+import pandas._tseries as lib
 
 class TestTseriesUtil(unittest.TestCase):
 
@@ -26,7 +27,7 @@ def test_getMergeVec(self):
         old = Index([1, 5, 10])
         new = Index(range(12))
 
-        filler, mask = tseries.getFillVec(old, new, old.indexMap,
+        filler, mask = lib.getFillVec(old, new, old.indexMap,
                                           new.indexMap, None)
 
         expect_filler = [-1, 0, -1, -1, -1, 1, -1, -1, -1, -1, 2, -1]
@@ -39,7 +40,7 @@ def test_getMergeVec(self):
         # corner case
         old = Index([1, 4])
         new = Index(range(5, 10))
-        filler, mask = tseries.getFillVec(old, new, old.indexMap,
+        filler, mask = lib.getFillVec(old, new, old.indexMap,
                                           new.indexMap, None)
 
         expect_filler = [-1, -1, -1, -1, -1]
@@ -51,7 +52,7 @@ def test_backfill(self):
         old = Index([1, 5, 10])
         new = Index(range(12))
 
-        filler, mask = tseries.getFillVec(old, new, old.indexMap,
+        filler, mask = lib.getFillVec(old, new, old.indexMap,
                                           new.indexMap, 'BACKFILL')
 
         expect_filler = [0, 0, 1, 1, 1, 1, 2, 2, 2, 2, 2, -1]
@@ -64,7 +65,7 @@ def test_backfill(self):
         # corner case
         old = Index([1, 4])
         new = Index(range(5, 10))
-        filler, mask = tseries.getFillVec(old, new, old.indexMap,
+        filler, mask = lib.getFillVec(old, new, old.indexMap,
                                           new.indexMap, 'BACKFILL')
 
         expect_filler = [-1, -1, -1, -1, -1]
@@ -76,7 +77,7 @@ def test_pad(self):
         old = Index([1, 5, 10])
         new = Index(range(12))
 
-        filler, mask = tseries.getFillVec(old, new, old.indexMap,
+        filler, mask = lib.getFillVec(old, new, old.indexMap,
                                           new.indexMap, 'PAD')
 
         expect_filler = [-1, 0, 0, 0, 0, 1, 1, 1, 1, 1, 2, 2]
@@ -89,13 +90,47 @@ def test_pad(self):
         # corner case
         old = Index([5, 10])
         new = Index(range(5))
-        filler, mask = tseries.getFillVec(old, new, old.indexMap,
+        filler, mask = lib.getFillVec(old, new, old.indexMap,
                                           new.indexMap, 'PAD')
 
         expect_filler = [-1, -1, -1, -1, -1]
         expect_mask = np.zeros(5, dtype=bool)
         self.assert_(np.array_equal(filler, expect_filler))
         self.assert_(np.array_equal(mask, expect_mask))
 
+def test_inner_join_indexer():
+    a = np.array([1, 2, 3, 4, 5])
+    b = np.array([0, 3, 5, 7, 9])
+
+    index, ares, bres = lib.inner_join_indexer(a, b)
+
+    index_exp = np.array([3, 5], dtype=np.int64)
+    assert_almost_equal(index, index_exp)
+
+    aexp = np.array([2, 4])
+    bexp = np.array([1, 2])
+    assert_almost_equal(ares, aexp)
+    assert_almost_equal(bres, bexp)
+
+def test_outer_join_indexer():
+    a = np.array([1, 2, 3, 4, 5])
+    b = np.array([0, 3, 5, 7, 9])
+
+    index, ares, bres = lib.outer_join_indexer(a, b)
+
+    index_exp = np.array([0, 1, 2, 3, 4, 5, 7, 9], dtype=np.int64)
+    assert_almost_equal(index, index_exp)
+
+    aexp = np.array([-1, 0, 1, 2, 3, 4, -1, -1], dtype=np.int32)
+    bexp = np.array([0, -1, -1, 1, -1, 2, 3, 4])
+    assert_almost_equal(ares, aexp)
+    assert_almost_equal(bres, bexp)
+
 class TestMoments(unittest.TestCase):
     pass
+
+if __name__ == '__main__':
+    import nose
+    nose.runmodule(argv=[__file__,'-vvs','-x','--pdb', '--pdb-failure'],
+                   exit=False)
+

scripts/bench_join.py

@@ -78,3 +78,7 @@ def do_left_join_frame(a, b):
     a.index._indexMap = None
     b.index._indexMap = None
     return a.join(b, how='left')
+
+# a = np.array([1, 2, 3, 4, 5], dtype=np.int64)
+# b = np.array([0, 3, 5, 7, 9], dtype=np.int64)
+# print lib.inner_join_indexer(a, b)