ClickHouse

1
#include "config.h"
2

3
#if USE_ULID
4

5
#include <Columns/ColumnFixedString.h>
6
#include <Columns/ColumnString.h>
7
#include <Columns/ColumnsDateTime.h>
8
#include <DataTypes/DataTypeFixedString.h>
9
#include <DataTypes/DataTypeString.h>
10
#include <Functions/extractTimeZoneFromFunctionArguments.h>
11
#include <Functions/FunctionFactory.h>
12
#include <Functions/FunctionHelpers.h>
13
#include <Functions/IFunction.h>
14
#include <Interpreters/Context.h>
15

16
#include <ulid.h>
17

18

19
namespace DB
20
{
21

22
namespace ErrorCodes
23
{
24
    extern const int NUMBER_OF_ARGUMENTS_DOESNT_MATCH;
25
    extern const int BAD_ARGUMENTS;
26
    extern const int ILLEGAL_TYPE_OF_ARGUMENT;
27
    extern const int ILLEGAL_COLUMN;
28
}
29

30
class FunctionULIDStringToDateTime : public IFunction
31
{
32
public:
33
    static constexpr size_t ULID_LENGTH = 26;
34
    static constexpr UInt32 DATETIME_SCALE = 3;
35

36
    static constexpr auto name = "ULIDStringToDateTime";
37

38
    static FunctionPtr create(ContextPtr /*context*/)
39
    {
40
        return std::make_shared<FunctionULIDStringToDateTime>();
41
    }
42

43
    String getName() const override { return name; }
44

45
    bool isVariadic() const override { return true; }
46
    size_t getNumberOfArguments() const override { return 0; }
47

48
    bool isSuitableForShortCircuitArgumentsExecution(const DataTypesWithConstInfo & /*arguments*/) const override { return true; }
49

50
    DataTypePtr getReturnTypeImpl(const ColumnsWithTypeAndName & arguments) const override
51
    {
52
        if (arguments.empty() || arguments.size() > 2)
53
            throw Exception(
54
                ErrorCodes::NUMBER_OF_ARGUMENTS_DOESNT_MATCH,
55
                "Wrong number of arguments for function {}: should be 1 or 2",
56
                getName());
57

58
        const auto * arg_fixed_string = checkAndGetDataType<DataTypeFixedString>(arguments[0].type.get());
59
        const auto * arg_string = checkAndGetDataType<DataTypeString>(arguments[0].type.get());
60

61
        if (!arg_string && !(arg_fixed_string && arg_fixed_string->getN() == ULID_LENGTH))
62
            throw Exception(
63
                ErrorCodes::ILLEGAL_TYPE_OF_ARGUMENT,
64
                "Illegal type {} of argument of function {}. Must be String or FixedString(26).",
65
                arguments[0].type->getName(),
66
                getName());
67

68
        String timezone;
69
        if (arguments.size() == 2)
70
        {
71
            timezone = extractTimeZoneNameFromColumn(arguments[1].column.get(), arguments[1].name);
72

73
            if (timezone.empty())
74
                throw Exception(ErrorCodes::ILLEGAL_TYPE_OF_ARGUMENT,
75
                    "Function {} supports a 2nd argument (optional) that must be a valid time zone",
76
                    getName());
77
        }
78

79
        return std::make_shared<DataTypeDateTime64>(DATETIME_SCALE, timezone);
80
    }
81

82
    bool useDefaultImplementationForConstants() const override { return true; }
83

84
    ColumnPtr executeImpl(const ColumnsWithTypeAndName & arguments, const DataTypePtr &, size_t input_rows_count) const override
85
    {
86
        auto col_res = ColumnDateTime64::create(input_rows_count, DATETIME_SCALE);
87
        auto & vec_res = col_res->getData();
88

89
        const ColumnPtr column = arguments[0].column;
90

91
        const auto * column_fixed_string = checkAndGetColumn<ColumnFixedString>(column.get());
92
        const auto * column_string = checkAndGetColumn<ColumnString>(column.get());
93

94
        if (column_fixed_string)
95
        {
96
            if (column_fixed_string->getN() != ULID_LENGTH)
97
                throw Exception(
98
                    ErrorCodes::ILLEGAL_COLUMN,
99
                    "Illegal column {} of argument of function {}, expected String or FixedString({})",
100
                    arguments[0].name, getName(), ULID_LENGTH
101
                );
102

103
            const auto & vec_src = column_fixed_string->getChars();
104

105
            for (size_t i = 0; i < input_rows_count; ++i)
106
            {
107
                DateTime64 time = decode(vec_src.data() + i * ULID_LENGTH);
108
                vec_res[i] = time;
109
            }
110
        }
111
        else if (column_string)
112
        {
113
            const auto & vec_src = column_string->getChars();
114
            const auto & offsets_src = column_string->getOffsets();
115

116
            size_t src_offset = 0;
117

118
            for (size_t i = 0; i < input_rows_count; ++i)
119
            {
120
                DateTime64 time = 0;
121

122
                size_t string_size = offsets_src[i] - src_offset;
123
                if (string_size != ULID_LENGTH + 1)
124
                    throw Exception(
125
                        ErrorCodes::ILLEGAL_COLUMN,
126
                        "Illegal column {} of argument of function {}, ULID must be {} characters long",
127
                        arguments[0].name, getName(), ULID_LENGTH
128
                    );
129

130
                time = decode(vec_src.data() + src_offset);
131

132
                src_offset += string_size;
133
                vec_res[i] = time;
134
            }
135
        }
136
        else
137
            throw Exception(
138
                ErrorCodes::ILLEGAL_COLUMN,
139
                "Illegal column {} of argument of function {}, expected String or FixedString({})",
140
                arguments[0].name, getName(), ULID_LENGTH
141
            );
142

143
        return col_res;
144
    }
145

146
    static DateTime64 decode(const UInt8 * data)
147
    {
148
        unsigned char buffer[16];
149
        int ret = ulid_decode(buffer, reinterpret_cast<const char *>(data));
150
        if (ret != 0)
151
            throw Exception(
152
                ErrorCodes::BAD_ARGUMENTS,
153
                "Cannot parse ULID {}",
154
                std::string_view(reinterpret_cast<const char *>(data), ULID_LENGTH)
155
            );
156

157
        /// Timestamp in milliseconds is the first 48 bits of the decoded ULID
158
        Int64 ms = 0;
159
        memcpy(reinterpret_cast<UInt8 *>(&ms) + 2, buffer, 6);
160

161
#    if __BYTE_ORDER__ == __ORDER_LITTLE_ENDIAN__
162
        ms = std::byteswap(ms);
163
#    endif
164

165
        return DecimalUtils::decimalFromComponents<DateTime64>(ms / intExp10(DATETIME_SCALE), ms % intExp10(DATETIME_SCALE), DATETIME_SCALE);
166
    }
167
};
168

169

170
REGISTER_FUNCTION(ULIDStringToDateTime)
171
{
172
    factory.registerFunction<FunctionULIDStringToDateTime>(FunctionDocumentation
173
        {
174
            .description=R"(
175
This function extracts the timestamp from a ULID and returns it as a DateTime64(3) typed value.
176
The function expects the ULID to be provided as the first argument, which can be either a String or a FixedString(26) data type.
177
An optional second argument can be passed to specify a timezone for the timestamp.
178
)",
179
            .examples{
180
                {"ulid", "SELECT ULIDStringToDateTime(generateULID())", ""},
181
                {"timezone", "SELECT ULIDStringToDateTime(generateULID(), 'Asia/Istanbul')", ""}},
182
            .categories{"ULID"}
183
        },
184
        FunctionFactory::CaseSensitive);
185
}
186

187
}
188

189
#endif
190