Repository: pandas-dev/pandas
Pull request: ENH: pass through schema keyword in to_parquet for pyarrow
- [X] closes #30189 
- [x] tests added / passed
- [X] passes `black pandas`
- [x] passes `git diff upstream/master -u -- "*.py" | flake8 --diff`
- [x] whatsnew entry


Problem statement:
No mechanism to enforce schema in to_parquet
#### Code Sample, a copy-pastable example if possible

```python
import pandas as pd
import dask as dd

df1 = pd.DataFrame({'col':  pd.Series([True, False], dtype='object)})
df2 = pd.DataFrame({'col':  pd.Series([None, None], dtype='object)})

df1.to_parquet('file1', engine='pyarrow')
df1.to_parquet('file2', engine='pyarrow')

out_df = dd.read_parquet(['file1', 'file2'], engine='pyarrow').compute()
```
#### Problem description

When writing parquet files using pandas there is no mechanism to enforce schema on the output file for the pyarrow engine; without this ability users are forced to fall back on pyarrows inferred schema which isn't always correct. This poses particular challenges when writing multiple files in sequence where partitions contain some or exclusively missing values.

In the example above `file1` is written with schema `bool` while `file2` is written with schema `null`. When attempting to load both partitions using something like dask this generates errors resulting from schema mismatch.

pyarrow exposes the functionality to solve this problem by passing a `schema` argument to `pa.parquet.Table` but all additional keyword arguments to `to_parquet` in pandas are passed to `pa.parquet.write_table`. 

```python
        if index is None:
            from_pandas_kwargs = {}
        else:
            from_pandas_kwargs = {"preserve_index": index}
        table = self.api.Table.from_pandas(df, **from_pandas_kwargs)
        if partition_cols is not None:
            self.api.parquet.write_to_dataset(
                table,
                path,
                compression=compression,
                coerce_timestamps=coerce_timestamps,
                partition_cols=partition_cols,
                **kwargs,
            )
        else:
            self.api.parquet.write_table(
                table,
                path,
                compression=compression,
                coerce_timestamps=coerce_timestamps,
                **kwargs,
            )
```

This could be solved by adding a `schema` argument to `to_parquet` which is then added to `from_pandas_kwargs` before invokation of `self.api.Table.from_pandas(df, **from_pandas_kwargs)`. 

#### Expected Output

|   |  col  |
|---|-------|
| 0 | True  |
| 1 | False |
| 0 | None  |
| 1 | None  |


#### Output of ``pd.show_versions()``


INSTALLED VERSIONS

commit           : None
python           : 3.7.4.final.0
python-bits      : 64
OS               : Darwin
OS-release       : 18.7.0
machine          : x86_64
processor        : i386
byteorder        : little
LC_ALL           : None
LANG             : en_US.UTF-8
LOCALE           : en_US.UTF-8

pandas           : 0.25.3
numpy            : 1.17.2
pytz             : 2019.3
dateutil         : 2.8.0
pip              : 19.2.3
setuptools       : 41.4.0
Cython           : 0.29.13
pytest           : 5.2.1
hypothesis       : None
sphinx           : 2.2.0
blosc            : None
feather          : None
xlsxwriter       : None
lxml.etree       : None
html5lib         : None
pymysql          : None
psycopg2         : 2.8.3 (dt dec pq3 ext lo64)
jinja2           : 2.10.3
IPython          : 7.8.0
pandas_datareader: None
bs4              : None
bottleneck       : None
fastparquet      : 0.3.2
gcsfs            : None
lxml.etree       : None
matplotlib       : 3.1.1
numexpr          : None
odfpy            : None
openpyxl         : None
pandas_gbq       : None
pyarrow          : 0.15.1
pytables         : None
s3fs             : None
scipy            : 1.3.1
sqlalchemy       : 1.3.9
tables           : None
xarray           : None
xlrd             : None
xlwt             : None
xlsxwriter       : None

Hints from issue discussion:
@ieaves Thanks for the clear report!

> This could be solved by adding a schema argument to to_parquet which is then added to from_pandas_kwargs before invokation of self.api.Table.from_pandas(df, **from_pandas_kwargs).

That seems a good idea to me. The changes in pyarrow 0.15 to more correctly follow the specified `schema` were exactly for the purpose to have better control over the final schema, so we should try to expose this in pandas' `to_parquet` as well.
Dask did something similar, they also pass through the `schema` argument now (https://github.com/dask/dask/pull/5150)

@ieaves Would you be interested in doing a PR for this? (don't feel obliged!) 



@jorisvandenbossche of course! I'll get a PR out in the next few days.

Commit message: black

Patch to review:
diff --git a/doc/source/whatsnew/v1.0.0.rst b/doc/source/whatsnew/v1.0.0.rst
index cd11a15e74b1f..469f445b901cd 100755
--- a/doc/source/whatsnew/v1.0.0.rst
+++ b/doc/source/whatsnew/v1.0.0.rst
@@ -792,7 +792,7 @@ I/O
 - Bug in :func:`read_json` where default encoding was not set to ``utf-8`` (:issue:`29565`)
 - Bug in :class:`PythonParser` where str and bytes were being mixed when dealing with the decimal field (:issue:`29650`)
 - :meth:`read_gbq` now accepts ``progress_bar_type`` to display progress bar while the data downloads. (:issue:`29857`)
--
+- :meth: `to_parquet` now appropriately handles the ``schema`` argument for user defined schemas in the pyarrow engine. (:issue: `30270`)
 
 Plotting
 ^^^^^^^^
diff --git a/pandas/io/parquet.py b/pandas/io/parquet.py
index a044cfcdf6a01..0621c02ce68cb 100644
--- a/pandas/io/parquet.py
+++ b/pandas/io/parquet.py
@@ -96,6 +96,10 @@ def write(
             from_pandas_kwargs = {}
         else:
             from_pandas_kwargs = {"preserve_index": index}
+
+        if "schema" in kwargs:
+            from_pandas_kwargs["schema"] = kwargs.pop("schema")
+
         table = self.api.Table.from_pandas(df, **from_pandas_kwargs)
         if partition_cols is not None:
             self.api.parquet.write_to_dataset(
diff --git a/pandas/tests/io/test_parquet.py b/pandas/tests/io/test_parquet.py
index a98c93c250070..45e5131d059d2 100644
--- a/pandas/tests/io/test_parquet.py
+++ b/pandas/tests/io/test_parquet.py
@@ -504,6 +504,14 @@ def test_empty_dataframe(self, pa):
         df = pd.DataFrame()
         check_round_trip(df, pa)
 
+    def test_write_with_schema(self, pa):
+        import pyarrow
+
+        df = pd.DataFrame({"x": [0, 1]})
+        schema = pyarrow.schema([pyarrow.field("x", type=pyarrow.bool_())])
+        out_df = df.astype(bool)
+        check_round_trip(df, pa, write_kwargs={"schema": schema}, expected=out_df)
+
     @pytest.mark.skip(reason="broken test")
     @td.skip_if_no("pyarrow", min_version="0.15.0")
     def test_additional_extension_arrays(self, pa):


Code review comments:
diff_hunk: @@ -96,6 +96,10 @@ def write(
             from_pandas_kwargs = {}
         else:
             from_pandas_kwargs = {"preserve_index": index}
+
+        if "schema" in kwargs:
path: pandas/io/parquet.py
text: @user1:
you could make this section simplier, something like

```
from_pandas_kwargs: Dict[str, Any] = {schema: kwargs.pop('schema', None)}
if index is not None:
    from_pandas_kwargs['preserve_index'] = index
```