[Python] Severe performance regression in isin() filter After pyarrow v18

[DanielElisenberg]

Describe the bug, including details regarding any error messages, version, and platform.

Severe performance regression in isin() filter After pyarrow v18

Introduction

We use pyarrow at work for statistical analysis of big populations. One of the key features we
use to achieve this is the incredibly fast predicate-pushdown-filters. We use these to only retrieve data for our
current working population.
We discovered a bug/performance-issue that seems to have been introduced in version 18 of pyarrow. This coincides
with the removal of numpy as a dependency. This seems to affect the isin-filter which handles arrays, so a link to
the removal of numpy seems feasible

Reproducible example

I have made a repository containing a reproducible example here.
It requires you to have python version 3.11.* installed and poetry. Run the example with the shell script at the root of the
repository:

./run_tests.sh

Explanation

Let's say we have a Parquet file that contains an observation for each unit in a population. It also contains the
start and stop time of this observation:

unit_id	observation	start	stop
1	23	123	456
2	23	123	456
...	...	...	...
50_000_000	23	123	456

The dataset includes 50 million units, each with a single observation. Now let's say I was currently working with
a population of the 10 million first units. I would then use predicate pushdown with the isin filter to read only
the relevant population:

from pyarrow import dataset

parquet_path = "big_dataset"
population = [unit_id for unit_id in range(1, 10_000_001)]
population_filter = dataset.field("unit_id").isin(population)
dataset.dateset(parquet_path).to_table(population_filter)

In earlier pyarrow versions, this was blazingly fast, but after v18, it has become up to 40 times slower. It is to the point
where we can't upgrade. Previously, our analytics services could handle large volumes of requests per minute, but now they freeze.

Stats

Running the code in the repository containing a reproducible example yields
these results on my local machine:

Generated in-memory table with 50,000,000 rows. Now writing...
✅ Done. Data written to: ../BIG_DATASET

=== PYARROW VERSION 15 ===
Retrieved 10,000,000 rows in 1.71 seconds.

=== PYARROW VERSION 16 ===
Retrieved 10,000,000 rows in 1.67 seconds.

=== PYARROW VERSION 17 ===
Retrieved 10,000,000 rows in 1.66 seconds.

=== PYARROW VERSION 18 ===
Retrieved 10,000,000 rows in 61.83 seconds.

=== PYARROW VERSION 19 ===
Retrieved 10,000,000 rows in 61.09 seconds.

=== PYARROW VERSION 20 ===
Retrieved 10,000,000 rows in 61.53 seconds.

Component(s)

Python

[raulcd]

I took a quick look at the repository, but the scripts are expecting the local dataset available, can you provide the dataset or a way of generating a dataset to reproduce the issue easily?

About the removal of numpy, I'll have to diagnose but we didn't remove numpy, we made it an optional runtime dependency, the code paths should use numpy if available which makes me hesitant on thinking that it is related, but could be.

Have you run any profiling that points you to that?

[DanielElisenberg]

I am sorry if I was not clear in the description 🙇🏻 The script to generate the variables is in v15/write.py. I did not commit the dataset as it is quite big, so it makes more sense to generate it with code.

I only noticed that if you look in each directory there is a poetry.lock file detailing the dependencies of the environment, and numpy disappears from v18 and onwards when doing a poetry add pyarrow==<version> (same as doing a pip install pyarrow==<version>). I have no other reason for connecting it to numpy than this simple observation.

EDIT: I have tried adding in numpy to the environment and rerunning the code as well, but it did not help.

[raulcd]

Sorry, I missed the write.py script on the first subfolder. I've been able to reproduce the issue locally comparing Arrow 17 and Arrow 18. After some bisect, I've been able to track down the issue to this specific commit 44b72d5 which is related to the is_in functionality, the PR:

• GH-43187: [C++] Support basic is_in predicate simplification #43761

The previous commit doesn't have the issue a87a8e0

pyarrow built locally at commit 44b72d5, has the issue:

[((44b72d5c25...))]$ python read.py 
=== PYARROW VERSION 17 ===
Retrieved 10,000,000 rows in 115.71 seconds.

pyarrow built locally at commit a87a8e0, does not have the issue:

[((a87a8e0efe...))]$ python read.py 
=== PYARROW VERSION 17 ===
Retrieved 10,000,000 rows in 3.22 seconds.

@bkietz @mapleFU @pitrou thoughts? this is a 50x performance penalty on this filter expression:

population = [i for i in range(1, 10_000_001)]
filter_expr = ds.field("unit_id").isin(population)

start = perf_counter()
table = dataset.to_table(filter=filter_expr)
end = perf_counter()

[pitrou]

In #43761, the simplification should simply be disabled if the value_set size is above a certain constant threshold (for example 50 items). This should be easy to do.

[rok]

It would be nice to have benchmark cover for this case. (I see we have some cover here but if I read it right it won't cover predicate pushdown).