2017-05-02 21:08:37 +00:00
|
|
|
#pragma once
|
|
|
|
|
|
|
|
#include <IO/WriteHelpers.h>
|
|
|
|
#include <IO/ReadHelpers.h>
|
|
|
|
|
|
|
|
#include <DataTypes/DataTypeArray.h>
|
|
|
|
#include <DataTypes/DataTypesNumber.h>
|
|
|
|
#include <DataTypes/DataTypeString.h>
|
|
|
|
|
|
|
|
#include <Columns/ColumnArray.h>
|
|
|
|
|
|
|
|
#include <Common/SpaceSaving.h>
|
|
|
|
|
2017-11-24 13:55:31 +00:00
|
|
|
#include <Common/FieldVisitors.h>
|
2017-05-02 21:08:37 +00:00
|
|
|
|
|
|
|
#include <AggregateFunctions/AggregateFunctionGroupArray.h>
|
|
|
|
|
|
|
|
|
|
|
|
namespace DB
|
|
|
|
{
|
|
|
|
|
|
|
|
|
|
|
|
// Allow NxK more space before calculating top K to increase accuracy
|
2017-05-12 20:38:03 +00:00
|
|
|
#define TOP_K_DEFAULT 10
|
2017-05-02 21:08:37 +00:00
|
|
|
#define TOP_K_LOAD_FACTOR 3
|
|
|
|
#define TOP_K_MAX_SIZE 0xFFFFFF
|
|
|
|
|
|
|
|
|
|
|
|
template <typename T>
|
|
|
|
struct AggregateFunctionTopKData
|
|
|
|
{
|
2017-05-10 00:17:45 +00:00
|
|
|
using Set = SpaceSaving
|
|
|
|
<
|
|
|
|
T,
|
|
|
|
HashCRC32<T>,
|
2017-05-12 20:38:03 +00:00
|
|
|
HashTableGrower<4>,
|
|
|
|
HashTableAllocatorWithStackMemory<sizeof(T) * (1 << 4)>
|
2017-05-10 00:17:45 +00:00
|
|
|
>;
|
2017-05-02 21:08:37 +00:00
|
|
|
Set value;
|
|
|
|
};
|
|
|
|
|
|
|
|
|
|
|
|
template <typename T>
|
|
|
|
class AggregateFunctionTopK
|
|
|
|
: public IUnaryAggregateFunction<AggregateFunctionTopKData<T>, AggregateFunctionTopK<T>>
|
|
|
|
{
|
|
|
|
private:
|
|
|
|
using State = AggregateFunctionTopKData<T>;
|
2017-10-14 06:39:48 +00:00
|
|
|
UInt64 threshold = TOP_K_DEFAULT;
|
|
|
|
UInt64 reserved = TOP_K_LOAD_FACTOR * threshold;
|
2017-05-02 21:08:37 +00:00
|
|
|
|
|
|
|
public:
|
|
|
|
String getName() const override { return "topK"; }
|
|
|
|
|
|
|
|
DataTypePtr getReturnType() const override
|
|
|
|
{
|
|
|
|
return std::make_shared<DataTypeArray>(std::make_shared<DataTypeNumber<T>>());
|
|
|
|
}
|
|
|
|
|
|
|
|
void setArgument(const DataTypePtr & argument)
|
|
|
|
{
|
|
|
|
}
|
|
|
|
|
|
|
|
void setParameters(const Array & params) override
|
|
|
|
{
|
|
|
|
if (params.size() != 1)
|
|
|
|
throw Exception("Aggregate function " + getName() + " requires exactly one parameter.", ErrorCodes::NUMBER_OF_ARGUMENTS_DOESNT_MATCH);
|
|
|
|
|
2017-10-14 06:39:48 +00:00
|
|
|
UInt64 k = applyVisitor(FieldVisitorConvertToNumber<UInt64>(), params[0]);
|
2017-05-02 21:08:37 +00:00
|
|
|
|
|
|
|
if (k > TOP_K_MAX_SIZE)
|
|
|
|
throw Exception("Too large parameter for aggregate function " + getName() + ". Maximum: " + toString(TOP_K_MAX_SIZE),
|
|
|
|
ErrorCodes::ARGUMENT_OUT_OF_BOUND);
|
|
|
|
|
|
|
|
threshold = k;
|
|
|
|
reserved = TOP_K_LOAD_FACTOR * k;
|
|
|
|
}
|
|
|
|
|
|
|
|
void addImpl(AggregateDataPtr place, const IColumn & column, size_t row_num, Arena *) const
|
|
|
|
{
|
|
|
|
auto & set = this->data(place).value;
|
2017-05-05 23:36:02 +00:00
|
|
|
if (set.capacity() != reserved)
|
2017-05-02 21:08:37 +00:00
|
|
|
set.resize(reserved);
|
|
|
|
set.insert(static_cast<const ColumnVector<T> &>(column).getData()[row_num]);
|
|
|
|
}
|
|
|
|
|
|
|
|
void merge(AggregateDataPtr place, ConstAggregateDataPtr rhs, Arena * arena) const override
|
|
|
|
{
|
|
|
|
this->data(place).value.merge(this->data(rhs).value);
|
|
|
|
}
|
|
|
|
|
|
|
|
void serialize(ConstAggregateDataPtr place, WriteBuffer & buf) const override
|
|
|
|
{
|
|
|
|
this->data(place).value.write(buf);
|
|
|
|
}
|
|
|
|
|
|
|
|
void deserialize(AggregateDataPtr place, ReadBuffer & buf, Arena *) const override
|
|
|
|
{
|
|
|
|
auto & set = this->data(place).value;
|
|
|
|
set.resize(reserved);
|
|
|
|
set.read(buf);
|
|
|
|
}
|
|
|
|
|
|
|
|
void insertResultInto(ConstAggregateDataPtr place, IColumn & to) const override
|
|
|
|
{
|
|
|
|
ColumnArray & arr_to = static_cast<ColumnArray &>(to);
|
|
|
|
ColumnArray::Offsets_t & offsets_to = arr_to.getOffsets();
|
|
|
|
|
|
|
|
const typename State::Set & set = this->data(place).value;
|
2017-10-08 22:53:38 +00:00
|
|
|
auto result_vec = set.topK(threshold);
|
|
|
|
size_t size = result_vec.size();
|
2017-05-02 21:08:37 +00:00
|
|
|
|
|
|
|
offsets_to.push_back((offsets_to.size() == 0 ? 0 : offsets_to.back()) + size);
|
|
|
|
|
|
|
|
typename ColumnVector<T>::Container_t & data_to = static_cast<ColumnVector<T> &>(arr_to.getData()).getData();
|
|
|
|
size_t old_size = data_to.size();
|
|
|
|
data_to.resize(old_size + size);
|
|
|
|
|
|
|
|
size_t i = 0;
|
2017-10-08 22:53:38 +00:00
|
|
|
for (auto it = result_vec.begin(); it != result_vec.end(); ++it, ++i)
|
2017-05-02 21:08:37 +00:00
|
|
|
data_to[old_size + i] = it->key;
|
|
|
|
}
|
2017-09-17 20:22:39 +00:00
|
|
|
|
|
|
|
const char * getHeaderFilePath() const override { return __FILE__; }
|
2017-05-02 21:08:37 +00:00
|
|
|
};
|
|
|
|
|
|
|
|
|
|
|
|
/// Generic implementation, it uses serialized representation as object descriptor.
|
|
|
|
struct AggregateFunctionTopKGenericData
|
|
|
|
{
|
2017-05-10 00:17:45 +00:00
|
|
|
using Set = SpaceSaving
|
|
|
|
<
|
|
|
|
StringRef,
|
|
|
|
StringRefHash,
|
2017-05-12 20:38:03 +00:00
|
|
|
HashTableGrower<4>,
|
|
|
|
HashTableAllocatorWithStackMemory<sizeof(StringRef) * (1 << 4)>
|
2017-05-10 00:17:45 +00:00
|
|
|
>;
|
2017-05-02 21:08:37 +00:00
|
|
|
|
|
|
|
Set value;
|
|
|
|
};
|
|
|
|
|
|
|
|
/** Template parameter with true value should be used for columns that store their elements in memory continuously.
|
|
|
|
* For such columns topK() can be implemented more efficently (especially for small numeric arrays).
|
|
|
|
*/
|
|
|
|
template <bool is_plain_column = false>
|
|
|
|
class AggregateFunctionTopKGeneric : public IUnaryAggregateFunction<AggregateFunctionTopKGenericData, AggregateFunctionTopKGeneric<is_plain_column>>
|
|
|
|
{
|
|
|
|
private:
|
|
|
|
using State = AggregateFunctionTopKGenericData;
|
|
|
|
DataTypePtr input_data_type;
|
2017-10-14 06:39:48 +00:00
|
|
|
UInt64 threshold = TOP_K_DEFAULT;
|
|
|
|
UInt64 reserved = TOP_K_LOAD_FACTOR * threshold;
|
2017-05-02 21:08:37 +00:00
|
|
|
|
|
|
|
static void deserializeAndInsert(StringRef str, IColumn & data_to);
|
|
|
|
|
|
|
|
public:
|
|
|
|
String getName() const override { return "topK"; }
|
|
|
|
|
|
|
|
void setArgument(const DataTypePtr & argument)
|
|
|
|
{
|
|
|
|
input_data_type = argument;
|
|
|
|
}
|
|
|
|
|
|
|
|
void setParameters(const Array & params) override
|
|
|
|
{
|
|
|
|
if (params.size() != 1)
|
|
|
|
throw Exception("Aggregate function " + getName() + " requires exactly one parameter.", ErrorCodes::NUMBER_OF_ARGUMENTS_DOESNT_MATCH);
|
|
|
|
|
2017-10-14 06:39:48 +00:00
|
|
|
UInt64 k = applyVisitor(FieldVisitorConvertToNumber<UInt64>(), params[0]);
|
2017-05-02 21:08:37 +00:00
|
|
|
|
|
|
|
if (k > TOP_K_MAX_SIZE)
|
|
|
|
throw Exception("Too large parameter for aggregate function " + getName() + ". Maximum: " + toString(TOP_K_MAX_SIZE),
|
|
|
|
ErrorCodes::ARGUMENT_OUT_OF_BOUND);
|
|
|
|
|
|
|
|
threshold = k;
|
|
|
|
reserved = TOP_K_LOAD_FACTOR * k;
|
|
|
|
}
|
|
|
|
|
|
|
|
DataTypePtr getReturnType() const override
|
|
|
|
{
|
|
|
|
return std::make_shared<DataTypeArray>(input_data_type->clone());
|
|
|
|
}
|
|
|
|
|
|
|
|
bool allocatesMemoryInArena() const override
|
|
|
|
{
|
|
|
|
return true;
|
|
|
|
}
|
|
|
|
|
|
|
|
void serialize(ConstAggregateDataPtr place, WriteBuffer & buf) const override
|
|
|
|
{
|
|
|
|
this->data(place).value.write(buf);
|
|
|
|
}
|
|
|
|
|
|
|
|
void deserialize(AggregateDataPtr place, ReadBuffer & buf, Arena * arena) const override
|
|
|
|
{
|
|
|
|
auto & set = this->data(place).value;
|
2017-09-28 07:01:54 +00:00
|
|
|
set.clear();
|
2017-10-08 22:10:13 +00:00
|
|
|
set.resize(reserved);
|
2017-05-02 21:08:37 +00:00
|
|
|
|
2017-10-08 22:09:00 +00:00
|
|
|
// Specialized here because there's no deserialiser for StringRef
|
2017-05-02 21:08:37 +00:00
|
|
|
size_t count = 0;
|
|
|
|
readVarUInt(count, buf);
|
2017-10-08 22:09:00 +00:00
|
|
|
for (size_t i = 0; i < count; ++i)
|
|
|
|
{
|
2017-06-26 00:11:32 +00:00
|
|
|
auto ref = readStringBinaryInto(*arena, buf);
|
2017-05-02 21:08:37 +00:00
|
|
|
UInt64 count, error;
|
|
|
|
readVarUInt(count, buf);
|
|
|
|
readVarUInt(error, buf);
|
2017-06-26 00:11:32 +00:00
|
|
|
set.insert(ref, count, error);
|
|
|
|
arena->rollback(ref.size);
|
2017-05-02 21:08:37 +00:00
|
|
|
}
|
2017-09-28 07:01:54 +00:00
|
|
|
|
|
|
|
set.readAlphaMap(buf);
|
2017-05-02 21:08:37 +00:00
|
|
|
}
|
|
|
|
|
|
|
|
void addImpl(AggregateDataPtr place, const IColumn & column, size_t row_num, Arena * arena) const
|
|
|
|
{
|
|
|
|
auto & set = this->data(place).value;
|
2017-10-08 22:09:00 +00:00
|
|
|
if (set.capacity() != reserved)
|
|
|
|
{
|
2017-05-02 21:08:37 +00:00
|
|
|
set.resize(reserved);
|
|
|
|
}
|
|
|
|
|
2017-05-10 00:17:45 +00:00
|
|
|
StringRef str_serialized = column.getDataAt(row_num);
|
2017-06-26 00:11:32 +00:00
|
|
|
set.insert(str_serialized);
|
2017-05-02 21:08:37 +00:00
|
|
|
}
|
|
|
|
|
|
|
|
void merge(AggregateDataPtr place, ConstAggregateDataPtr rhs, Arena * arena) const override
|
|
|
|
{
|
|
|
|
this->data(place).value.merge(this->data(rhs).value);
|
|
|
|
}
|
|
|
|
|
|
|
|
void insertResultInto(ConstAggregateDataPtr place, IColumn & to) const override
|
|
|
|
{
|
|
|
|
ColumnArray & arr_to = static_cast<ColumnArray &>(to);
|
|
|
|
ColumnArray::Offsets_t & offsets_to = arr_to.getOffsets();
|
|
|
|
IColumn & data_to = arr_to.getData();
|
|
|
|
|
2017-10-08 22:53:38 +00:00
|
|
|
auto result_vec = this->data(place).value.topK(threshold);
|
|
|
|
offsets_to.push_back((offsets_to.size() == 0 ? 0 : offsets_to.back()) + result_vec.size());
|
2017-05-02 21:08:37 +00:00
|
|
|
|
2017-10-08 22:53:38 +00:00
|
|
|
for (auto & elem : result_vec)
|
2017-05-02 21:08:37 +00:00
|
|
|
{
|
|
|
|
deserializeAndInsert(elem.key, data_to);
|
|
|
|
}
|
|
|
|
}
|
2017-09-17 20:22:39 +00:00
|
|
|
|
|
|
|
const char * getHeaderFilePath() const override { return __FILE__; }
|
2017-05-02 21:08:37 +00:00
|
|
|
};
|
|
|
|
|
|
|
|
|
|
|
|
template <>
|
|
|
|
inline void AggregateFunctionTopKGeneric<false>::deserializeAndInsert(StringRef str, IColumn & data_to)
|
|
|
|
{
|
|
|
|
data_to.deserializeAndInsertFromArena(str.data);
|
|
|
|
}
|
|
|
|
|
|
|
|
template <>
|
|
|
|
inline void AggregateFunctionTopKGeneric<true>::deserializeAndInsert(StringRef str, IColumn & data_to)
|
|
|
|
{
|
|
|
|
data_to.insertData(str.data, str.size);
|
|
|
|
}
|
|
|
|
|
|
|
|
|
2017-05-12 20:38:03 +00:00
|
|
|
#undef TOP_K_DEFAULT
|
2017-05-02 21:08:37 +00:00
|
|
|
#undef TOP_K_MAX_SIZE
|
|
|
|
#undef TOP_K_LOAD_FACTOR
|
|
|
|
|
|
|
|
}
|