Change numeric.c according to postgres 14 (#1753)

ydb-platform · Feb 9, 2024 · 9db23bc · 9db23bc
1 parent de2ba86
commit 9db23bc
Show file tree

Hide file tree

Showing 2 changed files with 104 additions and 26 deletions.
diff --git a/ydb/library/yql/parser/pg_wrapper/postgresql/src/backend/utils/adt/numeric.c b/ydb/library/yql/parser/pg_wrapper/postgresql/src/backend/utils/adt/numeric.c
@@ -4089,58 +4089,86 @@ int64_to_numeric(int64 val)
 }
 
 /*
- * Convert val1/(10**val2) to numeric.  This is much faster than normal
+ * Convert val1/(10**log10val2) to numeric.  This is much faster than normal
  * numeric division.
  */
 Numeric
 int64_div_fast_to_numeric(int64 val1, int log10val2)
 {
 	Numeric		res;
 	NumericVar	result;
-	int64		saved_val1 = val1;
+	int			rscale;
 	int			w;
 	int			m;
 
+	init_var(&result);
+
+	/* result scale */
+	rscale = log10val2 < 0 ? 0 : log10val2;
+
 	/* how much to decrease the weight by */
 	w = log10val2 / DEC_DIGITS;
-	/* how much is left */
+	/* how much is left to divide by */
 	m = log10val2 % DEC_DIGITS;
+	if (m < 0)
+	{
+		m += DEC_DIGITS;
+		w--;
+	}
 
 	/*
-	 * If there is anything left, multiply the dividend by what's left, then
-	 * shift the weight by one more.
+	 * If there is anything left to divide by (10^m with 0 < m < DEC_DIGITS),
+	 * multiply the dividend by 10^(DEC_DIGITS - m), and shift the weight by
+	 * one more.
 	 */
 	if (m > 0)
 	{
-		static __thread int	pow10[] = {1, 10, 100, 1000};
+#if DEC_DIGITS == 4
+		static __thread int pow10[] = {1, 10, 100, 1000};
+#elif DEC_DIGITS == 2
+		static __thread int pow10[] = {1, 10};
+#elif DEC_DIGITS == 1
+		static __thread int pow10[] = {1};
+#else
+#error unsupported NBASE
+#endif
+		int64		factor = pow10[DEC_DIGITS - m];
+		int64		new_val1;
 
 		StaticAssertStmt(lengthof(pow10) == DEC_DIGITS, "mismatch with DEC_DIGITS");
-		if (unlikely(pg_mul_s64_overflow(val1, pow10[DEC_DIGITS - m], &val1)))
+
+		if (unlikely(pg_mul_s64_overflow(val1, factor, &new_val1)))
 		{
-			/*
-			 * If it doesn't fit, do the whole computation in numeric the slow
-			 * way.  Note that va1l may have been overwritten, so use
-			 * saved_val1 instead.
-			 */
-			int			val2 = 1;
-
-			for (int i = 0; i < log10val2; i++)
-				val2 *= 10;
-			res = numeric_div_opt_error(int64_to_numeric(saved_val1), int64_to_numeric(val2), NULL);
-			res = DatumGetNumeric(DirectFunctionCall2(numeric_round,
-													  NumericGetDatum(res),
-													  Int32GetDatum(log10val2)));
-			return res;
+#ifdef HAVE_INT128
+			/* do the multiplication using 128-bit integers */
+			int128		tmp;
+
+			tmp = (int128) val1 * (int128) factor;
+
+			int128_to_numericvar(tmp, &result);
+#else
+			/* do the multiplication using numerics */
+			NumericVar	tmp;
+
+			init_var(&tmp);
+
+			int64_to_numericvar(val1, &result);
+			int64_to_numericvar(factor, &tmp);
+			mul_var(&result, &tmp, &result, 0);
+
+			free_var(&tmp);
+#endif
 		}
+		else
+			int64_to_numericvar(new_val1, &result);
+
 		w++;
 	}
-
-	init_var(&result);
-
-	int64_to_numericvar(val1, &result);
+	else
+		int64_to_numericvar(val1, &result);
 
 	result.weight -= w;
-	result.dscale += w * DEC_DIGITS - (DEC_DIGITS - m);
+	result.dscale = rscale;
 
 	res = make_result(&result);
 

diff --git a/ydb/library/yql/parser/pg_wrapper/ut/arrow_ut.cpp b/ydb/library/yql/parser/pg_wrapper/ut/arrow_ut.cpp
@@ -110,6 +110,31 @@ Y_UNIT_TEST(PgConvertNumericDecimal128Scale1) {
     checkResult<false>(expected, result, &reader, numeric_out);
 }
 
+Y_UNIT_TEST(PgConvertNumericDecimal128ScaleNegative) {
+    TArenaMemoryContext arena;
+
+    int32_t precision = 8;
+    int32_t scale = -3;
+    std::shared_ptr<arrow::DataType> type(new arrow::Decimal128Type(precision, scale));
+    arrow::Decimal128Builder builder(type);
+
+    const char* expected[] = {
+        "12345678000", "-12345678000", nullptr
+    };
+
+    ARROW_OK(builder.Append(arrow::Decimal128::FromString("12345678").ValueOrDie()));
+    ARROW_OK(builder.Append(arrow::Decimal128::FromString("-12345678").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;
 
@@ -160,6 +185,31 @@ Y_UNIT_TEST(PgConvertNumericDecimal128Scale3) {
     checkResult<false>(expected, result, &reader, numeric_out);
 }
 
+Y_UNIT_TEST(PgConvertNumericDecimal128Scale4) {
+    TArenaMemoryContext arena;
+
+    int32_t precision = 7;
+    int32_t scale = 4;
+    std::shared_ptr<arrow::DataType> type(new arrow::Decimal128Type(precision, scale));
+    arrow::Decimal128Builder builder(type);
+
+    const char* expected[] = {
+        "123.4567", "-123.4567", nullptr
+    };
+
+    ARROW_OK(builder.Append(arrow::Decimal128::FromString("123.4567").ValueOrDie()));
+    ARROW_OK(builder.Append(arrow::Decimal128::FromString("-123.4567").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;