Skip to content

Commit

Permalink
fix(pandas): handle struct columns with NA elements
Browse files Browse the repository at this point in the history
  • Loading branch information
cpcloud authored and kszucs committed Jun 7, 2022
1 parent 5886a9a commit 9a7c510
Showing 1 changed file with 16 additions and 3 deletions.
19 changes: 16 additions & 3 deletions ibis/backends/pandas/execution/structs.py
Original file line number Diff line number Diff line change
@@ -1,7 +1,7 @@
"""Pandas backend execution of struct fields and literals."""

import collections
import operator
import functools

import pandas as pd
from pandas.core.groupby import SeriesGroupBy
Expand All @@ -15,17 +15,30 @@ def execute_node_struct_field_dict(op, data, **kwargs):
return data[op.field]


@execute_node.register(ops.StructField, type(None))
def execute_node_struct_field_none(op, data, **kwargs):
return None


@execute_node.register(ops.StructField, pd.Series)
def execute_node_struct_field_series(op, data, **kwargs):
field = op.field
return data.map(operator.itemgetter(field)).rename(field)
return data.map(functools.partial(_safe_getter, field=field)).rename(field)


def _safe_getter(value, field: str):
try:
return value[field]
except TypeError:
return value


@execute_node.register(ops.StructField, SeriesGroupBy)
def execute_node_struct_field_series_group_by(op, data, **kwargs):
field = op.field

return (
data.obj.map(operator.itemgetter(field))
data.obj.map(functools.partial(_safe_getter, field=field))
.rename(field)
.groupby(data.grouper.groupings)
)

0 comments on commit 9a7c510

Please sign in to comment.