Merge dtype warning (dask#4917)

mcsoini · jcrist · commit f9cd802066fa · 2019-06-13T11:03:55.000-05:00
* add test covering the merge column dtype mismatch warning * for various merge types: checks that the resulting dataframe has either no nans or that a UserWarning has been thrown * Add warning for mismatches between column data types * fixes issue dask#4574 * Warning is thrown if the on-columns of left and right have different dtypes * flake8 fixes * fixes * use asciitable for warning string
diff --git a/dask/dataframe/multi.py b/dask/dataframe/multi.py
@@ -69,7 +69,8 @@
 from .io import from_pandas
 from . import methods
 from .shuffle import shuffle, rearrange_by_divisions
-from .utils import strip_unknown_categories, is_dataframe_like, is_series_like
+from .utils import (strip_unknown_categories, is_dataframe_like,
+                    is_series_like, asciitable)
 
 
 def align_partitions(*dfs):
@@ -325,6 +326,30 @@ def single_partition_join(left, right, **kwargs):
     return new_dd_object(graph, name, meta, divisions)
 
 
+def warn_dtype_mismatch(left, right, left_on, right_on):
+    """ Checks for merge column dtype mismatches and throws a warning (#4574)
+    """
+
+    if not isinstance(left_on, list):
+        left_on = [left_on]
+    if not isinstance(right_on, list):
+        right_on = [right_on]
+
+    if (all(col in left.columns for col in left_on) and
+            all(col in right.columns for col in right_on)):
+        dtype_mism = [((lo, ro), left.dtypes[lo], right.dtypes[ro])
+                      for lo, ro in zip(left_on, right_on)
+                      if not left.dtypes[lo] is right.dtypes[ro]]
+
+        if dtype_mism:
+            col_tb = asciitable(('Merge columns', 'left dtype', 'right dtype'),
+                                dtype_mism)
+
+            warnings.warn(('Merging dataframes with merge column data '
+                           'type mismatches: \n{}\nCast dtypes explicitly to '
+                           'avoid unexpected results.').format(col_tb))
+
+
 @wraps(pd.merge)
 def merge(left, right, how='inner', on=None, left_on=None, right_on=None,
           left_index=False, right_index=False, suffixes=('_x', '_y'),
@@ -413,6 +438,9 @@ def merge(left, right, how='inner', on=None, left_on=None, right_on=None,
                               empty_index_dtype=meta.index.dtype)
     # Catch all hash join
     else:
+        if left_on and right_on:
+            warn_dtype_mismatch(left, right, left_on, right_on)
+
         return hash_join(left, left.index if left_index else left_on,
                          right, right.index if right_index else right_on,
                          how, npartitions, suffixes, shuffle=shuffle,
diff --git a/dask/dataframe/tests/test_multi.py b/dask/dataframe/tests/test_multi.py
@@ -345,6 +345,42 @@ def test_concat_different_dtypes(join):
     assert len(set(map(str, dtypes_list))) == 1  # all the same
 
 
+@pytest.mark.parametrize('how', ['inner', 'outer', 'left', 'right'])
+@pytest.mark.parametrize('on_index', [True, False])
+def test_merge_columns_dtypes(how, on_index):
+    # tests results of merges with merge columns having different dtypes;
+    # asserts that either the merge was successful or the corresponding warning is raised
+    # addresses issue #4574
+
+    df1 = pd.DataFrame({"A": list(np.arange(5).astype(float)) * 2,
+                        "B": list(np.arange(5)) * 2})
+    df2 = pd.DataFrame({"A": np.arange(5), "B": np.arange(5)})
+
+    a = dd.from_pandas(df1, 2)  # merge column "A" is float
+    b = dd.from_pandas(df2, 2)  # merge column "A" is int
+
+    on = ["A"]
+    left_index = right_index = on_index
+
+    if on_index:
+        a = a.set_index("A")
+        b = b.set_index("A")
+        on = None
+
+    with pytest.warns(None) as record:
+        result = dd.merge(a, b, on=on, how=how,
+                          left_index=left_index, right_index=right_index)
+
+    warned = any('merge column data type mismatches' in str(r) for r in record)
+
+    # result type depends on merge operation -> convert to pandas
+    result = result if isinstance(result, pd.DataFrame) else result.compute()
+
+    has_nans = result.isna().values.any()
+
+    assert (has_nans and warned) or not has_nans
+
+
 @pytest.mark.parametrize('how', ['inner', 'outer', 'left', 'right'])
 @pytest.mark.parametrize('shuffle', ['disk', 'tasks'])
 def test_merge(how, shuffle):
