blob: e2e9038cd8b5192ea3f7627b4a465048bf6c7024 (
plain) (
blame)
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
|
#include <memory>
#include <Columns/ColumnString.h>
#include <Core/Block.h>
#include <DataTypes/DataTypeString.h>
#include <Functions/FunctionFactory.h>
#include <Functions/IFunction.h>
#include <Storages/MergeTree/MergeTreePartition.h>
namespace DB
{
namespace ErrorCodes
{
extern const int NUMBER_OF_ARGUMENTS_DOESNT_MATCH;
}
/** partitionId(x, y, ...) is a function that computes partition ids of arguments.
* The function is slow and should not be called for large amount of rows.
*/
class FunctionPartitionId : public IFunction
{
public:
static constexpr auto name = "partitionId";
static FunctionPtr create(ContextPtr) { return std::make_shared<FunctionPartitionId>(); }
String getName() const override { return name; }
bool isVariadic() const override { return true; }
size_t getNumberOfArguments() const override { return 0; }
bool isInjective(const ColumnsWithTypeAndName & /*sample_columns*/) const override { return true; }
bool useDefaultImplementationForNulls() const override { return true; }
bool useDefaultImplementationForConstants() const override { return true; }
bool isSuitableForShortCircuitArgumentsExecution(const DataTypesWithConstInfo & /*arguments*/) const override { return true; }
DataTypePtr getReturnTypeImpl(const DataTypes & arguments) const override
{
if (arguments.empty())
throw Exception(ErrorCodes::NUMBER_OF_ARGUMENTS_DOESNT_MATCH, "Function {} requires at least one argument.", getName());
return std::make_shared<DataTypeString>();
}
ColumnPtr executeImpl(const ColumnsWithTypeAndName & arguments, const DataTypePtr &, size_t input_rows_count) const override
{
Block sample_block(arguments);
size_t size = arguments.size();
auto result_column = ColumnString::create();
for (size_t j = 0; j < input_rows_count; ++j)
{
Row row(size);
for (size_t i = 0; i < size; ++i)
arguments[i].column->get(j, row[i]);
MergeTreePartition partition(std::move(row));
result_column->insert(partition.getID(sample_block));
}
return result_column;
}
};
REGISTER_FUNCTION(PartitionId)
{
factory.registerFunction<FunctionPartitionId>();
}
}
|
