Repository: dask/dask

Issue:
Inconsistent Schema with partitioned pyarrow files
When writing partitoned Datasets to parquet files, and it so happens that there are is a partition where one column has no values, inconsistent schemas are written to the parquet file so that reading is prevented later:

```python
import dask as da
import dask.dataframe as dd
import pandas as pd
import numpy as np
import os
import shutil

PATH_DASK = '/tmp/dask.pa'

if os.path.exists(PATH_DASK):
    shutil.rmtree(PATH_DASK)

arrays = np.array([np.array([0, 1, 2]), np.array([3, 4]), np.nan, np.nan])
strings = np.array(['a', 'b', np.nan, np.nan])

df = pd.DataFrame([0, 0, 1, 1], columns=['partition_column'])
df['arrays'] = pd.Series(arrays)
df['strings'] = pd.Series(strings)

da_df = dd.from_pandas(df, npartitions=2)
da_df.to_parquet(PATH_DASK, partition_on=['partition_column'])
dd.read_parquet(PATH_DASK)
```

Fails with:
```
ValueError: Schema in partition[partition_column=1] /tmp/dask.pa/partition_column=1/4cd3552267d84d66a5ec8764f869a13b.parquet was different. 
arrays: null
strings: string
__index_level_0__: int64
metadata
--------
...

vs

arrays: list<item: int64>
  child 0, item: int64
strings: string
__index_level_0__: int64
metadata
--------
```

I found #3454 relating to https://issues.apache.org/jira/browse/ARROW-2860. However the Pyarrow Issue only relates to reading files with different schemas. However Pyarrow in it self doesn't even suffer this Bug, as Pyarrow itself would't even write those different schemas in the first place:

```python
import pyarrow as pa
import pyarrow.parquet as pq
import pandas as pd
import os
import shutil

PATH_PYARROW = '/tmp/pyarrow.pa'

if os.path.exists(PATH_PYARROW):
    shutil.rmtree(PATH_PYARROW)

arrays = np.array([np.array([0, 1, 2]), np.array([3, 4]), np.nan, np.nan])
strings = np.array([np.nan, np.nan, 'a', 'b'])

df = pd.DataFrame([0, 0, 1, 1], columns=['partition_column'])
df['arrays'] = pd.Series(arrays)
df['strings'] = pd.Series(strings)

table = pa.Table.from_pandas(df)

pq.write_to_dataset(table, root_path=PATH_PYARROW, partition_cols=['partition_column'])

dataset = pq.ParquetDataset(PATH_PYARROW)
df_read = dataset.read().to_pandas()
df_read
```
Works perfectly fine. Even reading the Pyarrow Files with Dask now works perfectly fine:
`dd.read_parquet(PATH_PYARROW).compute()`

This is simply due to the reason that Pyarrow still uses the same schema accross partitions even if ther is missing data:
```python
import os
files = [val for sublist in [[os.path.join(i[0], j) for j in i[2]] for i in os.walk(PATH_PYARROW)] for val in sublist]
for file in files:
    dataset_part = pq.ParquetDataset(file)    
    print(file)
    print(dataset_part.schema)
    print(dataset_part.read().to_pandas())
```
Outputs:
```python
/tmp/pyarrow.pa/partition_column=0/f32bf0bfedcc4312829e7c4371e544e6.parquet
<pyarrow._parquet.ParquetSchema object at 0x7fdd1c103eb8>
arrays.list.item: INT64
strings: BYTE_ARRAY UTF8
__index_level_0__: INT64
 
      arrays strings
0  [0, 1, 2]     nan
1     [3, 4]     nan
/tmp/pyarrow.pa/partition_column=1/044dd125e0b5472db23bff9a954a1bff.parquet
<pyarrow._parquet.ParquetSchema object at 0x7fdd1c1034e0>
arrays.list.item: INT64
strings: BYTE_ARRAY UTF8
__index_level_0__: INT64
 
  arrays strings
2   None       a
3   None       b
```

So summing it up: In Pyarrow the `pyarrow.parquet.write_to_dataset` wrapper around `pyarrow.parquet.write_table` takes care that the schema in individual files doesn't get screwed up. Dask blindly uses `pyarrow.parquet.write_table` on each partition and hence ends up with a wrong schema.

This means that even if https://issues.apache.org/jira/browse/ARROW-2860 is fixed, loading of these wrong schemas might work, but they are still written wrongly to the file in the first place.

I believe that fixing this in dask is probably going to be hard but might still be worth a try. Maybe one simple but not ideal solution would be to to stick to the schema the first written partition produces and raise an error otherwise. At least we should make sure, that after all partitions are written to file, the schemas match. Otherwise people will save their files without an error raised and notice later on that opening is no longer possible.


Discussion:
cc @xhochy @pitrou thoughts on the best way to handle this?
I'm not a Parquet expert so it's hard to say for sure. Intuitively, I'd say all files in a Parquet partition had better have the same schema, but perhaps there's a specification somewhere that spells it out?
To answer a bit more concretely, ideally you may use the [Array.cast](https://arrow.apache.org/docs/python/generated/pyarrow.Array.html#pyarrow.Array.cast) method to homogenize datatypes accross columns.

However, it doesn't support all types (yet?):
```python
>>> a = pa.array([None, None])                                                            
>>> a                                                                                     
<pyarrow.lib.NullArray object at 0x7fd924b7d228>
2 nulls
>>> a.cast(pa.int64())                                                                    
<pyarrow.lib.Int64Array object at 0x7fd924af03b8>
[
  null,
  null
]
>>> a.cast(pa.list_(pa.int64()))                                                          
Traceback (most recent call last):
  File "<ipython-input-4-a5483e3f21af>", line 1, in <module>
    a.cast(pa.list_(pa.int64()))
  File "pyarrow/array.pxi", line 413, in pyarrow.lib.Array.cast
  File "pyarrow/error.pxi", line 89, in pyarrow.lib.check_status
ArrowNotImplementedError: No cast implemented from null to list<item: int64>

```

It is certainly much much better to have the same schema throughout. In some circumstances and some parquet frameworks, schema evolution is allowed, but that is not in the official spec.
That said, we could add such a facility in PyArrow. Homogenizing arrays may be needed for other purposes at some point. @wesm
We've already had several bug reports about this, here is one:

https://issues.apache.org/jira/browse/ARROW-2860

Here is my comment from 15th September

> There's other issues reported about this. We should handle schema normalization across components of a dataset. I don't have time to work on this for 0.11 but I hope that someone else can take a look soon

In general, we need more help on Parquet support in Apache Arrow.  I think this issue can be completely solved in Python code without touching any C++
That being said, we should probably implement concatenation with schema normalization in C++ so it doesn't have to be redone for R and Ruby. Ruby already has Parquet support and R is not far behind
I've just run into the same problem. I understand that I can use pyarrow.parquet.write_to_dataset as a workaround, right?
But I would like to do it properly using dask API. Is someone working in this issue? Can I help?
@mrocklin @wesm Any movement on this issue?  We are running into it as well.  Please let me know if you need help.
Not yet. It would be great if someone could write a patch for this or contact me about paying for development work on the project. Absent these things I don't know when it will get fixed 
I was having this problem too and the following, not very elegant, patch fixes it for my data:

```patch
From 396492823960544fbf7632450562908d516df069 Mon Sep 17 00:00:00 2001
From: Sarah Bird <birdsarah@mozilla.com>
Date: Sat, 20 Apr 2019 03:12:51 -0500
Subject: [PATCH 1/1] First hack to pass in schema

---
 dask/dataframe/io/parquet.py | 9 +++++++--
 1 file changed, 7 insertions(+), 2 deletions(-)

diff --git a/dask/dataframe/io/parquet.py b/dask/dataframe/io/parquet.py
index 81e2a846..8015dba9 100644
--- a/dask/dataframe/io/parquet.py
+++ b/dask/dataframe/io/parquet.py
@@ -958,7 +958,7 @@ def _read_pyarrow_parquet_piece(fs, piece, columns, index_cols, is_series,
 
 _pyarrow_write_table_kwargs = {'row_group_size', 'version', 'use_dictionary',
                                'compression', 'use_deprecated_int96_timestamps',
-                               'coerce_timestamps', 'flavor', 'chunk_size'}
+                               'coerce_timestamps', 'flavor', 'chunk_size', 'pyarrow_schema'}
 
 _pyarrow_write_metadata_kwargs = {'version', 'use_deprecated_int96_timestamps',
                                   'coerce_timestamps'}
@@ -1000,7 +1000,12 @@ def _write_partition_pyarrow(df, path, fs, filename, write_index,
                              partition_on, metadata_path=None, **kwargs):
     import pyarrow as pa
     from pyarrow import parquet
-    t = pa.Table.from_pandas(df, preserve_index=write_index)
+
+    if 'pyarrow_schema' in kwargs:
+        pyarrow_schema = kwargs.pop('pyarrow_schema')
+        t = pa.Table.from_pandas(df, preserve_index=write_index, schema=pyarrow_schema)
+    else:
+        t = pa.Table.from_pandas(df, preserve_index=write_index)
 
     if partition_on:
         parquet.write_to_dataset(t, path, partition_cols=partition_on,
-- 
2.20.1
```

I then manually create a pyarrow schema and pass it in:

```python
df.to_parquet(
    path='out.parquet',
    engine='pyarrow',
    pyarrow_schema=pa.schema([
        ('field_1', pa.string()),
        ('field_2', pa.int64()),
    ])
```
Thanks @birdsarah !  This isn't too terrible.  Do others think it would be reasonable fix this in Dask until we can get a solution at the lower level?  

Wes, I'm working on making open source contributions possible for my team.  Unfortunately, it is an uphill battle, however I am making progress.  Stay tuned!  I have folks, including myself, that would love to help out!
Probably the thing to do is produce the pyarrow schema from the meta? Or does that not produce a useful schema for an empty dataframe, is that the problem? 
Being allowed to explicitly specify the schema does seem like a useful option in any case.
> Probably the thing to do is produce the pyarrow schema from the meta

@martindurant I think that sounds right. 

I'm not having an empty dataframe problem so I can't speak to that. The problem I'm having is that the schema is inferrred each time from each partition of data and attempts to specify the meta have been ignored. I have a string field that's sometimes all null for a given partition, then my individual partitions are written with different schemas and I can't read back.

I was surprised to see the meta not being used, but I also didn't know how to make a pyarrow schema from the meta data, which is why i'm hard coding the pyarrow schema. I'm already making the meta data schema and I'd prefer to not make the schema twice. 

That said, I had specified the meta data. Are there were cases where you'd get to line ~1000 above and the df would not have _meta on it.
> I'm not having an empty dataframe problem

So although a dataframe should always have one, the meta is usually empty, so calling pyarrow on it directly to make a schema would produce exactly the wrong thing :|  I suppose you could always make the schema from a partition of the data that you know to be sound.
fastparquet has this problem too, and you need to, in the general case, specify how you want to interpret object columns.
>  in the general case, specify how you want to interpret object columns.

Good point. Maybe it's possible to go in the other direction then. I'm working with my initial data processing pipeline. Here's the code:

```python
dfs = [delayed(process_file)(file_name) for file_name in file_names]
df = dd.from_delayed(dfs, meta=meta)
out = df.to_parquet(
    path='out.parquet',
    engine='pyarrow',
    compression='snappy',
    pyarrow_schema=pyarrow_schema
)
```

Seems like there's potential to give a pyarrow schema to meta, as it's actually more specific, and then anything that needs to be cast as an object could be.

Obviously there's no need to optimize to my use case, just sharing.
That way around is certainly possible
```
schema=pa.schema([
         ('field_1', pa.string()),
         ('field_2', pa.int64()),
     ])
pa.Table.from_pandas(pd.DataFrame(columns=schema.fields, schema=schema)).to_pandas()
```
Generally, an arrow schema versus a pandas dtypes object do not map perfectly to one-another, I'm not sure you can do better; but still it may be useful to pass the arrow schema down functions like to_parquet which need direct use of them.
Thanks. Do you know how to handle categoricals?
No, I really don't, sorry :| Surely someone from arrow can help with that one.
>  Do you know how to handle categoricals?

@birdsarah do you mean in light of manually creating the schema (as this is the workaround you are doing above)?

In arrow, Categoricals map to DictionaryType:

```
In [23]: cat = pd.Series(pd.Categorical(['a', 'b', 'a']))

In [24]: a = pyarrow.array(cat)

In [25]: a
Out[25]: 
<pyarrow.lib.DictionaryArray object at 0x7f0c56993ca8>

-- dictionary:
  [
    "a",
    "b"
  ]
-- indices:
  [
    0,
    1,
    0
  ]

In [26]: a.type
Out[26]: DictionaryType(dictionary<values=string, indices=int8, ordered=0>)
```

So you can create the type manually to add this to the schema (from the datatype of the indices, and the unique categories):

```
In [30]: pa.dictionary(pa.int8(), pa.array(['a', 'b']))
Out[30]: DictionaryType(dictionary<values=string, indices=int8, ordered=0>)
```

I think you should be able to get the categories from the dtype of the empty `_meta` ?

> the meta is usually empty, so calling pyarrow on it directly to make a schema would produce exactly the wrong thing :|

Indeed, when passing an empty DataFrame to arrow, it will infer all object columns as 'null' type.
@martindurant 

> `pa.Table.from_pandas(pd.DataFrame(columns=schema.fields, schema=schema)).to_pandas()`

pandas doesn't accept schema keyword
I think it was meant to be 
```
pa.Table.from_pandas(pd.DataFrame(columns=schema.fields), schema=schema).to_pandas()
```
thank you!!!
Interestingly (maybe) that didn't work with timestamps because

```
~/miniconda/envs/daskyarn/lib/python3.6/site-packages/distributed/client.py in _gather(self, futures, errors, direct, local_worker)
   1496                             six.reraise(type(exception),
   1497                                         exception,
-> 1498                                         traceback)
   1499                     if errors == 'skip':
   1500                         bad_keys.add(key)

~/miniconda/envs/daskyarn/lib/python3.6/site-packages/six.py in reraise(tp, value, tb)
    690                 value = tp()
    691             if value.__traceback__ is not tb:
--> 692                 raise value.with_traceback(tb)
    693             raise value
    694         finally:

/mnt/yarn/usercache/hadoop/appcache/application_1556407391758_0007/container_1556407391758_0007_01_000026/environment/lib/python3.6/site-packages/dask/dataframe/utils.py in check_meta()

ValueError: Metadata mismatch found in `from_delayed`.

Partition type: `DataFrame`
+------------+---------------------+----------------+
| Column     | Found               | Expected       |
+------------+---------------------+----------------+
| time_stamp | datetime64[ns, UTC] | datetime64[ns] |
```
Pandas made `datetime64[ns, UTC] ` but the meta from `pa.Table.from_pandas(pd.DataFrame(columns=schema.fields), schema=schema).to_pandas()` was just `datetime64[ns] `.
The implicit timezone in parquet is UTC, I don't know if arrow follows this. The implicit timezone in pandas where none is given is also UTC, so the two above could be considered equivalent.
> The implicit timezone in pandas where none is given is also UTC, so the two above could be considered equivalent.

this is incorrect; no time zone means it’s a naive time; this is distinct from UTC
> The implicit timezone in parquet is UTC, I don't know if arrow follows this. 

Timestamps in Arrow are UTC-normalized if time-zone aware. If time-zone naive, they have the same kind of interpretation that pandas gives them (https://github.com/apache/arrow/blob/master/format/Schema.fbs#L171)
> Pandas made datetime64[ns, UTC] but the meta from pa.Table.from_pandas(pd.DataFrame(columns=schema.fields), schema=schema).to_pandas() was just datetime64[ns] .

@birdsarah can you give a bit more details how you ran into this? Has your original dask dataframe a tz-aware column, but the way to create an empty `meta` as above does not match that?
@jorisvandenbossche yes, it's a bit lost in this thread. I am currently manually making a meta dataframe and a *separate* pyarrow schema. I was lamenting this and @martindurant suggested making the meta dataframe from the pyarrow schema by running `pa.Table.from_pandas(pd.DataFrame(columns=fields), schema=schema).to_pandas()` on my pyarrow schema. But it failed at the timezone step.

My pyarrow schema looks like this:

```python
pyarrow_schema = pa.schema([
    ('argument_0', pa.string()),
    ('argument_1', pa.string()),
    ('argument_2', pa.string()),
    ('argument_3', pa.string()),
    ('argument_4', pa.string()),
    ('argument_5', pa.string()),
    ('argument_6', pa.string()),
    ('argument_7', pa.string()),
    ('arguments', pa.string()),
    ('arguments_len', pa.int64()),
    ('call_stack', pa.string()),
    ('crawl_id', pa.int64()),
    ('document_url', pa.string()),
    ('event_ordinal', pa.int64()),
    ('extension_session_uuid', pa.string()),
    ('frame_id', pa.int64()),
    ('func_name', pa.string()),
    ('in_iframe', pa.bool_()),
    ('incognito', pa.int32()),
    ('operation', pa.string()),
    ('page_scopedevent_ordinal', pa.float64()),
    ('script_col', pa.int64()),
    ('script_line', pa.int64()),
    ('script_loc_eval', pa.string()),
    ('script_url', pa.string()),
    ('symbol', pa.string()),
    ('tab_id', pa.int64()),
    ('time_stamp', pa.timestamp('ns')),
    ('top_level_url', pa.string()),
    ('value', pa.string()),
    ('value_1000', pa.string()),
    ('value_len', pa.int64()),
    ('visit_id', pa.int64()),
    ('window_id', pa.int64()), 
])
```
@birdsarah you create the time stamp field with `pa.timestamp('ns')`, but if you have tz-aware data, this would need to be `pa.timestamp('ns', 'UTC')`.

That fails when creating the empty pandas DataFrame:

```
In [58]: schema = pa.schema([('timestamps', pa.timestamp('us', 'UTC'))])

In [59]: pa.Table.from_pandas(pd.DataFrame(columns=schema.names),
    ...:                      schema=schema, preserve_index=False).to_pandas()
---------------------------------------------------------------------------
TypeError                                 Traceback (most recent call last)
<ipython-input-59-3e6fd251c01b> in <module>
----> 1 pa.Table.from_pandas(pd.DataFrame(columns=schema.names), schema=schema, preserve_index=False).to_pandas()

~/miniconda3/envs/dev37/lib/python3.7/site-packages/pyarrow/array.pxi in pyarrow.lib._PandasConvertible.to_pandas()

~/miniconda3/envs/dev37/lib/python3.7/site-packages/pyarrow/table.pxi in pyarrow.lib.Table._to_pandas()

~/miniconda3/envs/dev37/lib/python3.7/site-packages/pyarrow/pandas_compat.py in table_to_blockmanager(options, table, categories, ignore_metadata)
    600         column_indexes = pandas_metadata.get('column_indexes', [])
    601         index_descriptors = pandas_metadata['index_columns']
--> 602         table = _add_any_metadata(table, pandas_metadata)
    603         table, index = _reconstruct_index(table, index_descriptors,
    604                                           all_columns)

~/miniconda3/envs/dev37/lib/python3.7/site-packages/pyarrow/pandas_compat.py in _add_any_metadata(table, pandas_metadata)
    927                 col = table[idx]
    928                 converted = col.to_pandas()
--> 929                 tz = col_meta['metadata']['timezone']
    930                 tz_aware_type = pa.timestamp('ns', tz=tz)
    931                 with_metadata = pa.Array.from_pandas(converted.values,

TypeError: 'NoneType' object is not subscriptable
```

But that can be solved by specifying `ignore_metadata=True`
```
In [60]: pa.Table.from_pandas(pd.DataFrame(columns=schema.names),
    ...:                      schema=schema, preserve_index=False).to_pandas(ignore_metadata=True)
Out[60]: 
Empty DataFrame
Columns: [timestamps]
Index: []

In [61]: _.dtypes
Out[61]: 
timestamps    datetime64[ns, UTC]
dtype: object
```

The reason that it originally fails is because the pandas metadata that are created when converting the empty DataFrame to an arrow Table does not match the actual schema (and in case of timezone aware data, those metadata are used, but the expected timezone is not found)
Thanks. To get this issue back on track, I believe we still need a solution where dask lets you specify the schema or uses the provided metadata. I am currently running patched dask and it's not ideal.
> dask lets you specify the schema

Short term, this seems to be the only solution, given the ambiguity going back and forth from pandas. I think it may be reasonable to provide a `schema=` keyword to `to_parquet` for the pyarrow branch as one of the arbitrary kwargs that gets passed to the partition-writing tasks.
I can work up a PR. I'm not sure how I would write a test for the error described here though. Perhaps we can collaborate.
I guess make a dask dataframe with a zoned datetime column and no rows (or only NaT) in one partition.
The original issue here has nothing to do with datetime. That whole discussion only arised out of trying to build a shortcut and i would say it's separate if we're not going to try and infer the pyarrow schema from the already provided data to meta.
OK, just thought datetime was the most complex case; the simplest case (int?) would do.

>  if we're not going to try and infer the pyarrow schema 

Right, so it seems this isn't as simple to do as I had originally hoped. It it were straight-forward, no extra options would be needed.
There's also array cases and other datatypes which i feel like i don't know enough about pyarrow to make sure are properly covered. Providing your own pyarrow schema bypasses any attempts to be clever that could fail.
So trying to summarize the different threads from the discussion above:

-  Provide a `schema=` keyword to `to_parquet` for the pyarrow branch as one of the arbitrary kwargs that gets passed to the partition-writing tasks (and we might want to do something similar in pandas' `to_parquet`). 
   This can be used as a work-around to ensure consistent schema's across partitions / multiple writes.
- On the Arrow side, schema evolution / normalization can help with dealing with such situations *when reading* parquet files with inconsistent schema's (https://issues.apache.org/jira/browse/ARROW-2659 and related issues).
-  To prevent that dask writes parquet datasets with inconsistent schema's, it should somehow infer the schema for the whole dask.dataframe, so it can enforce that when writing each partition. This, however, is not that simple to solve (the empty `meta` dataframe alone is not enough to reliably create a schema):
   - In fastparquet, this is somewhat tackled with the `object_encoding` keyword, to specify how to interpret object dtype columns (as all other column dtypes can faithfully be mapped to arrow / parquet types)

Does that sound correct? Other items?
That sounds like a fair summary to me. The first point should not be required if we can tackle the third (or ameliorated by the second), although maybe it's possible to get such inconsistent schemas from elsewhere.
I can take a stab at the Dask related parts summarized above, but I wanted to check first. Is anyone else working on this?
I'm not. I don't currently have time, but would LOVE to see this happen. Thanks!

Resolving patch:
diff --git a/dask/dataframe/io/parquet/arrow.py b/dask/dataframe/io/parquet/arrow.py
--- a/dask/dataframe/io/parquet/arrow.py
+++ b/dask/dataframe/io/parquet/arrow.py
@@ -304,6 +304,7 @@ def write_partition(
         fmd=None,
         compression=None,
         index_cols=None,
+        schema=None,
         **kwargs,
     ):
         md_list = []
@@ -311,7 +312,7 @@ def write_partition(
         if index_cols:
             df = df.set_index(index_cols)
             preserve_index = True
-        t = pa.Table.from_pandas(df, preserve_index=preserve_index)
+        t = pa.Table.from_pandas(df, preserve_index=preserve_index, schema=schema)
         if partition_on:
             pq.write_to_dataset(
                 t,


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
@@ -938,6 +938,112 @@ def test_to_parquet_default_writes_nulls(tmpdir):
     assert table[1].null_count == 2
 
 
+def test_to_parquet_pyarrow_w_inconsistent_schema_by_partition_fails_by_default(tmpdir):
+    check_pyarrow()
+
+    df = pd.DataFrame(
+        {"partition_column": [0, 0, 1, 1], "strings": ["a", "b", None, None]}
+    )
+
+    ddf = dd.from_pandas(df, npartitions=2)
+    ddf.to_parquet(str(tmpdir), engine="pyarrow", partition_on=["partition_column"])
+
+    # Test that read fails because of default behavior when schema not provided
+    with pytest.raises(ValueError) as e_info:
+        dd.read_parquet(str(tmpdir), engine="pyarrow", gather_statistics=False)
+        assert e_info.message.contains("ValueError: Schema in partition")
+        assert e_info.message.contains("was different")
+
+
+def test_to_parquet_pyarrow_w_inconsistent_schema_by_partition_succeeds_w_manual_schema(
+    tmpdir
+):
+    check_pyarrow()
+
+    # Data types to test: strings, arrays, ints, timezone aware timestamps
+    in_arrays = [[0, 1, 2], [3, 4], np.nan, np.nan]
+    out_arrays = [[0, 1, 2], [3, 4], None, None]
+    in_strings = ["a", "b", np.nan, np.nan]
+    out_strings = ["a", "b", None, None]
+    tstamp = pd.Timestamp(1513393355, unit="s")
+    in_tstamps = [tstamp, tstamp, pd.NaT, pd.NaT]
+    out_tstamps = [
+        # Timestamps come out in numpy.datetime64 format
+        tstamp.to_datetime64(),
+        tstamp.to_datetime64(),
+        np.datetime64("NaT"),
+        np.datetime64("NaT"),
+    ]
+    timezone = "US/Eastern"
+    tz_tstamp = pd.Timestamp(1513393355, unit="s", tz=timezone)
+    in_tz_tstamps = [tz_tstamp, tz_tstamp, pd.NaT, pd.NaT]
+    out_tz_tstamps = [
+        # Timezones do not make it through a write-read cycle.
+        tz_tstamp.tz_convert(None).to_datetime64(),
+        tz_tstamp.tz_convert(None).to_datetime64(),
+        np.datetime64("NaT"),
+        np.datetime64("NaT"),
+    ]
+
+    df = pd.DataFrame(
+        {
+            "partition_column": [0, 0, 1, 1],
+            "arrays": in_arrays,
+            "strings": in_strings,
+            "tstamps": in_tstamps,
+            "tz_tstamps": in_tz_tstamps,
+        }
+    )
+
+    ddf = dd.from_pandas(df, npartitions=2)
+    schema = pa.schema(
+        [
+            ("arrays", pa.list_(pa.int64())),
+            ("strings", pa.string()),
+            ("tstamps", pa.timestamp("ns")),
+            ("tz_tstamps", pa.timestamp("ns", timezone)),
+            ("partition_column", pa.int64()),
+        ]
+    )
+    ddf.to_parquet(
+        str(tmpdir), engine="pyarrow", partition_on="partition_column", schema=schema
+    )
+    ddf_after_write = (
+        dd.read_parquet(str(tmpdir), engine="pyarrow", gather_statistics=False)
+        .compute()
+        .reset_index(drop=True)
+    )
+
+    # Check array support
+    arrays_after_write = ddf_after_write.arrays.values
+    for i in range(len(df)):
+        assert np.array_equal(arrays_after_write[i], out_arrays[i]), type(out_arrays[i])
+
+    # Check datetime support
+    tstamps_after_write = ddf_after_write.tstamps.values
+    for i in range(len(df)):
+        # Need to test NaT seperately
+        if np.isnat(tstamps_after_write[i]):
+            assert np.isnat(out_tstamps[i])
+        else:
+            assert tstamps_after_write[i] == out_tstamps[i]
+
+    # Check timezone aware datetime support
+    tz_tstamps_after_write = ddf_after_write.tz_tstamps.values
+    for i in range(len(df)):
+        # Need to test NaT seperately
+        if np.isnat(tz_tstamps_after_write[i]):
+            assert np.isnat(out_tz_tstamps[i])
+        else:
+            assert tz_tstamps_after_write[i] == out_tz_tstamps[i]
+
+    # Check string support
+    assert np.array_equal(ddf_after_write.strings.values, out_strings)
+
+    # Check partition column
+    assert np.array_equal(ddf_after_write.partition_column, df.partition_column)
+
+
 @write_read_engines_xfail
 def test_partition_on(tmpdir, write_engine, read_engine):
     tmpdir = str(tmpdir)