ARROW-10600: [Go] Implement Decimal256 (#13792)

Authored-by: Matt Topol <[email protected]> Signed-off-by: Matt Topol <[email protected]>
apache · Aug 5, 2022 · 6d1bc62 · 6d1bc62
1 parent 6a3fb97
commit 6d1bc62
Show file tree

Hide file tree

Showing 23 changed files with 1,772 additions and 79 deletions.
diff --git a/dev/archery/archery/integration/datagen.py b/dev/archery/archery/integration/datagen.py
@@ -1599,8 +1599,7 @@ def _temp_path():
 
         generate_decimal128_case(),
 
-        generate_decimal256_case()
-        .skip_category('Go')  # TODO(ARROW-7948): Decimal + Go
+        generate_decimal256_case()        
         .skip_category('JS'),
 
         generate_datetime_case(),

diff --git a/docs/source/status.rst b/docs/source/status.rst
@@ -46,7 +46,7 @@ Data Types
 +-------------------+-------+-------+-------+------------+-------+-------+-------+
 | Decimal128        | ✓     | ✓     | ✓     |            |  ✓    |  ✓    | ✓     |
 +-------------------+-------+-------+-------+------------+-------+-------+-------+
-| Decimal256        | ✓     | ✓     |       |            |  ✓    |       | ✓     |
+| Decimal256        | ✓     | ✓     | ✓     |            |  ✓    |       | ✓     |
 +-------------------+-------+-------+-------+------------+-------+-------+-------+
 | Date32/64         | ✓     | ✓     | ✓     | ✓          |  ✓    |  ✓    | ✓     |
 +-------------------+-------+-------+-------+------------+-------+-------+-------+

diff --git a/go/arrow/array/array.go b/go/arrow/array/array.go
@@ -114,10 +114,6 @@ var (
 	makeArrayFn [64]arrayConstructorFn
 )
 
-func unsupportedArrayType(data arrow.ArrayData) arrow.Array {
-	panic("unsupported data type: " + data.DataType().ID().String())
-}
-
 func invalidDataType(data arrow.ArrayData) arrow.Array {
 	panic("invalid data type: " + data.DataType().ID().String())
 }
@@ -166,7 +162,7 @@ func init() {
 		arrow.INTERVAL_MONTHS:         func(data arrow.ArrayData) arrow.Array { return NewMonthIntervalData(data) },
 		arrow.INTERVAL_DAY_TIME:       func(data arrow.ArrayData) arrow.Array { return NewDayTimeIntervalData(data) },
 		arrow.DECIMAL128:              func(data arrow.ArrayData) arrow.Array { return NewDecimal128Data(data) },
-		arrow.DECIMAL256:              unsupportedArrayType,
+		arrow.DECIMAL256:              func(data arrow.ArrayData) arrow.Array { return NewDecimal256Data(data) },
 		arrow.LIST:                    func(data arrow.ArrayData) arrow.Array { return NewListData(data) },
 		arrow.STRUCT:                  func(data arrow.ArrayData) arrow.Array { return NewStructData(data) },
 		arrow.SPARSE_UNION:            func(data arrow.ArrayData) arrow.Array { return NewSparseUnionData(data) },

diff --git a/go/arrow/array/array_test.go b/go/arrow/array/array_test.go
@@ -74,6 +74,7 @@ func TestMakeFromData(t *testing.T) {
 		{name: "month_interval", d: arrow.FixedWidthTypes.MonthInterval},
 		{name: "day_time_interval", d: arrow.FixedWidthTypes.DayTimeInterval},
 		{name: "decimal128", d: &testDataType{arrow.DECIMAL128}},
+		{name: "decimal256", d: &testDataType{arrow.DECIMAL256}},
 		{name: "month_day_nano_interval", d: arrow.FixedWidthTypes.MonthDayNanoInterval},
 
 		{name: "list", d: &testDataType{arrow.LIST}, child: []arrow.ArrayData{
@@ -122,9 +123,6 @@ func TestMakeFromData(t *testing.T) {
 		{name: "extension", d: &testDataType{arrow.EXTENSION}, expPanic: true, expError: "arrow/array: DataType for ExtensionArray must implement arrow.ExtensionType"},
 		{name: "extension", d: types.NewUUIDType()},
 
-		// unsupported types
-		{name: "decimal256", d: &testDataType{arrow.DECIMAL256}, expPanic: true, expError: "unsupported data type: DECIMAL256"},
-
 		// invalid types
 		{name: "invalid(-1)", d: &testDataType{arrow.Type(-1)}, expPanic: true, expError: "invalid data type: Type(-1)"},
 		{name: "invalid(63)", d: &testDataType{arrow.Type(63)}, expPanic: true, expError: "invalid data type: Type(63)"},

diff --git a/go/arrow/array/builder.go b/go/arrow/array/builder.go
@@ -299,6 +299,9 @@ func NewBuilder(mem memory.Allocator, dtype arrow.DataType) Builder {
 			return NewDecimal128Builder(mem, typ)
 		}
 	case arrow.DECIMAL256:
+		if typ, ok := dtype.(*arrow.Decimal256Type); ok {
+			return NewDecimal256Builder(mem, typ)
+		}
 	case arrow.LIST:
 		typ := dtype.(*arrow.ListType)
 		return NewListBuilder(mem, typ.Elem())

diff --git a/go/arrow/array/compare.go b/go/arrow/array/compare.go
@@ -276,6 +276,9 @@ func Equal(left, right arrow.Array) bool {
 	case *Decimal128:
 		r := right.(*Decimal128)
 		return arrayEqualDecimal128(l, r)
+	case *Decimal256:
+		r := right.(*Decimal256)
+		return arrayEqualDecimal256(l, r)
 	case *Date32:
 		r := right.(*Date32)
 		return arrayEqualDate32(l, r)
@@ -531,6 +534,9 @@ func arrayApproxEqual(left, right arrow.Array, opt equalOption) bool {
 	case *Decimal128:
 		r := right.(*Decimal128)
 		return arrayEqualDecimal128(l, r)
+	case *Decimal256:
+		r := right.(*Decimal256)
+		return arrayEqualDecimal256(l, r)
 	case *Date32:
 		r := right.(*Date32)
 		return arrayEqualDate32(l, r)

diff --git a/go/arrow/array/decimal128_test.go b/go/arrow/array/decimal128_test.go
@@ -77,7 +77,9 @@ func TestNewDecimal128Builder(t *testing.T) {
 
 	assert.Equal(t, want, a.Values(), "unexpected Decimal128Values")
 	assert.Equal(t, []byte{0xb7}, a.NullBitmapBytes()[:1]) // 4 bytes due to minBuilderCapacity
+	assert.Equal(t, 4, a.Data().Buffers()[0].Len(), "should be 4 bytes due to minBuilderCapacity")
 	assert.Len(t, a.Values(), 10, "unexpected length of Decimal128Values")
+	assert.Equal(t, 10*arrow.Decimal128SizeBytes, a.Data().Buffers()[1].Len())
 
 	a.Release()
 	ab.Append(decimal128.FromI64(7))
@@ -88,6 +90,7 @@ func TestNewDecimal128Builder(t *testing.T) {
 	assert.Equal(t, 0, a.NullN())
 	assert.Equal(t, []decimal128.Num{decimal128.FromI64(7), decimal128.FromI64(8)}, a.Values())
 	assert.Len(t, a.Values(), 2)
+	assert.Equal(t, 2*arrow.Decimal128SizeBytes, a.Data().Buffers()[1].Len())
 
 	a.Release()
 }