Merge pull request #50893 from rschu1ze/snowflake-crash

Fix LOGICAL_ERROR in snowflakeToDateTime*()
This commit is contained in:
Robert Schulze 2023-06-14 09:42:21 +02:00 committed by GitHub
commit edd29492c3
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
3 changed files with 55 additions and 27 deletions

View File

@ -5,6 +5,8 @@
#include <Functions/FunctionHelpers.h>
#include <DataTypes/DataTypeDateTime64.h>
#include <DataTypes/DataTypesNumber.h>
#include <Columns/ColumnConst.h>
#include <Columns/ColumnsDateTime.h>
#include <Columns/ColumnsNumber.h>
#include <Interpreters/Context.h>
@ -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<const ColumnUInt32 &>(col).getData();
const auto & src_data = typeid_cast<const ColumnUInt32 &>(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<const ColumnInt64 &>(col).getData();
for (size_t i = 0; i < input_rows_count; ++i)
if (const auto * src_column_non_const = typeid_cast<const ColumnInt64 *>(&src_column))
{
result_data[i] = static_cast<UInt32>(
((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<UInt32>(
((src_data[i] >> time_shift) + snowflake_epoch) / 1000);
}
else if (const auto * src_column_const = typeid_cast<const ColumnConst *>(&src_column))
{
Int64 src_val = src_column_const->getValue<Int64>();
for (size_t i = 0; i < input_rows_count; ++i)
res_data[i] = static_cast<UInt32>(
((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<const ColumnDecimal<DateTime64> &>(col).getData();
const auto & src_data = typeid_cast<const ColumnDecimal<DateTime64> &>(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<DateTime64>::create(input_rows_count, 3);
auto & result_data = res_column->getData();
auto & res_data = res_column->getData();
const auto & source_data = typeid_cast<const ColumnInt64 &>(col).getData();
for (size_t i = 0; i < input_rows_count; ++i)
if (const auto * src_column_non_const = typeid_cast<const ColumnInt64 *>(&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<const ColumnConst *>(&src_column))
{
Int64 src_val = src_column_const->getValue<Int64>();
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;
}
};

View File

@ -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

View File

@ -29,4 +29,14 @@ SELECT
snowflakeToDateTime(i64, tz) as dt,
toTypeName(dt),
snowflakeToDateTime64(i64, tz) as dt64,
toTypeName(dt64);
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;