diff --git a/ibis/backends/pandas/execution/structs.py b/ibis/backends/pandas/execution/structs.py index 0cec23a6ae3e..baf19c73ddd5 100644 --- a/ibis/backends/pandas/execution/structs.py +++ b/ibis/backends/pandas/execution/structs.py @@ -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 @@ -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) )