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 6e4e241c perf(arrow/array): compare union values in runs (#1244)
6e4e241c is described below
commit 6e4e241c4e6c85b149632b2070d0ef55b6eb982d
Author: Minh Vu <[email protected]>
AuthorDate: Fri Aug 28 23:46:31 2026 +0200
perf(arrow/array): compare union values in runs (#1244)
## Summary
- **Batch consecutive union values with the same type code.**
- Sparse unions compare one child slice per run.
- Dense unions batch only when both child offset sequences are
contiguous.
- Apply the same path to exact and approximate equality.
- Add sliced, null, skipped-offset, mismatch, and benchmark coverage.
## Benchmark
Local Apple M1 Pro run with 65,536 rows and 64-row type runs:
- Sparse `Equal`: **14.8 ms -> 0.95 ms**, **262,150 -> 4,102 allocs/op**
- Dense `Equal`: **13.9 ms -> 1.04 ms**, **262,150 -> 4,102 allocs/op**
- Alternating type codes stay roughly flat, as expected.
## Tests
- `go test ./...`
- `go vet ./arrow/array`
- `GOOS=linux GOARCH=386 go build ./arrow/array/...`
---
arrow/array/union.go | 55 +++++++--
arrow/array/union_equal_benchmark_test.go | 184 ++++++++++++++++++++++++++++++
arrow/array/union_equal_test.go | 151 ++++++++++++++++++++++++
3 files changed, 378 insertions(+), 12 deletions(-)
diff --git a/arrow/array/union.go b/arrow/array/union.go
index 061845b9..17138f8b 100644
--- a/arrow/array/union.go
+++ b/arrow/array/union.go
@@ -451,18 +451,20 @@ func arraySparseUnionEqual(l, r *SparseUnion) bool {
childIDs := l.unionType.ChildIDs()
leftCodes, rightCodes := l.RawTypeCodes(), r.RawTypeCodes()
- for i := 0; i < l.data.length; i++ {
+ for i := 0; i < l.data.length; {
typeID := leftCodes[i]
if typeID != rightCodes[i] {
return false
}
+ end := sparseUnionRunEnd(leftCodes, rightCodes, typeID, i,
l.data.length)
childNum := childIDs[typeID]
- eq := SliceEqual(l.children[childNum], int64(i), int64(i+1),
- r.children[childNum], int64(i), int64(i+1))
+ eq := SliceEqual(l.children[childNum], int64(i), int64(end),
+ r.children[childNum], int64(i), int64(end))
if !eq {
return false
}
+ i = end
}
return true
}
@@ -471,22 +473,32 @@ func arraySparseUnionApproxEqual(l, r *SparseUnion, opt
equalOption) bool {
childIDs := l.unionType.ChildIDs()
leftCodes, rightCodes := l.RawTypeCodes(), r.RawTypeCodes()
- for i := 0; i < l.data.length; i++ {
+ for i := 0; i < l.data.length; {
typeID := leftCodes[i]
if typeID != rightCodes[i] {
return false
}
+ end := sparseUnionRunEnd(leftCodes, rightCodes, typeID, i,
l.data.length)
childNum := childIDs[typeID]
- eq := sliceApproxEqual(l.children[childNum], int64(i),
int64(i+1),
- r.children[childNum], int64(i), int64(i+1), opt)
+ eq := sliceApproxEqual(l.children[childNum], int64(i),
int64(end),
+ r.children[childNum], int64(i), int64(end), opt)
if !eq {
return false
}
+ i = end
}
return true
}
+func sparseUnionRunEnd(leftCodes, rightCodes []arrow.UnionTypeCode, typeID
arrow.UnionTypeCode, start, length int) int {
+ end := start + 1
+ for end < length && leftCodes[end] == typeID && rightCodes[end] ==
typeID {
+ end++
+ }
+ return end
+}
+
// DenseUnion represents an array where each logical value is taken from
// a single child, at a specific offset. A buffer of 8-bit type ids
// indicates which child a given logical value is to be taken from and
@@ -702,18 +714,22 @@ func arrayDenseUnionEqual(l, r *DenseUnion) bool {
leftCodes, rightCodes := l.RawTypeCodes(), r.RawTypeCodes()
leftOffsets, rightOffsets := l.RawValueOffsets(), r.RawValueOffsets()
- for i := 0; i < l.data.length; i++ {
+ for i := 0; i < l.data.length; {
typeID := leftCodes[i]
if typeID != rightCodes[i] {
return false
}
+ end := denseUnionRunEnd(leftCodes, rightCodes, leftOffsets,
rightOffsets, typeID, i, l.data.length)
childNum := childIDs[typeID]
- eq := SliceEqual(l.children[childNum], int64(leftOffsets[i]),
int64(leftOffsets[i]+1),
- r.children[childNum], int64(rightOffsets[i]),
int64(rightOffsets[i]+1))
+ leftStart, leftEnd := int64(leftOffsets[i]),
int64(leftOffsets[end-1])+1
+ rightStart, rightEnd := int64(rightOffsets[i]),
int64(rightOffsets[end-1])+1
+ eq := SliceEqual(l.children[childNum], leftStart, leftEnd,
+ r.children[childNum], rightStart, rightEnd)
if !eq {
return false
}
+ i = end
}
return true
}
@@ -723,22 +739,37 @@ func arrayDenseUnionApproxEqual(l, r *DenseUnion, opt
equalOption) bool {
leftCodes, rightCodes := l.RawTypeCodes(), r.RawTypeCodes()
leftOffsets, rightOffsets := l.RawValueOffsets(), r.RawValueOffsets()
- for i := 0; i < l.data.length; i++ {
+ for i := 0; i < l.data.length; {
typeID := leftCodes[i]
if typeID != rightCodes[i] {
return false
}
+ end := denseUnionRunEnd(leftCodes, rightCodes, leftOffsets,
rightOffsets, typeID, i, l.data.length)
childNum := childIDs[typeID]
- eq := sliceApproxEqual(l.children[childNum],
int64(leftOffsets[i]), int64(leftOffsets[i]+1),
- r.children[childNum], int64(rightOffsets[i]),
int64(rightOffsets[i]+1), opt)
+ leftStart, leftEnd := int64(leftOffsets[i]),
int64(leftOffsets[end-1])+1
+ rightStart, rightEnd := int64(rightOffsets[i]),
int64(rightOffsets[end-1])+1
+ eq := sliceApproxEqual(l.children[childNum], leftStart, leftEnd,
+ r.children[childNum], rightStart, rightEnd, opt)
if !eq {
return false
}
+ i = end
}
return true
}
+func denseUnionRunEnd(leftCodes, rightCodes []arrow.UnionTypeCode,
leftOffsets, rightOffsets []int32, typeID arrow.UnionTypeCode, start, length
int) int {
+ end := start + 1
+ for end < length &&
+ leftCodes[end] == typeID && rightCodes[end] == typeID &&
+ int64(leftOffsets[end]) == int64(leftOffsets[end-1])+1 &&
+ int64(rightOffsets[end]) == int64(rightOffsets[end-1])+1 {
+ end++
+ }
+ return end
+}
+
// UnionBuilder is a convenience interface for building Union arrays of
// either Dense or Sparse mode.
type UnionBuilder interface {
diff --git a/arrow/array/union_equal_benchmark_test.go
b/arrow/array/union_equal_benchmark_test.go
new file mode 100644
index 00000000..7d04c562
--- /dev/null
+++ b/arrow/array/union_equal_benchmark_test.go
@@ -0,0 +1,184 @@
+// 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"
+ "testing"
+
+ "github.com/apache/arrow-go/v18/arrow"
+ "github.com/apache/arrow-go/v18/arrow/array"
+ "github.com/apache/arrow-go/v18/arrow/memory"
+)
+
+var benchmarkUnionEqualResult bool
+
+func BenchmarkUnionEqual(b *testing.B) {
+ const rows = 65536
+
+ for _, mode := range []string{"sparse", "dense"} {
+ for _, pattern := range []string{"one-type", "runs-64",
"alternating"} {
+ for _, mismatch := range []bool{false, true} {
+ for _, comparison := range []struct {
+ name string
+ fn func(arrow.Array, arrow.Array) bool
+ }{
+ {name: "equal", fn: array.Equal},
+ {name: "approx", fn: func(left, right
arrow.Array) bool {
+ return array.ApproxEqual(left,
right)
+ }},
+ } {
+ name := fmt.Sprintf("%s/%s/%s/%s",
mode, pattern, comparison.name, unionEqualBenchmarkMismatchName(mismatch))
+ b.Run(name, func(b *testing.B) {
+ left :=
makeUnionEqualBenchmarkArray(b, mode, pattern, rows, false)
+ right :=
makeUnionEqualBenchmarkArray(b, mode, pattern, rows, mismatch)
+ defer left.Release()
+ defer right.Release()
+
+ b.ReportAllocs()
+ b.SetBytes(int64(rows))
+ b.ResetTimer()
+ for i := 0; i < b.N; i++ {
+
benchmarkUnionEqualResult = comparison.fn(left, right)
+ }
+ })
+ }
+ }
+ }
+ }
+}
+
+func unionEqualBenchmarkMismatchName(mismatch bool) string {
+ if mismatch {
+ return "mismatch-last"
+ }
+ return "equal"
+}
+
+func makeUnionEqualBenchmarkArray(b *testing.B, mode, pattern string, rows
int, mismatch bool) arrow.Array {
+ b.Helper()
+
+ typeIDs := make([]int8, rows)
+ offsets := make([]int32, rows)
+ sparseInts := make([]int32, rows)
+ sparseStrings := make([]string, rows)
+ denseInts := make([]int32, 0, rows)
+ denseStrings := make([]string, 0, rows)
+ childOffsets := [2]int32{}
+
+ for i := 0; i < rows; i++ {
+ childID := unionEqualBenchmarkChildID(pattern, i)
+ typeIDs[i] = int8(childID)
+
+ if mode == "sparse" {
+ sparseInts[i] = int32(i)
+ sparseStrings[i] = fmt.Sprintf("value-%d", i)
+ if mismatch && i == rows-1 {
+ if childID == 0 {
+ sparseInts[i]++
+ } else {
+ sparseStrings[i] = "different"
+ }
+ }
+ continue
+ }
+
+ offsets[i] = childOffsets[childID]
+ if childID == 0 {
+ value := int32(i)
+ if mismatch && i == rows-1 {
+ value++
+ }
+ denseInts = append(denseInts, value)
+ } else {
+ value := fmt.Sprintf("value-%d", i)
+ if mismatch && i == rows-1 {
+ value = "different"
+ }
+ denseStrings = append(denseStrings, value)
+ }
+ childOffsets[childID]++
+ }
+
+ typeIDsArray := makeUnionEqualBenchmarkInt8Array(b, typeIDs)
+ defer typeIDsArray.Release()
+ if mode == "sparse" {
+ intArray := makeUnionEqualBenchmarkInt32Array(b, sparseInts)
+ defer intArray.Release()
+ stringArray := makeUnionEqualBenchmarkStringArray(b,
sparseStrings)
+ defer stringArray.Release()
+
+ result, err := array.NewSparseUnionFromArrays(typeIDsArray,
[]arrow.Array{intArray, stringArray})
+ if err != nil {
+ b.Fatal(err)
+ }
+ return result
+ }
+
+ offsetsArray := makeUnionEqualBenchmarkInt32Array(b, offsets)
+ defer offsetsArray.Release()
+ intArray := makeUnionEqualBenchmarkInt32Array(b, denseInts)
+ defer intArray.Release()
+ stringArray := makeUnionEqualBenchmarkStringArray(b, denseStrings)
+ defer stringArray.Release()
+
+ result, err := array.NewDenseUnionFromArrays(typeIDsArray,
offsetsArray, []arrow.Array{intArray, stringArray})
+ if err != nil {
+ b.Fatal(err)
+ }
+ return result
+}
+
+func unionEqualBenchmarkChildID(pattern string, index int) int {
+ switch pattern {
+ case "one-type":
+ return 0
+ case "runs-64":
+ return (index / 64) % 2
+ case "alternating":
+ return index % 2
+ default:
+ panic("unsupported union equality benchmark pattern")
+ }
+}
+
+func makeUnionEqualBenchmarkInt8Array(b *testing.B, values []int8) arrow.Array
{
+ b.Helper()
+ builder := array.NewInt8Builder(memory.DefaultAllocator)
+ builder.AppendValues(values, nil)
+ result := builder.NewInt8Array()
+ builder.Release()
+ return result
+}
+
+func makeUnionEqualBenchmarkInt32Array(b *testing.B, values []int32)
arrow.Array {
+ b.Helper()
+ builder := array.NewInt32Builder(memory.DefaultAllocator)
+ builder.AppendValues(values, nil)
+ result := builder.NewInt32Array()
+ builder.Release()
+ return result
+}
+
+func makeUnionEqualBenchmarkStringArray(b *testing.B, values []string)
arrow.Array {
+ b.Helper()
+ builder := array.NewStringBuilder(memory.DefaultAllocator)
+ builder.AppendValues(values, nil)
+ result := builder.NewStringArray()
+ builder.Release()
+ return result
+}
diff --git a/arrow/array/union_equal_test.go b/arrow/array/union_equal_test.go
new file mode 100644
index 00000000..e997b970
--- /dev/null
+++ b/arrow/array/union_equal_test.go
@@ -0,0 +1,151 @@
+// 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 (
+ "strings"
+ "testing"
+
+ "github.com/apache/arrow-go/v18/arrow"
+ "github.com/apache/arrow-go/v18/arrow/array"
+ "github.com/apache/arrow-go/v18/arrow/memory"
+ "github.com/stretchr/testify/assert"
+ "github.com/stretchr/testify/require"
+)
+
+func TestSparseUnionEqualRuns(t *testing.T) {
+ mem := memory.NewCheckedAllocator(memory.DefaultAllocator)
+ defer mem.AssertSize(t, 0)
+
+ left := newSparseUnionEqualTestArray(t, mem,
+ `[0, 0, 0, 1, 1, 0, 0, 1]`,
+ `[10, null, 12, 13, 14, 15, 16, 17]`,
+ `["a", "b", "c", "d", "e", "f", "g", "h"]`)
+ defer left.Release()
+ right := newSparseUnionEqualTestArray(t, mem,
+ `[0, 0, 0, 1, 1, 0, 0, 1]`,
+ `[10, null, 12, 13, 14, 15, 16, 17]`,
+ `["a", "b", "c", "d", "e", "f", "g", "h"]`)
+ defer right.Release()
+
+ assert.True(t, array.Equal(left, right))
+ assert.True(t, array.ApproxEqual(left, right))
+
+ leftSlice := array.NewSlice(left, 1, 7)
+ defer leftSlice.Release()
+ rightSlice := array.NewSlice(right, 1, 7)
+ defer rightSlice.Release()
+ assert.True(t, array.Equal(leftSlice, rightSlice))
+ assert.True(t, array.ApproxEqual(leftSlice, rightSlice))
+
+ different := newSparseUnionEqualTestArray(t, mem,
+ `[0, 0, 0, 1, 1, 0, 0, 1]`,
+ `[10, 99, 12, 13, 14, 15, 16, 17]`,
+ `["a", "b", "c", "d", "e", "f", "g", "h"]`)
+ defer different.Release()
+ assert.False(t, array.Equal(left, different))
+ assert.False(t, array.ApproxEqual(left, different))
+}
+
+func TestDenseUnionEqualRuns(t *testing.T) {
+ mem := memory.NewCheckedAllocator(memory.DefaultAllocator)
+ defer mem.AssertSize(t, 0)
+
+ left := newDenseUnionEqualTestArray(t, mem,
+ `[0, 0, 0, 1, 1, 0]`,
+ `[0, 2, 3, 0, 1, 4]`,
+ `[10, 99, 20, null, 40]`,
+ `["a", "b"]`)
+ defer left.Release()
+ right := newDenseUnionEqualTestArray(t, mem,
+ `[0, 0, 0, 1, 1, 0]`,
+ `[0, 2, 3, 0, 1, 4]`,
+ `[10, 88, 20, null, 40]`,
+ `["a", "b"]`)
+ defer right.Release()
+
+ assert.True(t, array.Equal(left, right))
+ assert.True(t, array.ApproxEqual(left, right))
+
+ leftSlice := array.NewSlice(left, 1, 5)
+ defer leftSlice.Release()
+ rightSlice := array.NewSlice(right, 1, 5)
+ defer rightSlice.Release()
+ assert.True(t, array.Equal(leftSlice, rightSlice))
+ assert.True(t, array.ApproxEqual(leftSlice, rightSlice))
+
+ leftWithDifferentOffsets := newDenseUnionEqualTestArray(t, mem,
+ `[0, 0, 0, 0]`,
+ `[0, 1, 2, 3]`,
+ `[10, 20, 30, 40]`,
+ `[]`)
+ defer leftWithDifferentOffsets.Release()
+ rightWithDifferentOffsets := newDenseUnionEqualTestArray(t, mem,
+ `[0, 0, 0, 0]`,
+ `[1, 3, 4, 5]`,
+ `[99, 10, 88, 20, 30, 40]`,
+ `[]`)
+ defer rightWithDifferentOffsets.Release()
+ assert.True(t, array.Equal(leftWithDifferentOffsets,
rightWithDifferentOffsets))
+ assert.True(t, array.ApproxEqual(leftWithDifferentOffsets,
rightWithDifferentOffsets))
+
+ different := newDenseUnionEqualTestArray(t, mem,
+ `[0, 0, 0, 1, 1, 0]`,
+ `[0, 2, 3, 0, 1, 4]`,
+ `[10, 88, 21, null, 40]`,
+ `["a", "b"]`)
+ defer different.Release()
+ assert.False(t, array.Equal(left, different))
+ assert.False(t, array.ApproxEqual(left, different))
+}
+
+func newSparseUnionEqualTestArray(t *testing.T, mem memory.Allocator, typeIDs,
ints, stringsJSON string) *array.SparseUnion {
+ t.Helper()
+ typeIDsArray := newUnionEqualTestArray(t, mem,
arrow.PrimitiveTypes.Int8, typeIDs)
+ defer typeIDsArray.Release()
+ intArray := newUnionEqualTestArray(t, mem, arrow.PrimitiveTypes.Int32,
ints)
+ defer intArray.Release()
+ stringArray := newUnionEqualTestArray(t, mem, arrow.BinaryTypes.String,
stringsJSON)
+ defer stringArray.Release()
+
+ result, err := array.NewSparseUnionFromArrays(typeIDsArray,
[]arrow.Array{intArray, stringArray})
+ require.NoError(t, err)
+ return result
+}
+
+func newDenseUnionEqualTestArray(t *testing.T, mem memory.Allocator, typeIDs,
offsets, ints, stringsJSON string) *array.DenseUnion {
+ t.Helper()
+ typeIDsArray := newUnionEqualTestArray(t, mem,
arrow.PrimitiveTypes.Int8, typeIDs)
+ defer typeIDsArray.Release()
+ offsetsArray := newUnionEqualTestArray(t, mem,
arrow.PrimitiveTypes.Int32, offsets)
+ defer offsetsArray.Release()
+ intArray := newUnionEqualTestArray(t, mem, arrow.PrimitiveTypes.Int32,
ints)
+ defer intArray.Release()
+ stringArray := newUnionEqualTestArray(t, mem, arrow.BinaryTypes.String,
stringsJSON)
+ defer stringArray.Release()
+
+ result, err := array.NewDenseUnionFromArrays(typeIDsArray,
offsetsArray, []arrow.Array{intArray, stringArray})
+ require.NoError(t, err)
+ return result
+}
+
+func newUnionEqualTestArray(t *testing.T, mem memory.Allocator, dtype
arrow.DataType, values string) arrow.Array {
+ t.Helper()
+ result, _, err := array.FromJSON(mem, dtype, strings.NewReader(values))
+ require.NoError(t, err)
+ return result
+}
