address some comments

Author: chelsea-lin

db_dtypes/json.py

@@ -19,13 +19,10 @@
 
 import numpy as np
 import pandas as pd
-from pandas.core.arrays.arrow.array import ArrowExtensionArray
-from pandas.core.arrays.masked import BaseMaskedArray
-from pandas.core.dtypes.common import is_dict_like, is_integer, is_list_like, is_scalar
-from pandas.core.dtypes.dtypes import ExtensionDtype
-from pandas.core.indexers import check_array_indexer, unpack_tuple_and_ellipses
+import pandas.arrays as arrays
+import pandas.core.dtypes.common as common
+import pandas.core.indexers as indexers
 import pyarrow as pa
-import pyarrow.compute as pc
 
 
 @pd.api.extensions.register_extension_dtype
@@ -63,78 +60,81 @@ def __from_arrow__(array: typing.Union[pa.Array, pa.ChunkedArray]) -> JSONArray:
         return JSONArray(array)
 
 
-class JSONArray(ArrowExtensionArray):
+class JSONArray(arrays.ArrowExtensionArray):
     """Extension array that handles BigQuery JSON data, leveraging a string-based
     pyarrow array for storage. It enables seamless conversion to JSON objects when
     accessing individual elements."""
 
     _dtype = JSONDtype()
 
     def __init__(self, values, dtype=None, copy=False) -> None:
-        if isinstance(values, (pa.Array, pa.ChunkedArray)) and pa.types.is_string(
-            values.type
-        ):
-            values = pc.cast(values, pa.large_string())
-
-        super().__init__(values)
         self._dtype = JSONDtype()
-
-        if not pa.types.is_large_string(self._pa_array.type) and not (
-            pa.types.is_dictionary(self._pa_array.type)
-            and pa.types.is_large_string(self._pa_array.type.value_type)
-        ):
-            raise ValueError(
-                "ArrowStringArray requires a PyArrow (chunked) array of "
-                "large_string type"
-            )
+        if isinstance(values, pa.Array):
+            self._pa_array = pa.chunked_array([values])
+        elif isinstance(values, pa.ChunkedArray):
+            self._pa_array = values
+        else:
+            raise ValueError(f"Unsupported type '{type(values)}' for JSONArray")
 
     @classmethod
     def _box_pa(
         cls, value, pa_type: pa.DataType | None = None
     ) -> pa.Array | pa.ChunkedArray | pa.Scalar:
         """Box value into a pyarrow Array, ChunkedArray or Scalar."""
         if isinstance(value, pa.Scalar) or not (
-            is_list_like(value) and not is_dict_like(value)
+            common.is_list_like(value) and not common.is_dict_like(value)
         ):
             return cls._box_pa_scalar(value, pa_type)
         return cls._box_pa_array(value, pa_type)
 
     @classmethod
     def _box_pa_scalar(cls, value, pa_type: pa.DataType | None = None) -> pa.Scalar:
         """Box value into a pyarrow Scalar."""
-        value = JSONArray._seralizate_json(value)
-        pa_scalar = super()._box_pa_scalar(value, pa_type)
-        if pa.types.is_string(pa_scalar.type) and pa_type is None:
-            pa_scalar = pc.cast(pa_scalar, pa.large_string())
+        if isinstance(value, pa.Scalar):
+            pa_scalar = value
+        if pd.isna(value):
+            pa_scalar = pa.scalar(None, type=pa_type)
+        else:
+            value = JSONArray._serialize_json(value)
+            pa_scalar = pa.scalar(value, type=pa_type, from_pandas=True)
+
+        if pa_type is not None and pa_scalar.type != pa_type:
+            pa_scalar = pa_scalar.cast(pa_type)
         return pa_scalar
 
     @classmethod
     def _box_pa_array(
         cls, value, pa_type: pa.DataType | None = None, copy: bool = False
     ) -> pa.Array | pa.ChunkedArray:
         """Box value into a pyarrow Array or ChunkedArray."""
-        if (
-            not isinstance(value, cls)
-            and not isinstance(value, (pa.Array, pa.ChunkedArray))
-            and not isinstance(value, BaseMaskedArray)
-        ):
-            value = [JSONArray._seralizate_json(x) for x in value]
-        pa_array = super()._box_pa_array(value, pa_type)
-        if pa.types.is_string(pa_array.type) and pa_type is None:
-            pa_array = pc.cast(pa_array, pa.large_string())
+        if isinstance(value, cls):
+            pa_array = value._pa_array
+        elif isinstance(value, (pa.Array, pa.ChunkedArray)):
+            pa_array = value
+        else:
+            try:
+                value = [JSONArray._serialize_json(x) for x in value]
+                pa_array = pa.array(value, type=pa_type, from_pandas=True)
+            except (pa.ArrowInvalid, pa.ArrowTypeError):
+                # GH50430: let pyarrow infer type, then cast
+                pa_array = pa.array(value, from_pandas=True)
+
+        if pa_type is not None and pa_array.type != pa_type:
+            pa_array = pa_array.cast(pa_type)
+
         return pa_array
 
     @classmethod
     def _from_sequence(cls, scalars, *, dtype=None, copy=False):
         """Construct a new ExtensionArray from a sequence of scalars."""
         result = []
         for scalar in scalars:
-            result.append(JSONArray._seralizate_json(scalar))
+            result.append(JSONArray._serialize_json(scalar))
         return cls(pa.array(result, type=pa.large_string(), from_pandas=True))
 
     @classmethod
     def _from_sequence_of_strings(
-        cls, strings, *, dtype: ExtensionDtype, copy: bool = False
+        cls, strings, *, dtype, copy: bool = False
     ) -> JSONArray:
         """Construct a new ExtensionArray from a sequence of strings."""
         return cls._from_sequence(strings, dtype=dtype, copy=copy)
@@ -152,7 +152,7 @@ def _from_factorized(cls, values, original):
         return cls._from_sequence(values, dtype=original.dtype)
 
     @staticmethod
-    def _seralizate_json(value):
+    def _serialize_json(value):
         """A static method that converts a JSON value into a string representation."""
         if isinstance(value, str) or pd.isna(value):
             return value
@@ -176,19 +176,19 @@ def dtype(self) -> JSONDtype:
 
     def __contains__(self, key) -> bool:
         """Return for `item in self`."""
-        return super().__contains__(JSONArray._seralizate_json(key))
+        return super().__contains__(JSONArray._serialize_json(key))
 
     def insert(self, loc: int, item) -> JSONArray:
         """
         Make new ExtensionArray inserting new item at location. Follows Python
         list.append semantics for negative values.
         """
-        val = JSONArray._seralizate_json(item)
+        val = JSONArray._serialize_json(item)
         return super().insert(loc, val)
 
     def __getitem__(self, item):
         """Select a subset of self."""
-        item = check_array_indexer(self, item)
+        item = indexers.check_array_indexer(self, item)
 
         if isinstance(item, np.ndarray):
             if not len(item):
@@ -203,9 +203,9 @@ def __getitem__(self, item):
                     "boolean arrays are valid indices."
                 )
         elif isinstance(item, tuple):
-            item = unpack_tuple_and_ellipses(item)
+            item = indexers.unpack_tuple_and_ellipses(item)
 
-        if is_scalar(item) and not is_integer(item):
+        if common.is_scalar(item) and not common.is_integer(item):
             # e.g. "foo" or 2.5
             # exception message copied from numpy
             raise IndexError(