Skip to content

Commit

Permalink
fix cover
Browse files Browse the repository at this point in the history
  • Loading branch information
chelsea-lin committed Jan 17, 2025
1 parent fdad61e commit 14a6dcd
Show file tree
Hide file tree
Showing 2 changed files with 8 additions and 11 deletions.
14 changes: 3 additions & 11 deletions db_dtypes/json.py
Original file line number Diff line number Diff line change
Expand Up @@ -96,9 +96,9 @@ def __init__(self, values) -> None:
else:
raise NotImplementedError(f"Unsupported pandas version: {pd.__version__}")

def __arrow_array__(self):
def __arrow_array__(self, type=None):
"""Convert to an arrow array. This is required for pyarrow extension."""
return self.pa_data
return pa.array(self.pa_data, type=JSONArrowType())

@classmethod
def _box_pa(
Expand Down Expand Up @@ -159,12 +159,7 @@ def _serialize_json(value):
def _deserialize_json(value):
"""A static method that converts a JSON string back into its original value."""
if not pd.isna(value):
# Attempt to interpret the value as a JSON object.
# If it's not valid JSON, treat it as a regular string.
try:
return json.loads(value)
except json.JSONDecodeError:
return value
return json.loads(value)
else:
return value

Expand Down Expand Up @@ -279,9 +274,6 @@ def __arrow_ext_serialize__(self) -> bytes:
def __arrow_ext_deserialize__(cls, storage_type, serialized) -> JSONArrowType:
return JSONArrowType()

def __hash__(self) -> int:
return hash(str(self))

def to_pandas_dtype(self):
return JSONDtype()

Expand Down
5 changes: 5 additions & 0 deletions tests/unit/test_json.py
Original file line number Diff line number Diff line change
Expand Up @@ -118,6 +118,11 @@ def test_as_numpy_array():
pd._testing.assert_equal(result, expected)


def test_json_arrow_array():
data = db_dtypes.JSONArray._from_sequence(JSON_DATA.values())
assert isinstance(data.__arrow_array__(), pa.ExtensionArray)


def test_json_arrow_storage_type():
arrow_json_type = db_dtypes.JSONArrowType()
assert arrow_json_type.extension_name == "dbjson"
Expand Down

0 comments on commit 14a6dcd

Please sign in to comment.