Repository: dask/dask

Issue:
DataFrame.assign behaves unpredictably with pd.Series
Calling `dataframe.assign(key=pd.Series())` causes the resulting dataframe to behave oddly. In particular, head() returns an empty dataframe and sample() is broken; however, tail() performs normally. 

For example:

```
>>> df.head()
...
[5 rows x 24 columns]

>>> df = df.assign(new_col=pd.Series(range(size)))
>>> df.head()
...
[0 rows x 25 columns]

>>> df.tail()
...
[5 rows x 25 columns]

>>> df.sample(0.1).compute()
ValueError: a must be greater than 0

Traceback
---------
  File "/.../python2.7/site-packages/dask/async.py", line 267, in execute_task
    result = _execute_task(task, data)
  File "/.../python2.7/site-packages/dask/async.py", line 249, in _execute_task
    return func(*args2)
  File "/.../python2.7/site-packages/pandas/core/generic.py", line 2572, in sample
    locs = rs.choice(axis_length, size=n, replace=replace, p=weights)
  File "mtrand.pyx", line 1391, in mtrand.RandomState.choice (numpy/random/mtrand/mtrand.c:15395)
```

Assigning an np.Array works fine.

DataFrame.assign behaves unpredictably with pd.Series
Calling `dataframe.assign(key=pd.Series())` causes the resulting dataframe to behave oddly. In particular, head() returns an empty dataframe and sample() is broken; however, tail() performs normally. 

For example:

```
>>> df.head()
...
[5 rows x 24 columns]

>>> df = df.assign(new_col=pd.Series(range(size)))
>>> df.head()
...
[0 rows x 25 columns]

>>> df.tail()
...
[5 rows x 25 columns]

>>> df.sample(0.1).compute()
ValueError: a must be greater than 0

Traceback
---------
  File "/.../python2.7/site-packages/dask/async.py", line 267, in execute_task
    result = _execute_task(task, data)
  File "/.../python2.7/site-packages/dask/async.py", line 249, in _execute_task
    return func(*args2)
  File "/.../python2.7/site-packages/pandas/core/generic.py", line 2572, in sample
    locs = rs.choice(axis_length, size=n, replace=replace, p=weights)
  File "mtrand.pyx", line 1391, in mtrand.RandomState.choice (numpy/random/mtrand/mtrand.c:15395)
```

Assigning an np.Array works fine.



Discussion:
We should probably raise an `TypeError` in this case.  We should probably only accept `dask.dataframe.Series` objects.

Did some more digging, looks like the base issue is with aligning partitions when one of the DataFrames' includes None. In my case, the first frame had a partition `(None, None)` and the series has a partition `(0, len(series))`. `multi.align_partitions` flattens the new division list to `[None, 0, len(series)]`, which throws a wrench in things.

A secondary issue is that, if align_partitions were to generate the list `[0, len(series)]`, `repartition_divisions` would break with "left side of the new division must be equal or smaller than old division", because None (left side of old division) < 0.

What's the best way to fix these? Should we force frames with None divisions to (0, len(df)) before aligning? Should repartition_divisions handle the case where the old division contains None, or throw a better exception?

The core issue here is that when receiving a `pd.Series` dask converts it to a dask series, and tries to partition the series the same way the dataframe is partitioned, so each dataframe partition could be merged with the right slice of the original `pd.Series`. This is done to optimize the data stored and shared with workers.

In your case dask cannot slice the pd.Series to fit to the dataframe's partitions simply because they're unknown.

I can think of two ways to handle this (others can probably think of something better):
1. Explicitly declining assignments of `dask.Series` (or anything else, actually?) when partitions don't fit. Suggesting the user should `repartition` or `set_index`. `None` partitions should never fit. This should probably be implemented by raising an exception in `align_partitions` (or `repartition`) if any df has a `None` partition.
2. Dropping a small portion of dask's usability here (since data is already in memory to begin with, there's only some profit from converting the `pd.Series` to `dd.Series`). This should not be the default behavior and raise a warning.

> Should we force frames with None divisions to (0, len(df)) before aligning?

The way to do this is by explicitly calling `dd.set_index` method. This is an expensive operation so it remains explicit. Keep in mind indices in both the dataframe and series should align, and that assign will use those indices. Choose your index properly.

We should probably raise an `TypeError` in this case.  We should probably only accept `dask.dataframe.Series` objects.

Did some more digging, looks like the base issue is with aligning partitions when one of the DataFrames' includes None. In my case, the first frame had a partition `(None, None)` and the series has a partition `(0, len(series))`. `multi.align_partitions` flattens the new division list to `[None, 0, len(series)]`, which throws a wrench in things.

A secondary issue is that, if align_partitions were to generate the list `[0, len(series)]`, `repartition_divisions` would break with "left side of the new division must be equal or smaller than old division", because None (left side of old division) < 0.

What's the best way to fix these? Should we force frames with None divisions to (0, len(df)) before aligning? Should repartition_divisions handle the case where the old division contains None, or throw a better exception?

The core issue here is that when receiving a `pd.Series` dask converts it to a dask series, and tries to partition the series the same way the dataframe is partitioned, so each dataframe partition could be merged with the right slice of the original `pd.Series`. This is done to optimize the data stored and shared with workers.

In your case dask cannot slice the pd.Series to fit to the dataframe's partitions simply because they're unknown.

I can think of two ways to handle this (others can probably think of something better):
1. Explicitly declining assignments of `dask.Series` (or anything else, actually?) when partitions don't fit. Suggesting the user should `repartition` or `set_index`. `None` partitions should never fit. This should probably be implemented by raising an exception in `align_partitions` (or `repartition`) if any df has a `None` partition.
2. Dropping a small portion of dask's usability here (since data is already in memory to begin with, there's only some profit from converting the `pd.Series` to `dd.Series`). This should not be the default behavior and raise a warning.

> Should we force frames with None divisions to (0, len(df)) before aligning?

The way to do this is by explicitly calling `dd.set_index` method. This is an expensive operation so it remains explicit. Keep in mind indices in both the dataframe and series should align, and that assign will use those indices. Choose your index properly.


Resolving patch:
diff --git a/dask/dataframe/multi.py b/dask/dataframe/multi.py
--- a/dask/dataframe/multi.py
+++ b/dask/dataframe/multi.py
@@ -65,7 +65,7 @@
 
 from ..base import tokenize
 from ..compatibility import apply
-from .core import (_Frame, Scalar, DataFrame, map_partitions,
+from .core import (_Frame, DataFrame, map_partitions,
                    Index, _maybe_from_pandas, new_dd_object)
 from .io import from_pandas
 from .shuffle import shuffle, rearrange_by_divisions
@@ -91,7 +91,6 @@ def align_partitions(*dfs):
     dfs: sequence of dd.DataFrame, dd.Series and dd.base.Scalar
         Sequence of dataframes to be aligned on their index
 
-
     Returns
     -------
     dfs: sequence of dd.DataFrame, dd.Series and dd.base.Scalar
@@ -105,6 +104,10 @@ def align_partitions(*dfs):
     dfs1 = [df for df in dfs if isinstance(df, _Frame)]
     if len(dfs) == 0:
         raise ValueError("dfs contains no DataFrame and Series")
+    if not all(df.known_divisions for df in dfs1):
+        raise ValueError("Not all divisions are known, can't align "
+                         "partitions. Please use `set_index` or "
+                         "`set_partition` to set the index.")
     divisions = list(unique(merge_sorted(*[df.divisions for df in dfs1])))
     dfs2 = [df.repartition(divisions, force=True)
             if isinstance(df, _Frame) else df for df in dfs]
@@ -129,22 +132,19 @@ def align_partitions(*dfs):
 
 
 def _maybe_align_partitions(args):
-    """ Align DataFrame blocks if divisions are different """
-    # passed to align_partitions
-    indexer, dasks = zip(*[x for x in enumerate(args)
-                           if isinstance(x[1], (_Frame, Scalar))])
+    """Align DataFrame blocks if divisions are different.
 
-    # to get current divisions
-    dfs = [df for df in dasks if isinstance(df, _Frame)]
-    if len(dfs) == 0:
-        # no need to align
+    Note that if all divisions are unknown, but have equal npartitions, then
+    they will be passed through unchanged. This is different than
+    `align_partitions`, which will fail if divisions aren't all known"""
+    dfs = [df for df in args if isinstance(df, _Frame)]
+    if not dfs:
         return args
 
     divisions = dfs[0].divisions
     if not all(df.divisions == divisions for df in dfs):
-        dasks, _, _ = align_partitions(*dasks)
-        for i, d in zip(indexer, dasks):
-            args[i] = d
+        dfs2 = iter(align_partitions(*dfs)[0])
+        return [a if not isinstance(a, _Frame) else next(dfs2) for a in args]
     return args
 
 


Test patch:
diff --git a/dask/dataframe/tests/test_dataframe.py b/dask/dataframe/tests/test_dataframe.py
--- a/dask/dataframe/tests/test_dataframe.py
+++ b/dask/dataframe/tests/test_dataframe.py
@@ -744,20 +744,40 @@ def test_index():
 
 
 def test_assign():
+    d_unknown = dd.from_pandas(full, npartitions=3, sort=False)
+    assert not d_unknown.known_divisions
     res = d.assign(c=1,
                    d='string',
-                   e=np.float64(1),
-                   f=d.a.sum(),
-                   g=d.a + 1,
-                   h=d.a + d.b)
+                   e=d.a.sum(),
+                   f=d.a + d.b)
+    res_unknown = d_unknown.assign(c=1,
+                                   d='string',
+                                   e=d_unknown.a.sum(),
+                                   f=d_unknown.a + d_unknown.b)
     sol = full.assign(c=1,
                       d='string',
-                      e=np.float64(1),
-                      f=full.a.sum(),
-                      g=full.a + 1,
-                      h=full.a + full.b)
+                      e=full.a.sum(),
+                      f=full.a + full.b)
     assert eq(res, sol)
-    pytest.raises(TypeError, lambda: d.assign(c=list(range(9))))
+    assert eq(res_unknown, sol)
+
+    res = d.assign(c=full.a + 1)
+    assert eq(res, full.assign(c=full.a + 1))
+
+    # divisions unknown won't work with pandas
+    with pytest.raises(ValueError):
+        d_unknown.assign(c=full.a + 1)
+
+    # unsupported type
+    with pytest.raises(TypeError):
+        d.assign(c=list(range(9)))
+
+    # Fails when assigning known divisions to unknown divisions
+    with pytest.raises(ValueError):
+        d_unknown.assign(foo=d.a)
+    # Fails when assigning unknown divisions to known divisions
+    with pytest.raises(ValueError):
+        d.assign(foo=d_unknown.a)
 
 
 def test_map():
diff --git a/dask/dataframe/tests/test_multi.py b/dask/dataframe/tests/test_multi.py
--- a/dask/dataframe/tests/test_multi.py
+++ b/dask/dataframe/tests/test_multi.py
@@ -2,7 +2,7 @@
 import numpy as np
 import pandas as pd
 from dask.dataframe.multi import (align_partitions, join_indexed_dataframes,
-        hash_join, concat_indexed_dataframes)
+        hash_join, concat_indexed_dataframes, _maybe_align_partitions)
 import pandas.util.testing as tm
 from dask.async import get_sync
 from dask.dataframe.utils import eq
@@ -59,7 +59,7 @@ def _check(a, b, aa, bb):
                      (dd.from_pandas(ldf, 2), dd.from_pandas(rdf, 2)),
                      (dd.from_pandas(ldf, 2), dd.from_pandas(rdf, 3)),
                      (dd.from_pandas(ldf, 3), dd.from_pandas(rdf, 2))]:
-        (lresult, rresult), div, parts = dd.multi.align_partitions(lhs, rhs)
+        (lresult, rresult), div, parts = align_partitions(lhs, rhs)
         assert eq(lresult, ldf)
         assert eq(rresult, rdf)
 
@@ -75,11 +75,78 @@ def _check(a, b, aa, bb):
                      (dd.from_pandas(ldf, 2), dd.from_pandas(rdf, 2)),
                      (dd.from_pandas(ldf, 2), dd.from_pandas(rdf, 3)),
                      (dd.from_pandas(ldf, 3), dd.from_pandas(rdf, 2))]:
-        (lresult, rresult), div, parts = dd.multi.align_partitions(lhs, rhs)
+        (lresult, rresult), div, parts = align_partitions(lhs, rhs)
         assert eq(lresult, ldf)
         assert eq(rresult, rdf)
 
 
+def test_align_partitions_unknown_divisions():
+    df = pd.DataFrame({'a': [1, 2, 3, 4, 5, 6, 7],
+                       'b': [7, 6, 5, 4, 3, 2, 1]})
+    # One known, one unknown
+    ddf = dd.from_pandas(df, npartitions=2)
+    ddf2 = dd.from_pandas(df, npartitions=2, sort=False)
+    assert not ddf2.known_divisions
+
+    with pytest.raises(ValueError):
+        align_partitions(ddf, ddf2)
+
+    # Both unknown
+    ddf = dd.from_pandas(df + 1, npartitions=2, sort=False)
+    ddf2 = dd.from_pandas(df, npartitions=2, sort=False)
+    assert not ddf.known_divisions
+    assert not ddf2.known_divisions
+
+    with pytest.raises(ValueError):
+        align_partitions(ddf, ddf2)
+
+
+def test__maybe_align_partitions():
+    df = pd.DataFrame({'a': [1, 2, 3, 4, 5, 6, 7],
+                       'b': [7, 6, 5, 4, 3, 2, 1]})
+    # Both known, same divisions
+    ddf = dd.from_pandas(df + 1, npartitions=2)
+    ddf2 = dd.from_pandas(df, npartitions=2)
+
+    a, b = _maybe_align_partitions([ddf, ddf2])
+    assert a is ddf
+    assert b is ddf2
+
+    # Both unknown, same divisions
+    ddf = dd.from_pandas(df + 1, npartitions=2, sort=False)
+    ddf2 = dd.from_pandas(df, npartitions=2, sort=False)
+    assert not ddf.known_divisions
+    assert not ddf2.known_divisions
+
+    a, b = _maybe_align_partitions([ddf, ddf2])
+    assert a is ddf
+    assert b is ddf2
+
+    # Both known, different divisions
+    ddf = dd.from_pandas(df + 1, npartitions=2)
+    ddf2 = dd.from_pandas(df, npartitions=3)
+
+    a, b = _maybe_align_partitions([ddf, ddf2])
+    assert a.divisions == b.divisions
+
+    # Both unknown, different divisions
+    ddf = dd.from_pandas(df + 1, npartitions=2, sort=False)
+    ddf2 = dd.from_pandas(df, npartitions=3, sort=False)
+    assert not ddf.known_divisions
+    assert not ddf2.known_divisions
+
+    with pytest.raises(ValueError):
+        _maybe_align_partitions([ddf, ddf2])
+
+    # One known, one unknown
+    ddf = dd.from_pandas(df, npartitions=2)
+    ddf2 = dd.from_pandas(df, npartitions=2, sort=False)
+    assert not ddf2.known_divisions
+
+    with pytest.raises(ValueError):
+        _maybe_align_partitions([ddf, ddf2])
+
+
 def test_join_indexed_dataframe_to_indexed_dataframe():
     A = pd.DataFrame({'x': [1, 2, 3, 4, 5, 6]},
                      index=[1, 2, 3, 4, 6, 7])
@@ -600,17 +667,17 @@ def test_cheap_single_partition_merge():
             a.merge(b, on='x', how='inner'))
 
 
-@pytest.mark.parametrize('lhs', [pd.DataFrame({'A': [1, 2, 3],
-                                               'B': list('abc'),
-                                               'C': 'foo',
-                                               'D': 1.0},
-                                              columns=list('DCBA'))])
-@pytest.mark.parametrize('rhs', [pd.DataFrame({'G': [4, 5],
-                                               'H': 6.0,
-                                               'I': 'bar',
-                                               'B': list('ab')},
-                                              columns=list('GHIB'))])
-def test_merge_maintains_columns(lhs, rhs):
+def test_merge_maintains_columns():
+    lhs = pd.DataFrame({'A': [1, 2, 3],
+                        'B': list('abc'),
+                        'C': 'foo',
+                        'D': 1.0},
+                        columns=list('DCBA'))
+    rhs = pd.DataFrame({'G': [4, 5],
+                        'H': 6.0,
+                        'I': 'bar',
+                        'B': list('ab')},
+                        columns=list('GHIB'))
     ddf = dd.from_pandas(lhs, npartitions=1)
     merged = dd.merge(ddf, rhs, on='B').compute()
     assert tuple(merged.columns) == ('D', 'C', 'B', 'A', 'G', 'H', 'I')