Added wide dates and decimals support for s3:parquet (#25186) (#26076)

Author: evanevanevanevannnn

ydb/library/yql/providers/s3/actors/ya.make

@@ -29,11 +29,9 @@ PEERDIR(
     ydb/core/base
     ydb/core/fq/libs/events
     ydb/library/yql/dq/actors/compute
-    yql/essentials/minikql/computation
     ydb/library/yql/providers/common/arrow
     ydb/library/yql/providers/common/arrow/interface
     ydb/library/yql/providers/common/http_gateway
-    yql/essentials/providers/common/schema/mkql
     ydb/library/yql/providers/common/token_accessor/client
     ydb/library/yql/providers/generic/pushdown
     ydb/library/yql/providers/s3/actors_factory
@@ -44,9 +42,12 @@ PEERDIR(
     ydb/library/yql/providers/s3/object_listers
     ydb/library/yql/providers/s3/proto
     ydb/library/yql/providers/s3/range_helpers
+    ydb/library/yql/udfs/common/clickhouse/client
+    yql/essentials/minikql
+    yql/essentials/minikql/computation
+    yql/essentials/providers/common/schema/mkql
     yql/essentials/public/issue
     yql/essentials/public/types
-    ydb/library/yql/udfs/common/clickhouse/client
     yql/essentials/utils
 )
 

ydb/library/yql/providers/s3/actors/yql_arrow_column_converters.cpp

@@ -12,6 +12,7 @@
 #include <yql/essentials/public/udf/arrow/block_builder.h>
 #include <yql/essentials/public/udf/arrow/block_item.h>
 #include <yql/essentials/public/udf/arrow/block_reader.h>
+#include <yql/essentials/public/udf/udf_data_type.h>
 #include <yql/essentials/utils/yql_panic.h>
 
 #include <arrow/api.h>
@@ -89,9 +90,9 @@ ui32 GetMultiplierForDatetime(arrow::TimeUnit::type unit) {
         case arrow::TimeUnit::SECOND:
             return 1;
         case arrow::TimeUnit::MILLI:
-            throw parquet::ParquetException(TStringBuilder() << "millisecond accuracy does not fit into the datetime");
+            return 1000;
         case arrow::TimeUnit::MICRO:
-            throw parquet::ParquetException(TStringBuilder() << "microsecond accuracy does not fit into the datetime");
+            return 1000000;
         case arrow::TimeUnit::NANO:
             throw parquet::ParquetException(TStringBuilder() << "nanosecond accuracy does not fit into the datetime");
     }
@@ -127,6 +128,35 @@ std::shared_ptr<arrow::Array> ArrowTypeAsYqlDatetime(const std::shared_ptr<arrow
     return builder.Build(true).make_array();
 }
 
+template <bool isOptional, typename TArrowType>
+std::shared_ptr<arrow::Array> ArrowTimestampAsYqlDatetime(const std::shared_ptr<arrow::DataType>& targetType, const std::shared_ptr<arrow::Array>& value, ui32 multiplier) {
+    ::NYql::NUdf::TFixedSizeArrayBuilder<TArrowType, isOptional> builder(NKikimr::NMiniKQL::TTypeInfoHelper(), targetType, *arrow::system_memory_pool(), value->length());
+    ::NYql::NUdf::TFixedSizeBlockReader<i64, isOptional> reader;
+    for (i64 i = 0; i < value->length(); ++i) {
+        const NUdf::TBlockItem item = reader.GetItem(*value->data(), i);
+        if constexpr (isOptional) {
+            if (!item) {
+                builder.Add(item);
+                continue;
+            }
+        } else if (!item) {
+            throw parquet::ParquetException(TStringBuilder() << "null value for datetime could not be represented in non-optional type");
+        }
+
+        const i64 baseValue = item.As<i64>();
+        if (baseValue < 0 && baseValue > static_cast<int64_t>(::NYql::NUdf::MAX_DATETIME)) {
+            throw parquet::ParquetException(TStringBuilder() << "datetime in parquet is out of range [0, " << ::NYql::NUdf::MAX_DATETIME << "]: " << baseValue);
+        }
+
+        if (baseValue % multiplier) {
+            throw parquet::ParquetException(TStringBuilder() << "datetime in parquet should have integer amount of seconds, have: " << baseValue * 1.0 / multiplier);
+        }
+        const TArrowType v = baseValue / static_cast<ui64>(multiplier);
+        builder.Add(NUdf::TBlockItem(static_cast<TArrowType>(v)));
+    }
+    return builder.Build(true).make_array();
+}
+
 template <bool isOptional>
 std::shared_ptr<arrow::Array> ArrowStringAsYqlDateTime(const std::shared_ptr<arrow::DataType>& targetType, const std::shared_ptr<arrow::Array>& value, const NDB::FormatSettings& formatSettings) {
     ::NYql::NUdf::TFixedSizeArrayBuilder<ui32, isOptional> builder(NKikimr::NMiniKQL::TTypeInfoHelper(), targetType, *arrow::system_memory_pool(), value->length());
@@ -330,8 +360,16 @@ TColumnConverter ArrowDate64AsYqlDatetime(const std::shared_ptr<arrow::DataType>
 TColumnConverter ArrowTimestampAsYqlDatetime(const std::shared_ptr<arrow::DataType>& targetType, bool isOptional, arrow::TimeUnit::type timeUnit) {
     return [targetType, isOptional, multiplier = GetMultiplierForDatetime(timeUnit)](const std::shared_ptr<arrow::Array>& value) {
         return isOptional
-                ? ArrowTypeAsYqlDatetime<true, i64>(targetType, value, multiplier)
-                : ArrowTypeAsYqlDatetime<false, i64>(targetType, value, multiplier);
+                ? ArrowTimestampAsYqlDatetime<true, ui32>(targetType, value, multiplier)
+                : ArrowTimestampAsYqlDatetime<false, ui32>(targetType, value, multiplier);
+    };
+}
+
+TColumnConverter ArrowTimestampAsYqlDatetime64(const std::shared_ptr<arrow::DataType>& targetType, bool isOptional, arrow::TimeUnit::type timeUnit) {
+    return [targetType, isOptional, multiplier = GetMultiplierForDatetime(timeUnit)](const std::shared_ptr<arrow::Array>& value) {
+        return isOptional
+                ? ArrowTimestampAsYqlDatetime<true, i64>(targetType, value, multiplier)
+                : ArrowTimestampAsYqlDatetime<false, i64>(targetType, value, multiplier);
     };
 }
 
@@ -624,6 +662,8 @@ TColumnConverter BuildCustomConverter(const std::shared_ptr<arrow::DataType>& or
             switch (slotItem) {
                 case NUdf::EDataSlot::Datetime:
                     return ArrowTimestampAsYqlDatetime(targetType, isOptional, timestampType.unit());
+                case NUdf::EDataSlot::Datetime64:
+                    return ArrowTimestampAsYqlDatetime64(targetType, isOptional, timestampType.unit());
                 case NUdf::EDataSlot::Timestamp:
                     return ArrowTimestampAsYqlTimestamp(targetType, isOptional, timestampType.unit());
                 case NUdf::EDataSlot::String:
@@ -654,7 +694,7 @@ TColumnConverter BuildCustomConverter(const std::shared_ptr<arrow::DataType>& or
                     ) {
                         return [](const std::shared_ptr<arrow::Array>& value) {
                             auto decimals = std::static_pointer_cast<arrow::Decimal128Array>(value);
-                            auto output = std::make_shared<arrow::FixedSizeBinaryArray>(arrow::fixed_size_binary(16), decimals->length(), decimals->values());
+                            auto output = std::make_shared<arrow::FixedSizeBinaryArray>(arrow::fixed_size_binary(16), decimals->length(), decimals->values(), decimals->null_bitmap(), decimals->null_count());
                             return output;
                         };
                     }
@@ -691,6 +731,39 @@ TColumnConverter YqlBlockTzDateToArrow(const std::string& columnName, const std:
     };
 }
 
+template <bool isOptional>
+TColumnConverter DecimalToArrowBaseConverter(const std::shared_ptr<arrow::DataType>& targetType) {
+    return [targetType](const std::shared_ptr<arrow::Array>& value) {
+        arrow::Decimal128Builder builder(targetType, arrow::default_memory_pool());
+        ::NYql::NUdf::TFixedSizeBlockReader<NYql::NDecimal::TInt128, isOptional> reader;
+
+        for (i64 i = 0; i < value->length(); ++i) {
+            NUdf::TBlockItem item = reader.GetItem(*value->data(), i);
+
+            if (!item) {
+                THROW_ARROW_NOT_OK(builder.AppendNull());
+                continue;
+            }
+
+            NYql::NDecimal::TInt128 val = item.GetInt128();
+            arrow::Decimal128 newValue((uint8_t*)(&val));
+            THROW_ARROW_NOT_OK(builder.Append(newValue));
+        }
+
+        std::shared_ptr<arrow::Array> array;
+        THROW_ARROW_NOT_OK(builder.Finish(&array));
+        return array;
+    };
+}
+
+TColumnConverter DecimalToArrowConverter(bool isOptional, const std::shared_ptr<arrow::DataType>& targetType) {
+    if (isOptional) {
+        return DecimalToArrowBaseConverter<true>(targetType);
+    } else {
+        return DecimalToArrowBaseConverter<false>(targetType);
+    }
+}
+
 }
 
 namespace NYql::NDq {
@@ -728,7 +801,9 @@ TColumnConverter BuildOutputColumnConverter(const std::string& columnName, NKiki
     YQL_ENSURE(ConvertArrowType(columnType, yqlArrowType), "Got unsupported yql block type: " << *columnType << " in column " << columnName);
     YQL_ENSURE(S3ConvertArrowOutputType(columnType, s3OutputType), "Got unsupported s3 output block type: " << *columnType << " in column " << columnName);
 
-    if (columnType->IsOptional()) {
+    bool isOptional = columnType->IsOptional();
+
+    if (isOptional) {
         columnType = AS_TYPE(TOptionalType, columnType)->GetItemType();
     }
     YQL_ENSURE(columnType->IsData(), "Allowed only data types for S3 output, but got: " << *columnType << " in column " << columnName);
@@ -752,13 +827,18 @@ TColumnConverter BuildOutputColumnConverter(const std::string& columnName, NKiki
         case NUdf::EDataSlot::Datetime:
         case NUdf::EDataSlot::Timestamp:
             return {};
+        case NUdf::EDataSlot::Date32:
+        case NUdf::EDataSlot::Datetime64:
+        case NUdf::EDataSlot::Timestamp64:
         case NUdf::EDataSlot::Utf8:
         case NUdf::EDataSlot::Json:
             return ArrowComputeConvertor(columnName, yqlArrowType, s3OutputType);
         case NUdf::EDataSlot::TzDate:
         case NUdf::EDataSlot::TzDatetime:
         case NUdf::EDataSlot::TzTimestamp:
             return YqlBlockTzDateToArrow(columnName, yqlArrowType);
+        case NUdf::EDataSlot::Decimal:
+            return DecimalToArrowConverter(isOptional, s3OutputType);
         default:
             YQL_ENSURE(false, "Got unsupported s3 output block type: " << *columnType << " in column " << columnName);
     }
@@ -815,7 +895,7 @@ std::shared_ptr<arrow::RecordBatch> ConvertArrowColumns(std::shared_ptr<arrow::R
 }
 
 // Type conversion same as in ClickHouseClient.SerializeFormat udf
-bool S3ConvertArrowOutputType(NUdf::EDataSlot slot, std::shared_ptr<arrow::DataType>& type) {
+bool S3ConvertArrowOutputType(NUdf::EDataSlot slot, std::shared_ptr<arrow::DataType>& type, TType* itemType) {
     switch (slot) {
         case NUdf::EDataSlot::Int8:
             type = arrow::int8();
@@ -835,11 +915,17 @@ bool S3ConvertArrowOutputType(NUdf::EDataSlot slot, std::shared_ptr<arrow::DataT
         case NUdf::EDataSlot::Int32:
             type = arrow::int32();
             return true;
+        case NUdf::EDataSlot::Date32:
+            type = arrow::date32();
+            return true;
         case NUdf::EDataSlot::Datetime:
         case NUdf::EDataSlot::TzDatetime:
         case NUdf::EDataSlot::Uint32:
             type = arrow::uint32();
             return true;
+        case NUdf::EDataSlot::Datetime64:
+            type = arrow::timestamp(arrow::TimeUnit::SECOND, "UTC");
+            return true;
         case NUdf::EDataSlot::Int64:
             type = arrow::int64();
             return true;
@@ -857,7 +943,17 @@ bool S3ConvertArrowOutputType(NUdf::EDataSlot slot, std::shared_ptr<arrow::DataT
         case NUdf::EDataSlot::Json:
             type = arrow::binary();
             return true;
+        case NUdf::EDataSlot::Decimal: {
+            if (itemType) {
+                auto [precision, scale] = static_cast<TDataDecimalType*>(itemType)->GetParams();
+                type = arrow::decimal128(precision, scale);
+            } else {
+                type = arrow::decimal128(22, 9);
+            }
+            return true;
+        }
         case NUdf::EDataSlot::Timestamp:
+        case NUdf::EDataSlot::Timestamp64:
         case NUdf::EDataSlot::TzTimestamp:
             type = arrow::timestamp(arrow::TimeUnit::MICRO, "UTC");
             return true;
@@ -880,7 +976,7 @@ bool S3ConvertArrowOutputType(TType* itemType, std::shared_ptr<arrow::DataType>&
         return false;
     }
 
-    return S3ConvertArrowOutputType(*slot, type);
+    return S3ConvertArrowOutputType(*slot, type, itemType);
 }
 
 void BuildOutputColumnConverters(const NKikimr::NMiniKQL::TStructType* outputStructType, std::vector<TColumnConverter>& columnConverters) {

ydb/library/yql/providers/s3/actors/yql_arrow_column_converters.h

@@ -35,7 +35,7 @@ std::shared_ptr<arrow::RecordBatch> ConvertArrowColumns(
     std::shared_ptr<arrow::RecordBatch> batch,
     std::vector<TColumnConverter>& columnConverters);
 
-bool S3ConvertArrowOutputType(NUdf::EDataSlot slot, std::shared_ptr<arrow::DataType>& type);
+bool S3ConvertArrowOutputType(NUdf::EDataSlot slot, std::shared_ptr<arrow::DataType>& type, NKikimr::NMiniKQL::TType* itemType = nullptr);
 bool S3ConvertArrowOutputType(NKikimr::NMiniKQL::TType* itemType, std::shared_ptr<arrow::DataType>& type);
 
 void BuildOutputColumnConverters(

ydb/library/yql/providers/s3/common/util.cpp

@@ -35,6 +35,20 @@ struct TTypeError {
     TString Error = "unsupported type";
 };
 
+std::optional<TTypeError> ValidateIoDataType(const TDataExprType* type, std::vector<EDataSlot> extraTypes = {}) {
+    const auto dataSlot = type->GetSlot();
+    if (IsDataTypeBigDate(dataSlot)) {
+        return TTypeError{type, "big dates is not supported"};
+    }
+    if (IsDataTypeNumeric(dataSlot) || IsDataTypeDateOrTzDate(dataSlot)) {
+        return std::nullopt;
+    }
+    if (IsIn({EDataSlot::Bool, EDataSlot::String, EDataSlot::Utf8, EDataSlot::Json, EDataSlot::Uuid}, dataSlot) || IsIn(extraTypes, dataSlot)) {
+        return std::nullopt;
+    }
+    return TTypeError{type};
+}
+
 std::optional<TTypeError> ValidateJsonListIoType(const TTypeAnnotationNode* type, std::function<std::optional<TTypeError>(const TTypeAnnotationNode*)> defaultHandler) {
     switch (type->GetKind()) {
         case ETypeAnnotationKind::Null:
@@ -103,6 +117,51 @@ std::optional<TTypeError> ValidateJsonListInputType(const TTypeAnnotationNode* t
     });
 }
 
+std::optional<TTypeError> ValidateParquetIoType(const TTypeAnnotationNode* type, bool underOptional = false) {
+    switch (type->GetKind()) {
+        case ETypeAnnotationKind::Data: {
+            const auto dataSlot = type->Cast<TDataExprType>()->GetSlot();
+            if (IsDataTypeNumeric(dataSlot) || IsDataTypeDateOrTzDate(dataSlot) || IsDataTypeDecimal(dataSlot) || IsDataTypeBigDate(dataSlot)) {
+                return std::nullopt;
+            }
+            if (IsIn({EDataSlot::Bool, EDataSlot::String, EDataSlot::Utf8, EDataSlot::Json, EDataSlot::Uuid}, dataSlot)) {
+                return std::nullopt;
+            }
+            return TTypeError{type};
+        }
+        case ETypeAnnotationKind::Optional: {
+            if (underOptional) {
+                return TTypeError{type, "double optional is not supported"};
+            }
+            return ValidateParquetIoType(type->Cast<TOptionalExprType>()->GetItemType(), true);
+        }
+        case ETypeAnnotationKind::List: {
+            if (underOptional) {
+                return TTypeError{type, "list under optional is not supported"};
+            }
+            return ValidateIoDataType(type->Cast<TDataExprType>());
+        }
+        case ETypeAnnotationKind::Tuple: {
+            if (underOptional) {
+                return TTypeError{type, "tuple under optional is not supported"};
+            }
+            for (const auto* item : type->Cast<TTupleExprType>()->GetItems()) {
+                if (const auto error = ValidateIoDataType(item->Cast<TDataExprType>())) {
+                    return error;
+                }
+            }
+            return std::nullopt;
+        }
+        case ETypeAnnotationKind::Pg: {
+            return std::nullopt;
+        }
+        default: {
+            break;
+        }
+    }
+    return TTypeError{type};
+}
+
 // Type compatible with Yson2.From udf
 std::optional<TTypeError> DefaultJsonListOutputTypeHandler(const TTypeAnnotationNode* type) {
     if (type->GetKind() == ETypeAnnotationKind::Variant) {
@@ -115,18 +174,8 @@ std::optional<TTypeError> ValidateJsonListOutputType(const TTypeAnnotationNode*
     return ValidateJsonListIoType(type, &DefaultJsonListOutputTypeHandler);
 }
 
-std::optional<TTypeError> ValidateIoDataType(const TDataExprType* type, std::vector<EDataSlot> extraTypes = {}) {
-    const auto dataSlot = type->GetSlot();
-    if (IsDataTypeBigDate(dataSlot)) {
-        return TTypeError{type, "big dates is not supported"};
-    }
-    if (IsDataTypeNumeric(dataSlot) || IsDataTypeDateOrTzDate(dataSlot)) {
-        return std::nullopt;
-    }
-    if (IsIn({EDataSlot::Bool, EDataSlot::String, EDataSlot::Utf8, EDataSlot::Json, EDataSlot::Uuid}, dataSlot) || IsIn(extraTypes, dataSlot)) {
-        return std::nullopt;
-    }
-    return TTypeError{type};
+std::optional<TTypeError> ValidateParquetOutputType(const TTypeAnnotationNode* type) {
+    return ValidateParquetIoType(type);
 }
 
 // Data type compatible with ClickHouseClient.ParseBlocks udf and S3 coro read actor
@@ -247,6 +296,10 @@ bool ValidateS3WriteSchema(TPositionHandle pos, std::string_view format, const T
         return ValidateIoSchema(pos, schemaStructRowType, "S3 json_list output format", ctx, &ValidateJsonListOutputType);
     }
 
+    if (format == "parquet"sv) {
+        return ValidateIoSchema(pos, schemaStructRowType, "S3 parquet output format", ctx, &ValidateParquetOutputType);
+    }
+
     return ValidateIoSchema(pos, schemaStructRowType, TStringBuilder() << "S3 " << format << " output format", ctx, [](const TTypeAnnotationNode* type) {
         return ValidateGenericIoType(type, &ValidateClickHouseUdfDataType);
     });