ClickHouse

1
#include <Functions/IFunction.h>
2
#include <Functions/FunctionFactory.h>
3
#include <Columns/ColumnsNumber.h>
4
#include <DataTypes/DataTypesNumber.h>
5
#include <atomic>
6

7

8
namespace DB
9
{
10
namespace
11
{
12

13
/** Incremental number of row within all columns passed to this function. */
14
class FunctionRowNumberInAllBlocks : public IFunction
15
{
16
private:
17
    mutable std::atomic<size_t> rows{0};
18

19
public:
20
    static constexpr auto name = "rowNumberInAllBlocks";
21
    static FunctionPtr create(ContextPtr)
22
    {
23
        return std::make_shared<FunctionRowNumberInAllBlocks>();
24
    }
25

26
    /// Get the name of the function.
27
    String getName() const override
28
    {
29
        return name;
30
    }
31

32
    bool isStateful() const override
33
    {
34
        return true;
35
    }
36

37
    size_t getNumberOfArguments() const override
38
    {
39
        return 0;
40
    }
41

42
    bool isDeterministic() const override
43
    {
44
        return false;
45
    }
46

47
    bool isDeterministicInScopeOfQuery() const override
48
    {
49
        return false;
50
    }
51

52
    bool isSuitableForShortCircuitArgumentsExecution(const DataTypesWithConstInfo & /*arguments*/) const override { return false; }
53

54
    DataTypePtr getReturnTypeImpl(const DataTypes & /*arguments*/) const override
55
    {
56
        return std::make_shared<DataTypeUInt64>();
57
    }
58

59
    ColumnPtr executeImplDryRun(const ColumnsWithTypeAndName &, const DataTypePtr &, size_t input_rows_count) const override
60
    {
61
        return ColumnUInt64::create(input_rows_count);
62
    }
63

64
    ColumnPtr executeImpl(const ColumnsWithTypeAndName &, const DataTypePtr &, size_t input_rows_count) const override
65
    {
66
        size_t current_row_number = rows.fetch_add(input_rows_count);
67

68
        auto column = ColumnUInt64::create();
69
        auto & data = column->getData();
70
        data.resize(input_rows_count);
71
        for (size_t i = 0; i < input_rows_count; ++i)
72
            data[i] = current_row_number + i;
73

74
        return column;
75
    }
76
};
77

78
}
79

80
REGISTER_FUNCTION(RowNumberInAllBlocks)
81
{
82
    factory.registerFunction<FunctionRowNumberInAllBlocks>();
83
}
84

85
}
86