to-csv works with single partition

mrocklin · Jul 22, 2015 · d7dd219 · d7dd219
1 parent f34977d
commit d7dd219
Show file tree

Hide file tree

Showing 2 changed files with 7 additions and 6 deletions.
diff --git a/dask/dataframe/io.py b/dask/dataframe/io.py
@@ -614,4 +614,4 @@ def to_csv(df, filename, **kwargs):
                              (tuple, [(df._name, i), filename]),
                              kwargs2))
 
-    get(merge(dsk, df.dask), (name, i), get=myget)
+    get(merge(dsk, df.dask), (name, df.npartitions - 1), get=myget)
diff --git a/dask/dataframe/tests/test_io.py b/dask/dataframe/tests/test_io.py
@@ -462,11 +462,12 @@ def test_read_hdf():
 def test_to_csv():
     df = pd.DataFrame({'x': ['a', 'b', 'c', 'd'],
                        'y': [1, 2, 3, 4]}, index=[1., 2., 3., 4.])
-    a = dd.from_pandas(df, 2)
 
-    with tmpfile('csv') as fn:
-        a.to_csv(fn, get=get_sync)
+    for npartitions in [1, 2]:
+        a = dd.from_pandas(df, npartitions)
+        with tmpfile('csv') as fn:
+            a.to_csv(fn, get=get_sync)
 
-        result = pd.read_csv(fn, index_col=0)
+            result = pd.read_csv(fn, index_col=0)
 
-        tm.assert_frame_equal(result, df)
+            tm.assert_frame_equal(result, df)