From cd7f431d8cd2808f8b49847fce0df0cc4270229b Mon Sep 17 00:00:00 2001 From: Igor Izvekov Date: Wed, 7 Jun 2023 18:20:43 +0300 Subject: [PATCH] GH-35965: [Go] Fix `Decimal256DictionaryBuilder` (#35966) ### Rationale for this change `Decimal256DictionaryBuilder` cannot append `decimal256` value and prefers `decimal128` instead. ``` cannot use v (variable of type decimal256.Num) as decimal128.Num value in argument to builder.Append ``` ### What changes are included in this PR? ### Are these changes tested? Yes ### Are there any user-facing changes? Yes * Closes: #35965 Authored-by: izveigor Signed-off-by: Matt Topol --- go/arrow/array/dictionary.go | 14 ++++++++++---- go/arrow/array/dictionary_test.go | 31 ++++++++++++++++++++++++++++++- 2 files changed, 40 insertions(+), 5 deletions(-) diff --git a/go/arrow/array/dictionary.go b/go/arrow/array/dictionary.go index 1cbaac5a2cab7..48f37e91d1012 100644 --- a/go/arrow/array/dictionary.go +++ b/go/arrow/array/dictionary.go @@ -28,6 +28,7 @@ import ( "github.com/apache/arrow/go/v13/arrow" "github.com/apache/arrow/go/v13/arrow/bitutil" "github.com/apache/arrow/go/v13/arrow/decimal128" + "github.com/apache/arrow/go/v13/arrow/decimal256" "github.com/apache/arrow/go/v13/arrow/float16" "github.com/apache/arrow/go/v13/arrow/internal/debug" "github.com/apache/arrow/go/v13/arrow/memory" @@ -45,12 +46,12 @@ import ( // // For example, the array: // -// ["foo", "bar", "foo", "bar", "foo", "bar"] +// ["foo", "bar", "foo", "bar", "foo", "bar"] // // with dictionary ["bar", "foo"], would have the representation of: // -// indices: [1, 0, 1, 0, 1, 0] -// dictionary: ["bar", "foo"] +// indices: [1, 0, 1, 0, 1, 0] +// dictionary: ["bar", "foo"] // // The indices in principle may be any integer type. type Dictionary struct { @@ -883,6 +884,11 @@ func getvalFn(arr arrow.Array) func(i int) interface{} { val := typedarr.Value(i) return (*(*[arrow.Decimal128SizeBytes]byte)(unsafe.Pointer(&val)))[:] } + case *Decimal256: + return func(i int) interface{} { + val := typedarr.Value(i) + return (*(*[arrow.Decimal256SizeBytes]byte)(unsafe.Pointer(&val)))[:] + } case *DayTimeInterval: return func(i int) interface{} { val := typedarr.Value(i) @@ -1373,7 +1379,7 @@ type Decimal256DictionaryBuilder struct { dictionaryBuilder } -func (b *Decimal256DictionaryBuilder) Append(v decimal128.Num) error { +func (b *Decimal256DictionaryBuilder) Append(v decimal256.Num) error { return b.appendValue((*(*[arrow.Decimal256SizeBytes]byte)(unsafe.Pointer(&v)))[:]) } func (b *Decimal256DictionaryBuilder) InsertDictValues(arr *Decimal256) (err error) { diff --git a/go/arrow/array/dictionary_test.go b/go/arrow/array/dictionary_test.go index 23b21e5aba7a5..230cc706a848d 100644 --- a/go/arrow/array/dictionary_test.go +++ b/go/arrow/array/dictionary_test.go @@ -27,6 +27,7 @@ import ( "github.com/apache/arrow/go/v13/arrow/array" "github.com/apache/arrow/go/v13/arrow/bitutil" "github.com/apache/arrow/go/v13/arrow/decimal128" + "github.com/apache/arrow/go/v13/arrow/decimal256" "github.com/apache/arrow/go/v13/arrow/memory" "github.com/apache/arrow/go/v13/internal/types" "github.com/stretchr/testify/assert" @@ -878,7 +879,7 @@ func TestFixedSizeBinaryDictionaryStringRoundTrip(t *testing.T) { assert.True(t, array.Equal(arr, arr1)) } -func TestDecimalDictionaryBuilderBasic(t *testing.T) { +func TestDecimal128DictionaryBuilderBasic(t *testing.T) { mem := memory.NewCheckedAllocator(memory.DefaultAllocator) defer mem.AssertSize(t, 0) @@ -906,6 +907,34 @@ func TestDecimalDictionaryBuilderBasic(t *testing.T) { assert.True(t, array.ArrayApproxEqual(expected, result)) } +func TestDecimal256DictionaryBuilderBasic(t *testing.T) { + mem := memory.NewCheckedAllocator(memory.DefaultAllocator) + defer mem.AssertSize(t, 0) + + test := []decimal256.Num{decimal256.FromI64(12), decimal256.FromI64(12), decimal256.FromI64(11), decimal256.FromI64(12)} + dictType := &arrow.DictionaryType{IndexType: &arrow.Int8Type{}, ValueType: &arrow.Decimal256Type{Precision: 2, Scale: 0}} + bldr := array.NewDictionaryBuilder(mem, dictType) + defer bldr.Release() + + builder := bldr.(*array.Decimal256DictionaryBuilder) + for _, v := range test { + assert.NoError(t, builder.Append(v)) + } + + result := bldr.NewDictionaryArray() + defer result.Release() + + indices, _, _ := array.FromJSON(mem, dictType.IndexType, strings.NewReader("[0, 0, 1, 0]")) + defer indices.Release() + dict, _, _ := array.FromJSON(mem, dictType.ValueType, strings.NewReader("[12, 11]")) + defer dict.Release() + + expected := array.NewDictionaryArray(dictType, indices, dict) + defer expected.Release() + + assert.True(t, array.ArrayApproxEqual(expected, result)) +} + func TestNullDictionaryBuilderBasic(t *testing.T) { mem := memory.NewCheckedAllocator(memory.DefaultAllocator) defer mem.AssertSize(t, 0)