Skip to content

Commit 6e4e241

Browse files
fallintoplacefallintoplace
authored
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/...`
1 parent 2120f2e commit 6e4e241

3 files changed

Lines changed: 378 additions & 12 deletions

File tree

‎arrow/array/union.go‎

Lines changed: 43 additions & 12 deletions
Original file line numberDiff line numberDiff line change
@@ -451,18 +451,20 @@ func arraySparseUnionEqual(l, r *SparseUnion) bool {
451451
childIDs := l.unionType.ChildIDs()
452452
leftCodes, rightCodes := l.RawTypeCodes(), r.RawTypeCodes()
453453

454-
for i := 0; i < l.data.length; i++ {
454+
for i := 0; i < l.data.length; {
455455
typeID := leftCodes[i]
456456
if typeID != rightCodes[i] {
457457
return false
458458
}
459459

460+
end := sparseUnionRunEnd(leftCodes, rightCodes, typeID, i, l.data.length)
460461
childNum := childIDs[typeID]
461-
eq := SliceEqual(l.children[childNum], int64(i), int64(i+1),
462-
r.children[childNum], int64(i), int64(i+1))
462+
eq := SliceEqual(l.children[childNum], int64(i), int64(end),
463+
r.children[childNum], int64(i), int64(end))
463464
if !eq {
464465
return false
465466
}
467+
i = end
466468
}
467469
return true
468470
}
@@ -471,22 +473,32 @@ func arraySparseUnionApproxEqual(l, r *SparseUnion, opt equalOption) bool {
471473
childIDs := l.unionType.ChildIDs()
472474
leftCodes, rightCodes := l.RawTypeCodes(), r.RawTypeCodes()
473475

474-
for i := 0; i < l.data.length; i++ {
476+
for i := 0; i < l.data.length; {
475477
typeID := leftCodes[i]
476478
if typeID != rightCodes[i] {
477479
return false
478480
}
479481

482+
end := sparseUnionRunEnd(leftCodes, rightCodes, typeID, i, l.data.length)
480483
childNum := childIDs[typeID]
481-
eq := sliceApproxEqual(l.children[childNum], int64(i), int64(i+1),
482-
r.children[childNum], int64(i), int64(i+1), opt)
484+
eq := sliceApproxEqual(l.children[childNum], int64(i), int64(end),
485+
r.children[childNum], int64(i), int64(end), opt)
483486
if !eq {
484487
return false
485488
}
489+
i = end
486490
}
487491
return true
488492
}
489493

494+
func sparseUnionRunEnd(leftCodes, rightCodes []arrow.UnionTypeCode, typeID arrow.UnionTypeCode, start, length int) int {
495+
end := start + 1
496+
for end < length && leftCodes[end] == typeID && rightCodes[end] == typeID {
497+
end++
498+
}
499+
return end
500+
}
501+
490502
// DenseUnion represents an array where each logical value is taken from
491503
// a single child, at a specific offset. A buffer of 8-bit type ids
492504
// indicates which child a given logical value is to be taken from and
@@ -702,18 +714,22 @@ func arrayDenseUnionEqual(l, r *DenseUnion) bool {
702714
leftCodes, rightCodes := l.RawTypeCodes(), r.RawTypeCodes()
703715
leftOffsets, rightOffsets := l.RawValueOffsets(), r.RawValueOffsets()
704716

705-
for i := 0; i < l.data.length; i++ {
717+
for i := 0; i < l.data.length; {
706718
typeID := leftCodes[i]
707719
if typeID != rightCodes[i] {
708720
return false
709721
}
710722

723+
end := denseUnionRunEnd(leftCodes, rightCodes, leftOffsets, rightOffsets, typeID, i, l.data.length)
711724
childNum := childIDs[typeID]
712-
eq := SliceEqual(l.children[childNum], int64(leftOffsets[i]), int64(leftOffsets[i]+1),
713-
r.children[childNum], int64(rightOffsets[i]), int64(rightOffsets[i]+1))
725+
leftStart, leftEnd := int64(leftOffsets[i]), int64(leftOffsets[end-1])+1
726+
rightStart, rightEnd := int64(rightOffsets[i]), int64(rightOffsets[end-1])+1
727+
eq := SliceEqual(l.children[childNum], leftStart, leftEnd,
728+
r.children[childNum], rightStart, rightEnd)
714729
if !eq {
715730
return false
716731
}
732+
i = end
717733
}
718734
return true
719735
}
@@ -723,22 +739,37 @@ func arrayDenseUnionApproxEqual(l, r *DenseUnion, opt equalOption) bool {
723739
leftCodes, rightCodes := l.RawTypeCodes(), r.RawTypeCodes()
724740
leftOffsets, rightOffsets := l.RawValueOffsets(), r.RawValueOffsets()
725741

726-
for i := 0; i < l.data.length; i++ {
742+
for i := 0; i < l.data.length; {
727743
typeID := leftCodes[i]
728744
if typeID != rightCodes[i] {
729745
return false
730746
}
731747

748+
end := denseUnionRunEnd(leftCodes, rightCodes, leftOffsets, rightOffsets, typeID, i, l.data.length)
732749
childNum := childIDs[typeID]
733-
eq := sliceApproxEqual(l.children[childNum], int64(leftOffsets[i]), int64(leftOffsets[i]+1),
734-
r.children[childNum], int64(rightOffsets[i]), int64(rightOffsets[i]+1), opt)
750+
leftStart, leftEnd := int64(leftOffsets[i]), int64(leftOffsets[end-1])+1
751+
rightStart, rightEnd := int64(rightOffsets[i]), int64(rightOffsets[end-1])+1
752+
eq := sliceApproxEqual(l.children[childNum], leftStart, leftEnd,
753+
r.children[childNum], rightStart, rightEnd, opt)
735754
if !eq {
736755
return false
737756
}
757+
i = end
738758
}
739759
return true
740760
}
741761

762+
func denseUnionRunEnd(leftCodes, rightCodes []arrow.UnionTypeCode, leftOffsets, rightOffsets []int32, typeID arrow.UnionTypeCode, start, length int) int {
763+
end := start + 1
764+
for end < length &&
765+
leftCodes[end] == typeID && rightCodes[end] == typeID &&
766+
int64(leftOffsets[end]) == int64(leftOffsets[end-1])+1 &&
767+
int64(rightOffsets[end]) == int64(rightOffsets[end-1])+1 {
768+
end++
769+
}
770+
return end
771+
}
772+
742773
// UnionBuilder is a convenience interface for building Union arrays of
743774
// either Dense or Sparse mode.
744775
type UnionBuilder interface {

‎arrow/array/union_equal_benchmark_test.go‎

Lines changed: 184 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,184 @@
1+
// Licensed to the Apache Software Foundation (ASF) under one
2+
// or more contributor license agreements. See the NOTICE file
3+
// distributed with this work for additional information
4+
// regarding copyright ownership. The ASF licenses this file
5+
// to you under the Apache License, Version 2.0 (the
6+
// "License"); you may not use this file except in compliance
7+
// with the License. You may obtain a copy of the License at
8+
//
9+
// http://www.apache.org/licenses/LICENSE-2.0
10+
//
11+
// Unless required by applicable law or agreed to in writing, software
12+
// distributed under the License is distributed on an "AS IS" BASIS,
13+
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
14+
// See the License for the specific language governing permissions and
15+
// limitations under the License.
16+
17+
package array_test
18+
19+
import (
20+
"fmt"
21+
"testing"
22+
23+
"github.com/apache/arrow-go/v18/arrow"
24+
"github.com/apache/arrow-go/v18/arrow/array"
25+
"github.com/apache/arrow-go/v18/arrow/memory"
26+
)
27+
28+
var benchmarkUnionEqualResult bool
29+
30+
func BenchmarkUnionEqual(b *testing.B) {
31+
const rows = 65536
32+
33+
for _, mode := range []string{"sparse", "dense"} {
34+
for _, pattern := range []string{"one-type", "runs-64", "alternating"} {
35+
for _, mismatch := range []bool{false, true} {
36+
for _, comparison := range []struct {
37+
name string
38+
fn func(arrow.Array, arrow.Array) bool
39+
}{
40+
{name: "equal", fn: array.Equal},
41+
{name: "approx", fn: func(left, right arrow.Array) bool {
42+
return array.ApproxEqual(left, right)
43+
}},
44+
} {
45+
name := fmt.Sprintf("%s/%s/%s/%s", mode, pattern, comparison.name, unionEqualBenchmarkMismatchName(mismatch))
46+
b.Run(name, func(b *testing.B) {
47+
left := makeUnionEqualBenchmarkArray(b, mode, pattern, rows, false)
48+
right := makeUnionEqualBenchmarkArray(b, mode, pattern, rows, mismatch)
49+
defer left.Release()
50+
defer right.Release()
51+
52+
b.ReportAllocs()
53+
b.SetBytes(int64(rows))
54+
b.ResetTimer()
55+
for i := 0; i < b.N; i++ {
56+
benchmarkUnionEqualResult = comparison.fn(left, right)
57+
}
58+
})
59+
}
60+
}
61+
}
62+
}
63+
}
64+
65+
func unionEqualBenchmarkMismatchName(mismatch bool) string {
66+
if mismatch {
67+
return "mismatch-last"
68+
}
69+
return "equal"
70+
}
71+
72+
func makeUnionEqualBenchmarkArray(b *testing.B, mode, pattern string, rows int, mismatch bool) arrow.Array {
73+
b.Helper()
74+
75+
typeIDs := make([]int8, rows)
76+
offsets := make([]int32, rows)
77+
sparseInts := make([]int32, rows)
78+
sparseStrings := make([]string, rows)
79+
denseInts := make([]int32, 0, rows)
80+
denseStrings := make([]string, 0, rows)
81+
childOffsets := [2]int32{}
82+
83+
for i := 0; i < rows; i++ {
84+
childID := unionEqualBenchmarkChildID(pattern, i)
85+
typeIDs[i] = int8(childID)
86+
87+
if mode == "sparse" {
88+
sparseInts[i] = int32(i)
89+
sparseStrings[i] = fmt.Sprintf("value-%d", i)
90+
if mismatch && i == rows-1 {
91+
if childID == 0 {
92+
sparseInts[i]++
93+
} else {
94+
sparseStrings[i] = "different"
95+
}
96+
}
97+
continue
98+
}
99+
100+
offsets[i] = childOffsets[childID]
101+
if childID == 0 {
102+
value := int32(i)
103+
if mismatch && i == rows-1 {
104+
value++
105+
}
106+
denseInts = append(denseInts, value)
107+
} else {
108+
value := fmt.Sprintf("value-%d", i)
109+
if mismatch && i == rows-1 {
110+
value = "different"
111+
}
112+
denseStrings = append(denseStrings, value)
113+
}
114+
childOffsets[childID]++
115+
}
116+
117+
typeIDsArray := makeUnionEqualBenchmarkInt8Array(b, typeIDs)
118+
defer typeIDsArray.Release()
119+
if mode == "sparse" {
120+
intArray := makeUnionEqualBenchmarkInt32Array(b, sparseInts)
121+
defer intArray.Release()
122+
stringArray := makeUnionEqualBenchmarkStringArray(b, sparseStrings)
123+
defer stringArray.Release()
124+
125+
result, err := array.NewSparseUnionFromArrays(typeIDsArray, []arrow.Array{intArray, stringArray})
126+
if err != nil {
127+
b.Fatal(err)
128+
}
129+
return result
130+
}
131+
132+
offsetsArray := makeUnionEqualBenchmarkInt32Array(b, offsets)
133+
defer offsetsArray.Release()
134+
intArray := makeUnionEqualBenchmarkInt32Array(b, denseInts)
135+
defer intArray.Release()
136+
stringArray := makeUnionEqualBenchmarkStringArray(b, denseStrings)
137+
defer stringArray.Release()
138+
139+
result, err := array.NewDenseUnionFromArrays(typeIDsArray, offsetsArray, []arrow.Array{intArray, stringArray})
140+
if err != nil {
141+
b.Fatal(err)
142+
}
143+
return result
144+
}
145+
146+
func unionEqualBenchmarkChildID(pattern string, index int) int {
147+
switch pattern {
148+
case "one-type":
149+
return 0
150+
case "runs-64":
151+
return (index / 64) % 2
152+
case "alternating":
153+
return index % 2
154+
default:
155+
panic("unsupported union equality benchmark pattern")
156+
}
157+
}
158+
159+
func makeUnionEqualBenchmarkInt8Array(b *testing.B, values []int8) arrow.Array {
160+
b.Helper()
161+
builder := array.NewInt8Builder(memory.DefaultAllocator)
162+
builder.AppendValues(values, nil)
163+
result := builder.NewInt8Array()
164+
builder.Release()
165+
return result
166+
}
167+
168+
func makeUnionEqualBenchmarkInt32Array(b *testing.B, values []int32) arrow.Array {
169+
b.Helper()
170+
builder := array.NewInt32Builder(memory.DefaultAllocator)
171+
builder.AppendValues(values, nil)
172+
result := builder.NewInt32Array()
173+
builder.Release()
174+
return result
175+
}
176+
177+
func makeUnionEqualBenchmarkStringArray(b *testing.B, values []string) arrow.Array {
178+
b.Helper()
179+
builder := array.NewStringBuilder(memory.DefaultAllocator)
180+
builder.AppendValues(values, nil)
181+
result := builder.NewStringArray()
182+
builder.Release()
183+
return result
184+
}

0 commit comments

Comments
 (0)