diff --git a/arrow/array/dictionary.go b/arrow/array/dictionary.go index fe7d66be8..dc05c0609 100644 --- a/arrow/array/dictionary.go +++ b/arrow/array/dictionary.go @@ -518,6 +518,12 @@ type IndexBuilder struct { UnsafeAppend func(int) } +func appendDictionaryIndices[T arrow.IntType | arrow.UintType](dst []T, indices []int) { + for i, idx := range indices { + dst[i] = T(idx) + } +} + func createIndexBuilder(mem memory.Allocator, dt arrow.FixedWidthDataType) (ret IndexBuilder, err error) { ret = IndexBuilder{Builder: NewBuilder(mem, dt)} switch dt.ID() { @@ -1181,57 +1187,42 @@ func (b *dictionaryBuilder) IndexBuilder() IndexBuilder { } func (b *dictionaryBuilder) AppendIndices(indices []int, valid []bool) { - b.length += len(indices) + if len(indices) != len(valid) && len(valid) != 0 { + panic("len(indices) != len(valid) && len(valid) != 0") + } + + if len(indices) == 0 { + return + } + + b.idxBuilder.Reserve(len(indices)) switch idxbldr := b.idxBuilder.Builder.(type) { case *Int8Builder: - vals := make([]int8, len(indices)) - for i, v := range indices { - vals[i] = int8(v) - } - idxbldr.AppendValues(vals, valid) + appendDictionaryIndices(idxbldr.rawData[idxbldr.length:], indices) + idxbldr.builder.unsafeAppendBoolsToBitmap(valid, len(indices)) case *Int16Builder: - vals := make([]int16, len(indices)) - for i, v := range indices { - vals[i] = int16(v) - } - idxbldr.AppendValues(vals, valid) + appendDictionaryIndices(idxbldr.rawData[idxbldr.length:], indices) + idxbldr.builder.unsafeAppendBoolsToBitmap(valid, len(indices)) case *Int32Builder: - vals := make([]int32, len(indices)) - for i, v := range indices { - vals[i] = int32(v) - } - idxbldr.AppendValues(vals, valid) + appendDictionaryIndices(idxbldr.rawData[idxbldr.length:], indices) + idxbldr.builder.unsafeAppendBoolsToBitmap(valid, len(indices)) case *Int64Builder: - vals := make([]int64, len(indices)) - for i, v := range indices { - vals[i] = int64(v) - } - idxbldr.AppendValues(vals, valid) + appendDictionaryIndices(idxbldr.rawData[idxbldr.length:], indices) + idxbldr.builder.unsafeAppendBoolsToBitmap(valid, len(indices)) case *Uint8Builder: - vals := make([]uint8, len(indices)) - for i, v := range indices { - vals[i] = uint8(v) - } - idxbldr.AppendValues(vals, valid) + appendDictionaryIndices(idxbldr.rawData[idxbldr.length:], indices) + idxbldr.builder.unsafeAppendBoolsToBitmap(valid, len(indices)) case *Uint16Builder: - vals := make([]uint16, len(indices)) - for i, v := range indices { - vals[i] = uint16(v) - } - idxbldr.AppendValues(vals, valid) + appendDictionaryIndices(idxbldr.rawData[idxbldr.length:], indices) + idxbldr.builder.unsafeAppendBoolsToBitmap(valid, len(indices)) case *Uint32Builder: - vals := make([]uint32, len(indices)) - for i, v := range indices { - vals[i] = uint32(v) - } - idxbldr.AppendValues(vals, valid) + appendDictionaryIndices(idxbldr.rawData[idxbldr.length:], indices) + idxbldr.builder.unsafeAppendBoolsToBitmap(valid, len(indices)) case *Uint64Builder: - vals := make([]uint64, len(indices)) - for i, v := range indices { - vals[i] = uint64(v) - } - idxbldr.AppendValues(vals, valid) + appendDictionaryIndices(idxbldr.rawData[idxbldr.length:], indices) + idxbldr.builder.unsafeAppendBoolsToBitmap(valid, len(indices)) } + b.length += len(indices) } func (b *dictionaryBuilder) DictionarySize() int { diff --git a/arrow/array/dictionary_append_indices_benchmark_test.go b/arrow/array/dictionary_append_indices_benchmark_test.go new file mode 100644 index 000000000..5ed517190 --- /dev/null +++ b/arrow/array/dictionary_append_indices_benchmark_test.go @@ -0,0 +1,91 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +package array_test + +import ( + "fmt" + "strconv" + "testing" + + "github.com/apache/arrow-go/v18/arrow" + "github.com/apache/arrow-go/v18/arrow/array" + "github.com/apache/arrow-go/v18/arrow/memory" +) + +func BenchmarkDictionaryBuilderAppendIndices(b *testing.B) { + const ( + length = 1 << 16 + cardinality = 1 << 8 + ) + + indices := make([]int, length) + valid := make([]bool, length) + for i := range indices { + indices[i] = i % cardinality + valid[i] = i%10 != 0 + } + + dictionaryValues := make([]string, cardinality) + for i := range dictionaryValues { + dictionaryValues[i] = strconv.Itoa(i) + } + + indexTypes := []arrow.DataType{ + arrow.PrimitiveTypes.Int8, + arrow.PrimitiveTypes.Uint8, + arrow.PrimitiveTypes.Int16, + arrow.PrimitiveTypes.Uint16, + arrow.PrimitiveTypes.Int32, + arrow.PrimitiveTypes.Uint32, + arrow.PrimitiveTypes.Int64, + arrow.PrimitiveTypes.Uint64, + } + + for _, indexType := range indexTypes { + indexType := indexType + b.Run(fmt.Sprintf("%s/non-null", indexType), func(b *testing.B) { + benchmarkDictionaryBuilderAppendIndices(b, indexType, indices, nil, dictionaryValues) + }) + b.Run(fmt.Sprintf("%s/nullable", indexType), func(b *testing.B) { + benchmarkDictionaryBuilderAppendIndices(b, indexType, indices, valid, dictionaryValues) + }) + } +} + +func benchmarkDictionaryBuilderAppendIndices(b *testing.B, indexType arrow.DataType, indices []int, valid []bool, dictionaryValues []string) { + mem := memory.NewGoAllocator() + dictBuilder := array.NewStringBuilder(mem) + dictBuilder.AppendValues(dictionaryValues, nil) + dictionary := dictBuilder.NewStringArray() + dictBuilder.Release() + defer dictionary.Release() + + builder := array.NewDictionaryBuilderWithDict(mem, &arrow.DictionaryType{ + IndexType: indexType, + ValueType: arrow.BinaryTypes.String, + }, dictionary) + defer builder.Release() + + b.ReportAllocs() + b.SetBytes(int64(len(indices) * indexType.(arrow.FixedWidthDataType).Bytes())) + b.ResetTimer() + for b.Loop() { + builder.AppendIndices(indices, valid) + arr := builder.NewDictionaryArray() + arr.Release() + } +} diff --git a/arrow/array/dictionary_test.go b/arrow/array/dictionary_test.go index 2324b743d..e63351891 100644 --- a/arrow/array/dictionary_test.go +++ b/arrow/array/dictionary_test.go @@ -1946,6 +1946,43 @@ func TestDictionaryAppendIndices(t *testing.T) { assert.Equal(t, fmt.Sprint(indices), arrIndices.String()) }) } + + valid := []bool{true, false, true, false, true, true, false, true} + for _, typ := range indexTypes { + t.Run(fmt.Sprintf("%s with validity", typ), func(t *testing.T) { + scoped := memory.NewCheckedAllocatorScope(mem) + defer scoped.CheckSize(t) + + dictType := &arrow.DictionaryType{ + IndexType: typ, ValueType: dict.DataType()} + bldr := array.NewDictionaryBuilderWithDict(mem, dictType, dict) + defer bldr.Release() + + bldr.AppendIndices(indices[:3], valid[:3]) + bldr.AppendIndices(indices[3:], valid[3:]) + + arr := bldr.NewDictionaryArray() + defer arr.Release() + + assert.EqualValues(t, len(indices), arr.Len()) + assert.EqualValues(t, 3, arr.NullN()) + for i, wantValid := range valid { + assert.Equal(t, !wantValid, arr.IsNull(i)) + assert.Equal(t, indices[i], arr.GetValueIndex(i)) + } + }) + } + + t.Run("validity length mismatch", func(t *testing.T) { + dictType := &arrow.DictionaryType{ + IndexType: arrow.PrimitiveTypes.Int32, ValueType: dict.DataType()} + bldr := array.NewDictionaryBuilderWithDict(mem, dictType, dict) + defer bldr.Release() + + assert.PanicsWithValue(t, "len(indices) != len(valid) && len(valid) != 0", func() { + bldr.AppendIndices([]int{0}, []bool{true, false}) + }) + }) } type panicAllocator struct {