jonasdedden commented on PR #51611:
URL: https://github.com/apache/arrow/pull/51611#issuecomment-6186182816
@zanmato1984 thanks so much for the investigation! There indeed was a lot of
performance optimization left to do, and I think now things are way faster and
don't have extreme memory overhead in edge cases.
**1. List array + scalar query:** batched `equal` + bitmap scan was the
simplest way to reuse `equal`'s type promotion; the only early exit was in the
bitmap scan. Each list is now searched directly. Short lists (with < 32 values)
are compared in full without data-dependent branches, which keeps them stable
when matches are frequent (early-exit loop lead to mispredicts in benchmarks).
Long lists stop after the first 16-value chunk holding a match.
**2. List array + query array:** the `take` was also a way to reuse `equal`.
The same per-list search now reads `queries[i]` directly, so the repeated
values (and the indices to build them) are gone.
**3. Scalar list + query array:** now uses `is_in` instead of broadcasting.
It aligns with `equal` once both sides are cast to `equal`'s common type, with
one difference: `is_in` tells -0.0 from 0.0, so floats get 0.0 added first.
Promotion still follows `equal`: values and list values are cast to the
types `equal` would compare them as, and a length-1 `equal` call on nulls
raises `equal`'s exact type and time zone errors.
Some benchmarks with 50M `int64` values in one `list<int64>` array:
| List length | Match position | Scalar query: before -> after | Query
array: before -> after |
|---:|---|---|---|
| 3 | first | 49 ms, 10 MB -> 34 ms, 4 MB | 218 ms, 804 MB -> 34 ms, 4 MB |
| 3 | last | 50 ms, 10 MB -> 33 ms, 4 MB | 218 ms, 804 MB -> 35 ms, 4 MB |
| 3 | none | 49 ms, 10 MB -> 34 ms, 4 MB | 219 ms, 804 MB -> 35 ms, 4 MB |
| 3 | random (30%) | 49 ms, 10 MB -> 36 ms, 4 MB | 217 ms, 804 MB -> 33 ms,
4 MB |
| 10 | none | 28 ms, 8 MB -> 19 ms, 1 MB | 196 ms, 801 MB -> 21 ms, 1 MB |
| 10 | random (30%) | 28 ms, 8 MB -> 20 ms, 1 MB | 198 ms, 801 MB -> 21 ms,
1 MB |
| 100 | first | 26 ms, 6 MB -> 9 ms, 0 MB | 202 ms, 800 MB -> 9 ms, 0 MB |
| 100 | none | 34 ms, 6 MB -> 19 ms, 0 MB | 210 ms, 800 MB -> 20 ms, 0 MB |
| 1,000 | first | 20 ms, 6 MB -> 1 ms, 0 MB | 184 ms, 800 MB -> 1 ms, 0 MB |
| 1,000 | middle | 20 ms, 6 MB -> 15 ms, 0 MB | 183 ms, 800 MB -> 15 ms, 0
MB |
| 1,000 | none | 21 ms, 6 MB -> 17 ms, 0 MB | 186 ms, 800 MB -> 18 ms, 0 MB |
| 1,000 | random (30%) | 20 ms, 6 MB -> 1 ms, 0 MB | 184 ms, 800 MB -> 1 ms,
0 MB |
With 10% null child values:
| List length | Match position | Scalar query: before -> after | Query
array: before -> after |
|---:|---|---|---|
| 3 | none | 53 ms -> 51 ms | 219 ms -> 51 ms |
| 3 | random (30%) | 52 ms -> 50 ms | 219 ms -> 50 ms |
| 1,000 | none | 22 ms -> 17 ms | 185 ms -> 18 ms |
| 1,000 | random (30%) | 20 ms -> 3 ms | 184 ms -> 3 ms |
Scalar list + query array (about half the queries match):
| List length | Queries | Before (broadcast) | After (`is_in`) |
|---:|---:|---|---|
| 10 | 1,000,000 | 156 ms, 245 MB | 9.2 ms, 1 MB |
| 1,000 | 1,000 | 4.1 ms, 24 MB | <0.1 ms, 0 MB |
| 1,000 | 100,000 | 428 ms, 2,400 MB | 1.4 ms, 0 MB |
| 100,000 | 1,000 | 427 ms, 2,400 MB | 0.9 ms, 6 MB |
The lookup table paid off even at the smallest size I tried; broadcasting
copies the list once per query, so its memory grows with list length times
queries.
<details>
<summary>Benchmark scripts</summary>
List array, scalar query vs query array, by match position (`python
bench_positions.py`):
```python
import json
import subprocess
import sys
import time
import numpy as np
import pyarrow as pa
import pyarrow.compute as pc
N_VALUES = 50_000_000
def make_input(length, where):
rng = np.random.default_rng(42)
values = rng.integers(2, 1 << 40, N_VALUES)
starts = np.arange(0, N_VALUES // length * length, length)
if where == "first":
values[starts] = 1
elif where == "middle":
values[starts + length // 2] = 1
elif where == "last":
values[starts + length - 1] = 1
elif where == "random":
values[rng.random(N_VALUES) < 0.3] = 1
offsets = np.arange(0, N_VALUES + 1, length, dtype=np.int32)
return pa.ListArray.from_arrays(offsets, pa.array(values))
def run_one(length, where, mode):
lists = make_input(length, where)
query = (pa.scalar(1, pa.int64()) if mode == "scalar"
else pa.array(np.full(len(lists), 1, dtype=np.int64)))
pool = pa.default_memory_pool()
baseline = pool.bytes_allocated()
times = []
for _ in range(5):
start = time.perf_counter()
pc.list_contains(lists, query)
times.append(time.perf_counter() - start)
return {"ms": min(times) * 1e3, "mb": (pool.max_memory() - baseline) /
1e6}
if __name__ == "__main__":
if len(sys.argv) == 4:
print(json.dumps(run_one(int(sys.argv[1]), sys.argv[2],
sys.argv[3])))
sys.exit()
for length in (3, 10, 100, 1000):
for where in ("first", "middle", "last", "none", "random"):
cells = []
for mode in ("scalar", "array"):
out = json.loads(subprocess.run(
[sys.executable, __file__, str(length), where, mode],
capture_output=True, text=True, check=True).stdout)
cells.append(f"{out['ms']:.0f} ms, {out['mb']:.0f} MB")
print(f"| {length} | {where} | " + " | ".join(cells) + " |",
flush=True)
```
With 10% null child values (uses `make_input` from above):
```python
import time
import numpy as np
import pyarrow as pa
import pyarrow.compute as pc
from bench_positions import make_input
for length in (3, 1000):
for where in ("none", "random"):
clean = make_input(length, where)
rng = np.random.default_rng(7)
values = clean.values.to_numpy()
values = pa.array(values, mask=rng.random(len(values)) < 0.1)
lists = pa.ListArray.from_arrays(clean.offsets, values)
queries = pa.array(np.full(len(lists), 1, dtype=np.int64))
cells = []
for q in (pa.scalar(1, pa.int64()), queries):
ts = []
for _ in range(5):
t = time.perf_counter()
pc.list_contains(lists, q)
ts.append(time.perf_counter() - t)
cells.append(f"{min(ts) * 1e3:.0f} ms")
print(f"| {length} | {where} | " + " | ".join(cells) + " |",
flush=True)
```
Scalar list + query array:
```python
import json
import subprocess
import sys
import time
import numpy as np
import pyarrow as pa
import pyarrow.compute as pc
SIZES = [(10, 1_000_000), (1_000, 1_000), (1_000, 100_000), (100_000, 1_000)]
def run_one(list_length, n_queries):
rng = np.random.default_rng(42)
values = rng.permutation(2 * list_length)[:list_length].astype(np.int64)
lst = pa.scalar(values, pa.list_(pa.int64()))
queries = pa.array(rng.integers(0, 2 * list_length, n_queries))
pool = pa.default_memory_pool()
baseline = pool.bytes_allocated()
times = []
for _ in range(5):
start = time.perf_counter()
pc.list_contains(lst, queries)
times.append(time.perf_counter() - start)
return {"ms": min(times) * 1e3, "mb": (pool.max_memory() - baseline) /
1e6}
if __name__ == "__main__":
if len(sys.argv) == 3:
print(json.dumps(run_one(int(sys.argv[1]), int(sys.argv[2]))))
sys.exit()
for list_length, n_queries in SIZES:
out = json.loads(subprocess.run(
[sys.executable, __file__, str(list_length), str(n_queries)],
capture_output=True, text=True, check=True).stdout)
print(f"| {list_length:,} | {n_queries:,} | {out['ms']:,.1f} ms,
{out['mb']:,.0f} MB |")
```
</details>
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
