Added decimal128 type support with small precision (#1428)

3 files changed, 147 insertions, 3 deletions

diff --git a/ydb/library/yql/parser/pg_wrapper/arrow.cpp b/ydb/library/yql/parser/pg_wrapper/arrow.cpp
index a855125206f..d89c8b881d0 100644
--- a/ydb/library/yql/parser/pg_wrapper/arrow.cpp
+++ b/ydb/library/yql/parser/pg_wrapper/arrow.cpp
@@ -177,6 +177,38 @@ Numeric PgFloatToNumeric(double item, ui64 scale, int digits) {
     }
 }
 
+std::shared_ptr<arrow::Array> PgDecimal128ConvertNumeric(const std::shared_ptr<arrow::Array>& value, int32_t precision, int32_t scale) {
+    TArenaMemoryContext arena;
+    const auto& data = value->data();
+    size_t length = data->length;
+    arrow::BinaryBuilder builder;
+
+    auto input = data->GetValues<arrow::Decimal128>(1);
+    for (size_t i = 0; i < length; ++i) {
+        if (value->IsNull(i)) {
+            ARROW_OK(builder.AppendNull());
+            continue;
+        }
+
+        Numeric v = PgDecimal128ToNumeric(input[i].high_bits(), input[i].low_bits(), precision, scale);
+       
+        auto datum = NumericGetDatum(v);
+        auto ptr = (char*)datum;
+        auto len = GetFullVarSize((const text*)datum);
+        NUdf::ZeroMemoryContext(ptr);
+        ARROW_OK(builder.Append(ptr - sizeof(void*), len + sizeof(void*)));  
+    }
+
+    std::shared_ptr<arrow::BinaryArray> ret;
+    ARROW_OK(builder.Finish(&ret));
+    return ret;
+}
+
+Numeric PgDecimal128ToNumeric(int64_t high_bits, uint64_t low_bits, int32_t precision, int32_t scale) {
+    Numeric res = int64_div_fast_to_numeric(low_bits, scale);
+    return res;
+}
+
 TColumnConverter BuildPgNumericColumnConverter(const std::shared_ptr<arrow::DataType>& originalType) {
     switch (originalType->id()) {
     case arrow::Type::INT16:
@@ -199,6 +231,14 @@ TColumnConverter BuildPgNumericColumnConverter(const std::shared_ptr<arrow::Data
         return [](const std::shared_ptr<arrow::Array>& value) {
             return PgConvertNumeric<double>(value);
         };
+    case arrow::Type::DECIMAL128: {
+        auto decimal128Ptr = std::static_pointer_cast<arrow::Decimal128Type>(originalType);
+        int32_t precision = decimal128Ptr->precision();
+        int32_t scale     = decimal128Ptr->scale();
+        return [precision, scale](const std::shared_ptr<arrow::Array>& value) {
+            return PgDecimal128ConvertNumeric(value, precision, scale);
+        };
+    }
     default:
         return {};
     }
diff --git a/ydb/library/yql/parser/pg_wrapper/arrow_impl.h b/ydb/library/yql/parser/pg_wrapper/arrow_impl.h
index 391e9c8792c..a9ea1e9c7d1 100644
--- a/ydb/library/yql/parser/pg_wrapper/arrow_impl.h
+++ b/ydb/library/yql/parser/pg_wrapper/arrow_impl.h
@@ -11,6 +11,7 @@ extern "C" {
 namespace NYql {
 
 Numeric PgFloatToNumeric(double item, ui64 scale, int digits);
+Numeric PgDecimal128ToNumeric(int64_t high_bits, uint64_t low_bits, int32_t precision, int32_t scale);
 TColumnConverter BuildPgColumnConverter(const std::shared_ptr<arrow::DataType>& originalType, NKikimr::NMiniKQL::TPgType* targetType);
 
 template<typename T>
@@ -27,9 +28,9 @@ std::shared_ptr<arrow::Array> PgConvertNumeric(const std::shared_ptr<arrow::Arra
         }
         T item = input[i];
         Numeric v;
-        if constexpr(std::is_same_v<T,double>) {
+        if constexpr(std::is_same_v<T, double>) {
             v = PgFloatToNumeric(item, 1000000000000LL, 12);
-        } else if constexpr(std::is_same_v<T,float>) {
+        } else if constexpr(std::is_same_v<T, float>) {
             v = PgFloatToNumeric(item, 1000000LL, 6);
         } else {
             v = int64_to_numeric(item);
@@ -46,5 +47,8 @@ std::shared_ptr<arrow::Array> PgConvertNumeric(const std::shared_ptr<arrow::Arra
     return ret;
 }
 
+
+std::shared_ptr<arrow::Array> PgDecimal128ConvertNumeric(const std::shared_ptr<arrow::Array>& value, int32_t precision, int32_t scale);
+
 }
 
diff --git a/ydb/library/yql/parser/pg_wrapper/ut/arrow_ut.cpp b/ydb/library/yql/parser/pg_wrapper/ut/arrow_ut.cpp
index eb8d8edb6e6..e8a7466822c 100644
--- a/ydb/library/yql/parser/pg_wrapper/ut/arrow_ut.cpp
+++ b/ydb/library/yql/parser/pg_wrapper/ut/arrow_ut.cpp
@@ -78,13 +78,113 @@ Y_UNIT_TEST(PgConvertNumericDouble) {
     auto result = PgConvertNumeric<double>(array);
     
     const char* expected[] = {
-        "1.1", "31.37", nullptr, "-1.337", "0"
+        "1.1", "31.37", nullptr, "-1.337", "0", "1.234111"
     };
 
     NYql::NUdf::TStringBlockReader<arrow::BinaryType, true> reader;
     checkResult<false>(expected, result, &reader, numeric_out);
 }
 
+Y_UNIT_TEST(PgConvertNumericDecimal128Scale1) {
+    TArenaMemoryContext arena;
+
+    int32_t precision = 6;
+    int32_t scale = 1;
+    std::shared_ptr<arrow::DataType> type(new arrow::Decimal128Type(precision, scale));
+    arrow::Decimal128Builder builder(type);
+
+    const char* expected[] = {
+        "12345.0", "-12345.0", nullptr
+    };
+
+    ARROW_OK(builder.Append(arrow::Decimal128::FromString("12345.0").ValueOrDie()));
+    ARROW_OK(builder.Append(arrow::Decimal128::FromString("-12345.0").ValueOrDie()));
+    ARROW_OK(builder.AppendNull());
+
+    std::shared_ptr<arrow::Array> array;
+    ARROW_OK(builder.Finish(&array));
+
+    auto result = PgDecimal128ConvertNumeric(array, precision, scale);
+
+    NYql::NUdf::TStringBlockReader<arrow::BinaryType, true> reader;
+    checkResult<false>(expected, result, &reader, numeric_out);
+}
+
+Y_UNIT_TEST(PgConvertNumericDecimal128Scale2) {
+    TArenaMemoryContext arena;
+
+    int32_t precision = 5;
+    int32_t scale = 2;
+    std::shared_ptr<arrow::DataType> type(new arrow::Decimal128Type(precision, scale));
+    arrow::Decimal128Builder builder(type);
+
+    const char* expected[] = {
+        "123.45", "-123.45", nullptr
+    };
+
+    ARROW_OK(builder.Append(arrow::Decimal128::FromString("123.45").ValueOrDie()));
+    ARROW_OK(builder.Append(arrow::Decimal128::FromString("-123.45").ValueOrDie()));
+    ARROW_OK(builder.AppendNull());
+
+    std::shared_ptr<arrow::Array> array;
+    ARROW_OK(builder.Finish(&array));
+
+    auto result = PgDecimal128ConvertNumeric(array, precision, scale);
+
+    NYql::NUdf::TStringBlockReader<arrow::BinaryType, true> reader;
+    checkResult<false>(expected, result, &reader, numeric_out);
+}
+
+Y_UNIT_TEST(PgConvertNumericDecimal128Scale3) {
+    TArenaMemoryContext arena;
+
+    int32_t precision = 3;
+    int32_t scale = 3;
+    std::shared_ptr<arrow::DataType> type(new arrow::Decimal128Type(precision, scale));
+    arrow::Decimal128Builder builder(type);
+
+    const char* expected[] = {
+        "0.123", "-0.123", nullptr
+    };
+
+    ARROW_OK(builder.Append(arrow::Decimal128::FromString("0.123").ValueOrDie()));
+    ARROW_OK(builder.Append(arrow::Decimal128::FromString("-0.123").ValueOrDie()));
+    ARROW_OK(builder.AppendNull());
+
+    std::shared_ptr<arrow::Array> array;
+    ARROW_OK(builder.Finish(&array));
+
+    auto result = PgDecimal128ConvertNumeric(array, precision, scale);
+
+    NYql::NUdf::TStringBlockReader<arrow::BinaryType, true> reader;
+    checkResult<false>(expected, result, &reader, numeric_out);
+}
+
+Y_UNIT_TEST(PgConvertNumericDecimal128Scale5) {
+    TArenaMemoryContext arena;
+
+    int32_t precision = 7;
+    int32_t scale = 5;
+    std::shared_ptr<arrow::DataType> type(new arrow::Decimal128Type(precision, scale));
+    arrow::Decimal128Builder builder(type);
+
+    const char* expected[] = {
+        "12.34567", "-12.34567", nullptr
+    };
+
+    ARROW_OK(builder.Append(arrow::Decimal128::FromReal(12.34567, precision, scale).ValueOrDie()));
+    ARROW_OK(builder.Append(arrow::Decimal128::FromReal(-12.34567, precision, scale).ValueOrDie()));
+    ARROW_OK(builder.AppendNull());
+
+    std::shared_ptr<arrow::Array> array;
+    ARROW_OK(builder.Finish(&array));
+
+    auto result = PgDecimal128ConvertNumeric(array, precision, scale);
+
+    NYql::NUdf::TStringBlockReader<arrow::BinaryType, true> reader;
+    checkResult<false>(expected, result, &reader, numeric_out);
+}
+
 Y_UNIT_TEST(PgConvertNumericInt) {
     TArenaMemoryContext arena;