diff --git a/arrow/array/data.go b/arrow/array/data.go index 47d83f56..5de367d5 100644 --- a/arrow/array/data.go +++ b/arrow/array/data.go @@ -187,13 +187,22 @@ func (d *Data) Dictionary() arrow.ArrayData { return d.dictionary } // SetDictionary allows replacing the dictionary for this particular Data object func (d *Data) SetDictionary(dict arrow.ArrayData) { - if d.dictionary != nil { - d.dictionary.Release() - d.dictionary = nil + var newDict *Data + if dict != nil { + var ok bool + newDict, ok = dict.(*Data) + if !ok { + panic("arrow/array: dictionary data must be *array.Data") + } } - if dict.(*Data) != nil { - dict.Retain() - d.dictionary = dict.(*Data) + + if newDict != nil { + newDict.Retain() + } + oldDict := d.dictionary + d.dictionary = newDict + if oldDict != nil { + oldDict.Release() } } diff --git a/arrow/array/data_test.go b/arrow/array/data_test.go index 95eb3848..54cc3930 100644 --- a/arrow/array/data_test.go +++ b/arrow/array/data_test.go @@ -24,6 +24,7 @@ import ( "github.com/apache/arrow-go/v18/arrow/memory" "github.com/apache/arrow-go/v18/internal/utils/maphash" "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" ) func TestDataReset(t *testing.T) { @@ -101,6 +102,57 @@ func TestDataResetClearsDictionary(t *testing.T) { mem.AssertSize(t, 0) } +func TestDataSetDictionaryWithSamePointerRetainsDictionary(t *testing.T) { + mem := memory.NewCheckedAllocator(memory.DefaultAllocator) + defer mem.AssertSize(t, 0) + + values := memory.NewResizableBuffer(mem) + values.Resize(4) + copy(values.Bytes(), arrow.Int32Traits.CastToBytes([]int32{42})) + dictData := NewData(arrow.PrimitiveTypes.Int32, 1, []*memory.Buffer{nil, values}, nil, 0, 0) + values.Release() + + dictType := &arrow.DictionaryType{ + IndexType: arrow.PrimitiveTypes.Int8, + ValueType: arrow.PrimitiveTypes.Int32, + } + data := NewDataWithDictionary(dictType, 1, []*memory.Buffer{nil, memory.NewBufferBytes([]byte{0})}, 0, 0, dictData) + defer data.Release() + dictData.Release() + + data.SetDictionary(dictData) + require.Same(t, dictData, data.Dictionary()) + require.NotNil(t, dictData.Buffers()[1]) + + valuesArray := NewInt32Data(dictData) + defer valuesArray.Release() + require.Equal(t, int32(42), valuesArray.Value(0)) +} + +func TestDataSetDictionaryNilClearsDictionary(t *testing.T) { + mem := memory.NewCheckedAllocator(memory.DefaultAllocator) + defer mem.AssertSize(t, 0) + + values := memory.NewResizableBuffer(mem) + values.Resize(4) + dictData := NewData(arrow.PrimitiveTypes.Int32, 1, []*memory.Buffer{nil, values}, nil, 0, 0) + values.Release() + + dictType := &arrow.DictionaryType{ + IndexType: arrow.PrimitiveTypes.Int8, + ValueType: arrow.PrimitiveTypes.Int32, + } + data := NewDataWithDictionary(dictType, 1, []*memory.Buffer{nil, memory.NewBufferBytes([]byte{0})}, 0, 0, dictData) + defer data.Release() + dictData.Release() + + require.NotNil(t, data.Dictionary()) + + data.SetDictionary(nil) + + require.Nil(t, data.Dictionary()) +} + func TestHashIncludesDataMetadataAndBuffers(t *testing.T) { seed := maphash.MakeSeed() hash := func(data arrow.ArrayData) uint64 {