Repository: dask/dask

Issue:
ENH: Support merging Dask DataFrames on a combination of columns and index levels
## Background
Pandas version 0.22.0 will have support for merging DataFrames on a combination of columns and index levels. This was accomplished by allowing the `on`, `left_on`, and `right_on` parameters to `df.merge` to accept index levels as well as column labels. 

- Feature discussion: pandas-dev/pandas#14355
- PR: pandas-dev/pandas#17484

## Proposal
I would like to update Dask DataFrame to support this feature as well (when used with pandas >= 0.22.0)

This will allow us to avoid shuffling when merging DataFrames (with known divisions) if the shared index level name is included in the `on` parameter list (by itself or with additional columns).

I'll have a PR up for discussion in a few days, but please share any comments or concerns in the meantime. @mrocklin @TomAugspurger 

Thanks!


Discussion:
This sounds great to me.  I'm also happy whenever anyone starts to get engaged in the deeper parts of dask.dataframe.  Please let me know if anything is unclear.

Resolving patch:
diff --git a/dask/dataframe/core.py b/dask/dataframe/core.py
--- a/dask/dataframe/core.py
+++ b/dask/dataframe/core.py
@@ -2846,7 +2846,7 @@ def _select_columns_or_index(self, columns_or_index):
                             if isinstance(columns_or_index, list)
                             else [columns_or_index])
 
-        column_names = [n for n in columns_or_index if self._is_column_label(n)]
+        column_names = [n for n in columns_or_index if self._is_column_label_reference(n)]
 
         selected_df = self[column_names]
         if self._contains_index_name(columns_or_index):
@@ -2855,34 +2855,38 @@ def _select_columns_or_index(self, columns_or_index):
 
         return selected_df
 
-    def _is_column_label(self, c):
+    def _is_column_label_reference(self, key):
         """
-        Test whether a value matches the label of a column in the DataFrame
+        Test whether a key is a column label reference
+
+        To be considered a column label reference, `key` must match the name of at
+        least one column.
         """
-        return (not is_dask_collection(c) and
-                (np.isscalar(c) or isinstance(c, tuple)) and
-                c in self.columns)
+        return (not is_dask_collection(key) and
+                (np.isscalar(key) or isinstance(key, tuple)) and
+                key in self.columns)
 
-    def _is_index_label(self, i):
+    def _is_index_level_reference(self, key):
         """
-        Test whether a value matches the label of the index of the DataFrame
+        Test whether a key is an index level reference
+
+        To be considered an index level reference, `key` must match the index name
+        and must NOT match the name of any column.
         """
         return (self.index.name is not None and
-                not is_dask_collection(i) and
-                (np.isscalar(i) or isinstance(i, tuple)) and
-                i == self.index.name)
+                not is_dask_collection(key) and
+                (np.isscalar(key) or isinstance(key, tuple)) and
+                key == self.index.name and
+                key not in self.columns)
 
     def _contains_index_name(self, columns_or_index):
         """
-        Test whether the input contains the label of the index of the DataFrame
+        Test whether the input contains a reference to the index of the DataFrame
         """
         if isinstance(columns_or_index, list):
-            return (any(self._is_index_label(n) and
-                        not self._is_column_label(n)
-                        for n in columns_or_index))
+            return any(self._is_index_level_reference(n) for n in columns_or_index)
         else:
-            return (self._is_index_label(columns_or_index) and
-                    not self._is_column_label(columns_or_index))
+            return self._is_index_level_reference(columns_or_index)
 
 
 # bind operators
diff --git a/dask/dataframe/multi.py b/dask/dataframe/multi.py
--- a/dask/dataframe/multi.py
+++ b/dask/dataframe/multi.py
@@ -203,7 +203,8 @@ def require(divisions, parts, required=None):
 
 
 def merge_indexed_dataframes(lhs, rhs, how='left', lsuffix='', rsuffix='',
-                             indicator=False):
+                             indicator=False, left_on=None, right_on=None,
+                             left_index=True, right_index=True):
     """ Join two partitioned dataframes along their index """
 
     (lhs, rhs), divisions, parts = align_partitions(lhs, rhs)
@@ -213,7 +214,8 @@ def merge_indexed_dataframes(lhs, rhs, how='left', lsuffix='', rsuffix='',
     right_empty = rhs._meta
 
     name = 'join-indexed-' + tokenize(lhs, rhs, how, lsuffix, rsuffix,
-                                      indicator)
+                                      indicator, left_on, right_on,
+                                      left_index, right_index)
 
     dsk = dict()
     for i, (a, b) in enumerate(parts):
@@ -222,12 +224,14 @@ def merge_indexed_dataframes(lhs, rhs, how='left', lsuffix='', rsuffix='',
         if b is None and how in ('left', 'outer'):
             b = right_empty
 
-        dsk[(name, i)] = (methods.merge, a, b, how, None, None, True, True,
+        dsk[(name, i)] = (methods.merge, a, b, how, left_on, right_on,
+                          left_index, right_index,
                           indicator, (lsuffix, rsuffix), left_empty,
                           right_empty)
 
     meta = pd.merge(lhs._meta_nonempty, rhs._meta_nonempty, how=how,
-                    left_index=True, right_index=True,
+                    left_index=left_index, right_index=right_index,
+                    left_on=left_on, right_on=right_on,
                     suffixes=(lsuffix, rsuffix), indicator=indicator)
     return new_dd_object(toolz.merge(lhs.dask, rhs.dask, dsk),
                          name, meta, divisions)
@@ -301,7 +305,8 @@ def single_partition_join(left, right, **kwargs):
         dsk = {(name, i): (apply, pd.merge, [left_key, right_key], kwargs)
                for i, right_key in enumerate(right.__dask_keys__())}
 
-        if kwargs.get('right_index'):
+        if kwargs.get('right_index') or right._contains_index_name(
+                kwargs.get('right_on')):
             divisions = right.divisions
         else:
             divisions = [None for _ in right.divisions]
@@ -311,7 +316,8 @@ def single_partition_join(left, right, **kwargs):
         dsk = {(name, i): (apply, pd.merge, [left_key, right_key], kwargs)
                for i, left_key in enumerate(left.__dask_keys__())}
 
-        if kwargs.get('left_index'):
+        if kwargs.get('left_index') or left._contains_index_name(
+                kwargs.get('left_on')):
             divisions = left.divisions
         else:
             divisions = [None for _ in left.divisions]
@@ -360,14 +366,22 @@ def merge(left, right, how='inner', on=None, left_on=None, right_on=None,
         right = from_pandas(right, npartitions=1)  # turn into DataFrame
 
     # Both sides are now dd.DataFrame or dd.Series objects
+    merge_indexed_left = (left_index or left._contains_index_name(
+        left_on)) and left.known_divisions
+
+    merge_indexed_right = (right_index or right._contains_index_name(
+        right_on)) and right.known_divisions
 
     # Both sides indexed
-    if (left_index and left.known_divisions and
-            right_index and right.known_divisions):  # Do indexed join
+    if merge_indexed_left and merge_indexed_right:  # Do indexed join
         return merge_indexed_dataframes(left, right, how=how,
                                         lsuffix=suffixes[0],
                                         rsuffix=suffixes[1],
-                                        indicator=indicator)
+                                        indicator=indicator,
+                                        left_on=left_on,
+                                        right_on=right_on,
+                                        left_index=left_index,
+                                        right_index=right_index)
 
     # Single partition on one side
     elif (left.npartitions == 1 and how in ('inner', 'right') or
@@ -386,11 +400,11 @@ def merge(left, right, how='inner', on=None, left_on=None, right_on=None,
                         left_on=left_on, right_on=right_on,
                         left_index=left_index, right_index=right_index,
                         suffixes=suffixes, indicator=indicator)
-        if left_index and left.known_divisions:
+        if merge_indexed_left and left.known_divisions:
             right = rearrange_by_divisions(right, right_on, left.divisions,
                                            max_branch, shuffle=shuffle)
             left = left.clear_divisions()
-        elif right_index and right.known_divisions:
+        elif merge_indexed_right and right.known_divisions:
             left = rearrange_by_divisions(left, left_on, right.divisions,
                                           max_branch, shuffle=shuffle)
             right = right.clear_divisions()
diff --git a/dask/dataframe/shuffle.py b/dask/dataframe/shuffle.py
--- a/dask/dataframe/shuffle.py
+++ b/dask/dataframe/shuffle.py
@@ -202,7 +202,7 @@ def shuffle(df, index, shuffle=None, npartitions=None, max_branch=32,
     shuffle_tasks
     """
     if not isinstance(index, _Frame):
-        index = df[index]
+        index = df._select_columns_or_index(index)
     partitions = index.map_partitions(partitioning_index,
                                       npartitions=npartitions or df.npartitions,
                                       meta=pd.Series([0]))


Test patch:
diff --git a/dask/dataframe/tests/test_merge_column_and_index.py b/dask/dataframe/tests/test_merge_column_and_index.py
new file mode 100644
--- /dev/null
+++ b/dask/dataframe/tests/test_merge_column_and_index.py
@@ -0,0 +1,182 @@
+import dask.dataframe as dd
+import numpy as np
+import pandas as pd
+import pytest
+
+from dask.dataframe.utils import assert_eq, PANDAS_VERSION
+
+
+# Fixtures
+# ========
+@pytest.fixture
+def df_left():
+    # Create frame with 10 partitions
+    # Frame has 11 distinct idx values
+    partition_sizes = np.array([3, 4, 2, 5, 3, 2, 5, 9, 4, 7, 4])
+    idx = [i for i, s in enumerate(partition_sizes) for _ in range(s)]
+    k = [i for s in partition_sizes for i in range(s)]
+    vi = range(len(k))
+
+    return pd.DataFrame(dict(
+        idx=idx,
+        k=k,
+        v1=vi
+    )).set_index(['idx'])
+
+
+@pytest.fixture
+def df_right():
+    # Create frame with 10 partitions
+    # Frame has 11 distinct idx values
+    partition_sizes = np.array([4, 2, 5, 3, 2, 5, 9, 4, 7, 4, 8])
+    idx = [i for i, s in enumerate(partition_sizes) for _ in range(s)]
+    k = [i for s in partition_sizes for i in range(s)]
+    vi = range(len(k))
+
+    return pd.DataFrame(dict(
+        idx=idx,
+        k=k,
+        v1=vi
+    )).set_index(['idx'])
+
+
+@pytest.fixture
+def ddf_left(df_left):
+    # Create frame with 10 partitions
+    # Skip division on 2 so there is one mismatch with ddf_right
+    return dd.repartition(df_left, [0, 1, 3, 4, 5, 6, 7, 8, 9, 10, 11])
+
+
+@pytest.fixture
+def ddf_left_unknown(ddf_left):
+    return ddf_left.clear_divisions()
+
+
+@pytest.fixture
+def ddf_left_single(df_left):
+    return dd.from_pandas(df_left, npartitions=1, sort=False)
+
+
+@pytest.fixture
+def ddf_right(df_right):
+    # Create frame with 10 partitions
+    # Skip division on 3 so there is one mismatch with ddf_left
+    return dd.repartition(df_right, [0, 1, 2, 4, 5, 6, 7, 8, 9, 10, 11])
+
+
+@pytest.fixture
+def ddf_right_unknown(ddf_right):
+    return ddf_right.clear_divisions()
+
+
+@pytest.fixture
+def ddf_right_single(df_right):
+    return dd.from_pandas(df_right, npartitions=1, sort=False)
+
+
+@pytest.fixture(params=['inner', 'left', 'right', 'outer'])
+def how(request):
+    return request.param
+
+
+@pytest.fixture(params=[
+    'idx',
+    ['idx'],
+    ['idx', 'k'],
+    ['k', 'idx']])
+def on(request):
+    return request.param
+
+
+# Tests
+# =====
+@pytest.mark.skipif(PANDAS_VERSION < '0.22.0',
+                    reason="Need pandas col+index merge support (pandas-dev/pandas#14355)")
+def test_merge_known_to_known(df_left, df_right, ddf_left, ddf_right, on, how):
+    # Compute expected
+    expected = df_left.merge(df_right, on=on, how=how)
+
+    # Perform merge
+    result = ddf_left.merge(ddf_right, on=on, how=how, shuffle='tasks')
+
+    # Assertions
+    assert_eq(result, expected)
+    assert_eq(result.divisions, tuple(range(12)))
+    assert len(result.__dask_graph__()) < 80
+
+
+@pytest.mark.skipif(PANDAS_VERSION < '0.22.0',
+                    reason="Need pandas col+index merge support (pandas-dev/pandas#14355)")
+@pytest.mark.parametrize('how', ['inner', 'left'])
+def test_merge_known_to_single(df_left, df_right, ddf_left, ddf_right_single, on, how):
+    # Compute expected
+    expected = df_left.merge(df_right, on=on, how=how)
+
+    # Perform merge
+    result = ddf_left.merge(ddf_right_single, on=on, how=how, shuffle='tasks')
+
+    # Assertions
+    assert_eq(result, expected)
+    assert_eq(result.divisions, ddf_left.divisions)
+    assert len(result.__dask_graph__()) < 30
+
+
+@pytest.mark.skipif(PANDAS_VERSION < '0.22.0',
+                    reason="Need pandas col+index merge support (pandas-dev/pandas#14355)")
+@pytest.mark.parametrize('how', ['inner', 'right'])
+def test_merge_single_to_known(df_left, df_right, ddf_left_single, ddf_right, on, how):
+    # Compute expected
+    expected = df_left.merge(df_right, on=on, how=how)
+
+    # Perform merge
+    result = ddf_left_single.merge(ddf_right, on=on, how=how, shuffle='tasks')
+
+    # Assertions
+    assert_eq(result, expected)
+    assert_eq(result.divisions, ddf_right.divisions)
+    assert len(result.__dask_graph__()) < 30
+
+
+@pytest.mark.skipif(PANDAS_VERSION < '0.22.0',
+                    reason="Need pandas col+index merge support (pandas-dev/pandas#14355)")
+def test_merge_known_to_unknown(df_left, df_right, ddf_left, ddf_right_unknown, on, how):
+    # Compute expected
+    expected = df_left.merge(df_right, on=on, how=how)
+
+    # Perform merge
+    result = ddf_left.merge(ddf_right_unknown, on=on, how=how, shuffle='tasks')
+
+    # Assertions
+    assert_eq(result, expected)
+    assert_eq(result.divisions, tuple(None for _ in range(11)))
+    assert len(result.__dask_graph__()) >= 400
+
+
+@pytest.mark.skipif(PANDAS_VERSION < '0.22.0',
+                    reason="Need pandas col+index merge support (pandas-dev/pandas#14355)")
+def test_merge_unknown_to_known(df_left, df_right, ddf_left_unknown, ddf_right, on, how):
+    # Compute expected
+    expected = df_left.merge(df_right, on=on, how=how)
+
+    # Perform merge
+    result = ddf_left_unknown.merge(ddf_right, on=on, how=how, shuffle='tasks')
+
+    # Assertions
+    assert_eq(result, expected)
+    assert_eq(result.divisions, tuple(None for _ in range(11)))
+    assert len(result.__dask_graph__()) > 400
+
+
+@pytest.mark.skipif(PANDAS_VERSION < '0.22.0',
+                    reason="Need pandas col+index merge support (pandas-dev/pandas#14355)")
+def test_merge_unknown_to_unknown(df_left, df_right, ddf_left_unknown, ddf_right_unknown, on, how):
+    # Compute expected
+    expected = df_left.merge(df_right, on=on, how=how)
+
+    # Merge unknown to unknown
+    result = ddf_left_unknown.merge(ddf_right_unknown, on=on, how=how, shuffle='tasks')
+
+    # Assertions
+    assert_eq(result, expected)
+    assert_eq(result.divisions, tuple(None for _ in range(11)))
+    assert len(result.__dask_graph__()) > 400