[Python] Underscores in partition (string) values are dropped when reading dataset

[asfimport]

When reading a partitioned dataset, in which the partition column contains string values with underscores, pyarrow seems to be ignoring the underscores in the resulting values.

For example if I write and then read a dataset as follows:

import pyarrow as pa
import pandas as pd

df = pd.DataFrame({
    "year_week": ["2019_2", "2019_3"],
    "value": [1, 2]
})

table = pa.Table.from_pandas(df.head())
pq.write_to_dataset(table, 'test', partition_cols=["year_week"])

table2 = pq.ParquetDataset('test').read()

The resulting 'year_week' column in table 2 has lost the underscores:

table2[1] # Gives:

<Column name='year_week' type=DictionaryType(dictionary<values=int64, indices=int32, ordered=0>)>
[

  -- dictionary:
    [
      20192,
      20193
    ]
  -- indices:
    [
      0
    ],

  -- dictionary:
    [
      20192,
      20193
    ]
  -- indices:
    [
      1
    ]
]

Is this intentional behaviour or is this a bug in arrow?

Reporter: Julian de Ruiter
Assignee: Joris Van den Bossche / @jorisvandenbossche

Related issues:

• [Python] Datatypes not preserved for partition fields in roundtrip to partitioned parquet dataset (relates to)
• [Python][Dataset] Support using dataset API in pyarrow.parquet with a minimal ParquetDataset shim (depends upon)

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

[asfimport]

Joris Van den Bossche / @jorisvandenbossche:
Thanks for the report!

The problem is that we try to convert the keys to integer, and if that fails just preserve them as strings.
That is done here

arrow/python/pyarrow/parquet.py

Lines 659 to 663 in 961927a

	try:
	integer_keys = [int(x) for x in self.keys]
	dictionary = lib.array(integer_keys)
	except ValueError:
	dictionary = lib.array(self.keys)

        # Only integer and string partition types are supported right now
        try:
            integer_keys = [int(x) for x in self.keys]
            dictionary = lib.array(integer_keys)
        except ValueError:
            dictionary = lib.array(self.keys)

and apparently, Python will convert a string with an underscore to an integer ...

In [3]: int("2019_1")                                                                                                                                                                                              
Out[3]: 20191

I think this is because in recent Python versions underscores are allowed in integer literals (eg to separate thousands).
We could special case this and first check if there is an underscore in the string before trying to convert to integers, but that's a big ugly.

[asfimport]

Robin Kåveland:
Maybe it's an option to check all(key.isdigit() for key in self.keys) ? Might also be ugly, but apparently:

 

>>> '1231_2'.isdigit()
False
>>> int('123_2')
1232
>>> '1231'.isdigit()
True

 

This breaks stuff like base 16 partition keys, but at least the original data is still unchanged, so the user can coerce themselves.

[asfimport]

Julian de Ruiter:
I think isdigit might be a good solution, it would also guard against silently casting (for example) floats to integers. It also less ugly than explicitly testing for an underscore.

[asfimport]

Julian de Ruiter:
What about negative values though? isdigit returns false for '-10' for example, whilst int('-10') returns an int.

[asfimport]

Robin Kåveland:
That's just a slight modification to the expression above:

all(key.isdigit() or (key.startswith('-') and key[1:].isdigit()) for key in self.keys)

But it's starting to feel like a bad idea to attempt to do this coercing at all. Maybe it's better to force the user to deal with what type the partition key should have? If it were to be interpreted as something that looks a bit like pd.Categorical, it would be relatively cheap to read it into memory and deal with it after the user has read the file?

[asfimport]

Julian de Ruiter:
Ideally you would want to preserve the data type of the partition columns, but that's going to be hard to do properly if you also want to include other types such as dates, etc. Maybe it would be most consistent to just read partitioned columns as a categorical type and let the user handle the details, as this makes the code less magic than what happens now.

[asfimport]

Joris Van den Bossche / @jorisvandenbossche:
This now works with the new Datasets API:

In [2]: import pyarrow.dataset as ds

In [3]: dataset = ds.dataset("test/", format="parquet", partitioning="hive")                                                                                                           

In [4]: dataset.schema                                                                                                                                                                                             
Out[4]: 
value: int64
year_week: string

In [5]: dataset.to_table().to_pandas()                                                                                                                                                                             
Out[5]: 
   value year_week
0      1    2019_2
1      2    2019_3

So once we start using this new code in the parquet module (ARROW-8039), this issue should get resolved.

[asfimport]

Joris Van den Bossche / @jorisvandenbossche:
With the new API, you can also specify the exact schema of the partition keys in case the default inference is not what you want (eg specify string in case it looks like integers but you want to keep it as string)

For example:

part = ds.partitioning(pa.schema([("year_week", pa.int64())]), flavor="hive")   
dataset = ds.dataset("test", format="parquet", partitioning=part)      

(for the actual example code this doesn't work, because "2019_2" is no longer parsable as int now, and then it gives all nulls instead of raising an error, going to open a separate issue about that)

[asfimport]

Wes McKinney / @wesm:
Test added in 57b5082