[Python] Introspect partition keys and values in fragments.

[coady]

Describe the enhancement requested

It's not possible to programmatically determine the values of partition keys in a fragment. Fragments have a partition_expression attribute, but the Expression type doesn't allow any further introspection. I don't want to have to parse the string representation of the expression.

In []: dataset.partitioning.schema
Out[]: 
year: int32
month: int32

In []: fragment = next(dataset.get_fragments())

In []: fragment.partition_expression
Out[]: <pyarrow.compute.Expression ((year == 2013) and (month == 1))>

My broader use case is more performant (speed and memory) aggregation of partitioned data. Using pc._group_by requires loaded arrays, so it ignores that the data is already partitioned. And iterating get_fragments is crippled if one can't identify the fragment.

Component(s)

Python

[lidavidm]

There is _get_partition_keys which is private, but has been kept around. I think there were past discussions about promoting it to a public API.

[westonpace]

That's an interesting broader use case. Can you add a new GH issue for the broader ask of a simpler way to do aggregations on partitions?

[jorisvandenbossche]

There is _get_partition_keys which is private, but has been kept around. I think there were past discussions about promoting it to a public API.

Yes, we should probably do that (we did the same for pyarrow.parquet.filters_to_expression). I think we can re-scope this issue to do that.

Also @coady feel free to already use the current "private" method. It's private because it was thought to not be really user-facing, we know it is used (eg dask uses it as well), so we promise some stability for it.

[jorisvandenbossche]

Something related that might be worth to mention (not that it solves your exact use case here though): there is also a dataset.partitioning.dictionaries if you want to inspect all possible values for all fragments that a certain partitioning field can be.

[jorisvandenbossche]

Opened #33862 to make it public.