diff --git a/src/Functions/FunctionSnowflake.h b/src/Functions/FunctionSnowflake.h index ce3a48269b4..b57e21e9a50 100644 --- a/src/Functions/FunctionSnowflake.h +++ b/src/Functions/FunctionSnowflake.h @@ -5,6 +5,8 @@ #include #include #include +#include +#include #include #include @@ -53,22 +55,19 @@ public: ColumnPtr executeImpl(const ColumnsWithTypeAndName & arguments, const DataTypePtr &, size_t input_rows_count) const override { const auto & src = arguments[0]; - const auto & col = *src.column; + const auto & src_column = *src.column; auto res_column = ColumnInt64::create(input_rows_count); - auto & result_data = res_column->getData(); + auto & res_data = res_column->getData(); - const auto & source_data = typeid_cast(col).getData(); + const auto & src_data = typeid_cast(src_column).getData(); for (size_t i = 0; i < input_rows_count; ++i) - { - result_data[i] = (Int64(source_data[i]) * 1000 - snowflake_epoch) << time_shift; - } + res_data[i] = (Int64(src_data[i]) * 1000 - snowflake_epoch) << time_shift; return res_column; } }; - class FunctionSnowflakeToDateTime : public IFunction { private: @@ -105,18 +104,28 @@ public: ColumnPtr executeImpl(const ColumnsWithTypeAndName & arguments, const DataTypePtr &, size_t input_rows_count) const override { const auto & src = arguments[0]; - const auto & col = *src.column; + const auto & src_column = *src.column; auto res_column = ColumnUInt32::create(input_rows_count); - auto & result_data = res_column->getData(); + auto & res_data = res_column->getData(); - const auto & source_data = typeid_cast(col).getData(); - - for (size_t i = 0; i < input_rows_count; ++i) + if (const auto * src_column_non_const = typeid_cast(&src_column)) { - result_data[i] = static_cast( - ((source_data[i] >> time_shift) + snowflake_epoch) / 1000); + const auto & src_data = src_column_non_const->getData(); + for (size_t i = 0; i < input_rows_count; ++i) + res_data[i] = static_cast( + ((src_data[i] >> time_shift) + snowflake_epoch) / 1000); } + else if (const auto * src_column_const = typeid_cast(&src_column)) + { + Int64 src_val = src_column_const->getValue(); + for (size_t i = 0; i < input_rows_count; ++i) + res_data[i] = static_cast( + ((src_val >> time_shift) + snowflake_epoch) / 1000); + } + else + throw Exception(ErrorCodes::ILLEGAL_TYPE_OF_ARGUMENT, "Illegal argument for function {}", name); + return res_column; } }; @@ -147,16 +156,14 @@ public: ColumnPtr executeImpl(const ColumnsWithTypeAndName & arguments, const DataTypePtr &, size_t input_rows_count) const override { const auto & src = arguments[0]; - const auto & col = *src.column; + const auto & src_column = *src.column; auto res_column = ColumnInt64::create(input_rows_count); - auto & result_data = res_column->getData(); + auto & res_data = res_column->getData(); - const auto & source_data = typeid_cast &>(col).getData(); + const auto & src_data = typeid_cast &>(src_column).getData(); for (size_t i = 0; i < input_rows_count; ++i) - { - result_data[i] = (source_data[i] - snowflake_epoch) << time_shift; - } + res_data[i] = (src_data[i] - snowflake_epoch) << time_shift; return res_column; } @@ -199,17 +206,26 @@ public: ColumnPtr executeImpl(const ColumnsWithTypeAndName & arguments, const DataTypePtr &, size_t input_rows_count) const override { const auto & src = arguments[0]; - const auto & col = *src.column; + const auto & src_column = *src.column; auto res_column = ColumnDecimal::create(input_rows_count, 3); - auto & result_data = res_column->getData(); + auto & res_data = res_column->getData(); - const auto & source_data = typeid_cast(col).getData(); - - for (size_t i = 0; i < input_rows_count; ++i) + if (const auto * src_column_non_const = typeid_cast(&src_column)) { - result_data[i] = (source_data[i] >> time_shift) + snowflake_epoch; + const auto & src_data = src_column_non_const->getData(); + for (size_t i = 0; i < input_rows_count; ++i) + res_data[i] = (src_data[i] >> time_shift) + snowflake_epoch; } + else if (const auto * src_column_const = typeid_cast(&src_column)) + { + Int64 src_val = src_column_const->getValue(); + for (size_t i = 0; i < input_rows_count; ++i) + res_data[i] = (src_val >> time_shift) + snowflake_epoch; + } + else + throw Exception(ErrorCodes::ILLEGAL_TYPE_OF_ARGUMENT, "Illegal argument for function {}", name); + return res_column; } }; diff --git a/tests/queries/0_stateless/01942_snowflakeToDateTime.reference b/tests/queries/0_stateless/01942_snowflakeToDateTime.reference index bed18023f6a..83fae3ef809 100644 --- a/tests/queries/0_stateless/01942_snowflakeToDateTime.reference +++ b/tests/queries/0_stateless/01942_snowflakeToDateTime.reference @@ -1,3 +1,5 @@ const column UTC 1426860704886947840 2021-08-15 10:57:56 DateTime(\'UTC\') 2021-08-15 10:57:56.492 DateTime64(3, \'UTC\') Asia/Shanghai 1426860704886947840 2021-08-15 18:57:56 DateTime(\'Asia/Shanghai\') 2021-08-15 18:57:56.492 DateTime64(3, \'Asia/Shanghai\') +1 +1 diff --git a/tests/queries/0_stateless/01942_snowflakeToDateTime.sql b/tests/queries/0_stateless/01942_snowflakeToDateTime.sql index f6f171afabf..0092eca848c 100644 --- a/tests/queries/0_stateless/01942_snowflakeToDateTime.sql +++ b/tests/queries/0_stateless/01942_snowflakeToDateTime.sql @@ -29,4 +29,14 @@ SELECT snowflakeToDateTime(i64, tz) as dt, toTypeName(dt), snowflakeToDateTime64(i64, tz) as dt64, - toTypeName(dt64); \ No newline at end of file + toTypeName(dt64); + + +DROP TABLE IF EXISTS tab; +CREATE TABLE tab(val Int64, tz String) engine=Log; +INSERT INTO tab VALUES (42, 'Asia/Singapore'); + +SELECT 1 FROM tab WHERE snowflakeToDateTime(42::Int64, tz) != now() SETTINGS allow_nonconst_timezone_arguments = 1; +SELECT 1 FROM tab WHERE snowflakeToDateTime64(42::Int64, tz) != now() SETTINGS allow_nonconst_timezone_arguments = 1; + +DROP TABLE tab;