Skip to content

Commit

Permalink
fix(pandas-format): convert map keys
Browse files Browse the repository at this point in the history
  • Loading branch information
cpcloud committed Dec 18, 2023
1 parent e98fa3c commit bb92e9f
Showing 1 changed file with 44 additions and 15 deletions.
59 changes: 44 additions & 15 deletions ibis/formats/pandas.py
Original file line number Diff line number Diff line change
@@ -1,7 +1,6 @@
from __future__ import annotations

import contextlib
import json
import warnings

import numpy as np
Expand Down Expand Up @@ -240,6 +239,9 @@ def convert_Struct_element(cls, dtype):
converters = tuple(map(cls.get_element_converter, dtype.types))

def convert(values, names=dtype.names, converters=converters):
if values is None:
return values

items = values.items() if isinstance(values, dict) else zip(names, values)
return {
k: converter(v) if v is not None else v
Expand All @@ -250,19 +252,24 @@ def convert(values, names=dtype.names, converters=converters):

@classmethod
def convert_JSON_element(cls, _):
def try_json(x):
if x is None:
return x
import json

def convert(value):
if value is None:
return value
try:
return json.loads(x)
return json.loads(value)
except (TypeError, json.JSONDecodeError):
return x
return value

return try_json
return convert

@classmethod
def convert_Timestamp_element(cls, dtype):
def converter(value, dtype=dtype):
if value is None:
return value

with contextlib.suppress(AttributeError):
value = value.item()

Expand All @@ -276,23 +283,45 @@ def converter(value, dtype=dtype):
@classmethod
def convert_Array_element(cls, dtype):
convert_value = cls.get_element_converter(dtype.value_type)
return lambda values: [
convert_value(value) if value is not None else value for value in values
]

def convert(values):
if values is None:
return values

return [
convert_value(value) if value is not None else value for value in values
]

return convert

@classmethod
def convert_Map_element(cls, dtype):
convert_key = cls.get_element_converter(dtype.key_type)
convert_value = cls.get_element_converter(dtype.value_type)
return lambda row: {
key: convert_value(value) if value is not None else value
for key, value in dict(row).items()
}

def convert(raw_row):
if raw_row is None:
return raw_row

row = dict(raw_row)
return dict(
zip(map(convert_key, row.keys()), map(convert_value, row.values()))
)

return convert

@classmethod
def convert_UUID_element(cls, _):
from uuid import UUID

return lambda v: v if isinstance(v, UUID) else UUID(v)
def convert(value):
if value is None:
return value
elif isinstance(value, UUID):
return value
return UUID(value)

return convert


class DaskData(PandasData):
Expand Down

0 comments on commit bb92e9f

Please sign in to comment.