[Python] [Perf] Speed up null bitmap creation in NumPy to Arrow conversion #55443

Describe the enhancement requested

Hello there coming from xarray I noticed the pyarrow conversion from numpy to pyarrow was slow for datetime64 because pyarrow is scanning NaT values to drop them as null.

The conversion from NumPy to Arrow is slow in two cases:

• datetime64 and timedelta64 arrays. pyarrow always scans these arrays for NaT.
• Arrays with a mask= argument.

In the two cases, most of the time goes into making the validity bitmap. The values buffer is zero-copy.

Example

import timeit

import numpy as np
import pyarrow as pa

n = 25_920_000
values = np.random.rand(n)
dates = np.arange(n).astype("datetime64[s]")
mask = np.zeros(n, dtype=bool)

for label, fn in {
    "float64 (zero-copy)     ": lambda: pa.array(values),
    "float64 mask=           ": lambda: pa.array(values, mask=mask),
    "float64 from_pandas=True": lambda: pa.array(values, from_pandas=True),
    "datetime64              ": lambda: pa.array(dates),
}.items():
    t = min(timeit.repeat(fn, number=1, repeat=5))
    print(f"{label} {t * 1e3:.1f} ms")

pyarrow 26.0.1
float64 (zero-copy)      0.0 ms
float64 mask=            34.0 ms
float64 from_pandas=True 45.1 ms
datetime64               45.7 ms

For comparison, np.isnat(dates).any() takes approximately 8 ms on the same array.

Cause

ValuesToBitmap and MaskToBitmap functions has a loop with a branch for each element. The loop writes one bit at a time with SetBit or ClearBit. The compiler cannot vectorize this loop.

Solution

Update both functions to vectorize null_count and batching Bitmap setting.

• Count the nulls with a loop that has no branch for vectorization
• If there are no nulls, set all bits with bit_util::SetBitsTo.
• If there are nulls, write the bitmap one byte at a time with GenerateBitsUnrolled.

Component(s)

Python