feat: add conversion between snowflake id and datetime

src/Functions/FunctionSnowflake.h

@@ -0,0 +1,208 @@
+#pragma once
+
+#include <Functions/extractTimeZoneFromFunctionArguments.h>
+#include <Functions/IFunction.h>
+#include <Functions/FunctionHelpers.h>
+#include <DataTypes/DataTypeDateTime64.h>
+#include <DataTypes/DataTypesNumber.h>
+#include <Columns/ColumnsNumber.h>
+
+#include <common/arithmeticOverflow.h>
+
+
+namespace DB
+{
+
+namespace ErrorCodes
+{
+    extern const int NUMBER_OF_ARGUMENTS_DOESNT_MATCH;
+    extern const int ILLEGAL_TYPE_OF_ARGUMENT;
+}
+
+
+class FunctionDateTimeToSnowflake : public IFunction
+{
+private:
+    const char * name;
+public:
+    FunctionDateTimeToSnowflake( const char * name_)
+        : name(name_)
+    {
+    }
+
+    String getName() const override { return name; }
+    size_t getNumberOfArguments() const override { return 1; }
+    bool isVariadic() const override { return false; }
+    bool useDefaultImplementationForConstants() const override { return true; }
+
+    DataTypePtr getReturnTypeImpl(const ColumnsWithTypeAndName & arguments) const override
+    {
+        if (!isDateTime(arguments[0].type))
+            throw Exception(ErrorCodes::ILLEGAL_TYPE_OF_ARGUMENT, "The only argument for function {} must be DateTime", name);
+
+        return std::make_shared<DataTypeInt64>();
+    }
+
+    ColumnPtr executeImpl(const ColumnsWithTypeAndName & arguments, const DataTypePtr &, size_t input_rows_count) const override
+    {
+        const auto & src = arguments[0];
+        const auto & col = *src.column;
+
+        auto res_column = ColumnInt64::create(input_rows_count);
+        auto & result_data = res_column->getData();
+
+        const auto & source_data = typeid_cast<const ColumnUInt32 &>(col).getData();
+        for (size_t i = 0; i < input_rows_count; ++i)
+        {
+            result_data[i] = (int64_t(source_data[i])*1000-1288834974657)<<22;
+        }
+
+        return res_column;
+    }
+};
+
+
+class FunctionSnowflakeToDateTime : public IFunction
+{
+private:
+    const char * name;
+public:
+    FunctionSnowflakeToDateTime(const char * name_)
+        : name(name_)
+    {
+    }
+
+    String getName() const override { return name; }
+    size_t getNumberOfArguments() const override { return 0; }
+    bool isVariadic() const override { return true; }
+    bool useDefaultImplementationForConstants() const override { return true; }
+
+    DataTypePtr getReturnTypeImpl(const ColumnsWithTypeAndName & arguments) const override
+    {
+        
+        if (arguments.size() < 1 || arguments.size() > 2)
+            throw Exception(ErrorCodes::NUMBER_OF_ARGUMENTS_DOESNT_MATCH, "Function {} takes one or two arguments", name);
+
+        if (!typeid_cast<const DataTypeInt64 *>(arguments[0].type.get()))
+            throw Exception(ErrorCodes::ILLEGAL_TYPE_OF_ARGUMENT, "The first argument for function {} must be Int64", name);
+
+        std::string timezone;
+        if (arguments.size() == 2)
+            timezone = extractTimeZoneNameFromFunctionArguments(arguments, 1, 0);
+
+        return std::make_shared<DataTypeDateTime>(timezone);
+    }
+
+    ColumnPtr executeImpl(const ColumnsWithTypeAndName & arguments, const DataTypePtr &, size_t input_rows_count) const override
+    {
+        const auto & src = arguments[0];
+        const auto & col = *src.column;
+
+        auto res_column = ColumnUInt32::create(input_rows_count);
+        auto & result_data = res_column->getData();
+
+        const auto & source_data = typeid_cast<const ColumnInt64 &>(col).getData();
+
+        for (size_t i = 0; i < input_rows_count; ++i)
+        {
+            result_data[i] = ((source_data[i]>>22)+1288834974657)/1000;
+        }
+            
+        return res_column;
+    }
+};
+
+
+class FunctionDateTime64ToSnowflake : public IFunction
+{
+private:
+    const char * name;
+public:
+    FunctionDateTime64ToSnowflake( const char * name_)
+        : name(name_)
+    {
+    }
+
+    String getName() const override { return name; }
+    size_t getNumberOfArguments() const override { return 1; }
+    bool isVariadic() const override { return false; }
+    bool useDefaultImplementationForConstants() const override { return true; }
+
+    DataTypePtr getReturnTypeImpl(const ColumnsWithTypeAndName & arguments) const override
+    {
+        if (!isDateTime64(arguments[0].type))
+            throw Exception(ErrorCodes::ILLEGAL_TYPE_OF_ARGUMENT, "The only argument for function {} must be DateTime64", name);
+
+        return std::make_shared<DataTypeInt64>();
+    }
+
+    ColumnPtr executeImpl(const ColumnsWithTypeAndName & arguments, const DataTypePtr &, size_t input_rows_count) const override
+    {
+        const auto & src = arguments[0];
+        const auto & col = *src.column;
+
+        auto res_column = ColumnInt64::create(input_rows_count);
+        auto & result_data = res_column->getData();
+
+        const auto & source_data = typeid_cast<const ColumnDecimal<DateTime64> &>(col).getData();
+        for (size_t i = 0; i < input_rows_count; ++i)
+        {
+            result_data[i] = (source_data[i]-1288834974657)<<22;
+        }
+
+        return res_column;
+    }
+};
+
+
+class FunctionSnowflakeToDateTime64 : public IFunction
+{
+private:
+    const char * name;
+public:
+    FunctionSnowflakeToDateTime64(const char * name_)
+        : name(name_)
+    {
+    }
+
+    String getName() const override { return name; }
+    size_t getNumberOfArguments() const override { return 0; }
+    bool isVariadic() const override { return true; }
+    bool useDefaultImplementationForConstants() const override { return true; }
+
+    DataTypePtr getReturnTypeImpl(const ColumnsWithTypeAndName & arguments) const override
+    {
+        
+        if (arguments.size() < 1 || arguments.size() > 2)
+            throw Exception(ErrorCodes::NUMBER_OF_ARGUMENTS_DOESNT_MATCH, "Function {} takes one or two arguments", name);
+
+        if (!typeid_cast<const DataTypeInt64 *>(arguments[0].type.get()))
+            throw Exception(ErrorCodes::ILLEGAL_TYPE_OF_ARGUMENT, "The first argument for function {} must be Int64", name);
+
+        std::string timezone;
+        if (arguments.size() == 2)
+            timezone = extractTimeZoneNameFromFunctionArguments(arguments, 1, 0);
+
+        return std::make_shared<DataTypeDateTime64>(3, timezone);
+    }
+
+    ColumnPtr executeImpl(const ColumnsWithTypeAndName & arguments, const DataTypePtr &, size_t input_rows_count) const override
+    {
+        const auto & src = arguments[0];
+        const auto & col = *src.column;
+
+        auto res_column = ColumnDecimal<DateTime64>::create(input_rows_count, 3);
+        auto & result_data = res_column->getData();
+
+        const auto & source_data = typeid_cast<const ColumnInt64 &>(col).getData();
+
+        for (size_t i = 0; i < input_rows_count; ++i)
+        {
+            result_data[i] = (source_data[i]>>22)+1288834974657;
+        }
+            
+        return res_column;
+    }
+};
+
+}

src/Functions/dateTime64ToSnowflake.cpp

@@ -0,0 +1,14 @@
+#include <Functions/FunctionSnowflake.h>
+#include <Functions/FunctionFactory.h>
+
+namespace DB
+{
+
+void registerDateTime64ToSnowflake(FunctionFactory & factory)
+{
+    factory.registerFunction("dateTime64ToSnowflake",
+        [](ContextPtr){ return std::make_unique<FunctionToOverloadResolverAdaptor>(
+            std::make_shared<FunctionDateTime64ToSnowflake>("dateTime64ToSnowflake")); });
+}
+
+}

src/Functions/dateTimeToSnowflake.cpp

@@ -0,0 +1,14 @@
+#include <Functions/FunctionSnowflake.h>
+#include <Functions/FunctionFactory.h>
+
+namespace DB
+{
+
+void registerDateTimeToSnowflake(FunctionFactory & factory)
+{
+    factory.registerFunction("dateTimeToSnowflake",
+        [](ContextPtr){ return std::make_unique<FunctionToOverloadResolverAdaptor>(
+            std::make_shared<FunctionDateTimeToSnowflake>("dateTimeToSnowflake")); });
+}
+
+}

src/Functions/registerFunctions.cpp

@@ -51,6 +51,7 @@ void registerFunctionBitHammingDistance(FunctionFactory & factory);
 void registerFunctionTupleHammingDistance(FunctionFactory & factory);
 void registerFunctionsStringHash(FunctionFactory & factory);
 void registerFunctionValidateNestedArraySizes(FunctionFactory & factory);
+void registerFunctionsSnowflake(FunctionFactory & factory);
 #if !defined(ARCADIA_BUILD)
 void registerFunctionBayesAB(FunctionFactory &);
 #endif
@@ -115,6 +116,7 @@ void registerFunctions()
     registerFunctionTupleHammingDistance(factory);
     registerFunctionsStringHash(factory);
     registerFunctionValidateNestedArraySizes(factory);
+    registerFunctionsSnowflake(factory);
 
 #if !defined(ARCADIA_BUILD)
     registerFunctionBayesAB(factory);

src/Functions/registerFunctionsSnowflake.cpp

@@ -0,0 +1,22 @@
+namespace DB
+{
+
+class FunctionFactory;
+
+void registerDateTimeToSnowflake(FunctionFactory &);
+void registerSnowflakeToDateTime(FunctionFactory &);
+
+void registerDateTime64ToSnowflake(FunctionFactory &);
+void registerSnowflakeToDateTime64(FunctionFactory &);
+
+
+void registerFunctionsSnowflake(FunctionFactory & factory)
+{
+    registerDateTimeToSnowflake(factory);
+    registerSnowflakeToDateTime(factory);
+    
+    registerDateTime64ToSnowflake(factory);
+    registerSnowflakeToDateTime64(factory);
+}
+
+}

src/Functions/snowflakeToDateTime.cpp

@@ -0,0 +1,14 @@
+#include <Functions/FunctionSnowflake.h>
+#include <Functions/FunctionFactory.h>
+
+namespace DB
+{
+
+void registerSnowflakeToDateTime(FunctionFactory & factory)
+{
+    factory.registerFunction("snowflakeToDateTime",
+        [](ContextPtr){ return std::make_unique<FunctionToOverloadResolverAdaptor>(
+            std::make_shared<FunctionSnowflakeToDateTime>("snowflakeToDateTime")); });
+}
+
+}

src/Functions/snowflakeToDateTime64.cpp

@@ -0,0 +1,14 @@
+#include <Functions/FunctionSnowflake.h>
+#include <Functions/FunctionFactory.h>
+
+namespace DB
+{
+
+void registerSnowflakeToDateTime64(FunctionFactory & factory)
+{
+    factory.registerFunction("snowflakeToDateTime64",
+        [](ContextPtr){ return std::make_unique<FunctionToOverloadResolverAdaptor>(
+            std::make_shared<FunctionSnowflakeToDateTime64>("snowflakeToDateTime64")); });
+}
+
+}