BUG: Aggregation with mixed reduction datatypes (array, scalar, etc) fails on Dask backend

[timothydijamco]

Problem

Using .aggregate() with a reduction that outputs an array + a reduction that outputs a scalar fails on the Dask backend.

@udf.reduction(
    input_type=[dt.int64],
    output_type=dt.int64,
)
def sum_udf(v):
    return np.sum(v)

@udf.reduction(
    input_type=[dt.int64],
    output_type=dt.Array(dt.int64),
)
def collect_udf(v):
    return np.array(v)

t = t.aggregate(
    sum_col=sum_udf(t['int_col']),
    arr_col=collect_udf(t['int_col'])
)
t.execute()

~/miniconda3/envs/ibis-dev/lib/python3.9/site-packages/dask/dataframe/multi.py in concat(dfs, axis, join, interleave_partitions, ignore_unknown_divisions, **kwargs)
   1138             return concat_unindexed_dataframes(dfs, **kwargs)
   1139         else:
-> 1140             raise ValueError(
   1141                 "Unable to concatenate DataFrame with unknown "
   1142                 "division specifying axis=1"

ValueError: Unable to concatenate DataFrame with unknown division specifying axis=1

When executing .aggregate(), Dask computes the individual reductions, then concatenates them together to form a single DataFrame.

In this case, Dask does not want to concatenate the results of the two reductions because one has known divisions, and the other has unknown divisions:

1. The non-array reduction UDF results in a delayed object with unknown divisions. This is due to the way we are using dd.from_delayed in coerce_to_output: util.py#L128-L130
2. The array reduction UDF results in a delayed object with known divisions, due to coerce_to_output using dd.from_pandas: util.py#L134-L136.

Solution

In this PR, we change our use of dd.from_delayed in coerce_to_output to manually specify the divisions. Since the divisions are now always known, Dask will be able to concatenate the results of the individual reductions together.

[datapythonista]

lgtm, just this is probably worth a release note

[timothydijamco]

Thanks, added a release note

[jreback]

thanks @timothydijamco