Repository: dask/dask

Issue:
Support writing and reading parquet partitions
I'm saving partitioned data out by month to parquet with fastparquet 0.0.6, then reading in with dask 0.14.3 and pandas 0.20.1. I'm updating only the partitions as needed each time I need to update the data, then I'm trying to read it all with dask, to further operate on the entire thing and eventually store it in a location for tableau to pick it up.

It looks like the only option at the moment is to manually manage the different parquet partitions, but it would be nice if I could do something like spark supports, where you can do something like:

### Ideal End Goal

``` python
data_path = '/path/to/full/data/'

new_data = get_new_chunk_of_data()

new_data.write_parquet(data_path, mode='append', partition_on=['year', 'month', 'day'])
# or
new_data.write_parquet(data_path, mode='replace', partition_on=['year', 'month', 'day'])
```

For now, I'm trying to just handle the appending/replacing of each partition manually, then reading the entire thing with dask afterwards. So I often have a structure like:

```
/path/to/full/data/
/path/to/full/data/year=2017/
/path/to/full/data/year=2017/month=1/
/path/to/full/data/year=2017/month=1/day=1/
```

And I want to do: `dd.read_parquet('/path/to/full/data/')`, where dask generates a dataframe with columns of all those that exist within the parquet file, plus year, month, day. Right now it appears to error out like `[<paritition_name>] not in index` on this function call `meta = _meta_from_dtypes(all_columns, pf.columns, dtypes)`, line 102 of `dask/dataframe/io/parquet.py`.

### Simple Example

Here is a simple example of what I'm trying to do:

``` python
import dask.dataframe as dd
import pandas as pd
from fastparquet import write

# create dataframe for a single partition
dates = pd.date_range('2017-01-01', '2017-01-31')
col1 = range(0, len(dates))
col2 = list(reversed(col1))
df = pd.DataFrame({'event_date': dates, 'metric_1': col1, 'metric_2': col2})

# add column I will partition on
df['year_month'] = '2017-01'

# write with fastparquet using partitioning
write('./tmp/', df, partition_on=['year_month'], file_scheme='hive', write_index=False)
```

#### Resulting in

``` python
> glob.glob('./tmp/*')
['./tmp/_common_metadata', './tmp/_metadata', './tmp/year_month=2017-01']
```

#### Write a second partition

``` python
# pretending we actually have new data here
df['year_month'] = '2017-02'
write('./tmp/', df, partition_on=['year_month'], file_scheme='hive', write_index=False)

> glob.glob('./tmp/*')
['./tmp/_common_metadata',
 './tmp/_metadata',
 './tmp/year_month=2017-01',
 './tmp/year_month=2017-02']
```

#### Exceptions on Reading with Dask

```python
df2 = dd.read_parquet('./tmp/')
```

Errors with `ValueError: Multiple possible indexes exist: [u'event_date', u'metric_1', u'metric_2'].  Please select one with index='index-name'`

```python
df2 = dd.read_parquet('./tmp/', index='event_date')
```

Errors with `KeyError: "[u'year_month'] not in index"`

**Note:** reading with fastparquet does work, but then I can only convert it to a pandas dataframe.


Discussion:
Would you mind checking if https://github.com/dask/dask/pull/2354 fixes this for you? You might try `index=False` to make things simple.
I don't think there's explicit consideration yet that the index might be a partitioned value, but we certainly should. I expect to work on categorical and partitioned reading/writing shortly.

Resolving patch:
diff --git a/dask/dataframe/io/parquet.py b/dask/dataframe/io/parquet.py
--- a/dask/dataframe/io/parquet.py
+++ b/dask/dataframe/io/parquet.py
@@ -2,7 +2,7 @@
 
 import numpy as np
 import pandas as pd
-from toolz import first, partial
+from toolz import partial
 
 from ..core import DataFrame, Series
 from ..utils import UNKNOWN_CATEGORIES
@@ -64,24 +64,12 @@ def _read_fastparquet(fs, paths, myopen, columns=None, filters=None,
            not(fastparquet.api.filter_out_stats(rg, filters, pf.schema)) and
            not(fastparquet.api.filter_out_cats(rg, filters))]
 
-    # Find an index among the partially sorted columns
-    minmax = fastparquet.api.sorted_partitioned_columns(pf)
-
     if index is False:
         index_col = None
-    elif len(minmax) == 1:
-        index_col = first(minmax)
-    elif len(minmax) > 1:
-        if index:
-            index_col = index
-        elif 'index' in minmax:
-            index_col = 'index'
-        else:
-            raise ValueError("Multiple possible indexes exist: %s.  "
-                             "Please select one with index='index-name'"
-                             % sorted(minmax))
+    elif index is None:
+        index_col = pf._get_index()
     else:
-        index_col = None
+        index_col = index
 
     if columns is None:
         all_columns = tuple(pf.columns + list(pf.cats))
@@ -99,7 +87,8 @@ def _read_fastparquet(fs, paths, myopen, columns=None, filters=None,
         categories = pf.categories
     dtypes = pf._dtypes(categories)
 
-    meta = _meta_from_dtypes(all_columns, pf.columns, dtypes)
+    meta = _meta_from_dtypes(all_columns, tuple(pf.columns + list(pf.cats)),
+                             dtypes)
 
     for cat in categories:
         meta[cat] = pd.Series(pd.Categorical([],
@@ -118,7 +107,14 @@ def _read_fastparquet(fs, paths, myopen, columns=None, filters=None,
            for i, rg in enumerate(rgs)}
 
     if index_col:
-        divisions = list(minmax[index_col]['min']) + [minmax[index_col]['max'][-1]]
+        minmax = fastparquet.api.sorted_partitioned_columns(pf)
+        if index_col in minmax:
+            divisions = (list(minmax[index_col]['min']) +
+                         [minmax[index_col]['max'][-1]])
+            divisions = [divisions[i] for i, rg in enumerate(pf.row_groups)
+                         if rg in rgs] + [divisions[-1]]
+        else:
+            divisions = (None,) * (len(rgs) + 1)
     else:
         divisions = (None,) * (len(rgs) + 1)
 
@@ -290,9 +286,9 @@ def read_parquet(path, columns=None, filters=None, categories=None, index=None,
                              categories=categories, index=index)
 
 
-def to_parquet(path, df, compression=None, write_index=None, has_nulls=None,
+def to_parquet(path, df, compression=None, write_index=None, has_nulls=True,
                fixed_text=None, object_encoding=None, storage_options=None,
-               append=False, ignore_divisions=False):
+               append=False, ignore_divisions=False, partition_on=None):
     """Store Dask.dataframe to Parquet files
 
     Notes
@@ -311,11 +307,11 @@ def to_parquet(path, df, compression=None, write_index=None, has_nulls=None,
     write_index : boolean
         Whether or not to write the index.  Defaults to True *if* divisions are
         known.
-    has_nulls : bool, list or None
+    has_nulls : bool, list or 'infer'
         Specifies whether to write NULLs information for columns. If bools,
-        apply to all columns, if list, use for only the named columns, if None,
-        use only for columns which don't have a sentinel NULL marker (currently
-        object columns only).
+        apply to all columns, if list, use for only the named columns, if
+        'infer', use only for columns which don't have a sentinel NULL marker
+        (currently object columns only).
     fixed_text : dict {col: int}
         For column types that are written as bytes (bytes, utf8 strings, or
         json and bson-encoded objects), if a column is included here, the
@@ -333,6 +329,10 @@ def to_parquet(path, df, compression=None, write_index=None, has_nulls=None,
     ignore_divisions: bool (False)
         If False raises error when previous divisions overlap with the new
         appended divisions. Ignored if append=False.
+    partition_on: list
+        Construct directory-based partitioning by splitting on these fields'
+        values. Each dask partition will result in one or more datafiles,
+        there will be no global groupby.
 
     This uses the fastparquet project:
     http://fastparquet.readthedocs.io/en/latest
@@ -348,6 +348,7 @@ def to_parquet(path, df, compression=None, write_index=None, has_nulls=None,
 
     """
     import fastparquet
+    partition_on = partition_on or []
 
     fs, paths, myopen = get_fs_paths_myopen(path, None, 'wb',
                                             **(storage_options or {}))
@@ -361,15 +362,19 @@ def to_parquet(path, df, compression=None, write_index=None, has_nulls=None,
         index_col = df.columns[0]
     else:
         ignore_divisions = True
+        index_col = []
 
     object_encoding = object_encoding or 'utf8'
     if object_encoding == 'infer' or (isinstance(object_encoding, dict) and
                                       'infer' in object_encoding.values()):
         raise ValueError('"infer" not allowed as object encoding, '
                          'because this required data in memory.')
-    fmd = fastparquet.writer.make_metadata(df._meta, has_nulls=has_nulls,
-                                           fixed_text=fixed_text,
-                                           object_encoding=object_encoding)
+    if set(partition_on) - set(df.columns):
+        raise ValueError('Partitioning on non-existent column')
+    fmd = fastparquet.writer.make_metadata(
+        df._meta, has_nulls=has_nulls, fixed_text=fixed_text,
+        object_encoding=object_encoding, index_cols=[index_col],
+        ignore_columns=partition_on)
 
     if append:
         pf = fastparquet.api.ParquetFile(path, open_with=myopen, sep=sep)
@@ -410,18 +415,29 @@ def to_parquet(path, df, compression=None, write_index=None, has_nulls=None,
                  for i in range(i_offset, len(partitions) + i_offset)]
     outfiles = [sep.join([path, fn]) for fn in filenames]
 
-    writes = [delayed(fastparquet.writer.make_part_file)(
-              myopen(outfile, 'wb'), partition, fmd.schema,
-              compression=compression)
-              for outfile, partition in zip(outfiles, partitions)]
+    if partition_on:
+        writes = [delayed(fastparquet.writer.partition_on_columns)(
+            partition, partition_on, path, filename, fmd, sep,
+            compression, fs.open, fs.mkdirs)
+            for filename, partition in zip(filenames, partitions)
+        ]
+    else:
+        writes = [delayed(fastparquet.writer.make_part_file)(
+                  myopen(outfile, 'wb'), partition, fmd.schema,
+                  compression=compression)
+                  for outfile, partition in zip(outfiles, partitions)]
 
     out = delayed(writes).compute()
 
     for fn, rg in zip(filenames, out):
         if rg is not None:
-            for chunk in rg.columns:
-                chunk.file_path = fn
-            fmd.row_groups.append(rg)
+            if isinstance(rg, list):
+                for r in rg:
+                    fmd.row_groups.append(r)
+            else:
+                for chunk in rg.columns:
+                    chunk.file_path = fn
+                fmd.row_groups.append(rg)
 
     if len(fmd.row_groups) == 0:
         raise ValueError("All partitions were empty")


Test patch:
diff --git a/dask/dataframe/io/tests/test_parquet.py b/dask/dataframe/io/tests/test_parquet.py
--- a/dask/dataframe/io/tests/test_parquet.py
+++ b/dask/dataframe/io/tests/test_parquet.py
@@ -389,7 +389,6 @@ def test_timestamp_index():
 
 
 @pytest.mark.skipif(not pyarrow, reason='pyarrow not found')
-@pytest.mark.xfail(reason="to_parquet does not write nulls by default")
 def test_to_parquet_default_writes_nulls():
     import pyarrow.parquet as pq
 
@@ -399,4 +398,43 @@ def test_to_parquet_default_writes_nulls():
     with tmpfile() as fn:
         ddf.to_parquet(fn)
         table = pq.read_table(fn)
-        assert table[0].null_count == 2
+        assert table[1].null_count == 2
+
+
+def test_partition_on(tmpdir):
+    tmpdir = str(tmpdir)
+    df = pd.DataFrame({'a': np.random.choice(['A', 'B', 'C'], size=100),
+                       'b': np.random.random(size=100),
+                       'c': np.random.randint(1, 5, size=100)})
+    d = dd.from_pandas(df, npartitions=2)
+    d.to_parquet(tmpdir, partition_on=['a'])
+    out = dd.read_parquet(tmpdir, engine='fastparquet').compute()
+    for val in df.a.unique():
+        assert set(df.b[df.a == val]) == set(out.b[out.a == val])
+
+
+def test_filters(fn):
+
+    df = pd.DataFrame({'at': ['ab', 'aa', 'ba', 'da', 'bb']})
+    ddf = dd.from_pandas(df, npartitions=1)
+
+    # Ok with 1 partition and filters
+    ddf.repartition(npartitions=1, force=True).to_parquet(fn, write_index=False)
+    ddf2 = dd.read_parquet(fn, index=False,
+                           filters=[('at', '==', 'aa')]).compute()
+    assert_eq(ddf2, ddf)
+
+    # with >1 partition and no filters
+    ddf.repartition(npartitions=2, force=True).to_parquet(fn)
+    dd.read_parquet(fn).compute()
+    assert_eq(ddf2, ddf)
+
+    # with >1 partition and filters using base fastparquet
+    ddf.repartition(npartitions=2, force=True).to_parquet(fn)
+    df2 = fastparquet.ParquetFile(fn).to_pandas(filters=[('at', '==', 'aa')])
+    assert len(df2) > 0
+
+    # with >1 partition and filters
+    ddf.repartition(npartitions=2, force=True).to_parquet(fn)
+    dd.read_parquet(fn, filters=[('at', '==', 'aa')]).compute()
+    assert len(ddf2) > 0