YT-15805: Support tz types in arrow writer.

* Changelog entry
Type: feature
Component: proxy

Support timezone types in the arrow writer.
commit_hash:727b817a870c8a6b6256b3f6f31a6e29b7614e60

8 files changed, 457 insertions, 96 deletions

diff --git a/yt/yt/client/formats/config.cpp b/yt/yt/client/formats/config.cpp
index eff67b828da..89f0a2d27fa 100644
--- a/yt/yt/client/formats/config.cpp
+++ b/yt/yt/client/formats/config.cpp
@@ -360,4 +360,12 @@ void TYamlFormatConfig::Register(TRegistrar registrar)
 
 ////////////////////////////////////////////////////////////////////////////////
 
+void TArrowFormatConfig::Register(TRegistrar registrar)
+{
+    registrar.Parameter("enable_tz_index", &TThis::EnableTzIndex)
+        .Default(false);
+}
+
+////////////////////////////////////////////////////////////////////////////////
+
 } // namespace NYT::NFormats
diff --git a/yt/yt/client/formats/config.h b/yt/yt/client/formats/config.h
index 69f5e2fd90d..2016db008c3 100644
--- a/yt/yt/client/formats/config.h
+++ b/yt/yt/client/formats/config.h
@@ -419,4 +419,19 @@ DEFINE_REFCOUNTED_TYPE(TYamlFormatConfig)
 
 ////////////////////////////////////////////////////////////////////////////////
 
+struct TArrowFormatConfig
+    : public NYTree::TYsonStruct
+{
+    //! Return the timezone as index.
+    bool EnableTzIndex;
+
+    REGISTER_YSON_STRUCT(TArrowFormatConfig);
+
+    static void Register(TRegistrar registrar);
+};
+
+DEFINE_REFCOUNTED_TYPE(TArrowFormatConfig)
+
+////////////////////////////////////////////////////////////////////////////////
+
 } // namespace NYT::NFormats
diff --git a/yt/yt/client/formats/public.h b/yt/yt/client/formats/public.h
index bf531859ad4..baf2ae038b2 100644
--- a/yt/yt/client/formats/public.h
+++ b/yt/yt/client/formats/public.h
@@ -78,6 +78,7 @@ DECLARE_REFCOUNTED_STRUCT(TProtobufFormatConfig)
 DECLARE_REFCOUNTED_STRUCT(TWebJsonFormatConfig)
 DECLARE_REFCOUNTED_STRUCT(TSkiffFormatConfig)
 DECLARE_REFCOUNTED_STRUCT(TYamlFormatConfig)
+DECLARE_REFCOUNTED_STRUCT(TArrowFormatConfig)
 
 DECLARE_REFCOUNTED_STRUCT(IYamrConsumer)
 
diff --git a/yt/yt/client/table_client/logical_type.cpp b/yt/yt/client/table_client/logical_type.cpp
index a9b525a6fb7..0d71fc16c68 100644
--- a/yt/yt/client/table_client/logical_type.cpp
+++ b/yt/yt/client/table_client/logical_type.cpp
@@ -1942,19 +1942,15 @@ bool IsComparable(const TLogicalTypePtr& type)
 
 bool IsTzType(const TLogicalTypePtr& logicalType)
 {
-    switch (logicalType->GetMetatype()) {
-        case ELogicalMetatype::Simple:
-            switch (logicalType->AsSimpleTypeRef().GetElement()) {
-                case ESimpleLogicalValueType::TzDate:
-                case ESimpleLogicalValueType::TzDatetime:
-                case ESimpleLogicalValueType::TzTimestamp:
-                case ESimpleLogicalValueType::TzDate32:
-                case ESimpleLogicalValueType::TzDatetime64:
-                case ESimpleLogicalValueType::TzTimestamp64:
-                    return true;
-                default:
-                    return false;
-            }
+    auto simpleType = CastToV1Type(logicalType).first;
+    switch (simpleType) {
+        case ESimpleLogicalValueType::TzDate:
+        case ESimpleLogicalValueType::TzDatetime:
+        case ESimpleLogicalValueType::TzTimestamp:
+        case ESimpleLogicalValueType::TzDate32:
+        case ESimpleLogicalValueType::TzDatetime64:
+        case ESimpleLogicalValueType::TzTimestamp64:
+            return true;
         default:
             return false;
     }
diff --git a/yt/yt/client/table_client/row_base.h b/yt/yt/client/table_client/row_base.h
index 77af4fc1b41..0f447b04797 100644
--- a/yt/yt/client/table_client/row_base.h
+++ b/yt/yt/client/table_client/row_base.h
@@ -107,6 +107,22 @@ inline bool IsIntegralTypeSigned(ESimpleLogicalValueType type)
     }
 }
 
+inline int IsTzTypeSigned(ESimpleLogicalValueType type)
+{
+    switch (type) {
+        case ESimpleLogicalValueType::TzDate:
+        case ESimpleLogicalValueType::TzDatetime:
+        case ESimpleLogicalValueType::TzTimestamp:
+            return false;
+        case ESimpleLogicalValueType::TzDate32:
+        case ESimpleLogicalValueType::TzDatetime64:
+        case ESimpleLogicalValueType::TzTimestamp64:
+            return true;
+        default:
+            YT_ABORT();
+    }
+}
+
 inline bool IsIntegralType(ESimpleLogicalValueType type)
 {
     switch (type) {
@@ -169,6 +185,23 @@ inline int GetIntegralTypeBitWidth(ESimpleLogicalValueType type)
     }
 }
 
+inline int GetTzTypeBitWidth(ESimpleLogicalValueType type)
+{
+    switch (type) {
+        case ESimpleLogicalValueType::TzDate:
+            return 16;
+        case ESimpleLogicalValueType::TzDatetime:
+        case ESimpleLogicalValueType::TzDate32:
+            return 32;
+        case ESimpleLogicalValueType::TzTimestamp:
+        case ESimpleLogicalValueType::TzDatetime64:
+        case ESimpleLogicalValueType::TzTimestamp64:
+            return 64;
+        default:
+            YT_ABORT();
+    }
+}
+
 inline int GetIntegralTypeByteSize(ESimpleLogicalValueType type)
 {
     switch (type) {
diff --git a/yt/yt/library/formats/arrow_writer.cpp b/yt/yt/library/formats/arrow_writer.cpp
index 6a268b6a0cb..b186d1580b4 100644
--- a/yt/yt/library/formats/arrow_writer.cpp
+++ b/yt/yt/library/formats/arrow_writer.cpp
@@ -13,6 +13,8 @@
 
 #include <yt/yt/library/column_converters/column_converter.h>
 
+#include <yt/yt/library/tz_types/tz_types.h>
+
 #include <yt/yt/core/concurrency/async_stream.h>
 #include <yt/yt/core/concurrency/public.h>
 
@@ -31,6 +33,9 @@ namespace NYT::NFormats {
 using namespace NColumnConverters;
 using namespace NComplexTypes;
 using namespace NTableClient;
+using namespace NTzTypes;
+
+using namespace org::apache::arrow;
 
 constinit const auto Logger = FormatsLogger;
 
@@ -57,18 +62,78 @@ flatbuffers::Offset<flatbuffers::String> SerializeString(
     return flatbufBuilder->CreateString(str.data(), str.length());
 }
 
-std::tuple<org::apache::arrow::flatbuf::Type, flatbuffers::Offset<void>> SerializeColumnType(
+std::tuple<flatbuf::Type, flatbuffers::Offset<void>, std::vector<flatbuffers::Offset<flatbuf::Field>>> SerializeTzType(
     flatbuffers::FlatBufferBuilder* flatbufBuilder,
-    const TColumnSchema& schema)
+    ESimpleLogicalValueType type,
+    const TArrowFormatConfigPtr& arrowConfig)
+{
+    std::vector<flatbuffers::Offset<flatbuf::Field>> childrenOffset;
+
+    // Make timestamp field.
+
+    auto timestampOffset = flatbuf::CreateInt(
+        *flatbufBuilder,
+        GetTzTypeBitWidth(type),
+        IsTzTypeSigned(type)).Union();
+
+    auto timestampField = flatbuf::CreateField(
+        *flatbufBuilder,
+        SerializeString(flatbufBuilder, "Timestamp"),
+        /*nullable*/ false,
+        flatbuf::Type::Int,
+        timestampOffset);
+
+    childrenOffset.push_back(timestampField);
+
+    if (arrowConfig->EnableTzIndex) {
+        // Make tz index field.
+        auto tzIndexOffset = flatbuf::CreateInt(
+            *flatbufBuilder,
+            /*bitWidth*/ 16,
+            /*is_signed*/ false).Union();
+
+        auto tzIndexField = flatbuf::CreateField(
+            *flatbufBuilder,
+            SerializeString(flatbufBuilder, "TzIndex"),
+            /*nullable*/ false,
+            flatbuf::Type::Int,
+            tzIndexOffset);
+
+        childrenOffset.push_back(std::move(tzIndexField));
+    } else {
+        // Make tz name field.
+        auto tzNameOffset = flatbuf::CreateBinary(*flatbufBuilder).Union();
+
+        auto tzNameField = flatbuf::CreateField(
+            *flatbufBuilder,
+            SerializeString(flatbufBuilder, "TzName"),
+            /*nullable*/ false,
+            flatbuf::Type::Binary,
+            tzNameOffset);
+
+        childrenOffset.push_back(std::move(tzNameField));
+    }
+
+    return std::tuple(
+        flatbuf::Type::Struct_,
+        flatbuf::CreateStruct_(*flatbufBuilder).Union(),
+        std::move(childrenOffset));
+}
+
+std::tuple<flatbuf::Type, flatbuffers::Offset<void>, std::vector<flatbuffers::Offset<flatbuf::Field>>> SerializeColumnType(
+    flatbuffers::FlatBufferBuilder* flatbufBuilder,
+    const TColumnSchema& schema,
+    const TArrowFormatConfigPtr& arrowConfig)
 {
     auto simpleType = CastToV1Type(schema.LogicalType()).first;
     switch (simpleType) {
         case ESimpleLogicalValueType::Null:
         case ESimpleLogicalValueType::Void:
             return std::tuple(
-                org::apache::arrow::flatbuf::Type::Null,
-                org::apache::arrow::flatbuf::CreateNull(*flatbufBuilder)
-                    .Union());
+                flatbuf::Type::Null,
+                flatbuf::CreateNull(*flatbufBuilder)
+                    .Union(),
+                std::vector<flatbuffers::Offset<flatbuf::Field>>());
 
         case ESimpleLogicalValueType::Int64:
         case ESimpleLogicalValueType::Uint64:
@@ -79,81 +144,99 @@ std::tuple<org::apache::arrow::flatbuf::Type, flatbuffers::Offset<void>> Seriali
         case ESimpleLogicalValueType::Int32:
         case ESimpleLogicalValueType::Uint32:
             return std::tuple(
-                org::apache::arrow::flatbuf::Type::Int,
-                org::apache::arrow::flatbuf::CreateInt(
+                flatbuf::Type::Int,
+                flatbuf::CreateInt(
                     *flatbufBuilder,
                     GetIntegralTypeBitWidth(simpleType),
                     IsIntegralTypeSigned(simpleType))
-                    .Union());
+                    .Union(),
+                std::vector<flatbuffers::Offset<flatbuf::Field>>());
 
         case ESimpleLogicalValueType::Interval:
             return std::tuple(
-                org::apache::arrow::flatbuf::Type::Int,
-                org::apache::arrow::flatbuf::CreateInt(
+                flatbuf::Type::Int,
+                flatbuf::CreateInt(
                     *flatbufBuilder,
-                    64,
-                    true)
-                    .Union());
+                    /*bitWidth*/ 64,
+                    /*is_signed*/ true)
+                    .Union(),
+                std::vector<flatbuffers::Offset<flatbuf::Field>>());
 
         case ESimpleLogicalValueType::Date:
             return std::tuple(
-                org::apache::arrow::flatbuf::Type::Date,
-                org::apache::arrow::flatbuf::CreateDate(
+                flatbuf::Type::Date,
+                flatbuf::CreateDate(
                     *flatbufBuilder,
-                    org::apache::arrow::flatbuf::DateUnit::DAY)
-                    .Union());
+                    flatbuf::DateUnit::DAY)
+                    .Union(),
+                std::vector<flatbuffers::Offset<flatbuf::Field>>());
 
         case ESimpleLogicalValueType::Datetime:
             return std::tuple(
-                org::apache::arrow::flatbuf::Type::Date,
-                org::apache::arrow::flatbuf::CreateDate(
+                flatbuf::Type::Date,
+                flatbuf::CreateDate(
                     *flatbufBuilder,
-                    org::apache::arrow::flatbuf::DateUnit::MILLISECOND)
-                    .Union());
+                    flatbuf::DateUnit::MILLISECOND)
+                    .Union(),
+                std::vector<flatbuffers::Offset<flatbuf::Field>>());
 
         case ESimpleLogicalValueType::Timestamp:
             return std::tuple(
-                org::apache::arrow::flatbuf::Type::Timestamp,
-                org::apache::arrow::flatbuf::CreateTimestamp(
+                flatbuf::Type::Timestamp,
+                flatbuf::CreateTimestamp(
                     *flatbufBuilder,
-                    org::apache::arrow::flatbuf::TimeUnit::MICROSECOND)
-                    .Union());
+                    flatbuf::TimeUnit::MICROSECOND)
+                    .Union(),
+                std::vector<flatbuffers::Offset<flatbuf::Field>>());
 
         case ESimpleLogicalValueType::Double:
             return std::tuple(
-                org::apache::arrow::flatbuf::Type::FloatingPoint,
-                org::apache::arrow::flatbuf::CreateFloatingPoint(
+                flatbuf::Type::FloatingPoint,
+                flatbuf::CreateFloatingPoint(
                     *flatbufBuilder,
-                    org::apache::arrow::flatbuf::Precision::DOUBLE)
-                    .Union());
+                    flatbuf::Precision::DOUBLE)
+                    .Union(),
+                std::vector<flatbuffers::Offset<flatbuf::Field>>());
 
         case ESimpleLogicalValueType::Float:
             return std::tuple(
-                org::apache::arrow::flatbuf::Type::FloatingPoint,
-                org::apache::arrow::flatbuf::CreateFloatingPoint(
+                flatbuf::Type::FloatingPoint,
+                flatbuf::CreateFloatingPoint(
                     *flatbufBuilder,
-                    org::apache::arrow::flatbuf::Precision::SINGLE)
-                    .Union());
+                    flatbuf::Precision::SINGLE)
+                    .Union(),
+                std::vector<flatbuffers::Offset<flatbuf::Field>>());
 
         case ESimpleLogicalValueType::Boolean:
             return std::tuple(
-                org::apache::arrow::flatbuf::Type::Bool,
-                org::apache::arrow::flatbuf::CreateBool(*flatbufBuilder)
-                    .Union());
+                flatbuf::Type::Bool,
+                flatbuf::CreateBool(*flatbufBuilder)
+                    .Union(),
+                std::vector<flatbuffers::Offset<flatbuf::Field>>());
 
         case ESimpleLogicalValueType::String:
         case ESimpleLogicalValueType::Any:
             return std::tuple(
-                org::apache::arrow::flatbuf::Type::Binary,
-                org::apache::arrow::flatbuf::CreateBinary(*flatbufBuilder)
-                    .Union());
+                flatbuf::Type::Binary,
+                flatbuf::CreateBinary(*flatbufBuilder)
+                    .Union(),
+                std::vector<flatbuffers::Offset<flatbuf::Field>>());
+
+        case ESimpleLogicalValueType::TzDate:
+        case ESimpleLogicalValueType::TzDatetime:
+        case ESimpleLogicalValueType::TzTimestamp:
+        case ESimpleLogicalValueType::TzDate32:
+        case ESimpleLogicalValueType::TzDatetime64:
+        case ESimpleLogicalValueType::TzTimestamp64:
+            return SerializeTzType(flatbufBuilder, simpleType, arrowConfig);
 
         case ESimpleLogicalValueType::Utf8:
         case ESimpleLogicalValueType::Json:
             return std::tuple(
-                org::apache::arrow::flatbuf::Type::Utf8,
-                org::apache::arrow::flatbuf::CreateUtf8(*flatbufBuilder)
-                    .Union());
+                flatbuf::Type::Utf8,
+                flatbuf::CreateUtf8(*flatbufBuilder)
+                    .Union(),
+                std::vector<flatbuffers::Offset<flatbuf::Field>>());
 
         default:
             THROW_ERROR_EXCEPTION("Column %v has type %Qlv that is not currently supported by Arrow encoder",
@@ -232,6 +315,13 @@ bool IsRleButNotDictionaryEncodedStringLikeColumn(const TBatchColumn& column)
         !column.Rle->ValueColumn->Dictionary;
 }
 
+bool IsRleButNotDictionaryEncodedTzColumn(const TBatchColumn& column)
+{
+    return IsTzType(column.Type) &&
+        column.Rle &&
+        !column.Rle->ValueColumn->Dictionary;
+}
+
 bool IsRleAndDictionaryEncodedColumn(const TBatchColumn& column)
 {
     return column.Rle &&
@@ -242,7 +332,8 @@ bool IsDictionaryEncodedColumn(const TBatchColumn& column)
 {
     return column.Dictionary ||
         IsRleAndDictionaryEncodedColumn(column) ||
-        IsRleButNotDictionaryEncodedStringLikeColumn(column);
+        IsRleButNotDictionaryEncodedStringLikeColumn(column) ||
+        IsRleButNotDictionaryEncodedTzColumn(column);
 }
 
 
@@ -277,8 +368,8 @@ struct TRecordBatchSerializationContext final
     flatbuffers::FlatBufferBuilder* const FlatbufBuilder;
 
     i64 CurrentBodyOffset = 0;
-    std::vector<org::apache::arrow::flatbuf::FieldNode> FieldNodes;
-    std::vector<org::apache::arrow::flatbuf::Buffer> Buffers;
+    std::vector<flatbuf::FieldNode> FieldNodes;
+    std::vector<flatbuf::Buffer> Buffers;
     std::vector<TRecordBatchBodyPart> Parts;
 };
 
@@ -834,6 +925,169 @@ void SerializeStringLikeColumn(
         });
 }
 
+template<ESimpleLogicalValueType type>
+void SerializeTzColumnImpl(
+    const TTypedBatchColumn& typedColumn,
+    TRecordBatchSerializationContext* context,
+    const TArrowFormatConfigPtr& config)
+{
+    const auto* column = typedColumn.Column;
+    YT_VERIFY(column->Values);
+    YT_VERIFY(column->Values->BaseValue == 0);
+    YT_VERIFY(column->Values->BitWidth == 32);
+    YT_VERIFY(column->Values->ZigZagEncoded);
+    YT_VERIFY(column->Strings);
+    YT_VERIFY(column->Strings->AvgLength);
+    YT_VERIFY(!column->Rle);
+
+    auto startIndex = column->StartIndex;
+    auto endIndex = startIndex + column->ValueCount;
+    auto stringData = column->Strings->Data;
+    auto avgLength = *column->Strings->AvgLength;
+
+    auto offsets = column->GetTypedValues<ui32>();
+    auto startOffset = DecodeStringOffset(offsets, avgLength, startIndex);
+    auto endOffset = DecodeStringOffset(offsets, avgLength, endIndex);
+    auto stringsSize = endOffset - startOffset;
+    std::vector<ui32> tzOffsets(column->ValueCount + 1);
+
+    YT_LOG_DEBUG("Adding tz-type column (ColumnId: %v, StartIndex: %v, ValueCount: %v, StartOffset: %v, EndOffset: %v, StringsSize: %v)",
+        column->Id,
+        column->StartIndex,
+        column->ValueCount,
+        startOffset,
+        endOffset,
+        stringsSize);
+
+    SerializeColumnPrologue(typedColumn, context);
+
+    DecodeStringOffsets(
+        offsets,
+        avgLength,
+        startIndex,
+        endIndex,
+        TMutableRange(tzOffsets));
+
+    auto addEmptyBitmap = [=] () {
+        context->AddFieldNode(
+            column->ValueCount,
+            0);
+
+        context->AddBuffer(
+            0,
+            [=] (TMutableRef /*dstRef*/) {
+            });
+    };
+
+    addEmptyBitmap();
+
+    constexpr ESimpleLogicalValueType UnderlyingDateType = GetUnderlyingDateType<type>();
+    using TInt = TUnderlyingTimestampIntegerType<UnderlyingDateType>;
+
+    // Writing timestamp values.
+    context->AddBuffer(
+        sizeof(TInt) * column->ValueCount,
+        [=] (TMutableRef dstRef) {
+            auto currentStringData = stringData.Data();
+            auto dstValues = GetTypedValues<TInt>(dstRef);
+
+            auto* currentOutput = dstValues.Begin();
+            for (int rowOffset = 0; rowOffset < column->ValueCount; ++rowOffset) {
+                if (tzOffsets[rowOffset] != tzOffsets[rowOffset + 1]) {
+                    auto tzItem = ParseTzValue<TInt>(std::string_view(
+                        currentStringData + tzOffsets[rowOffset],
+                        currentStringData + tzOffsets[rowOffset + 1]));
+                    *currentOutput = tzItem.first;
+                }
+                ++currentOutput;
+            }
+        });
+
+    addEmptyBitmap();
+
+    if (config->EnableTzIndex) {
+        // Writing timezone indexes
+        context->AddBuffer(
+            sizeof(ui16) * column->ValueCount,
+            [=] (TMutableRef dstRef) {
+                auto currentStringData = stringData.Data();
+                auto dstValues = GetTypedValues<ui16>(dstRef);
+
+                auto* currentOutput = dstValues.Begin();
+                for (int rowOffset = 0; rowOffset < column->ValueCount; ++rowOffset) {
+                    if (tzOffsets[rowOffset] != tzOffsets[rowOffset + 1]) {
+                        auto tzItem = ParseTzValue<TInt>(std::string_view(
+                            currentStringData + tzOffsets[rowOffset],
+                            currentStringData + tzOffsets[rowOffset + 1]));
+                        *currentOutput = GetTzIndex(tzItem.second);
+                    }
+                    ++currentOutput;
+                }
+            });
+    } else {
+        // Writing timezone names.
+        TStringBuilder builder;
+        std::vector<i32> nameOffsets;
+        nameOffsets.reserve(column->ValueCount);
+        i32 tzStringsSize = 0;
+        auto currentStringData = stringData.Data();
+        for (int rowOffset = 0; rowOffset < column->ValueCount; ++rowOffset) {
+            nameOffsets.push_back(tzStringsSize);
+            if (tzOffsets[rowOffset] != tzOffsets[rowOffset + 1]) {
+                auto tzItem = ParseTzValue<TInt>(std::string_view(
+                    currentStringData + tzOffsets[rowOffset],
+                    currentStringData + tzOffsets[rowOffset + 1]));
+                tzStringsSize += tzItem.second.size();
+                builder.AppendString(tzItem.second);
+            }
+        }
+        nameOffsets.push_back(tzStringsSize);
+        auto tzStringsBuffer = builder.Flush();
+
+        context->AddBuffer(
+            sizeof(i32) * (column->ValueCount + 1),
+            [=] (TMutableRef dstRef) {
+                ::memcpy(
+                    dstRef.Begin(),
+                    nameOffsets.data(),
+                    nameOffsets.size() * sizeof(i32));
+            });
+
+        context->AddBuffer(
+            tzStringsSize,
+            [=] (TMutableRef dstRef) {
+                ::memcpy(
+                    dstRef.Begin(),
+                    tzStringsBuffer.data(),
+                    tzStringsSize);
+            });
+    }
+}
+
+void SerializeTzColumn(
+    const TTypedBatchColumn& typedColumn,
+    ESimpleLogicalValueType simpleType,
+    TRecordBatchSerializationContext* context,
+    const TArrowFormatConfigPtr& config)
+{
+    switch (simpleType) {
+#define XX(ytType)                                                                                      \
+    case ESimpleLogicalValueType::ytType: {                                                             \
+        return SerializeTzColumnImpl<ESimpleLogicalValueType::ytType>(typedColumn, context, config);    \
+    }
+    XX(TzDate)
+    XX(TzDatetime)
+    XX(TzTimestamp)
+    XX(TzDate32)
+    XX(TzDatetime64)
+    XX(TzTimestamp64)
+#undef XX
+
+    default:
+        YT_ABORT();
+    }
+}
+
 void SerializeBooleanColumn(
     const TTypedBatchColumn& typedColumn,
     TRecordBatchSerializationContext* context)
@@ -871,11 +1125,14 @@ void SerializeNullColumn(
 
 void SerializeColumn(
     const TTypedBatchColumn& typedColumn,
-    TRecordBatchSerializationContext* context)
+    TRecordBatchSerializationContext* context,
+    const TArrowFormatConfigPtr& config)
 {
     const auto* column = typedColumn.Column;
 
-    if (IsRleButNotDictionaryEncodedStringLikeColumn(*typedColumn.Column)) {
+    if (IsRleButNotDictionaryEncodedStringLikeColumn(*typedColumn.Column) ||
+        IsRleButNotDictionaryEncodedTzColumn(*typedColumn.Column))
+    {
         SerializeRleButNotDictionaryEncodedStringLikeColumn(typedColumn, context);
         return;
     }
@@ -893,6 +1150,8 @@ void SerializeColumn(
     auto simpleType = CastToV1Type(typedColumn.Type).first;
     if (IsIntegralType(simpleType)) {
         SerializeIntegerColumn(typedColumn, simpleType, context);
+    } else if (IsTzType(typedColumn.Type)) {
+        SerializeTzColumn(typedColumn, simpleType, context, config);
     } else if (simpleType == ESimpleLogicalValueType::Interval) {
         SerializeIntegerColumn(typedColumn, simpleType, context);
     }  else if (simpleType == ESimpleLogicalValueType::Date) {
@@ -923,19 +1182,20 @@ void SerializeColumn(
 auto SerializeRecordBatch(
     flatbuffers::FlatBufferBuilder* flatbufBuilder,
     int length,
-    TRange<TTypedBatchColumn> typedColumns)
+    TRange<TTypedBatchColumn> typedColumns,
+    const TArrowFormatConfigPtr& config)
 {
     auto context = New<TRecordBatchSerializationContext>(flatbufBuilder);
 
     for (const auto& typedColumn : typedColumns) {
-        SerializeColumn(typedColumn, context.Get());
+        SerializeColumn(typedColumn, context.Get(), config);
     }
 
     auto fieldNodesOffset = flatbufBuilder->CreateVectorOfStructs(context->FieldNodes);
 
     auto buffersOffset = flatbufBuilder->CreateVectorOfStructs(context->Buffers);
 
-    auto recordBatchOffset = org::apache::arrow::flatbuf::CreateRecordBatch(
+    auto recordBatchOffset = flatbuf::CreateRecordBatch(
         *flatbufBuilder,
         length,
         fieldNodesOffset,
@@ -962,6 +1222,7 @@ class TArrowWriter
 {
 public:
     TArrowWriter(
+        TArrowFormatConfigPtr config,
         TNameTablePtr nameTable,
         const std::vector<TTableSchemaPtr>& tableSchemas,
         const std::vector<std::optional<std::vector<std::string>>>& columns,
@@ -975,6 +1236,7 @@ public:
             enableContextSaving,
             std::move(controlAttributesConfig),
             keyColumnCount)
+        , ArrowConfig_(std::move(config))
     {
         YT_VERIFY(tableSchemas.size() > 0);
         YT_VERIFY(columns.size() == tableSchemas.size() || columns.size() == 0);
@@ -1141,6 +1403,8 @@ private:
     }
 
 private:
+    const TArrowFormatConfigPtr ArrowConfig_;
+
     int TableCount_ = 0;
     bool IsFirstBatch_ = true;
     i64 PrevTableIndex_ = 0;
@@ -1283,13 +1547,13 @@ private:
     }
 
     void RegisterMessage(
-        [[maybe_unused]] org::apache::arrow::flatbuf::MessageHeader type,
+        [[maybe_unused]] flatbuf::MessageHeader type,
         flatbuffers::FlatBufferBuilder&& flatbufBuilder,
         i64 bodySize = 0,
         std::function<void(TMutableRef)> bodyWriter = nullptr)
     {
         YT_LOG_DEBUG("Message registered (Type: %v, MessageSize: %v, BodySize: %v)",
-            org::apache::arrow::flatbuf::EnumNameMessageHeader(type),
+            flatbuf::EnumNameMessageHeader(type),
             flatbufBuilder.GetSize(),
             bodySize);
 
@@ -1305,7 +1569,7 @@ private:
         flatbuffers::FlatBufferBuilder flatbufBuilder;
 
         int arrowDictionaryIdCounter = 0;
-        std::vector<flatbuffers::Offset<org::apache::arrow::flatbuf::Field>> fieldOffsets;
+        std::vector<flatbuffers::Offset<flatbuf::Field>> fieldOffsets;
         for (int columnIndex = 0; columnIndex < std::ssize(TypedColumns_); columnIndex++) {
             const auto& typedColumn = TypedColumns_[columnIndex];
             auto iterSchema = ColumnSchemas_[tableIndex].find(typedColumn.Column->Id);
@@ -1313,58 +1577,59 @@ private:
             auto columnSchema = iterSchema->second;
             auto nameOffset = SerializeString(&flatbufBuilder, columnSchema.Name());
 
-            auto [typeType, typeOffset] = SerializeColumnType(&flatbufBuilder, columnSchema);
+            auto [typeType, typeOffset, childrenFields] = SerializeColumnType(&flatbufBuilder, columnSchema, ArrowConfig_);
 
-            flatbuffers::Offset<org::apache::arrow::flatbuf::DictionaryEncoding> dictionaryEncodingOffset;
-            auto index_type_offset = org::apache::arrow::flatbuf::CreateInt(flatbufBuilder, 32, false);
+            flatbuffers::Offset<flatbuf::DictionaryEncoding> dictionaryEncodingOffset;
+            auto indexTypeOffset = flatbuf::CreateInt(flatbufBuilder, /*bitWidth*/ 32,  /*is_signed*/ false);
 
             if (IsDictionaryEncodedColumn(*typedColumn.Column)) {
-                dictionaryEncodingOffset = org::apache::arrow::flatbuf::CreateDictionaryEncoding(
+                dictionaryEncodingOffset = flatbuf::CreateDictionaryEncoding(
                     flatbufBuilder,
                     arrowDictionaryIdCounter++,
-                    index_type_offset);
+                    indexTypeOffset);
             }
 
-            auto fieldOffset = org::apache::arrow::flatbuf::CreateField(
+            auto fieldOffset = flatbuf::CreateField(
                 flatbufBuilder,
                 nameOffset,
                 columnSchema.LogicalType()->IsNullable(),
                 typeType,
                 typeOffset,
-                dictionaryEncodingOffset);
+                dictionaryEncodingOffset,
+                flatbufBuilder.CreateVector(childrenFields.data(), childrenFields.size()));
 
             fieldOffsets.push_back(fieldOffset);
         }
 
         auto fieldsOffset = flatbufBuilder.CreateVector(fieldOffsets);
 
-        std::vector<flatbuffers::Offset<org::apache::arrow::flatbuf::KeyValue>> customMetadata;
+        std::vector<flatbuffers::Offset<flatbuf::KeyValue>> customMetadata;
 
         if (TableCount_ > 1) {
-            auto keyValueOffsett = org::apache::arrow::flatbuf::CreateKeyValue(
+            auto keyValueOffsett = flatbuf::CreateKeyValue(
                 flatbufBuilder,
                 flatbufBuilder.CreateString("TableId"),
                 flatbufBuilder.CreateString(std::to_string(tableIndex)));
             customMetadata.push_back(keyValueOffsett);
         }
 
-        auto schemaOffset = org::apache::arrow::flatbuf::CreateSchema(
+        auto schemaOffset = flatbuf::CreateSchema(
             flatbufBuilder,
-            org::apache::arrow::flatbuf::Endianness::Little,
+            flatbuf::Endianness::Little,
             fieldsOffset,
             flatbufBuilder.CreateVector(customMetadata));
 
-        auto messageOffset = org::apache::arrow::flatbuf::CreateMessage(
+        auto messageOffset = flatbuf::CreateMessage(
             flatbufBuilder,
-            org::apache::arrow::flatbuf::MetadataVersion::V4,
-            org::apache::arrow::flatbuf::MessageHeader::Schema,
+            flatbuf::MetadataVersion::V4,
+            flatbuf::MessageHeader::Schema,
             schemaOffset.Union(),
             0);
 
         flatbufBuilder.Finish(messageOffset);
 
         RegisterMessage(
-            org::apache::arrow::flatbuf::MessageHeader::Schema,
+            flatbuf::MessageHeader::Schema,
             std::move(flatbufBuilder));
     }
 
@@ -1404,7 +1669,9 @@ private:
                     columnIndex,
                     typedColumn.Column->Dictionary->DictionaryId,
                     typedColumn.Column->Dictionary->ValueColumn);
-            } else if (IsRleButNotDictionaryEncodedStringLikeColumn(*typedColumn.Column)) {
+            } else if (IsRleButNotDictionaryEncodedStringLikeColumn(*typedColumn.Column) ||
+                IsRleButNotDictionaryEncodedTzColumn(*typedColumn.Column))
+            {
                 YT_LOG_DEBUG("Adding dictionary batch for RLE but not dictionary-encoded string-like column (ColumnId: %v)",
                     typedColumn.Column->Id);
                 prepareDictionaryBatch(
@@ -1431,24 +1698,25 @@ private:
         auto [recordBatchOffset, bodySize, bodyWriter] = SerializeRecordBatch(
             &flatbufBuilder,
             typedColumn.Column->ValueCount,
-            TRange({typedColumn}));
+            TRange({typedColumn}),
+            ArrowConfig_);
 
-        auto dictionaryBatchOffset = org::apache::arrow::flatbuf::CreateDictionaryBatch(
+        auto dictionaryBatchOffset = flatbuf::CreateDictionaryBatch(
             flatbufBuilder,
             arrowDictionaryId,
             recordBatchOffset);
 
-        auto messageOffset = org::apache::arrow::flatbuf::CreateMessage(
+        auto messageOffset = flatbuf::CreateMessage(
             flatbufBuilder,
-            org::apache::arrow::flatbuf::MetadataVersion::V4,
-            org::apache::arrow::flatbuf::MessageHeader::DictionaryBatch,
+            flatbuf::MetadataVersion::V4,
+            flatbuf::MessageHeader::DictionaryBatch,
             dictionaryBatchOffset.Union(),
             bodySize);
 
         flatbufBuilder.Finish(messageOffset);
 
         RegisterMessage(
-            org::apache::arrow::flatbuf::MessageHeader::DictionaryBatch,
+            flatbuf::MessageHeader::DictionaryBatch,
             std::move(flatbufBuilder),
             bodySize,
             std::move(bodyWriter));
@@ -1461,19 +1729,20 @@ private:
         auto [recordBatchOffset, bodySize, bodyWriter] = SerializeRecordBatch(
             &flatbufBuilder,
             RowCount_,
-            TypedColumns_);
+            TypedColumns_,
+            ArrowConfig_);
 
-        auto messageOffset = org::apache::arrow::flatbuf::CreateMessage(
+        auto messageOffset = flatbuf::CreateMessage(
             flatbufBuilder,
-            org::apache::arrow::flatbuf::MetadataVersion::V4,
-            org::apache::arrow::flatbuf::MessageHeader::RecordBatch,
+            flatbuf::MetadataVersion::V4,
+            flatbuf::MessageHeader::RecordBatch,
             recordBatchOffset.Union(),
             bodySize);
 
         flatbufBuilder.Finish(messageOffset);
 
         RegisterMessage(
-            org::apache::arrow::flatbuf::MessageHeader::RecordBatch,
+            flatbuf::MessageHeader::RecordBatch,
             std::move(flatbufBuilder),
             bodySize,
             std::move(bodyWriter));
@@ -1535,15 +1804,17 @@ private:
 ////////////////////////////////////////////////////////////////////////////////
 
 ISchemalessFormatWriterPtr CreateWriterForArrow(
+    TArrowFormatConfigPtr config,
     TNameTablePtr nameTable,
-    const std::vector<TTableSchemaPtr>& schemas,
+    const std::vector<NTableClient::TTableSchemaPtr>& schemas,
     const std::vector<std::optional<std::vector<std::string>>>& columns,
     NConcurrency::IAsyncOutputStreamPtr output,
     bool enableContextSaving,
     TControlAttributesConfigPtr controlAttributesConfig,
     int keyColumnCount)
 {
-    auto result = New<TArrowWriter>(
+    return New<TArrowWriter>(
+        std::move(config),
         std::move(nameTable),
         schemas,
         columns,
@@ -1551,8 +1822,33 @@ ISchemalessFormatWriterPtr CreateWriterForArrow(
         enableContextSaving,
         std::move(controlAttributesConfig),
         keyColumnCount);
+}
 
-    return result;
+ISchemalessFormatWriterPtr CreateWriterForArrow(
+    const NYTree::IAttributeDictionary& attributes,
+    TNameTablePtr nameTable,
+    const std::vector<TTableSchemaPtr>& schemas,
+    const std::vector<std::optional<std::vector<std::string>>>& columns,
+    NConcurrency::IAsyncOutputStreamPtr output,
+    bool enableContextSaving,
+    TControlAttributesConfigPtr controlAttributesConfig,
+    int keyColumnCount)
+{
+    TArrowFormatConfigPtr arrowConfig;
+    try {
+        arrowConfig = ConvertTo<TArrowFormatConfigPtr>(attributes);
+    } catch (const std::exception& ex) {
+        THROW_ERROR_EXCEPTION(NFormats::EErrorCode::InvalidFormat, "Failed to parse config for arrow format") << ex;
+    }
+    return CreateWriterForArrow(
+        std::move(arrowConfig),
+        std::move(nameTable),
+        schemas,
+        columns,
+        std::move(output),
+        enableContextSaving,
+        std::move(controlAttributesConfig),
+        keyColumnCount);
 }
 
 ////////////////////////////////////////////////////////////////////////////////
diff --git a/yt/yt/library/formats/arrow_writer.h b/yt/yt/library/formats/arrow_writer.h
index 020090152cc..46240293bdd 100644
--- a/yt/yt/library/formats/arrow_writer.h
+++ b/yt/yt/library/formats/arrow_writer.h
@@ -13,6 +13,17 @@ namespace NYT::NFormats {
 ////////////////////////////////////////////////////////////////////////////////
 
 ISchemalessFormatWriterPtr CreateWriterForArrow(
+    TArrowFormatConfigPtr attributes,
+    NTableClient::TNameTablePtr nameTable,
+    const std::vector<NTableClient::TTableSchemaPtr>& schemas,
+    const std::vector<std::optional<std::vector<std::string>>>& columns,
+    NConcurrency::IAsyncOutputStreamPtr output,
+    bool enableContextSaving,
+    TControlAttributesConfigPtr controlAttributesConfig,
+    int keyColumnCount);
+
+ISchemalessFormatWriterPtr CreateWriterForArrow(
+    const NYTree::IAttributeDictionary& attributes,
     NTableClient::TNameTablePtr nameTable,
     const std::vector<NTableClient::TTableSchemaPtr>& schemas,
     const std::vector<std::optional<std::vector<std::string>>>& columns,
diff --git a/yt/yt/library/formats/format.cpp b/yt/yt/library/formats/format.cpp
index ff1791dd2e9..6b1733e6bf6 100644
--- a/yt/yt/library/formats/format.cpp
+++ b/yt/yt/library/formats/format.cpp
@@ -349,6 +349,7 @@ ISchemalessFormatWriterPtr CreateStaticTableWriterForFormat(
                 keyColumnCount);
         case EFormatType::Arrow:
             return CreateWriterForArrow(
+                format.Attributes(),
                 nameTable,
                 tableSchemas,
                 columns,