diff --git a/src/Functions/FunctionSnowflake.h b/src/Functions/FunctionSnowflake.h new file mode 100644 index 00000000000..cf3a91b8e69 --- /dev/null +++ b/src/Functions/FunctionSnowflake.h @@ -0,0 +1,208 @@ +#pragma once + +#include +#include +#include +#include +#include +#include + +#include + + +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(); + } + + 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(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(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(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(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(); + } + + 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 &>(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(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(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::create(input_rows_count, 3); + auto & result_data = res_column->getData(); + + const auto & source_data = typeid_cast(col).getData(); + + for (size_t i = 0; i < input_rows_count; ++i) + { + result_data[i] = (source_data[i]>>22)+1288834974657; + } + + return res_column; + } +}; + +} diff --git a/src/Functions/dateTime64ToSnowflake.cpp b/src/Functions/dateTime64ToSnowflake.cpp new file mode 100644 index 00000000000..87e35c25371 --- /dev/null +++ b/src/Functions/dateTime64ToSnowflake.cpp @@ -0,0 +1,14 @@ +#include +#include + +namespace DB +{ + +void registerDateTime64ToSnowflake(FunctionFactory & factory) +{ + factory.registerFunction("dateTime64ToSnowflake", + [](ContextPtr){ return std::make_unique( + std::make_shared("dateTime64ToSnowflake")); }); +} + +} diff --git a/src/Functions/dateTimeToSnowflake.cpp b/src/Functions/dateTimeToSnowflake.cpp new file mode 100644 index 00000000000..246f35cc1dc --- /dev/null +++ b/src/Functions/dateTimeToSnowflake.cpp @@ -0,0 +1,14 @@ +#include +#include + +namespace DB +{ + +void registerDateTimeToSnowflake(FunctionFactory & factory) +{ + factory.registerFunction("dateTimeToSnowflake", + [](ContextPtr){ return std::make_unique( + std::make_shared("dateTimeToSnowflake")); }); +} + +} diff --git a/src/Functions/registerFunctions.cpp b/src/Functions/registerFunctions.cpp index 7e8f35bc0c4..9b1a7faebbe 100644 --- a/src/Functions/registerFunctions.cpp +++ b/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); diff --git a/src/Functions/registerFunctionsSnowflake.cpp b/src/Functions/registerFunctionsSnowflake.cpp new file mode 100644 index 00000000000..f0c2feddfb5 --- /dev/null +++ b/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); +} + +} diff --git a/src/Functions/snowflakeToDateTime.cpp b/src/Functions/snowflakeToDateTime.cpp new file mode 100644 index 00000000000..37f5e07512f --- /dev/null +++ b/src/Functions/snowflakeToDateTime.cpp @@ -0,0 +1,14 @@ +#include +#include + +namespace DB +{ + +void registerSnowflakeToDateTime(FunctionFactory & factory) +{ + factory.registerFunction("snowflakeToDateTime", + [](ContextPtr){ return std::make_unique( + std::make_shared("snowflakeToDateTime")); }); +} + +} diff --git a/src/Functions/snowflakeToDateTime64.cpp b/src/Functions/snowflakeToDateTime64.cpp new file mode 100644 index 00000000000..ef9502a224e --- /dev/null +++ b/src/Functions/snowflakeToDateTime64.cpp @@ -0,0 +1,14 @@ +#include +#include + +namespace DB +{ + +void registerSnowflakeToDateTime64(FunctionFactory & factory) +{ + factory.registerFunction("snowflakeToDateTime64", + [](ContextPtr){ return std::make_unique( + std::make_shared("snowflakeToDateTime64")); }); +} + +}