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 PARAMETER_OUT_OF_BOUND;
extern const int NUMBER_OF_ARGUMENTS_DOESNT_MATCH; 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: private:
AggregateFunctionPtr nested_func; AggregateFunctionPtr nested_func;
size_t nested_size_of_data = 0; size_t nested_size_of_data = 0;
size_t num_arguments;
AggregateFunctionForEachData & ensureAggregateData(AggregateDataPtr place, size_t new_size, Arena & arena) const AggregateFunctionForEachData & ensureAggregateData(AggregateDataPtr place, size_t new_size, Arena & arena) const
{ {
@ -95,7 +97,7 @@ private:
public: public:
AggregateFunctionForEach(AggregateFunctionPtr nested_, const DataTypes & arguments) AggregateFunctionForEach(AggregateFunctionPtr nested_, const DataTypes & arguments)
: nested_func(nested_) : nested_func(nested_), num_arguments(arguments.size())
{ {
nested_size_of_data = nested_func->sizeOfData(); 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 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 ColumnArray & first_array_column = static_cast<const ColumnArray &>(*columns[0]);
const IColumn::Offsets & offsets = first_array_column.getOffsets(); 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 begin = row_num == 0 ? 0 : offsets[row_num - 1];
size_t end = offsets[row_num]; 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; char * nested_state = state.array_of_aggregate_datas;
for (size_t i = begin; i < end; ++i) 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; 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);