Skip to content

Commit 12596f4

Browse files
author
kinash-varvara
authored
Add decimal128 type support with big precision (#1607)
1 parent b399caf commit 12596f4

File tree

3 files changed

+72
-4
lines changed

3 files changed

+72
-4
lines changed

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

Lines changed: 21 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -1,5 +1,6 @@
11
#include "arrow.h"
22
#include "arrow_impl.h"
3+
#include "ydb/library/yql/minikql/defs.h"
34
#include <ydb/library/yql/parser/pg_wrapper/interface/arrow.h>
45
#include <ydb/library/yql/parser/pg_wrapper/interface/utils.h>
56
#include <ydb/library/yql/minikql/mkql_node_cast.h>
@@ -183,14 +184,17 @@ std::shared_ptr<arrow::Array> PgDecimal128ConvertNumeric(const std::shared_ptr<a
183184
size_t length = data->length;
184185
arrow::BinaryBuilder builder;
185186

187+
bool error;
188+
Numeric high_bits_mul = numeric_mul_opt_error(int64_to_numeric(int64_t(1) << 62), int64_to_numeric(4), &error);
189+
186190
auto input = data->GetValues<arrow::Decimal128>(1);
187191
for (size_t i = 0; i < length; ++i) {
188192
if (value->IsNull(i)) {
189193
ARROW_OK(builder.AppendNull());
190194
continue;
191195
}
192196

193-
Numeric v = PgDecimal128ToNumeric(input[i].high_bits(), input[i].low_bits(), precision, scale);
197+
Numeric v = PgDecimal128ToNumeric(input[i], precision, scale, high_bits_mul);
194198

195199
auto datum = NumericGetDatum(v);
196200
auto ptr = (char*)datum;
@@ -204,8 +208,22 @@ std::shared_ptr<arrow::Array> PgDecimal128ConvertNumeric(const std::shared_ptr<a
204208
return ret;
205209
}
206210

207-
Numeric PgDecimal128ToNumeric(int64_t high_bits, uint64_t low_bits, int32_t precision, int32_t scale) {
208-
Numeric res = int64_div_fast_to_numeric(low_bits, scale);
211+
Numeric PgDecimal128ToNumeric(arrow::Decimal128 value, int32_t precision, int32_t scale, Numeric high_bits_mul) {
212+
uint64_t low_bits = value.low_bits();
213+
int64 high_bits = value.high_bits();
214+
215+
if (low_bits > INT64_MAX){
216+
high_bits += 1;
217+
}
218+
219+
bool error;
220+
Numeric low_bits_res = int64_div_fast_to_numeric(low_bits, scale);
221+
Numeric high_bits_res = numeric_mul_opt_error(int64_div_fast_to_numeric(high_bits, scale), high_bits_mul, &error);
222+
MKQL_ENSURE(error == false, "Bad numeric multiplication.");
223+
224+
Numeric res = numeric_add_opt_error(high_bits_res, low_bits_res, &error);
225+
MKQL_ENSURE(error == false, "Bad numeric addition.");
226+
209227
return res;
210228
}
211229

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

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -11,7 +11,7 @@ extern "C" {
1111
namespace NYql {
1212

1313
Numeric PgFloatToNumeric(double item, ui64 scale, int digits);
14-
Numeric PgDecimal128ToNumeric(int64_t high_bits, uint64_t low_bits, int32_t precision, int32_t scale);
14+
Numeric PgDecimal128ToNumeric(arrow::Decimal128 val, int32_t precision, int32_t scale, Numeric high_bits_mul);
1515
TColumnConverter BuildPgColumnConverter(const std::shared_ptr<arrow::DataType>& originalType, NKikimr::NMiniKQL::TPgType* targetType);
1616

1717
template<typename T>

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

Lines changed: 50 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -185,6 +185,56 @@ Y_UNIT_TEST(PgConvertNumericDecimal128Scale5) {
185185
checkResult<false>(expected, result, &reader, numeric_out);
186186
}
187187

188+
Y_UNIT_TEST(PgConvertNumericDecimal128BigScale3) {
189+
TArenaMemoryContext arena;
190+
191+
int32_t precision = 20;
192+
int32_t scale = 3;
193+
std::shared_ptr<arrow::DataType> type(new arrow::Decimal128Type(precision, scale));
194+
arrow::Decimal128Builder builder(type);
195+
196+
const char* expected[] = {
197+
"36893488147419103.245", "-36893488147419103.245", nullptr
198+
};
199+
200+
ARROW_OK(builder.Append(arrow::Decimal128::FromString("36893488147419103.245").ValueOrDie()));
201+
ARROW_OK(builder.Append(arrow::Decimal128::FromString("-36893488147419103.245").ValueOrDie()));
202+
ARROW_OK(builder.AppendNull());
203+
204+
std::shared_ptr<arrow::Array> array;
205+
ARROW_OK(builder.Finish(&array));
206+
207+
auto result = PgDecimal128ConvertNumeric(array, precision, scale);
208+
209+
NYql::NUdf::TStringBlockReader<arrow::BinaryType, true> reader;
210+
checkResult<false>(expected, result, &reader, numeric_out);
211+
}
212+
213+
Y_UNIT_TEST(PgConvertNumericDecimal128BigScale1) {
214+
TArenaMemoryContext arena;
215+
216+
int32_t precision = 26;
217+
int32_t scale = 1;
218+
std::shared_ptr<arrow::DataType> type(new arrow::Decimal128Type(precision, scale));
219+
arrow::Decimal128Builder builder(type);
220+
221+
const char* expected[] = {
222+
"3868562622766813359059763.2", "-3868562622766813359059763.2", nullptr
223+
};
224+
225+
ARROW_OK(builder.Append(arrow::Decimal128::FromString("3868562622766813359059763.2").ValueOrDie()));
226+
ARROW_OK(builder.Append(arrow::Decimal128::FromString("-3868562622766813359059763.2").ValueOrDie()));
227+
ARROW_OK(builder.AppendNull());
228+
229+
std::shared_ptr<arrow::Array> array;
230+
ARROW_OK(builder.Finish(&array));
231+
232+
auto result = PgDecimal128ConvertNumeric(array, precision, scale);
233+
234+
NYql::NUdf::TStringBlockReader<arrow::BinaryType, true> reader;
235+
checkResult<false>(expected, result, &reader, numeric_out);
236+
}
237+
188238
Y_UNIT_TEST(PgConvertNumericInt) {
189239
TArenaMemoryContext arena;
190240

0 commit comments

Comments
 (0)