[Python] parquet.write_to_dataset creates empty partitions for non-observed dictionary items (categories)

[asfimport]

Hello,

it looks like, views with selection along categorical column are not properly respected.

For the following dummy dataframe:

 

d = pd.date_range('1990-01-01', freq='D', periods=10000)
vals = pd.np.random.randn(len(d), 4)
x = pd.DataFrame(vals, index=d, columns=['A', 'B', 'C', 'D'])
x['Year'] = x.index.year

The slice by Year is saved to partitioned parquet properly:

table = pa.Table.from_pandas(x[x.Year==1990], preserve_index=False)
pq.write_to_dataset(table, root_path='test_a.parquet', partition_cols=['Year'])

However, if we convert Year to pandas.Categorical - it will save the whole original dataframe, not only slice of Year=1990:

x['Year'] = x['Year'].astype('category')

table = pa.Table.from_pandas(x[x.Year==1990], preserve_index=False)
pq.write_to_dataset(table, root_path='test_b.parquet', partition_cols=['Year'])

 

 

Reporter: Vladimir

_{Note: This issue was originally created as ARROW-7617. Please see the migration documentation for further details.}

[asfimport]

Vladimir:
Apologies for the false alarm.

It is not saving data, just creating empty partitions for all categories. Maybe still undesired behavior, but not the bug as it was reported.

[asfimport]

Joris Van den Bossche / @jorisvandenbossche:
[~Filimonov] thanks for the report!

It's indeed writing all partitions, and this is because all the unique categories are still present in the type (the categories in pandas, the dictionary items in arrow), and then because of pandas' groupby behaviour of by default preserving all categories (used here: pq.write_to_dataset(table, root_path='test_b.parquet', partition_cols=['Year'])). Pandas has a observed=True option in groupby (see https://pandas.pydata.org/pandas-docs/stable/reference/api/pandas.DataFrame.groupby.html) to ensure only categories that are actually present in the data are used to group, which could be used here. PR certainly welcome!

[asfimport]

Joris Van den Bossche / @jorisvandenbossche:
Reopened and changed the title to reflect the actual issue

[asfimport]

Antoine Pitrou / @pitrou:
@arw2019 Do you still want to be assigned this issue?

[asfimport]

Andrew Wieteska / @arw2019:
@pitrou  unassigned myself