Merge afaa987 into efea0b8

ydb-platform · Apr 2, 2024 · 055a18b · 055a18b
2 parents efea0b8 + afaa987
commit 055a18b
Show file tree

Hide file tree

Showing 12 changed files with 554 additions and 98 deletions.
diff --git a/ydb/library/yql/minikql/computation/mkql_block_reader.cpp b/ydb/library/yql/minikql/computation/mkql_block_reader.cpp
@@ -195,6 +195,37 @@ class TTupleBlockItemConverter : public IBlockItemConverter {
     mutable TVector<TBlockItem> Items;
 };
 
+template <typename TTzDate, bool Nullable>
+class TTzDateBlockItemConverter : public IBlockItemConverter {
+public:
+    using TLayout = NYql::NUdf::TDataType<TTzDate>::TLayout;
+
+    NUdf::TUnboxedValuePod MakeValue(TBlockItem item, const THolderFactory& holderFactory) const final {
+        Y_UNUSED(holderFactory);
+        if constexpr (Nullable) {
+            if (!item) {
+                return {};
+            }
+        }
+
+        NUdf::TUnboxedValuePod value {item.Get<TLayout>()};
+        value.SetTimezoneId(item.GetTimezoneId());
+        return value;
+    }
+
+    TBlockItem MakeItem(const NUdf::TUnboxedValuePod& value) const final {
+        if constexpr (Nullable) {
+            if (!value) {
+                return {};
+            }
+        }
+
+        TBlockItem item {value.Get<TLayout>()};
+        item.SetTimezoneId(value.GetTimezoneId());
+        return item;
+    }
+};
+
 class TExternalOptionalBlockItemConverter : public IBlockItemConverter {
 public:
     TExternalOptionalBlockItemConverter(std::unique_ptr<IBlockItemConverter>&& inner)
@@ -229,6 +260,8 @@ struct TConverterTraits {
     template <typename TStringType, bool Nullable, NUdf::EDataSlot TOriginal = NUdf::EDataSlot::String, NUdf::EPgStringType PgString = NUdf::EPgStringType::None>
     using TStrings = TStringBlockItemConverter<TStringType, Nullable, PgString>;
     using TExtOptional = TExternalOptionalBlockItemConverter;
+    template<typename TTzDate, bool Nullable>
+    using TTzDateConverter = TTzDateBlockItemConverter<TTzDate, Nullable>;
 
     static std::unique_ptr<TResult> MakePg(const NUdf::TPgTypeDescription& desc, const NUdf::IPgBuilder* pgBuilder) {
         if (desc.PassByValue) {
@@ -258,6 +291,15 @@ struct TConverterTraits {
             return std::make_unique<TResourceBlockItemConverter<false>>();
         }
     }
+
+    template<typename TTzDate>
+    static std::unique_ptr<TResult> MakeTzDate(bool isOptional) {
+        if (isOptional) {
+            return std::make_unique<TTzDateConverter<TTzDate, true>>();
+        } else {
+            return std::make_unique<TTzDateConverter<TTzDate, false>>();
+        }
+    }
 };
 
 } // namespace

diff --git a/ydb/library/yql/minikql/computation/mkql_block_transport.cpp b/ydb/library/yql/minikql/computation/mkql_block_transport.cpp
@@ -516,6 +516,12 @@ struct TSerializerTraits {
         Y_UNUSED(isOptional);
         ythrow yexception() << "Serializer not implemented for block resources";
     }
+
+    template<typename TTzDate>
+    static std::unique_ptr<TResult> MakeTzDate(bool isOptional) {
+        Y_UNUSED(isOptional);
+        ythrow yexception() << "Serializer not implemented for block resources";
+    }
 };
 
 struct TDeserializerTraits {
@@ -540,6 +546,12 @@ struct TDeserializerTraits {
         Y_UNUSED(isOptional);
         ythrow yexception() << "Deserializer not implemented for block resources";
     }
+
+    template<typename TTzDate>
+    static std::unique_ptr<TResult> MakeTzDate(bool isOptional) {
+        Y_UNUSED(isOptional);
+        ythrow yexception() << "Deserializer not implemented for block resources";
+    }
 };
 
 } // namespace

diff --git a/ydb/library/yql/minikql/datetime/datetime.h b/ydb/library/yql/minikql/datetime/datetime.h
@@ -21,6 +21,7 @@ struct TTMStorage {
     unsigned int Second : 6;
     unsigned int Microsecond : 20;
     unsigned int TimezoneId : 16;
+    ui8 Reserved[2];
 
     TTMStorage() {
         Zero(*this);

diff --git a/ydb/library/yql/minikql/mkql_type_builder.cpp b/ydb/library/yql/minikql/mkql_type_builder.cpp
@@ -1451,6 +1451,33 @@ bool ConvertArrowType(NUdf::EDataSlot slot, std::shared_ptr<arrow::DataType>& ty
     case NUdf::EDataSlot::Json:
         type = arrow::utf8();
         return true;
+    case NUdf::EDataSlot::TzDate: {
+        auto&& [dateType, timezoneType] = MakeTzDateArrowFieldTypes<NUdf::EDataSlot::TzDate>();
+        std::vector<std::shared_ptr<arrow::Field>> fields {
+            std::make_shared<arrow::Field>("date", std::move(dateType)),
+            std::make_shared<arrow::Field>("timezoneId", std::move(timezoneType)),
+        };
+        type = std::make_shared<arrow::StructType>(fields);
+        return true;
+    }
+    case NUdf::EDataSlot::TzDatetime: {
+        auto&& [dateType, timezoneType] = MakeTzDateArrowFieldTypes<NUdf::EDataSlot::TzDate>();
+        std::vector<std::shared_ptr<arrow::Field>> fields {
+            std::make_shared<arrow::Field>("datetime", std::move(dateType)),
+            std::make_shared<arrow::Field>("timezoneId", std::move(timezoneType)),
+        };
+        type = std::make_shared<arrow::StructType>(fields);
+        return true;
+    }
+    case NUdf::EDataSlot::TzTimestamp: {
+        auto&& [dateType, timezoneType] = MakeTzDateArrowFieldTypes<NUdf::EDataSlot::TzDate>();
+        std::vector<std::shared_ptr<arrow::Field>> fields {
+            std::make_shared<arrow::Field>("timestamp", std::move(dateType)),
+            std::make_shared<arrow::Field>("timezoneId", std::move(timezoneType)),
+        };
+        type = std::make_shared<arrow::StructType>(fields);
+        return true;
+    }
     default:
         return false;
     }
@@ -2402,6 +2429,8 @@ struct TComparatorTraits {
     template <typename TStringType, bool Nullable, NUdf::EDataSlot TOriginal = NUdf::EDataSlot::String>
     using TStrings = NUdf::TStringBlockItemComparator<TStringType, Nullable>;
     using TExtOptional = NUdf::TExternalOptionalBlockItemComparator;
+    template <typename T, bool Nullable>
+    using TTzDateComparator = NUdf::TTzDateBlockItemComparator<T, Nullable>;
 
     static std::unique_ptr<TResult> MakePg(const NUdf::TPgTypeDescription& desc, const NUdf::IPgBuilder* pgBuilder) {
         Y_UNUSED(pgBuilder);
@@ -2412,6 +2441,15 @@ struct TComparatorTraits {
         Y_UNUSED(isOptional);
         ythrow yexception() << "Comparator not implemented for block resources: ";
     }
+
+    template<typename TTzDate>
+    static std::unique_ptr<TResult> MakeTzDate(bool isOptional) {
+        if (isOptional) {
+            return std::make_unique<TTzDateComparator<TTzDate, true>>();
+        } else {
+            return std::make_unique<TTzDateComparator<TTzDate, false>>();
+        }
+    }
 };
 
 struct THasherTraits {
@@ -2423,6 +2461,8 @@ struct THasherTraits {
     template <typename TStringType, bool Nullable, NUdf::EDataSlot TOriginal = NUdf::EDataSlot::String>
     using TStrings = NUdf::TStringBlockItemHasher<TStringType, Nullable>;
     using TExtOptional = NUdf::TExternalOptionalBlockItemHasher;
+    template <typename T, bool Nullable>
+    using TTzDateHasher = NYql::NUdf::TTzDateBlockItemHasher<T, Nullable>;
 
     static std::unique_ptr<TResult> MakePg(const NUdf::TPgTypeDescription& desc, const NUdf::IPgBuilder* pgBuilder) {
         Y_UNUSED(pgBuilder);
@@ -2433,6 +2473,15 @@ struct THasherTraits {
         Y_UNUSED(isOptional);
         ythrow yexception() << "Hasher not implemented for block resources";
     }
+
+    template<typename TTzDate>
+    static std::unique_ptr<TResult> MakeTzDate(bool isOptional) {
+        if (isOptional) {
+            return std::make_unique<TTzDateHasher<TTzDate, true>>();
+        } else {
+            return std::make_unique<TTzDateHasher<TTzDate, false>>();
+        }
+    }
 };
 
 NUdf::IBlockItemComparator::TPtr TBlockTypeHelper::MakeComparator(NUdf::TType* type) const {

diff --git a/ydb/library/yql/minikql/mkql_type_builder.h b/ydb/library/yql/minikql/mkql_type_builder.h
@@ -32,6 +32,22 @@ inline size_t CalcBlockLen(size_t maxBlockItemSize) {
 bool ConvertArrowType(TType* itemType, std::shared_ptr<arrow::DataType>& type);
 bool ConvertArrowType(NUdf::EDataSlot slot, std::shared_ptr<arrow::DataType>& type);
 
+template<NUdf::EDataSlot slot>
+std::pair<std::shared_ptr<arrow::DataType>, std::shared_ptr<arrow::DataType>> MakeTzDateArrowFieldTypes() {
+    const auto make_arrow_date_type = [] () {
+        if constexpr (slot == NUdf::EDataSlot::TzDate) {
+            return arrow::uint16();
+        }
+        if constexpr (slot == NUdf::EDataSlot::TzDatetime) {
+            return arrow::uint32();
+        }
+        return arrow::uint64();
+    };
+
+    const auto timezoneType = arrow::uint16();
+    return { make_arrow_date_type(), timezoneType };
+}
+
 class TArrowType : public NUdf::IArrowType {
 public:
     TArrowType(const std::shared_ptr<arrow::DataType>& type)

diff --git a/ydb/library/yql/providers/yt/comp_nodes/dq/arrow_converter.cpp b/ydb/library/yql/providers/yt/comp_nodes/dq/arrow_converter.cpp
@@ -438,6 +438,12 @@ struct TYsonBlockReaderTraits {
         Y_UNUSED(isOptional);
         ythrow yexception() << "Yson reader not implemented for block resources";
     }   
+
+    template<typename TTzDate>
+    static std::unique_ptr<TResult> MakeTzDate(bool isOptional) {
+        Y_UNUSED(isOptional);
+        ythrow yexception() << "Yson reader not implemented for block tz dates";
+    }   
 };
 
 template<bool IsDictionary>