ArrayCumSumLimitedImpl

This commit is contained in:
daoready 2018-08-23 21:06:51 +03:00
parent d2c4696bc6
commit 48cd3341a0
2 changed files with 117 additions and 0 deletions

View File

@ -17,6 +17,7 @@ void registerFunctionsHigherOrder(FunctionFactory & factory)
factory.registerFunction<FunctionArraySort>();
factory.registerFunction<FunctionArrayReverseSort>();
factory.registerFunction<FunctionArrayCumSum>();
factory.registerFunction<FunctionArrayCumSumLimited>();
}
}

View File

@ -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<DataTypeUInt8>(&*expression_return) ||
checkDataType<DataTypeUInt16>(&*expression_return) ||
checkDataType<DataTypeUInt32>(&*expression_return) ||
checkDataType<DataTypeUInt64>(&*expression_return))
return std::make_shared<DataTypeArray>(std::make_shared<DataTypeUInt64>());
if (checkDataType<DataTypeInt8>(&*expression_return) ||
checkDataType<DataTypeInt16>(&*expression_return) ||
checkDataType<DataTypeInt32>(&*expression_return) ||
checkDataType<DataTypeInt64>(&*expression_return))
return std::make_shared<DataTypeArray>(std::make_shared<DataTypeInt64>());
if (checkDataType<DataTypeFloat32>(&*expression_return) ||
checkDataType<DataTypeFloat64>(&*expression_return))
return std::make_shared<DataTypeArray>(std::make_shared<DataTypeFloat64>());
throw Exception("arrayCumSumLimited cannot add values of type " + expression_return->getName(), ErrorCodes::ILLEGAL_TYPE_OF_ARGUMENT);
}
template <typename Element, typename Result>
static bool executeType(const ColumnPtr & mapped, const ColumnArray & array, ColumnPtr & res_ptr)
{
const ColumnVector<Element> * column = checkAndGetColumn<ColumnVector<Element>>(&*mapped);
if (!column)
{
const ColumnConst * column_const = checkAndGetColumnConst<ColumnVector<Element>>(&*mapped);
if (!column_const)
return false;
const Element x = column_const->template getValue<Element>();
const IColumn::Offsets & offsets = array.getOffsets();
auto res_nested = ColumnVector<Result>::create();
typename ColumnVector<Result>::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<Element>::Container & data = column->getData();
auto res_nested = ColumnVector<Result>::create();
typename ColumnVector<Result>::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<Float32,Float64>(mapped, array, res) ||
executeType<Float64,Float64>(mapped, array, res))
return res;
else
throw Exception("Unexpected column for arrayCumSumLimited: " + mapped->getName());
}
};
template <typename Impl, typename Name>
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<ArrayMapImpl, NameArrayMap>;
using FunctionArrayFilter = FunctionArrayMapped<ArrayFilterImpl, NameArrayFilter>;
@ -970,5 +1085,6 @@ using FunctionArrayFirstIndex = FunctionArrayMapped<ArrayFirstIndexImpl, NameArr
using FunctionArraySort = FunctionArrayMapped<ArraySortImpl<true>, NameArraySort>;
using FunctionArrayReverseSort = FunctionArrayMapped<ArraySortImpl<false>, NameArrayReverseSort>;
using FunctionArrayCumSum = FunctionArrayMapped<ArrayCumSumImpl, NameArrayCumSum>;
using FunctionArrayCumSumLimited = FunctionArrayMapped<ArrayCumSumLimitedImpl, NameArrayCumSumLimited>;
}