huggingface / huggingface/datasets
PandasArrayExtensionArray.take coerces fill_value before checking whether anything is filled, breaking boolean masks on integer array columns
- Dominant language
- Python
- Stars
- 22k
- Forks
- 3.4k
- Avg merge
- 5d 7h
- Merged PRs (30d)
- 17
Description
### 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
```python
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:
```python
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:
```python
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 #8464)
- Python version: 3.11.9
- Platform: Windows 11
- PyArrow version: 25.0.1
- Pandas version: 3.0.5
- NumPy version: 2.4.6
Contributor guide
Research direction
Read src/datasets/features/features.py around PandasArrayExtensionArray.take, then run the reduced integer-array example from the issue, noting that it depends on the behavior enabled by #8464. Done means boolean selection with allow_fill=True no longer coerces an unused NaN for integer arrays, while the existing behavior for indices containing -1 remains unchanged.
Written by the indexing model from the issue text.
Assessment
- Tech stack
- numpy, pandas, python
- Domain
- data
- Issue type
- Bug
- Difficulty
- 2/5
- Estimated time
- 1-3 hours
- Activity status
- Quiet
- Clarity
- Clearly specified
- Newbie friendliness
- 72/100