How do I apply a filter on a map type column in a Pyarrow table while loading?

I have a file written in the Deltalake/Parquet format which has a map as one of the columns. The map stores various properties of the row entry in a "property_name": "property_value" format. I'd like to filter on a particular property stored in this map column, preferably before loading the table into memory using predicate pushdown if available.

This was my attempt to solve the problem via accessing the key-value data through nested fields:

1. Code to create example parquet file:

import pyarrow.parquet as pq
import pyarrow as pa
import pandas as pd
data = {'Name': ['Name1', 'Name2', 'Name3'],
        'Trial_Map': [{'a': 'a1', 'b': 'b1'}, {'a': 'a2', 'b': 'b2'}, {'a': 'a1', 'b': 'b3'}]}
df = pd.DataFrame(data)
schema = pa.schema([
    ('Name', pa.string()),
    ('Trial_Map', pa.map_(pa.string(), pa.string()))
])
table = pa.Table.from_pandas(df, schema=schema)
writer = pq.ParquetWriter("example.parquet", table.schema)
writer.write_table(table)
writer.close()

2. Code that attempts filtering while reading file to table:

import pyarrow.parquet as pq
import pyarrow.dataset as ds
condition = ((ds.field("Trial_Map", "Trial_Map", "key") == "a") &
             (ds.field("Trial_Map", "Trial_Map", "value") == "a1"))
table = pq.read_table("example.parquet", filters=condition)
print(table.schema)

However, this code gave me the following error:
pyarrow.lib.ArrowNotImplementedError: Function 'struct_field' has no kernel matching input types (map<string, string ('Trial_Map')>)

I would appreciate any help in solving this error/pointing me to any other methods of performing this pre-loading filter. Thanks for your time!

PyArrow's predicate pushdown feature allows you to filter data during the Parquet file reading process, reducing the amount of data loaded into memory. However, it seems that PyArrow does not currently support predicate pushdown for nested fields such as map columns.

To work around this limitation, you can consider using Apache Spark with PySpark, which provides more advanced functionality for querying and manipulating nested data structures. Spark supports predicate pushdown for complex nested types like maps, allowing you to filter data before loading it into memory.

Here's an example of how you can achieve the desired filtering using PySpark:

from pyspark.sql import SparkSession
from pyspark.sql.functions import col, explode
spark = SparkSession.builder.getOrCreate()
# Read the Parquet file into a DataFrame
df = spark.read.parquet("example.parquet")
# Explode the map column into individual key-value pairs
df_exploded = df.select("Name", explode("Trial_Map").alias("Trial_Map"))
# Apply the filter condition on the exploded DataFrame
filtered_df = df_exploded.filter((col("Trial_Map.key") == "a") & (col("Trial_Map.value") == "a1"))
# Show the filtered results
filtered_df.show()
# You can also convert the DataFrame back to a Parquet file if needed
filtered_df.write.parquet("filtered_example.parquet")

In this code, we first read the Parquet file into a DataFrame using PySpark. Then, we explode the map column into individual key-value pairs using the explode function. After that, we apply the desired filter condition on the exploded DataFrame using the filter function. Finally, we show the filtered results and optionally write them back to a Parquet file.

Using PySpark gives you more flexibility in handling complex data types and performing advanced operations on nested fields like maps.

PyArrow currently doesn't support directly selecting the values for a certain key using a nested field referenced (as you were trying with ds.field("Trial_Map", "key")), but there is a compute function that allows selecting those values, i.e. "map_lookup".

If we can assume that each key occurs only once in each map element (i.e. no duplicates per row), we can define a filter like this:

import pyarrow.compute as pc
map_filter = pc.map_lookup(pc.field("Trial_Map"), pa.scalar("a"), "first") == "a1"

Using it on your example, you can see it working in practice:

>>> pq.read_table("example.parquet").to_pandas()
    Name           Trial_Map
0  Name1  [(a, a1), (b, b1)]
1  Name2  [(a, a2), (b, b2)]
2  Name3  [(a, a1), (b, b3)]
>>> pq.read_table("example.parquet", filters=map_filter).to_pandas()
    Name           Trial_Map
0  Name1  [(a, a1), (b, b1)]
1  Name3  [(a, a1), (b, b3)]

The reason that this only works when there is only a single "a" key per map element is because I have to specify "first" to get the first value for key "a" for each element. If there are multiple ones, you can specify "all", but then you get a ListArray as result.
Applying the compute function directly on the table to illustrate this:

>>> pc.map_lookup(table["Trial_Map"].chunk(0), pa.scalar("a"), "first") 
<pyarrow.lib.StringArray object at 0x7f099d2b8040>
[
  "a1",
  "a2",
  "a1"
]
>>> pc.map_lookup(table["Trial_Map"].chunk(0), pa.scalar("a"), "all")
<pyarrow.lib.ListArray object at 0x7f099d3134c0>
[
  [
    "a1"
  ],
  [
    "a2"
  ],
  [
    "a1"
  ]
]

And then if we have this ListArray, the element-wise equality == "a1" doesn't work out of the box (there is an open enhancement request about this to add a function to check if a list contains some value: https://github.com/apache/arrow/issues/33295)