This is an automated email from the ASF dual-hosted git repository.
zeroshade pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/arrow-go.git
The following commit(s) were added to refs/heads/main by this push:
new 0e8c7d13 perf(compute): use typed memo insertion for numeric hashes
(#1251)
0e8c7d13 is described below
commit 0e8c7d132276a46703859b7a492892e496b8158d
Author: Minh Vu <[email protected]>
AuthorDate: Mon Aug 31 20:24:09 2026 +0200
perf(compute): use typed memo insertion for numeric hashes (#1251)
## Summary
- Use typed memo insertion in the shared numeric hash path.
- Keep the existing normalized uint8, uint16, uint32, and uint64
mappings.
- Avoid the per-value interface call to `MemoTable.GetOrInsert`.
- Add coverage for int32, int64, float32, float64, and both null modes.
- Add a numeric `dictionary_encode` benchmark.
## Benchmark
Apple M1 Pro. Compared with main at `6b039a76`. The benchmark uses
65,535 values, 100 unique values, 500ms per sample, and 5 samples. The
table shows median results.
| type | main | branch | change | allocs/op |
| --- | ---: | ---: | ---: | ---: |
| int32 | 1.168 ms | 0.989 ms | 15.3% faster | 56 -> 56 |
| int64 | 1.138 ms | 1.016 ms | 10.8% faster | 56 -> 56 |
| float32 | 1.662 ms | 1.114 ms | 33.0% faster | 64,935 -> 56 |
| float64 | 1.707 ms | 1.251 ms | 26.7% faster | 64,935 -> 56 |
Command:
```text
go test ./arrow/compute -run '^$' -bench
'^BenchmarkDictionaryEncodeNumeric$' -benchmem -benchtime=500ms -count=5
```
## Checks
- `go test ./arrow/compute/...`
- `go test -race ./arrow/compute/internal/kernels ./arrow/compute -run
'^(TestDictionaryEncodeNumericTypes|TestDictionaryEncode|TestDictionaryEncodeArrayWithSmallExecChunkSize|TestDictionaryEncodeResizesMemoTable|TestDictionaryEncodeStateResetAfterFinalize)$'
-count=1`
- `go vet ./arrow/compute/internal/kernels`
- `git diff --check`
---
arrow/compute/internal/kernels/vector_hash.go | 5 +-
arrow/compute/vector_hash_numeric_bench_test.go | 104 +++++++++++++++++++
arrow/compute/vector_hash_numeric_test.go | 126 ++++++++++++++++++++++++
3 files changed, 233 insertions(+), 2 deletions(-)
diff --git a/arrow/compute/internal/kernels/vector_hash.go
b/arrow/compute/internal/kernels/vector_hash.go
index 00f48173..ecad7e98 100644
--- a/arrow/compute/internal/kernels/vector_hash.go
+++ b/arrow/compute/internal/kernels/vector_hash.go
@@ -356,12 +356,13 @@ func doAppendFixedSize(action Action, memo
hashing.MemoTable, arr *exec.ArraySpa
})
}
-func doAppendNumeric[T arrow.IntType | arrow.UintType |
arrow.FloatType](action Action, memo hashing.MemoTable, arr *exec.ArraySpan)
error {
+func doAppendNumeric[T uint8 | uint16 | uint32 | uint64](action Action, memo
hashing.MemoTable, arr *exec.ArraySpan) error {
arrData := exec.GetSpanValues[T](arr, 1)
shouldEncodeNulls := action.ShouldEncodeNulls()
+ typedMemo := memo.(hashing.TypedMemoTable[T])
return bitutils.VisitBitBlocksShort(arr.Buffers[0].Buf, arr.Offset,
arr.Len,
func(pos int64) error {
- idx, found, err := memo.GetOrInsert(arrData[pos])
+ idx, found, err := typedMemo.InsertOrGet(arrData[pos])
if err != nil {
return err
}
diff --git a/arrow/compute/vector_hash_numeric_bench_test.go
b/arrow/compute/vector_hash_numeric_bench_test.go
new file mode 100644
index 00000000..872d6a44
--- /dev/null
+++ b/arrow/compute/vector_hash_numeric_bench_test.go
@@ -0,0 +1,104 @@
+// 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"
+)
+
+func BenchmarkDictionaryEncodeNumeric(b *testing.B) {
+ const (
+ nvalues = 65535
+ nunique = 100
+ )
+
+ mem := memory.DefaultAllocator
+ ctx := compute.WithAllocator(context.Background(), mem)
+
+ b.Run("int32", func(b *testing.B) {
+ builder := array.NewInt32Builder(mem)
+ values := make([]int32, nvalues)
+ for i := range values {
+ values[i] = int32(i % nunique)
+ }
+ builder.AppendValues(values, nil)
+ input := builder.NewInt32Array()
+ builder.Release()
+ defer input.Release()
+ benchmarkDictionaryEncodeNumeric(b, ctx, input,
arrow.Int32SizeBytes)
+ })
+
+ b.Run("int64", func(b *testing.B) {
+ builder := array.NewInt64Builder(mem)
+ values := make([]int64, nvalues)
+ for i := range values {
+ values[i] = int64(i % nunique)
+ }
+ builder.AppendValues(values, nil)
+ input := builder.NewInt64Array()
+ builder.Release()
+ defer input.Release()
+ benchmarkDictionaryEncodeNumeric(b, ctx, input,
arrow.Int64SizeBytes)
+ })
+
+ b.Run("float32", func(b *testing.B) {
+ builder := array.NewFloat32Builder(mem)
+ values := make([]float32, nvalues)
+ for i := range values {
+ values[i] = float32(i % nunique)
+ }
+ builder.AppendValues(values, nil)
+ input := builder.NewFloat32Array()
+ builder.Release()
+ defer input.Release()
+ benchmarkDictionaryEncodeNumeric(b, ctx, input,
arrow.Float32SizeBytes)
+ })
+
+ b.Run("float64", func(b *testing.B) {
+ builder := array.NewFloat64Builder(mem)
+ values := make([]float64, nvalues)
+ for i := range values {
+ values[i] = float64(i % nunique)
+ }
+ builder.AppendValues(values, nil)
+ input := builder.NewFloat64Array()
+ builder.Release()
+ defer input.Release()
+ benchmarkDictionaryEncodeNumeric(b, ctx, input,
arrow.Float64SizeBytes)
+ })
+}
+
+func benchmarkDictionaryEncodeNumeric(b *testing.B, ctx context.Context, input
arrow.Array, valueSize int) {
+ b.ReportAllocs()
+ b.SetBytes(int64(input.Len() * valueSize))
+ b.ResetTimer()
+ for i := 0; i < b.N; i++ {
+ result, err := compute.DictionaryEncodeArray(ctx,
compute.DictionaryEncodeOptions{}, input)
+ if err != nil {
+ b.Fatal(err)
+ }
+ result.Release()
+ }
+}
diff --git a/arrow/compute/vector_hash_numeric_test.go
b/arrow/compute/vector_hash_numeric_test.go
new file mode 100644
index 00000000..db6cb541
--- /dev/null
+++ b/arrow/compute/vector_hash_numeric_test.go
@@ -0,0 +1,126 @@
+// 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"
+ "strings"
+ "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 TestDictionaryEncodeNumericTypes(t *testing.T) {
+ tests := []struct {
+ name string
+ typ arrow.DataType
+ input string
+ masked string
+ encoded string
+ }{
+ {
+ name: "int32",
+ typ: arrow.PrimitiveTypes.Int32,
+ input: "[-3, 1, -3, null, 2]",
+ masked: "[-3, 1, 2]",
+ encoded: "[-3, 1, null, 2]",
+ },
+ {
+ name: "int64",
+ typ: arrow.PrimitiveTypes.Int64,
+ input: "[-3, 1, -3, null, 2]",
+ masked: "[-3, 1, 2]",
+ encoded: "[-3, 1, null, 2]",
+ },
+ {
+ name: "float32",
+ typ: arrow.PrimitiveTypes.Float32,
+ input: "[-3.5, 1.25, -3.5, null, 2.75]",
+ masked: "[-3.5, 1.25, 2.75]",
+ encoded: "[-3.5, 1.25, null, 2.75]",
+ },
+ {
+ name: "float64",
+ typ: arrow.PrimitiveTypes.Float64,
+ input: "[-3.5, 1.25, -3.5, null, 2.75]",
+ masked: "[-3.5, 1.25, 2.75]",
+ encoded: "[-3.5, 1.25, null, 2.75]",
+ },
+ }
+
+ mem := memory.NewCheckedAllocator(memory.DefaultAllocator)
+ defer mem.AssertSize(t, 0)
+ ctx := compute.WithAllocator(context.Background(), mem)
+
+ for _, tc := range tests {
+ t.Run(tc.name, func(t *testing.T) {
+ input, _, err := array.FromJSON(mem, tc.typ,
strings.NewReader(tc.input))
+ require.NoError(t, err)
+ defer input.Release()
+
+ maskedExpected, _, err := array.FromJSON(mem, tc.typ,
strings.NewReader(tc.masked))
+ require.NoError(t, err)
+ defer maskedExpected.Release()
+
+ encodedExpected, _, err := array.FromJSON(mem, tc.typ,
strings.NewReader(tc.encoded))
+ require.NoError(t, err)
+ defer encodedExpected.Release()
+
+ for _, mode := range []struct {
+ name string
+ nullEncoding compute.NullEncodingBehavior
+ expectedDict arrow.Array
+ expectedIndex []int32
+ nullCount int
+ }{
+ {
+ name: "mask nulls",
+ nullEncoding: compute.NullEncodingMask,
+ expectedDict: maskedExpected,
+ expectedIndex: []int32{0, 1, 0, 0, 2},
+ nullCount: 1,
+ },
+ {
+ name: "encode nulls",
+ nullEncoding:
compute.NullEncodingEncode,
+ expectedDict: encodedExpected,
+ expectedIndex: []int32{0, 1, 0, 2, 3},
+ nullCount: 0,
+ },
+ } {
+ t.Run(mode.name, func(t *testing.T) {
+ result, err :=
compute.DictionaryEncodeArray(ctx, compute.DictionaryEncodeOptions{
+ NullEncoding: mode.nullEncoding,
+ }, input)
+ require.NoError(t, err)
+ defer result.Release()
+
+ encoded := result.(*array.Dictionary)
+ require.True(t,
array.Equal(mode.expectedDict, encoded.Dictionary()))
+ require.Equal(t, mode.expectedIndex,
encoded.Indices().(*array.Int32).Int32Values())
+ require.Equal(t, mode.nullCount,
encoded.NullN())
+ })
+ }
+ })
+ }
+}
