Repository: vaexio/vaex (python)

Issue:
Using ds.join twice with masked values / and inner join
When joining ds1 with ds3 (right join to get ds31),
ds31 = ds1.join(ds3, left_on='a', right_on='c', how='right'):

\# | c | z | a | x
-- | -- | -- | -- | --
0 | 'a' | 1 | 'a' | 1
1 | 'b' | 8 | 'b' | 2
2 | 'e' | 27 | masked | masked
3 | 'f' | 64 | masked | masked

for ds1:

\# | a | x
-- | -- | --
0 | 'a' | 1
1 | 'b' | 2
2 | 'c' | 3


for ds3:

\# | c | z | a | x
-- | -- | -- | -- | --
0 | 'a' | 1 | 'a' | 1
1 | 'b' | 8 | 'b' | 2
2 | 'e' | 27 | masked | masked
3 | 'f' | 64 | masked | masked

Now I want to join (left) ds31 with ds4:
ds314 = ds31.join(ds4, left_on='a', right_on='d', how='left', inplace='False')

---------------------------------------------------------------------------
TypeError                                 Traceback (most recent call last)
<ipython-input-178-7b64de6bfa10> in <module>()
----> 1 ds314 = ds31.join(ds4, left_on='a', right_on='d', how='left', inplace='False')

/net/gaia/data/users/breddels/python/anaconda3/lib/python3.6/site-packages/vaex/dataset.py in join(self, other, on, left_on, right_on, lsuffix, rsuffix, how, inplace)
   5168             right_values = right.evaluate(right_on)
   5169             # maps from the left_values to row #
-> 5170             index_left = dict(zip(left_values, range(N)))
   5171             # idem for right
   5172             index_other = dict(zip(right_values, range(N_other)))

TypeError: unhashable type: 'MaskedConstant'

for ds4:

\# | d | z4
-- | -- | --
0 | 'a' | 1
1 | 'b' | 16
2 | 'c' | 81
3 | 'd' | 256
4 | 'e' | 625



In other words, vaex tries to find matches in ds4 for the masked rows of ds31.
I would like to have a resulting Table in which all ds4 is matched to ds31. Values for ds314 should be masked (or Nan) if the corresponding rows in ds31 are masked.

The same error occurs when:
ds431 = ds4.join(ds31, left_on='d', right_on='a', how='left', inplace='False')

So: **ds.join can't handle 'masked' values.**

In addition: a feature for **inner join** would be useful too.







Pull request description:
fix(core,join): when doing join after join, we avoid two levels of indirection
Fixes #96

Patch:
diff --git a/packages/vaex-core/vaex/column.py b/packages/vaex-core/vaex/column.py
index aca8c3cf..38d58714 100644
--- a/packages/vaex-core/vaex/column.py
+++ b/packages/vaex-core/vaex/column.py
@@ -76,6 +76,30 @@ class ColumnIndexed(Column):
         max_index = self.indices.max()
         assert max_index < self.df._length_original
 
+    @staticmethod
+    def index(df, column, name, indices, direct_indices_map=None):
+        """Creates a new column indexed by indices which avoids nested indices
+
+        :param df: Dataframe where column comes from
+        :param column: Column object or numpy array
+        :param name: name of column
+        :param indices: ndarray with indices which rows to take
+        :param direct_indices_map: cache of the nested indices (pass a dict for better performance), defaults to None
+        :return: A column object which avoids two levels of indirection
+        :rtype: ColumnIndexed
+        """
+        direct_indices_map = direct_indices_map or {}
+        if isinstance(column, ColumnIndexed):
+            # TODO: think about what happpens when the indices are masked.. ?
+            if id(column.indices) not in direct_indices_map:
+                direct_indices = column.indices[indices]
+                direct_indices_map[id(column.indices)] = direct_indices
+            else:
+                direct_indices = direct_indices_map[id(column.indices)]
+            return ColumnIndexed(column.df, direct_indices, column.name)
+        else:
+            return ColumnIndexed(df, indices, name)
+
     def __len__(self):
         return len(self.indices)
 
diff --git a/packages/vaex-core/vaex/dataframe.py b/packages/vaex-core/vaex/dataframe.py
index 3ba90a45..bf390b6a 100644
--- a/packages/vaex-core/vaex/dataframe.py
+++ b/packages/vaex-core/vaex/dataframe.py
@@ -3674,16 +3674,7 @@ class DataFrame(object):
             if column is not None:
                 # we optimize this somewhere, so we don't do multiple
                 # levels of indirection
-                if isinstance(column, ColumnIndexed):
-                    # TODO: think about what happpens when the indices are masked.. ?
-                    if id(column.indices) not in direct_indices_map:
-                        direct_indices = column.indices[indices]
-                        direct_indices_map[id(column.indices)] = direct_indices
-                    else:
-                        direct_indices = direct_indices_map[id(column.indices)]
-                    df.columns[name] = ColumnIndexed(column.df, direct_indices, column.name)
-                else:
-                    df.columns[name] = ColumnIndexed(df_trimmed, indices, name)
+                df.columns[name] = ColumnIndexed.index(df_trimmed, column, name, indices, direct_indices_map)
         df._length_original = len(indices)
         df._length_unfiltered = df._length_original
         df._cached_filtered_length = None
@@ -4957,23 +4948,21 @@ class DataFrameLocal(DataFrame):
         :param inplace: {inplace}
         :return:
         """
-        ds = self if inplace else self.copy()
         inner = False
+        left = self
+        right = other
         if how == 'left':
-            left = ds
-            right = other
+            pass
         elif how == 'right':
-            left = other
-            right = ds
+            left, right = right, left
             lprefix, rprefix = rprefix, lprefix
             lsuffix, rsuffix = rsuffix, lsuffix
             left_on, right_on = right_on, left_on
         elif how == 'inner':
-            left = ds
-            right = other
             inner = True
         else:
             raise ValueError('join type not supported: {}, only left and right'.format(how))
+        left = left if inplace else left.copy()
 
         for name in right:
             if name in left and rprefix + name + rsuffix == lprefix + name + lsuffix:
@@ -4997,7 +4986,7 @@ class DataFrameLocal(DataFrame):
                 else:
                     left.add_column(right_name, right.columns[name])
         else:
-            df = ds
+            df = left
             # we index the right side, this assumes right is smaller in size
             index = right._index(right_on)
             lookup = np.zeros(left._length_original, dtype=np.int64)
@@ -5038,9 +5027,10 @@ class DataFrameLocal(DataFrame):
                 left_mask_matched = lookup != -1  # all the places where we found a match to the right
                 lookup = lookup[left_mask_matched]  # filter the lookup table to the right
                 left_indices_matched = np.where(left_mask_matched)[0]  # convert mask to indices for the left
-                left = left.take(left_indices_matched)
+                left = left.take(left_indices_matched, unfiltered=True)
             else:
                 lookup = np.ma.array(lookup, mask=lookup==-1)
+            direct_indices_map = {}  # for performance, keeps a cache of two levels of indirection of indices
             for name in right:
                 right_name = name
                 if name in left:
@@ -5049,7 +5039,8 @@ class DataFrameLocal(DataFrame):
                 if name in right.virtual_columns:
                     left.add_virtual_column(right_name, right.virtual_columns[name])
                 else:
-                    left.add_column(right_name, ColumnIndexed(right, lookup, name))
+                    column = ColumnIndexed.index(right, right.columns[name], name, lookup, direct_indices_map)
+                    left.add_column(right_name, column)
         return left
 
     def export(self, path, column_names=None, byteorder="=", shuffle=False, selection=False, progress=None, virtual=False, sort=None, ascending=True):


Test patch:
diff --git a/tests/join_test.py b/tests/join_test.py
index c3e4b506..dddc0a79 100644
--- a/tests/join_test.py
+++ b/tests/join_test.py
@@ -53,6 +53,11 @@ def test_left_a_b():
     assert df.evaluate('y').tolist() == [0, None, 2]
     assert df.evaluate('y_r').tolist() == [None, 1, None]
 
+def test_join_indexed():
+    df = df_a.join(other=df_b, left_on='a', right_on='b', rsuffix='_r')
+    df_X = df_a.join(df, left_on='a', right_on='b', rsuffix='_r')
+    assert df_X.evaluate('b').tolist() == ['A', 'B', None]
+
 
 def test_left_a_b_filtered():
     df_af = df_a[df_a.x > 0]
@@ -94,6 +99,7 @@ def test_right_x_x():
     assert df.evaluate('x_r').tolist() == [2, 1, 0]
     assert df.evaluate('y').tolist() == [2, None, 0]
     assert df.evaluate('y_r').tolist() == [None, 1, 2]
+    assert 'y_r' not in df_b
 
 
 def test_left_dup():
@@ -133,6 +139,7 @@ def test_inner_a_d():
     assert df.x1.tolist() == ['dog', 'cat']
     assert df.x2.tolist() == [3.1, 25.]
 
+
 @pytest.mark.skip(reason='full join not supported yet')
 def test_full_a_d():
     df = df_a.join(df_d, on='a', right_on='a', how='full')
@@ -164,3 +171,9 @@ def test_left_on_virtual_col():
     assert df.x1.tolist() == [None, 'dog', 'cat']
     assert df.x2.tolist() == [None, 3.1, 25.]
     assert df.a_right.tolist() == [None, 'B', 'C']
+
+
+def test_join_filtered_inner():
+    df_a_filtered = df_a[df_a.y > 0]
+    df_joined = df_a_filtered.join(other=df_b, on='x', how='inner', rsuffix='_', allow_duplication=True)
+    assert len(df_joined) == len(df_a)