Implemented -ForEach aggregate functions combinator for multiple arguments; added test for -ForEachIf and -IfForEach variants [#CLICKHOUSE-3511].

This commit is contained in:
Alexey Milovidov 2017-12-24 02:16:42 +03:00
parent 1ba80c6768
commit cc79fb6684
3 changed files with 22 additions and 5 deletions

View File

@ -18,6 +18,7 @@ namespace ErrorCodes
{
extern const int PARAMETER_OUT_OF_BOUND;
extern const int NUMBER_OF_ARGUMENTS_DOESNT_MATCH;
extern const int SIZES_OF_ARRAYS_DOESNT_MATCH;
}
@ -47,6 +48,7 @@ class AggregateFunctionForEach final : public IAggregateFunctionDataHelper<Aggre
private:
AggregateFunctionPtr nested_func;
size_t nested_size_of_data = 0;
size_t num_arguments;
AggregateFunctionForEachData & ensureAggregateData(AggregateDataPtr place, size_t new_size, Arena & arena) const
{
@ -95,7 +97,7 @@ private:
public:
AggregateFunctionForEach(AggregateFunctionPtr nested_, const DataTypes & arguments)
: nested_func(nested_)
: nested_func(nested_), num_arguments(arguments.size())
{
nested_size_of_data = nested_func->sizeOfData();
@ -136,20 +138,33 @@ public:
void add(AggregateDataPtr place, const IColumn ** columns, size_t row_num, Arena * arena) const override
{
const IColumn * nested[num_arguments];
for (size_t i = 0; i < num_arguments; ++i)
nested[i] = &static_cast<const ColumnArray &>(*columns[i]).getData();
const ColumnArray & first_array_column = static_cast<const ColumnArray &>(*columns[0]);
const IColumn::Offsets & offsets = first_array_column.getOffsets();
const IColumn * array_data = &first_array_column.getData();
size_t begin = row_num == 0 ? 0 : offsets[row_num - 1];
size_t end = offsets[row_num];
size_t array_size = end - begin;
/// Sanity check. NOTE We can implement specialization for a case with single argument, if the check will hurt performance.
for (size_t i = 1; i < num_arguments; ++i)
{
const ColumnArray & ith_column = static_cast<const ColumnArray &>(*columns[i]);
const IColumn::Offsets & ith_offsets = ith_column.getOffsets();
AggregateFunctionForEachData & state = ensureAggregateData(place, array_size, *arena);
if (ith_offsets[row_num] != end || (row_num != 0 && ith_offsets[row_num - 1] != begin))
throw Exception("Arrays passed to " + getName() + " aggregate function have different sizes", ErrorCodes::SIZES_OF_ARRAYS_DOESNT_MATCH);
}
AggregateFunctionForEachData & state = ensureAggregateData(place, end - begin, *arena);
char * nested_state = state.array_of_aggregate_datas;
for (size_t i = begin; i < end; ++i)
{
nested_func->add(nested_state, static_cast<const IColumn **>(&array_data), i, arena);
nested_func->add(nested_state, nested, i, arena);
nested_state += nested_size_of_data;
}
}

View File

@ -0,0 +1 @@
[5,7,9] [1,2,3] [5,2,9]

View File

@ -0,0 +1 @@
SELECT sumForEach(arr), sumForEachIf(arr, arr[1] = 1), sumIfForEach(arr, arrayMap(x -> x != 5, arr)) FROM (SELECT arrayJoin([[1, 2, 3], [4, 5, 6]]) AS arr);