def test_map_grouping_mixed(mixed_grouping_size):
grouping, size = mixed_grouping_size
def fn(x):
return x * 2 + 5
result = map_grouping(fn, grouping)
expected = make_grouping_by_index(
grouping, list(map(fn, flatten_grouping(grouping)))
)
assert expected == result
def make_schema_with_nones(grouping):
"""
Create a grouping by replacing all grouping scalars values with None
"""
return map_grouping(lambda _: None, grouping)
def test_map_grouping_list(list_grouping_size):
grouping, size = list_grouping_size
result = map_grouping(lambda x: x * 2 + 5, grouping)
expected = [g * 2 + 5 for g in grouping]
assert expected == result
def test_map_grouping_dict(dict_grouping_size):
grouping, size = dict_grouping_size
result = map_grouping(lambda x: x * 2 + 5, grouping)
expected = {k: v * 2 + 5 for k, v in grouping.items()}
assert expected == result
def test_map_grouping_scalar(scalar_grouping_size):
grouping, size = scalar_grouping_size
result = map_grouping(lambda x: x * 2 + 5, grouping)
expected = grouping * 2 + 5
assert expected == result