def test_aggregate_list_like(backend, alltypes, df, agg_fn):
"""Tests .aggregate() where the result of an aggregation is a list-like.
We expect the list / np.array to be treated as a scalar (in other words,
the resulting table expression should have one element, which is the
list / np.array).
"""
udf = reduction(input_type=[dt.double],
output_type=dt.Array(dt.double))(agg_fn)
expr = alltypes.aggregate(result_col=udf(alltypes.double_col))
result = expr.execute()
# Expecting a 1-row DataFrame
expected = pd.DataFrame({'result_col': [agg_fn(df.double_col)]})
backend.assert_frame_equal(result, expected)
def create_mean_struct_udf(result_formatter):
return reduction(
input_type=[dt.double, dt.double],
output_type=dt.Struct(['mean', 'mean_weight'], [dt.double, dt.double]),
)(_format_struct_udf_return_type(mean_struct, result_formatter))