diff --git a/arrow/compute/internal/kernels/vector_hash.go b/arrow/compute/internal/kernels/vector_hash.go index ecad7e98b..6eacba05e 100644 --- a/arrow/compute/internal/kernels/vector_hash.go +++ b/arrow/compute/internal/kernels/vector_hash.go @@ -295,12 +295,13 @@ func doAppendBinary[OffsetT int32 | int64](action Action, memo hashing.MemoTable offsets = exec.GetSpanOffsets[OffsetT](arr, 1) data = arr.Buffers[2].Buf shouldEncodeNulls = action.ShouldEncodeNulls() + typedMemo = memo.(hashing.TypedMemoTable[[]byte]) ) return bitutils.VisitBitBlocksShort(bitmap, arr.Offset, arr.Len, func(pos int64) error { v := data[offsets[pos]:offsets[pos+1]] - idx, found, err := memo.GetOrInsert(v) + idx, found, err := typedMemo.InsertOrGet(v) if err != nil { return err } @@ -328,12 +329,13 @@ func doAppendFixedSize(action Action, memo hashing.MemoTable, arr *exec.ArraySpa sz := int64(arr.Type.(arrow.FixedWidthDataType).Bytes()) arrData := arr.Buffers[1].Buf[arr.Offset*sz:] shouldEncodeNulls := action.ShouldEncodeNulls() + typedMemo := memo.(hashing.TypedMemoTable[[]byte]) return bitutils.VisitBitBlocksShort(arr.Buffers[0].Buf, arr.Offset, arr.Len, func(pos int64) error { // fixed size type memo table we use a binary memo table // so get the raw bytes - idx, found, err := memo.GetOrInsert(arrData[pos*sz : (pos+1)*sz]) + idx, found, err := typedMemo.InsertOrGet(arrData[pos*sz : (pos+1)*sz]) if err != nil { return err } diff --git a/arrow/compute/vector_hash_binary_bench_test.go b/arrow/compute/vector_hash_binary_bench_test.go new file mode 100644 index 000000000..55f816a50 --- /dev/null +++ b/arrow/compute/vector_hash_binary_bench_test.go @@ -0,0 +1,87 @@ +// 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. + +//go:build go1.24 + +package compute_test + +import ( + "context" + "encoding/binary" + "fmt" + "testing" + + "github.com/apache/arrow-go/v18/arrow" + "github.com/apache/arrow-go/v18/arrow/array" + "github.com/apache/arrow-go/v18/arrow/compute" + "github.com/apache/arrow-go/v18/arrow/memory" +) + +func BenchmarkDictionaryEncodeBinary(b *testing.B) { + const ( + nvalues = 1 << 16 + nunique = 100 + ) + + mem := memory.DefaultAllocator + ctx := compute.WithAllocator(context.Background(), mem) + + b.Run("string", func(b *testing.B) { + values := make([]string, nvalues) + bytes := 0 + for i := range values { + values[i] = fmt.Sprintf("value-%08d", i%nunique) + bytes += len(values[i]) + } + + builder := array.NewStringBuilder(mem) + builder.AppendValues(values, nil) + input := builder.NewStringArray() + builder.Release() + defer input.Release() + + benchmarkDictionaryEncodeBinary(b, ctx, input, bytes) + }) + + b.Run("fixed-size-binary-16", func(b *testing.B) { + values := make([][]byte, nvalues) + for i := range values { + values[i] = make([]byte, 16) + binary.LittleEndian.PutUint32(values[i], uint32(i%nunique)) + } + + builder := array.NewFixedSizeBinaryBuilder(mem, &arrow.FixedSizeBinaryType{ByteWidth: 16}) + builder.AppendValues(values, nil) + input := builder.NewFixedSizeBinaryArray() + builder.Release() + defer input.Release() + + benchmarkDictionaryEncodeBinary(b, ctx, input, nvalues*16) + }) +} + +func benchmarkDictionaryEncodeBinary(b *testing.B, ctx context.Context, input arrow.Array, bytes int) { + b.ReportAllocs() + b.SetBytes(int64(bytes)) + b.ResetTimer() + for b.Loop() { + result, err := compute.DictionaryEncodeArray(ctx, compute.DictionaryEncodeOptions{}, input) + if err != nil { + b.Fatal(err) + } + result.Release() + } +} diff --git a/arrow/compute/vector_hash_binary_test.go b/arrow/compute/vector_hash_binary_test.go new file mode 100644 index 000000000..4812d4dc4 --- /dev/null +++ b/arrow/compute/vector_hash_binary_test.go @@ -0,0 +1,66 @@ +// 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. + +//go:build go1.18 + +package compute_test + +import ( + "context" + "testing" + + "github.com/apache/arrow-go/v18/arrow" + "github.com/apache/arrow-go/v18/arrow/array" + "github.com/apache/arrow-go/v18/arrow/compute" + "github.com/apache/arrow-go/v18/arrow/memory" + "github.com/stretchr/testify/require" +) + +func TestDictionaryEncodeFixedSizeBinary(t *testing.T) { + mem := memory.NewCheckedAllocator(memory.DefaultAllocator) + defer mem.AssertSize(t, 0) + + values := [][]byte{ + {0, 0, 0, 1}, + {0, 0, 0, 2}, + {0, 0, 0, 1}, + {0, 0, 0, 3}, + {0, 0, 0, 2}, + {0, 0, 0, 3}, + } + valid := []bool{true, true, true, false, true, true} + + builder := array.NewFixedSizeBinaryBuilder(mem, &arrow.FixedSizeBinaryType{ByteWidth: 4}) + builder.AppendValues(values, valid) + input := builder.NewFixedSizeBinaryArray() + builder.Release() + defer input.Release() + + ctx := compute.WithAllocator(context.Background(), mem) + result, err := compute.DictionaryEncodeArray(ctx, compute.DictionaryEncodeOptions{}, input) + require.NoError(t, err) + defer result.Release() + + encoded := result.(*array.Dictionary) + require.Equal(t, []int32{0, 1, 0, 0, 1, 2}, encoded.Indices().(*array.Int32).Int32Values()) + require.Equal(t, 1, encoded.NullN()) + require.Equal(t, 3, encoded.Dictionary().Len()) + + dict := encoded.Dictionary().(*array.FixedSizeBinary) + require.Equal(t, values[0], dict.Value(0)) + require.Equal(t, values[1], dict.Value(1)) + require.Equal(t, values[3], dict.Value(2)) +}