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 db47815d perf(arrow/array): batch union null and empty appends (#1216)
db47815d is described below
commit db47815d919ffdec0f9433365180d1a881406262
Author: Minh Vu <[email protected]>
AuthorDate: Fri Aug 28 22:15:56 2026 +0200
perf(arrow/array): batch union null and empty appends (#1216)
## Summary
- **Batch sparse union appends** through the child bulk APIs.
- **Write repeated type IDs and dense offsets** directly into their
buffers.
- **Keep dense unions compact** by sharing one physical null or empty
child value.
- **Treat non-positive counts as no-ops** for both union modes.
- Add parity, edge-case, and benchmark coverage.
## Benchmarks
Apple M1 Pro, darwin/arm64, 65,536 rows. Times are rounded medians from
3 runs.
| Mode | Operation | Before | After | Before allocs | After allocs |
| --- | --- | ---: | ---: | ---: | ---: |
| Sparse | nulls | ~1.83 ms | ~0.35 ms | 44 | 34 |
| Sparse | empty | ~1.83 ms | ~0.33 ms | 44 | 34 |
| Dense | nulls | ~0.55 ms | ~0.14 ms | 51 | 29 |
| Dense | empty | ~0.39 ms | ~0.14 ms | ~51 | 29 |
The benchmark uses the same builder setup before and after. Small row
counts stay close to benchmark noise.
## Tests
- `go test ./arrow/array -count=1`
- `go test -race ./arrow/array -run
'Test(UnionBuilderBulkAppendNullsAndEmptyValues|UnionBuilderZeroBulkAppendDoesNotMutateChildren|UnionBuilderBulkAppendMatchesScalar)$'
-count=1`
- `go vet ./arrow/array`
- `PARQUET_TEST_DATA="$PWD/parquet-testing/data" go test -count=1 ./...`
---
arrow/array/union.go | 68 ++++++---
arrow/array/union_builder_bulk_test.go | 257 +++++++++++++++++++++++++++++++++
2 files changed, 306 insertions(+), 19 deletions(-)
diff --git a/arrow/array/union.go b/arrow/array/union.go
index e6137eae..061845b9 100644
--- a/arrow/array/union.go
+++ b/arrow/array/union.go
@@ -934,6 +934,34 @@ func (b *unionBuilder) newData() *Data {
return NewData(b.Type(), length, []*memory.Buffer{nil, typesBuffer},
childData, 0, 0)
}
+func unsafeAppendRepeatedInt8(b *int8BufferBuilder, value int8, n int) {
+ if n <= 0 {
+ return
+ }
+
+ end := b.length + n
+ if b.capacity < end {
+ b.resize(bitutil.NextPowerOf2(end))
+ }
+ memory.Set(b.bytes[b.length:end], byte(value))
+ b.length = end
+}
+
+func unsafeAppendRepeatedInt32(b *int32BufferBuilder, value int32, n int) {
+ if n <= 0 {
+ return
+ }
+
+ end := b.length + n*arrow.Int32SizeBytes
+ if b.capacity < end {
+ b.resize(bitutil.NextPowerOf2(end))
+ }
+ for i := b.length; i < end; i += arrow.Int32SizeBytes {
+ arrow.Int32Traits.PutValue(b.bytes[i:], value)
+ }
+ b.length = end
+}
+
// SparseUnionBuilder is used to build a Sparse Union array using the Append
// methods. You can also add new types to the union on the fly by using
// AppendChild.
@@ -1002,17 +1030,20 @@ func (b *SparseUnionBuilder) AppendNull() {
// AppendNulls is identical to calling AppendNull() n times, except
// it will pre-allocate with reserve for all the nulls beforehand.
func (b *SparseUnionBuilder) AppendNulls(n int) {
+ if n <= 0 {
+ return
+ }
+
firstChildCode := b.codes[0]
b.Reserve(n)
for _, c := range b.codes {
b.typeIDtoBuilder[c].Reserve(n)
}
- for i := 0; i < n; i++ {
- b.typesBuilder.AppendValue(firstChildCode)
- b.typeIDtoBuilder[firstChildCode].AppendNull()
- for _, c := range b.codes[1:] {
- b.typeIDtoBuilder[c].AppendEmptyValue()
- }
+
+ unsafeAppendRepeatedInt8(b.typesBuilder, firstChildCode, n)
+ b.typeIDtoBuilder[firstChildCode].AppendNulls(n)
+ for _, c := range b.codes[1:] {
+ b.typeIDtoBuilder[c].AppendEmptyValues(n)
}
}
@@ -1029,16 +1060,19 @@ func (b *SparseUnionBuilder) AppendEmptyValue() {
// AppendEmptyValues is identical to calling AppendEmptyValue() n times,
// except it pre-allocates first so it is more efficient.
func (b *SparseUnionBuilder) AppendEmptyValues(n int) {
+ if n <= 0 {
+ return
+ }
+
b.Reserve(n)
firstChildCode := b.codes[0]
for _, c := range b.codes {
b.typeIDtoBuilder[c].Reserve(n)
}
- for i := 0; i < n; i++ {
- b.typesBuilder.AppendValue(firstChildCode)
- for _, c := range b.codes {
- b.typeIDtoBuilder[c].AppendEmptyValue()
- }
+
+ unsafeAppendRepeatedInt8(b.typesBuilder, firstChildCode, n)
+ for _, c := range b.codes {
+ b.typeIDtoBuilder[c].AppendEmptyValues(n)
}
}
@@ -1249,10 +1283,8 @@ func (b *DenseUnionBuilder) AppendNulls(n int) {
firstChildCode := b.codes[0]
childBuilder := b.typeIDtoBuilder[firstChildCode]
b.Reserve(n)
- for i := 0; i < n; i++ {
- b.typesBuilder.AppendValue(firstChildCode)
- b.offsetsBuilder.AppendValue(int32(childBuilder.Len()))
- }
+ unsafeAppendRepeatedInt8(b.typesBuilder, firstChildCode, n)
+ unsafeAppendRepeatedInt32(b.offsetsBuilder, int32(childBuilder.Len()),
n)
// only append a single null to the child builder, the offsets all
refer to the same value
childBuilder.AppendNull()
}
@@ -1280,10 +1312,8 @@ func (b *DenseUnionBuilder) AppendEmptyValues(n int) {
firstChildCode := b.codes[0]
childBuilder := b.typeIDtoBuilder[firstChildCode]
b.Reserve(n)
- for i := 0; i < n; i++ {
- b.typesBuilder.AppendValue(firstChildCode)
- b.offsetsBuilder.AppendValue(int32(childBuilder.Len()))
- }
+ unsafeAppendRepeatedInt8(b.typesBuilder, firstChildCode, n)
+ unsafeAppendRepeatedInt32(b.offsetsBuilder, int32(childBuilder.Len()),
n)
// only append a single empty value to the child builder, the offsets
all
// refer to the same value
childBuilder.AppendEmptyValue()
diff --git a/arrow/array/union_builder_bulk_test.go
b/arrow/array/union_builder_bulk_test.go
new file mode 100644
index 00000000..8922871f
--- /dev/null
+++ b/arrow/array/union_builder_bulk_test.go
@@ -0,0 +1,257 @@
+// 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"
+ "github.com/stretchr/testify/assert"
+ "github.com/stretchr/testify/require"
+)
+
+type bulkUnionBuilderFactory func(memory.Allocator) array.UnionBuilder
+
+var bulkUnionBuilderFactories = []struct {
+ name string
+ new bulkUnionBuilderFactory
+}{
+ {"sparse", func(mem memory.Allocator) array.UnionBuilder {
+ return array.NewSparseUnionBuilder(mem,
bulkUnionType(arrow.SparseMode).(*arrow.SparseUnionType))
+ }},
+ {"dense", func(mem memory.Allocator) array.UnionBuilder {
+ return array.NewDenseUnionBuilder(mem,
bulkUnionType(arrow.DenseMode).(*arrow.DenseUnionType))
+ }},
+}
+
+func bulkUnionType(mode arrow.UnionMode) arrow.UnionType {
+ fields := []arrow.Field{
+ {Name: "str", Type: arrow.BinaryTypes.String, Nullable: true},
+ {Name: "i32", Type: arrow.PrimitiveTypes.Int32, Nullable: true},
+ {Name: "f64", Type: arrow.PrimitiveTypes.Float64, Nullable:
true},
+ }
+ codes := []arrow.UnionTypeCode{8, 13, 7}
+ return arrow.UnionOf(mode, fields, codes)
+}
+
+func TestUnionBuilderBulkAppendNullsAndEmptyValues(t *testing.T) {
+ for _, tc := range bulkUnionBuilderFactories {
+ t.Run(tc.name, func(t *testing.T) {
+ mem :=
memory.NewCheckedAllocator(memory.NewGoAllocator())
+ defer mem.AssertSize(t, 0)
+
+ builder := tc.new(mem)
+ defer builder.Release()
+
+ builder.AppendNulls(0)
+ builder.AppendEmptyValues(0)
+ builder.AppendNulls(-1)
+ builder.AppendEmptyValues(-1)
+ require.Zero(t, builder.Len())
+ require.Zero(t, builder.Cap())
+
+ appendBulkUnionPrefix(builder, 3)
+ builder.AppendNulls(5)
+ builder.AppendEmptyValues(4)
+
+ result := builder.NewArray().(array.Union)
+ defer result.Release()
+ require.NoError(t, result.ValidateFull())
+
+ assert.Equal(t, 12, result.Len())
+ assert.Equal(t, []arrow.UnionTypeCode{8, 13, 7, 8, 8,
8, 8, 8, 8, 8, 8, 8}, result.RawTypeCodes())
+ })
+ }
+}
+
+func TestUnionBuilderZeroBulkAppendDoesNotMutateChildren(t *testing.T) {
+ tests := []struct {
+ name string
+ new func(memory.Allocator) array.UnionBuilder
+ }{
+ {"sparse", func(mem memory.Allocator) array.UnionBuilder {
+ return array.NewEmptySparseUnionBuilder(mem)
+ }},
+ {"dense", func(mem memory.Allocator) array.UnionBuilder {
+ return array.NewEmptyDenseUnionBuilder(mem)
+ }},
+ }
+
+ for _, tc := range tests {
+ t.Run(tc.name, func(t *testing.T) {
+ mem :=
memory.NewCheckedAllocator(memory.NewGoAllocator())
+ defer mem.AssertSize(t, 0)
+
+ builder := tc.new(mem)
+ defer builder.Release()
+
+ builder.AppendNulls(0)
+ builder.AppendEmptyValues(0)
+ builder.AppendNulls(-1)
+ builder.AppendEmptyValues(-1)
+
+ assert.Zero(t, builder.Len())
+ assert.Zero(t, builder.Cap())
+ })
+ }
+}
+
+func TestUnionBuilderBulkAppendMatchesScalar(t *testing.T) {
+ starts := []int{0, 1, 2, 7, 8, 15}
+ batchSizes := []int{-1, 0, 1, 2, 7, 16, 17}
+ operations := []struct {
+ name string
+ bulk func(array.UnionBuilder, int)
+ scalar func(array.UnionBuilder, int)
+ }{
+ {
+ name: "nulls",
+ bulk: func(builder array.UnionBuilder, n int) {
+ builder.AppendNulls(n)
+ },
+ scalar: func(builder array.UnionBuilder, n int) {
+ for i := 0; i < n; i++ {
+ builder.AppendNull()
+ }
+ },
+ },
+ {
+ name: "empty_values",
+ bulk: func(builder array.UnionBuilder, n int) {
+ builder.AppendEmptyValues(n)
+ },
+ scalar: func(builder array.UnionBuilder, n int) {
+ for i := 0; i < n; i++ {
+ builder.AppendEmptyValue()
+ }
+ },
+ },
+ }
+
+ for _, factory := range bulkUnionBuilderFactories {
+ t.Run(factory.name, func(t *testing.T) {
+ for operationIndex, operation := range operations {
+ t.Run(operation.name, func(t *testing.T) {
+ reuseOperation :=
operations[1-operationIndex]
+ for _, start := range starts {
+ for _, batchSize := range
batchSizes {
+
t.Run(fmt.Sprintf("start_%d_batch_%d", start, batchSize), func(t *testing.T) {
+ mem :=
memory.NewCheckedAllocator(memory.NewGoAllocator())
+ defer
mem.AssertSize(t, 0)
+
+ bulk :=
factory.new(mem)
+ defer
bulk.Release()
+ scalar :=
factory.new(mem)
+ defer
scalar.Release()
+
+
appendBulkUnionPrefix(bulk, start)
+
appendBulkUnionPrefix(scalar, start)
+
operation.bulk(bulk, batchSize)
+
operation.scalar(scalar, batchSize)
+
assertUnionBuilderArrayParity(t, bulk, scalar)
+
+
appendBulkUnionPrefix(bulk, start)
+
appendBulkUnionPrefix(scalar, start)
+
reuseOperation.bulk(bulk, 9)
+
reuseOperation.scalar(scalar, 9)
+
assertUnionBuilderArrayParity(t, bulk, scalar)
+ })
+ }
+ }
+ })
+ }
+ })
+ }
+}
+
+func appendBulkUnionPrefix(builder array.UnionBuilder, n int) {
+ codes := []arrow.UnionTypeCode{8, 13, 7}
+ for i := 0; i < n; i++ {
+ childID := i % len(codes)
+ builder.Append(codes[childID])
+
+ switch childID {
+ case 0:
+
builder.Child(childID).(*array.StringBuilder).Append(fmt.Sprintf("value-%d", i))
+ case 1:
+
builder.Child(childID).(*array.Int32Builder).Append(int32(i))
+ case 2:
+
builder.Child(childID).(*array.Float64Builder).Append(float64(i))
+ }
+
+ if builder.Mode() == arrow.SparseMode {
+ for i := 0; i < len(codes); i++ {
+ if i != childID {
+ builder.Child(i).AppendEmptyValue()
+ }
+ }
+ }
+ }
+}
+
+func assertUnionBuilderArrayParity(t *testing.T, bulk, scalar
array.UnionBuilder) {
+ t.Helper()
+
+ assert.Equal(t, scalar.Len(), bulk.Len())
+
+ bulkArray := bulk.NewArray().(array.Union)
+ defer bulkArray.Release()
+ scalarArray := scalar.NewArray().(array.Union)
+ defer scalarArray.Release()
+
+ require.NoError(t, bulkArray.ValidateFull())
+ require.NoError(t, scalarArray.ValidateFull())
+ assert.True(t, array.Equal(bulkArray, scalarArray))
+}
+
+func BenchmarkUnionBuilderBulkAppend(b *testing.B) {
+ for _, tc := range bulkUnionBuilderFactories {
+ b.Run(tc.name, func(b *testing.B) {
+ for _, rows := range []int{1, 16, 1024, 65536} {
+ b.Run(fmt.Sprintf("rows_%d", rows), func(b
*testing.B) {
+ b.Run("nulls", func(b *testing.B) {
+
benchmarkUnionBuilderBulkAppend(b, tc.new, rows, false)
+ })
+ b.Run("empty", func(b *testing.B) {
+
benchmarkUnionBuilderBulkAppend(b, tc.new, rows, true)
+ })
+ })
+ }
+ })
+ }
+}
+
+func benchmarkUnionBuilderBulkAppend(b *testing.B, factory
bulkUnionBuilderFactory, rows int, empty bool) {
+ builder := factory(memory.DefaultAllocator)
+ defer builder.Release()
+ b.ReportAllocs()
+
+ for b.Loop() {
+ if empty {
+ builder.AppendEmptyValues(rows)
+ } else {
+ builder.AppendNulls(rows)
+ }
+
+ arr := builder.NewArray()
+ arr.Release()
+ }
+}
