apacheGH-34863: [Go] Pow method for Decimal DataTypes (apache#34864)

### Rationale for this change Closes: apache#34863 ### Are these changes tested? Yes ### Are there any user-facing changes? Yes * Closes: apache#34863 Authored-by: izveigor <[email protected]> Signed-off-by: Matt Topol <[email protected]>
wgtmac · Apr 4, 2023 · 6024678 · 6024678
1 parent c4057bd
commit 6024678
Show file tree

Hide file tree

Showing 4 changed files with 70 additions and 9 deletions.
diff --git a/go/arrow/decimal128/decimal128.go b/go/arrow/decimal128/decimal128.go
@@ -134,6 +134,11 @@ func (n Num) Div(rhs Num) (res, rem Num) {
 	return FromBigInt(out), FromBigInt(remainder)
 }
 
+func (n Num) Pow(rhs Num) Num {
+	b := n.BigInt()
+	return FromBigInt(b.Exp(b, rhs.BigInt(), nil))
+}
+
 func scalePositiveFloat64(v float64, prec, scale int32) (float64, error) {
 	var pscale float64
 	if scale >= -38 && scale <= 38 {
@@ -168,12 +173,16 @@ func fromPositiveFloat64(v float64, prec, scale int32) (Num, error) {
 // Aren't floating point values so much fun?
 //
 // example value to use:
-//    v := float32(1.8446746e+15)
+//
+//	v := float32(1.8446746e+15)
 //
 // You'll end up with a different values if you do:
-// 	  FromFloat64(float64(v), 20, 4)
+//
+//	FromFloat64(float64(v), 20, 4)
+//
 // vs
-//    FromFloat32(v, 20, 4)
+//
+//	FromFloat32(v, 20, 4)
 //
 // because float64(v) == 1844674629206016 rather than 1844674600000000
 func fromPositiveFloat32(v float32, prec, scale int32) (Num, error) {
@@ -297,7 +306,9 @@ func (n Num) HighBits() int64 { return n.hi }
 // Sign returns:
 //
 // -1 if x <  0
-//  0 if x == 0
+//
+//	0 if x == 0
+//
 // +1 if x >  0
 func (n Num) Sign() int {
 	if n == (Num{}) {

diff --git a/go/arrow/decimal128/decimal128_test.go b/go/arrow/decimal128/decimal128_test.go
@@ -206,6 +206,26 @@ func TestDiv(t *testing.T) {
 	}
 }
 
+func TestPow(t *testing.T) {
+	for _, tc := range []struct {
+		n    decimal128.Num
+		rhs  decimal128.Num
+		want decimal128.Num
+	}{
+		{decimal128.New(0, 2), decimal128.New(0, 3), decimal128.New(0, 8)},
+		{decimal128.New(0, 2), decimal128.New(0, 65), decimal128.New(2, 0)},
+		{decimal128.New(0, 1), decimal128.New(0, 0), decimal128.New(0, 1)},
+		{decimal128.New(0, 0), decimal128.New(0, 1), decimal128.New(0, 0)},
+	} {
+		t.Run("pow", func(t *testing.T) {
+			n := tc.n.Pow(tc.rhs)
+			if got, want := n, tc.want; got != want {
+				t.Fatalf("invalid value. got=%v, want=%v", got, want)
+			}
+		})
+	}
+}
+
 func BenchmarkDecimalToBigInt(b *testing.B) {
 	var (
 		bi *big.Int

diff --git a/go/arrow/decimal256/decimal256.go b/go/arrow/decimal256/decimal256.go
@@ -45,8 +45,8 @@ type Num struct {
 // the highest bits with the rest of the values in order down to the
 // lowest bits
 //
-//   ie: New(1, 2, 3, 4) returns with the elements in little-endian order
-//       {4, 3, 2, 1} but each value is still represented as the native endianness
+//	ie: New(1, 2, 3, 4) returns with the elements in little-endian order
+//	    {4, 3, 2, 1} but each value is still represented as the native endianness
 func New(x1, x2, x3, x4 uint64) Num {
 	return Num{[4]uint64{x4, x3, x2, x1}}
 }
@@ -112,6 +112,11 @@ func (n Num) Div(rhs Num) (res, rem Num) {
 	return FromBigInt(out), FromBigInt(remainder)
 }
 
+func (n Num) Pow(rhs Num) Num {
+	b := n.BigInt()
+	return FromBigInt(b.Exp(b, rhs.BigInt(), nil))
+}
+
 var pt5 = big.NewFloat(0.5)
 
 func FromString(v string, prec, scale int32) (n Num, err error) {
@@ -212,12 +217,16 @@ func FromFloat64(v float64, prec, scale int32) (Num, error) {
 // Aren't floating point values so much fun?
 //
 // example value to use:
-//    v := float32(1.8446746e+15)
+//
+//	v := float32(1.8446746e+15)
 //
 // You'll end up with a different values if you do:
-// 	  FromFloat64(float64(v), 20, 4)
+//
+//	FromFloat64(float64(v), 20, 4)
+//
 // vs
-//    FromFloat32(v, 20, 4)
+//
+//	FromFloat32(v, 20, 4)
 //
 // because float64(v) == 1844674629206016 rather than 1844674600000000
 func fromPositiveFloat32(v float32, prec, scale int32) (Num, error) {

diff --git a/go/arrow/decimal256/decimal256_test.go b/go/arrow/decimal256/decimal256_test.go
@@ -233,6 +233,27 @@ func TestDiv(t *testing.T) {
 	}
 }
 
+func TestPow(t *testing.T) {
+	for _, tc := range []struct {
+		n    decimal256.Num
+		rhs  decimal256.Num
+		want decimal256.Num
+	}{
+		{decimal256.New(0, 0, 0, 2), decimal256.New(0, 0, 0, 3), decimal256.New(0, 0, 0, 8)},
+		{decimal256.New(0, 0, 2, 0), decimal256.New(0, 0, 0, 3), decimal256.New(8, 0, 0, 0)},
+		{decimal256.New(0, 0, 2, 2), decimal256.New(0, 0, 0, 3), decimal256.New(8, 24, 24, 8)},
+		{decimal256.New(0, 0, 0, 1), decimal256.New(0, 0, 0, 0), decimal256.New(0, 0, 0, 1)},
+		{decimal256.New(0, 0, 0, 0), decimal256.New(0, 0, 0, 1), decimal256.New(0, 0, 0, 0)},
+	} {
+		t.Run("pow", func(t *testing.T) {
+			n := tc.n.Pow(tc.rhs)
+			if got, want := n, tc.want; got != want {
+				t.Fatalf("invalid value. got=%v, want=%v", got, want)
+			}
+		})
+	}
+}
+
 func TestDecimalToBigInt(t *testing.T) {
 	tests := []struct {
 		arr [4]uint64