Skip to content

Commit e17595e

Browse files
committed
fix: support correct numpy construction for dbjson dtype in pandas 1.5
1 parent 4b84e4a commit e17595e

File tree

1 file changed

+13
-0
lines changed

1 file changed

+13
-0
lines changed

db_dtypes/json.py

Lines changed: 13 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -229,3 +229,16 @@ def _reduce(
229229
if name in ["min", "max"]:
230230
raise TypeError("JSONArray does not support min/max reducntion.")
231231
super()._reduce(name, skipna=skipna, keepdims=keepdims, **kwargs)
232+
233+
def __array__(self, dtype=None, copy: bool | None = None) -> np.ndarray:
234+
"""Correctly construct numpy arrays when passed to `np.asarray()`."""
235+
pa_type = self.pa_data.type
236+
data = self
237+
if dtype is None:
238+
empty = pa.array([], type=pa_type).to_numpy(zero_copy_only=False)
239+
dtype = empty.dtype
240+
result = np.empty(len(data), dtype=dtype)
241+
mask = data.isna()
242+
result[mask] = pd.NA
243+
result[~mask] = data[~mask].pa_data.to_numpy()
244+
return result

0 commit comments

Comments
 (0)