Skip to content

Commit a2b2097

Browse files
committed
fix: support correct numpy construction for dbjson dtype in pandas 1.5
1 parent c5e9a10 commit a2b2097

File tree

1 file changed

+13
-0
lines changed

1 file changed

+13
-0
lines changed

db_dtypes/json.py

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

0 commit comments

Comments
 (0)