Implement double/int -> pg numeric conversion YQL-16767

ydb/library/yql/parser/pg_wrapper/arrow.cpp

@@ -13,6 +13,7 @@
 extern "C" {
 #include "utils/date.h"
 #include "utils/timestamp.h"
+#include "utils/numeric.h"
 }
 
 namespace NYql {
@@ -151,6 +152,88 @@ std::shared_ptr<arrow::Array> PgConvertString(const std::shared_ptr<arrow::Array
     return ret;
 }
 
+Numeric PgFloatToNumeric(double item, ui64 scale, int digits) {
+    double intPart, fracPart;
+    bool error;
+    fracPart = modf(item, &intPart);
+    i64 fracInt = round(fracPart * scale);
+
+    // scale compaction: represent 711.56000 as 711.56
+    while (digits > 0 && fracInt % 10 == 0) {
+        fracInt /= 10;
+        digits -= 1;
+    }
+
+    if (digits == 0) {
+        return int64_to_numeric(intPart);
+    } else {
+        return numeric_add_opt_error(
+            int64_to_numeric(intPart),
+            int64_div_fast_to_numeric(fracInt, digits),
+            &error);
+    }
+}
+
+template<typename T>
+std::shared_ptr<arrow::Array> PgConvertNumeric(const std::shared_ptr<arrow::Array>& value) {
+    TArenaMemoryContext arena;
+    const auto& data = value->data();
+    size_t length = data->length;
+    arrow::BinaryBuilder builder;
+    auto input = data->GetValues<T>(1);
+    for (size_t i = 0; i < length; ++i) {
+        if (value->IsNull(i)) {
+            builder.AppendNull();
+            continue;
+        }
+        T item = input[i];
+        Numeric v;
+        if constexpr(std::is_same_v<T,double>) {
+            v = PgFloatToNumeric(item, 1000000000000LL, 12);
+        } else if constexpr(std::is_same_v<T,float>) {
+            v = PgFloatToNumeric(item, 1000000LL, 6);
+        } else {
+            v = int64_to_numeric(item);
+        }
+        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;
+}
+
+TColumnConverter BuildPgNumericColumnConverter(const std::shared_ptr<arrow::DataType>& originalType) {
+    switch (originalType->id()) {
+    case arrow::Type::INT16:
+        return [](const std::shared_ptr<arrow::Array>& value) {
+            return PgConvertNumeric<i16>(value);
+        };
+    case arrow::Type::INT32:
+        return [](const std::shared_ptr<arrow::Array>& value) {
+            return PgConvertNumeric<i32>(value);
+        };
+    case arrow::Type::INT64:
+        return [](const std::shared_ptr<arrow::Array>& value) {
+            return PgConvertNumeric<i64>(value);
+        };
+    case arrow::Type::FLOAT:
+        return [](const std::shared_ptr<arrow::Array>& value) {
+            return PgConvertNumeric<float>(value);
+        };
+    case arrow::Type::DOUBLE:
+        return [](const std::shared_ptr<arrow::Array>& value) {
+            return PgConvertNumeric<double>(value);
+        };
+    default:
+        return {};
+    }
+}
+
 template <typename T, typename F>
 TColumnConverter BuildPgFixedColumnConverter(const std::shared_ptr<arrow::DataType>& originalType, const F& f) {
     auto primaryType = NKikimr::NMiniKQL::GetPrimitiveDataType<T>();
@@ -200,6 +283,9 @@ TColumnConverter BuildPgColumnConverter(const std::shared_ptr<arrow::DataType>&
     case FLOAT8OID: {
         return BuildPgFixedColumnConverter<double>(originalType, [](auto value){ return Float8GetDatum(value); });
     }
+    case NUMERICOID: {
+        return BuildPgNumericColumnConverter(originalType);
+    }
     case BYTEAOID:
     case VARCHAROID:
     case TEXTOID:

ydb/library/yql/parser/pg_wrapper/arrow.h

@@ -17,6 +17,7 @@ extern "C" {
 #include "fmgr.h"
 #include "catalog/pg_type_d.h"
 #include "catalog/pg_collation_d.h"
+#include "utils/numeric.h"
 }
 
 #include "utils.h"
@@ -1298,4 +1299,9 @@ TExecFunc FindExec(Oid oid);
 
 const NPg::TAggregateDesc& ResolveAggregation(const TString& name, NKikimr::NMiniKQL::TTupleType* tupleType, const std::vector<ui32>& argsColumns, NKikimr::NMiniKQL::TType* returnType);
 
+Numeric PgFloatToNumeric(double item, ui64 scale, int digits);
+
+template<typename T>
+std::shared_ptr<arrow::Array> PgConvertNumeric(const std::shared_ptr<arrow::Array>& value);
+
 }

ydb/library/yql/parser/pg_wrapper/ut/arrow_ut.cpp

@@ -0,0 +1,110 @@
+#include <arrow/api.h>
+#include <arrow/array.h>
+
+#include <library/cpp/testing/unittest/registar.h>
+
+#include "arrow.h"
+
+extern "C" {
+#include "utils/numeric.h"
+#include "utils/fmgrprotos.h"
+}
+
+namespace NYql {
+
+Y_UNIT_TEST_SUITE(TArrowUtilsTests) {
+
+Y_UNIT_TEST(TestPgFloatToNumeric) {
+    TArenaMemoryContext arena;
+    auto n = PgFloatToNumeric(711.56, 1000000000000LL, 12);
+    auto value = TString(DatumGetCString(DirectFunctionCall1(numeric_out, NumericGetDatum(n))));
+    UNIT_ASSERT_VALUES_EQUAL(value, "711.56");
+
+    n = PgFloatToNumeric(-711.56, 1000000000000LL, 12);
+    value = TString(DatumGetCString(DirectFunctionCall1(numeric_out, NumericGetDatum(n))));
+    UNIT_ASSERT_VALUES_EQUAL(value, "-711.56");
+
+    n = PgFloatToNumeric(711.56f, 100000LL, 5);
+    value = TString(DatumGetCString(DirectFunctionCall1(numeric_out, NumericGetDatum(n))));
+    UNIT_ASSERT_VALUES_EQUAL(value, "711.56");
+
+    n = PgFloatToNumeric(-711.56f, 100000LL, 5);
+    value = TString(DatumGetCString(DirectFunctionCall1(numeric_out, NumericGetDatum(n))));
+    UNIT_ASSERT_VALUES_EQUAL(value, "-711.56");
+}
+
+
+Y_UNIT_TEST(PgConvertNumericDouble) {
+    TArenaMemoryContext arena;
+
+    arrow::DoubleBuilder builder;
+    builder.Append(1.1);
+    builder.Append(31.37);
+    builder.AppendNull();
+    builder.Append(-1.337);
+    builder.Append(0.0);
+
+    std::shared_ptr<arrow::Array> array;
+    builder.Finish(&array);
+
+    auto result = PgConvertNumeric<double>(array);
+    const auto& data = result->data();
+
+    const char* expected[] = {
+        "1.1", "31.37", nullptr, "-1.337", "0"
+    };
+
+    NUdf::TStringBlockReader<arrow::BinaryType, true> reader;
+    for (int i = 0; i < 5; i++) {
+        auto item = reader.GetItem(*data, i);
+        if (!item) {
+            UNIT_ASSERT(expected[i] == nullptr);
+        } else {
+            const char* addr = item.AsStringRef().Data() + sizeof(void*);
+            UNIT_ASSERT_VALUES_EQUAL(
+                TString(DatumGetCString(DirectFunctionCall1(numeric_out, (Datum)addr))),
+                expected[i]
+            );
+        }
+    }
+}
+
+Y_UNIT_TEST(PgConvertNumericInt) {
+    TArenaMemoryContext arena;
+
+    arrow::Int64Builder builder;
+    builder.Append(11);
+    builder.Append(3137);
+    builder.AppendNull();
+    builder.Append(-1337);
+    builder.Append(0);
+
+    std::shared_ptr<arrow::Array> array;
+    builder.Finish(&array);
+
+    auto result = PgConvertNumeric<i64>(array);
+    const auto& data = result->data();
+
+    const char* expected[] = {
+        "11", "3137", nullptr, "-1337", "0"
+    };
+
+    NUdf::TStringBlockReader<arrow::BinaryType, true> reader;
+    for (int i = 0; i < 5; i++) {
+        auto item = reader.GetItem(*data, i);
+        if (!item) {
+            UNIT_ASSERT(expected[i] == nullptr);
+        } else {
+            const char* addr = item.AsStringRef().Data() + sizeof(void*);
+            UNIT_ASSERT_VALUES_EQUAL(
+                TString(DatumGetCString(DirectFunctionCall1(numeric_out, (Datum)addr))),
+                expected[i]
+            );
+        }
+    }
+}
+
+} // Y_UNIT_TEST_SUITE(TArrowUtilsTests)
+
+} // namespace NYql
+

ydb/library/yql/parser/pg_wrapper/ut/ya.make

@@ -8,6 +8,7 @@ NO_COMPILER_WARNINGS()
 INCLUDE(../cflags.inc)
 
 SRCS(
+    arrow_ut.cpp
     codegen_ut.cpp
     error_ut.cpp
     parser_ut.cpp