perf(arrow/array): append dictionary indices directly

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 {

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()
+	}
+}

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 {