PandasArrayExtensionArray.take coerces fill_value before checking whether anything is filled, breaking boolean masks on integer array columns

[shashvat-singham]

Describe the bug

PandasArrayExtensionArray.take coerces fill_value to the array's value type before checking whether any element will actually be filled:

https://github.com/huggingface/datasets/blob/main/src/datasets/features/features.py#L965-L972

if allow_fill:
    fill_value = (
        self.dtype.na_value if fill_value is None else np.asarray(fill_value, dtype=self.dtype.value_type)
    )
    mask = indices == -1

Pandas calls take(indices, allow_fill=True, fill_value=nan) for an ordinary boolean mask, even when indices holds no -1 and nothing needs filling. For an integer-valued array np.asarray(nan, dtype="int32") raises, so row selection fails:

ValueError: cannot convert float NaN to integer

float64 is unaffected (nan is representable) and bool is unaffected (np.asarray(nan, dtype=bool) is True), so this only bites integer Array2D/Array3D/Array4D/Array5D columns.

Note this is currently masked by #8375 — execution raises AttributeError in dtype comparison before ever reaching take. It becomes reachable once #8464 lands, which is where I ran into it.

Steps to reproduce the bug

On top of #8464:

import pandas as pd
import datasets

for dtype, dummy in [("float64", 1.0), ("int32", 1), ("int64", 1), ("bool", True)]:
    features = datasets.Features({"foo": datasets.Array2D(dtype=dtype, shape=(2, 2))})
    ds = datasets.Dataset.from_dict({"foo": [[[dummy] * 2] * 2] * 2}, features=features)
    df = ds._data.to_pandas()
    try:
        print(f"{dtype:>8}: OK   -> {df[pd.Series([True, False])].shape}")
    except Exception as e:
        print(f"{dtype:>8}: FAIL -> {type(e).__name__}: {e}")

 float64: OK   -> (1, 1)
   int32: FAIL -> ValueError: cannot convert float NaN to integer
   int64: FAIL -> ValueError: cannot convert float NaN to integer
    bool: OK   -> (1, 1)

Reduced to the array itself, showing the fill is never needed:

import numpy as np
from datasets.features.features import PandasArrayExtensionArray

arr = PandasArrayExtensionArray(np.array([[[1, 1], [1, 1]], [[2, 2], [2, 2]]], dtype="int32"))

arr.take(np.array([0]), allow_fill=False)                      # [[[1, 1], [1, 1]]]
arr.take(np.array([0]), allow_fill=True, fill_value=np.nan)    # ValueError

No -1 is present in either call, so both should return the same thing.

Expected behavior

take should only resolve fill_value when it is actually going to be used, i.e. when mask.any(). With no -1 in indices, allow_fill=True and allow_fill=False should agree, and boolean masking should work for integer array columns as it already does for float64 and bool.

Note on scope

Deferring the coercion to the two places that consume it fixes the reported failure and is a small change. It does leave open what an integer array should actually fill with when a -1 genuinely is present — there is no integer NA, so self.dtype.na_value (nan) cannot be stored either. Options I see:

1. Defer the coercion only (fixes masking; a real fill on an int array still raises, as it does today).
2. Additionally promote the result to a float dtype when a fill is genuinely required.
3. Raise a clearer, explicit error for that case.

Happy to open a PR for (1) since it is self-contained, if that is the direction you'd prefer.

Environment info

• datasets version: 5.0.2.dev0 (main @ 48b7ee7, plus Make PandasArrayExtensionDtype._metadata a tuple #8464)
• Python version: 3.11.9
• Platform: Windows 11
• PyArrow version: 25.0.1
• Pandas version: 3.0.5
• NumPy version: 2.4.6