diff --git a/dbms/src/Functions/FunctionsHigherOrder.cpp b/dbms/src/Functions/FunctionsHigherOrder.cpp index a78769b061a..555181f0ebc 100644 --- a/dbms/src/Functions/FunctionsHigherOrder.cpp +++ b/dbms/src/Functions/FunctionsHigherOrder.cpp @@ -17,6 +17,7 @@ void registerFunctionsHigherOrder(FunctionFactory & factory) factory.registerFunction(); factory.registerFunction(); factory.registerFunction(); + factory.registerFunction(); } } diff --git a/dbms/src/Functions/FunctionsHigherOrder.h b/dbms/src/Functions/FunctionsHigherOrder.h index 3443d240f91..39d88517dc6 100644 --- a/dbms/src/Functions/FunctionsHigherOrder.h +++ b/dbms/src/Functions/FunctionsHigherOrder.h @@ -706,6 +706,120 @@ struct ArrayCumSumImpl }; +struct ArrayCumSumLimitedImpl +{ + static bool needBoolean() { return false; } + static bool needExpression() { return false; } + static bool needOneArray() { return false; } + + static DataTypePtr getReturnType(const DataTypePtr & expression_return, const DataTypePtr & /*array_element*/) + { + if (checkDataType(&*expression_return) || + checkDataType(&*expression_return) || + checkDataType(&*expression_return) || + checkDataType(&*expression_return)) + return std::make_shared(std::make_shared()); + + if (checkDataType(&*expression_return) || + checkDataType(&*expression_return) || + checkDataType(&*expression_return) || + checkDataType(&*expression_return)) + return std::make_shared(std::make_shared()); + + if (checkDataType(&*expression_return) || + checkDataType(&*expression_return)) + return std::make_shared(std::make_shared()); + + throw Exception("arrayCumSumLimited cannot add values of type " + expression_return->getName(), ErrorCodes::ILLEGAL_TYPE_OF_ARGUMENT); + } + + + template + static bool executeType(const ColumnPtr & mapped, const ColumnArray & array, ColumnPtr & res_ptr) + { + const ColumnVector * column = checkAndGetColumn>(&*mapped); + + if (!column) + { + const ColumnConst * column_const = checkAndGetColumnConst>(&*mapped); + + if (!column_const) + return false; + + const Element x = column_const->template getValue(); + const IColumn::Offsets & offsets = array.getOffsets(); + + auto res_nested = ColumnVector::create(); + typename ColumnVector::Container & res_values = res_nested->getData(); + res_values.resize(column_const->size()); + + size_t pos = 0; + for (size_t i = 0; i < offsets.size(); ++i) + { + // skip empty arrays + if (pos < offsets[i]) + { + res_values[pos++] = x; + for (; pos < offsets[i]; ++pos) + { + res_values[pos] = res_values[pos - 1] + x; + res_values[pos] = 0 if res_values[pos]<0; + } + } + } + + res_ptr = ColumnArray::create(std::move(res_nested), array.getOffsetsPtr()); + return true; + } + + const IColumn::Offsets & offsets = array.getOffsets(); + const typename ColumnVector::Container & data = column->getData(); + + auto res_nested = ColumnVector::create(); + typename ColumnVector::Container & res_values = res_nested->getData(); + res_values.resize(data.size()); + + size_t pos = 0; + for (size_t i = 0; i < offsets.size(); ++i) + { + // skip empty arrays + if (pos < offsets[i]) + { + res_values[pos] = data[pos]; + for (++pos; pos < offsets[i]; ++pos) + { + res_values[pos] = res_values[pos - 1] + data[pos]; + res_values[pos] = 0 if res_values[pos]<0; + } + } + } + res_ptr = ColumnArray::create(std::move(res_nested), array.getOffsetsPtr()); + return true; + + } + + static ColumnPtr execute(const ColumnArray & array, ColumnPtr mapped) + { + ColumnPtr res; + + if (executeType< UInt8 , UInt64>(mapped, array, res) || + executeType< UInt16, UInt64>(mapped, array, res) || + executeType< UInt32, UInt64>(mapped, array, res) || + executeType< UInt64, UInt64>(mapped, array, res) || + executeType< Int8 , Int64>(mapped, array, res) || + executeType< Int16, Int64>(mapped, array, res) || + executeType< Int32, Int64>(mapped, array, res) || + executeType< Int64, Int64>(mapped, array, res) || + executeType(mapped, array, res) || + executeType(mapped, array, res)) + return res; + else + throw Exception("Unexpected column for arrayCumSumLimited: " + mapped->getName()); + } + +}; + + template class FunctionArrayMapped : public IFunction { @@ -958,6 +1072,7 @@ struct NameArrayFirstIndex { static constexpr auto name = "arrayFirstIndex"; }; struct NameArraySort { static constexpr auto name = "arraySort"; }; struct NameArrayReverseSort { static constexpr auto name = "arrayReverseSort"; }; struct NameArrayCumSum { static constexpr auto name = "arrayCumSum"; }; +struct NameArrayCumSumLimited { static constexpr auto name = "arrayCumSumLimited"; }; using FunctionArrayMap = FunctionArrayMapped; using FunctionArrayFilter = FunctionArrayMapped; @@ -970,5 +1085,6 @@ using FunctionArrayFirstIndex = FunctionArrayMapped, NameArraySort>; using FunctionArrayReverseSort = FunctionArrayMapped, NameArrayReverseSort>; using FunctionArrayCumSum = FunctionArrayMapped; +using FunctionArrayCumSumLimited = FunctionArrayMapped; }