(arrow-go) branch main updated: perf(arrow): reduce metadata sorting overhead (#1361)

[ze...@apache.org]

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 9296b3fa perf(arrow): reduce metadata sorting overhead (#1361)
9296b3fa is described below

commit 9296b3fac5a1daf2ce1f8a124a5227449fc1bdcb
Author: Minh Vu <vu...@gmail.com>
AuthorDate: Tue Oct 13 16:30:29 2026 +0200

    perf(arrow): reduce metadata sorting overhead (#1361)
    
    **What**
    - Reduce metadata equality overhead for small metadata.
    
    **Why**
    - Small sorts spend time calling the comparator.
    
    **Implementation**
    - Use direct insertion sort for up to 12 pairs.
    - Keep the current sorter for larger metadata.
    - Keep small-sort locals out of the general sorting frame.
    
    Go 1.26.3, Apple M1 Pro, 6 interleaved samples of 100ms. Compared with
    `c1c8ed0a`.
    
    | Case | Base ns/op | PR ns/op | Change |
    | --- | ---: | ---: | ---: |
    | 8 pairs, matching | 123.6 | 112.9 | -8.62% |
    | 8 pairs, reversed | 233.9 | 192.8 | -17.57% |
    | 12 pairs, reversed | 460.5 | 375.1 | -18.53% |
    | 128 pairs, middle swap | 2,794 | 2,823 | ~ |
    | 2,049 pairs, middle swap | 40,600 | 40,680 | ~ |
    
    `~` means no significant change. Allocations are unchanged.
---
 arrow/metadata_equal_benchmark_test.go | 128 ++++++++++++++++++++++
 arrow/metadata_equal_test.go           | 191 +++++++++++++++++++++++++++++++++
 arrow/schema.go                        |  36 +++++++
 3 files changed, 355 insertions(+)

diff --git a/arrow/metadata_equal_benchmark_test.go b/arrow/metadata_equal_benchmark_test.go
new file mode 100644
index 00000000..0e790225
--- /dev/null
+++ b/arrow/metadata_equal_benchmark_test.go
@@ -0,0 +1,128 @@
+// 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 arrow
+
+import (
+	"fmt"
+	"slices"
+	"strings"
+	"testing"
+)
+
+var metadataEqualBenchmarkResult bool
+
+func metadataEqualityBenchmarkPair(size int, pattern string) (Metadata, Metadata) {
+	keys, values := make([]string, size), make([]string, size)
+	for i := range keys {
+		keys[i] = fmt.Sprintf("key_%04d", i)
+		values[i] = fmt.Sprintf("value_%04d", i)
+	}
+	left := NewMetadata(keys, values)
+	right := left.clone()
+	if strings.HasPrefix(pattern, "duplicate-") {
+		for i := range left.keys {
+			left.keys[i] = fmt.Sprintf("key_%04d", i/2)
+			left.values[i] = fmt.Sprintf("value_%04d", i/2)
+			right.keys[i], right.values[i] = left.keys[i], left.values[i]
+		}
+		pattern = strings.TrimPrefix(pattern, "duplicate-")
+	}
+	switch pattern {
+	case "matching-reversed":
+		slices.Reverse(left.keys)
+		slices.Reverse(left.values)
+		slices.Reverse(right.keys)
+		slices.Reverse(right.values)
+	case "different-value":
+		right.values[size-1] = "different"
+	case "different-key":
+		right.keys[size/2] = "different"
+	case "reversed":
+		slices.Reverse(right.keys)
+		slices.Reverse(right.values)
+	case "middle-swap":
+		pos := size/2 - 1
+		right.keys[pos], right.keys[pos+1] = right.keys[pos+1], right.keys[pos]
+		right.values[pos], right.values[pos+1] = right.values[pos+1], right.values[pos]
+	}
+	return left, right
+}
+
+func BenchmarkMetadataEqual(b *testing.B) {
+	for _, size := range []int{0, 1, 8, 12, 13, 128} {
+		for _, pattern := range []string{"matching", "matching-reversed", "different-value", "reversed", "different-key"} {
+			if size == 0 && pattern != "matching" || size == 1 && (pattern == "matching-reversed" || pattern == "reversed") {
+				continue
+			}
+			b.Run(fmt.Sprintf("size=%d/%s", size, pattern), func(b *testing.B) {
+				left, right := metadataEqualityBenchmarkPair(size, pattern)
+				if got := left.Equal(right); got != (pattern != "different-value" && pattern != "different-key") {
+					b.Fatalf("unexpected equality: %v", got)
+				}
+				b.ReportAllocs()
+				for b.Loop() {
+					metadataEqualBenchmarkResult = left.Equal(right)
+				}
+			})
+		}
+	}
+}
+
+func BenchmarkMetadataEqualDuplicateKeys(b *testing.B) {
+	for _, size := range []int{8, 12, 13, 128} {
+		for _, pattern := range []string{"matching", "reversed", "middle-swap"} {
+			b.Run(fmt.Sprintf("size=%d/%s", size, pattern), func(b *testing.B) {
+				left, right := metadataEqualityBenchmarkPair(size, "duplicate-"+pattern)
+				if !left.Equal(right) {
+					b.Fatal("metadata should be equal")
+				}
+				b.ReportAllocs()
+				for b.Loop() {
+					metadataEqualBenchmarkResult = left.Equal(right)
+				}
+			})
+		}
+	}
+}
+
+func BenchmarkSchemaEqualFieldMetadata(b *testing.B) {
+	const nfields = 16
+	for _, size := range []int{0, 1, 8, 32} {
+		for _, pattern := range []string{"matching", "reversed"} {
+			if size < 2 && pattern != "matching" {
+				continue
+			}
+			b.Run(fmt.Sprintf("size=%d/%s", size, pattern), func(b *testing.B) {
+				leftMetadata, rightMetadata := metadataEqualityBenchmarkPair(size, pattern)
+				leftFields, rightFields := make([]Field, nfields), make([]Field, nfields)
+				for i := range leftFields {
+					name := fmt.Sprintf("field_%d", i)
+					leftFields[i] = Field{Name: name, Type: PrimitiveTypes.Int64, Metadata: leftMetadata}
+					rightFields[i] = Field{Name: name, Type: PrimitiveTypes.Int64, Metadata: rightMetadata}
+				}
+				left, right := NewSchema(leftFields, nil), NewSchema(rightFields, nil)
+				if !left.Equal(right) {
+					b.Fatal("schemas should be equal")
+				}
+				b.ReportAllocs()
+				for b.Loop() {
+					metadataEqualBenchmarkResult = left.Equal(right)
+				}
+			})
+		}
+	}
+}
diff --git a/arrow/metadata_equal_test.go b/arrow/metadata_equal_test.go
new file mode 100644
index 00000000..5aa69350
--- /dev/null
+++ b/arrow/metadata_equal_test.go
@@ -0,0 +1,191 @@
+// 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 arrow
+
+import (
+	"fmt"
+	"slices"
+	"strings"
+	"testing"
+
+	"github.com/stretchr/testify/assert"
+)
+
+func TestMetadataEqual(t *testing.T) {
+	for _, tc := range []struct {
+		name        string
+		left, right Metadata
+		want        bool
+	}{
+		{"empty", Metadata{}, NewMetadata(nil, nil), true},
+		{"empty-map", Metadata{}, MetadataFrom(map[string]string{}), true},
+		{"single", NewMetadata([]string{"key"}, []string{"value"}), NewMetadata([]string{"key"}, []string{"value"}), true},
+		{"single-copied-strings", NewMetadata([]string{"key"}, []string{"value"}), NewMetadata([]string{strings.Clone("key")}, []string{strings.Clone("value")}), true},
+		{"single-long-key-difference", NewMetadata([]string{strings.Repeat("a", 4096) + "x"}, []string{"value"}), NewMetadata([]string{strings.Repeat("a", 4096) + "y"}, []string{"value"}), false},
+		{"single-long-value-difference", NewMetadata([]string{"key"}, []string{strings.Repeat("a", 4096) + "x"}), NewMetadata([]string{"key"}, []string{strings.Repeat("a", 4096) + "y"}), false},
+		{"different-lengths", Metadata{}, NewMetadata([]string{"key"}, []string{"value"}), false},
+		{"different-key", NewMetadata([]string{"key"}, []string{"value"}), NewMetadata([]string{"other"}, []string{"value"}), false},
+		{"different-value", NewMetadata([]string{"key"}, []string{"value"}), NewMetadata([]string{"key"}, []string{"other"}), false},
+		{"empty-key-value", NewMetadata([]string{""}, []string{""}), NewMetadata([]string{""}, []string{""}), true},
+		{"case-sensitive", NewMetadata([]string{"key"}, []string{"value"}), NewMetadata([]string{"Key"}, []string{"value"}), false},
+		{"unicode", NewMetadata([]string{"é", "键"}, []string{"值", ""}), NewMetadata([]string{"é", "键"}, []string{"值", ""}), true},
+		{"unicode-normalization", NewMetadata([]string{"é"}, []string{"value"}), NewMetadata([]string{"e\u0301"}, []string{"value"}), false},
+		{"embedded-zero", NewMetadata([]string{"key\x00"}, []string{"value\x00"}), NewMetadata([]string{"key\x00"}, []string{"value\x00"}), true},
+		{"matching-unsorted", NewMetadata([]string{"z", "a", "m"}, []string{"1", "2", "3"}), NewMetadata([]string{"z", "a", "m"}, []string{"1", "2", "3"}), true},
+		{"reordered", NewMetadata([]string{"z", "a", "m"}, []string{"1", "2", "3"}), NewMetadata([]string{"a", "m", "z"}, []string{"2", "3", "1"}), true},
+		{"reordered-different-value", NewMetadata([]string{"z", "a", "m"}, []string{"1", "2", "3"}), NewMetadata([]string{"a", "m", "z"}, []string{"2", "4", "1"}), false},
+		{"duplicate-keys", NewMetadata([]string{"key", "key"}, []string{"1", "2"}), NewMetadata([]string{"key", "key"}, []string{"1", "2"}), true},
+		{"duplicate-values-swapped", NewMetadata([]string{"key", "key"}, []string{"1", "2"}), NewMetadata([]string{"key", "key"}, []string{"2", "1"}), false},
+		{"duplicate-pairs-reordered", NewMetadata([]string{"key", "key", "a"}, []string{"v", "v", "a"}), NewMetadata([]string{"a", "key", "key"}, []string{"a", "v", "v"}), true},
+		{"duplicate-distinct-values-reordered", NewMetadata([]string{"b", "a", "b"}, []string{"first", "a", "second"}), NewMetadata([]string{"a", "b", "b"}, []string{"a", "first", "second"}), true},
+		{"duplicate-distinct-values-swapped", NewMetadata([]string{"b", "a", "b"}, []string{"first", "a", "second"}), NewMetadata([]string{"a", "b", "b"}, []string{"a", "second", "first"}), false},
+		{"duplicate-counts", NewMetadata([]string{"key", "key", "a"}, []string{"v", "v", "v"}), NewMetadata([]string{"key", "a", "a"}, []string{"v", "v", "v"}), false},
+	} {
+		t.Run(tc.name, func(t *testing.T) {
+			leftBefore := Metadata{keys: slices.Clone(tc.left.keys), values: slices.Clone(tc.left.values)}
+			rightBefore := Metadata{keys: slices.Clone(tc.right.keys), values: slices.Clone(tc.right.values)}
+			assert.Equal(t, tc.want, tc.left.Equal(tc.right))
+			assert.Equal(t, tc.want, tc.right.Equal(tc.left))
+			assert.Equal(t, leftBefore, tc.left)
+			assert.Equal(t, rightBefore, tc.right)
+		})
+	}
+}
+
+func TestMetadataEqualDuplicateKeys(t *testing.T) {
+	for _, size := range []int{1, 2, 3, 12, 13, 15, 16, 17, 31, 32, 33, 64, 127, 128, 129, 255, 256, 257, 1024, 2047, 2048, 2049} {
+		t.Run(fmt.Sprint(size), func(t *testing.T) {
+			keys, values := make([]string, size), make([]string, size)
+			for i := range keys {
+				keys[i] = fmt.Sprintf("key_%d", i%4)
+				values[i] = fmt.Sprintf("value_%d", i)
+			}
+			left, right := NewMetadata(keys, values), NewMetadata(keys, values)
+			assert.True(t, left.Equal(right))
+			right.Values()[size/2] = "changed"
+			assert.False(t, left.Equal(right))
+			assert.False(t, right.Equal(left))
+		})
+	}
+}
+
+func TestMetadataEqualReordered(t *testing.T) {
+	for _, size := range []int{2, 8, 11, 12, 13, 15, 16, 17, 31, 32, 33, 127, 128, 129, 256, 257, 1024, 2047, 2048, 2049} {
+		t.Run(fmt.Sprint(size), func(t *testing.T) {
+			keys, values := make([]string, size), make([]string, size)
+			for i := range keys {
+				keys[i] = fmt.Sprintf("key_%04d", i)
+				values[i] = fmt.Sprintf("value_%04d", i)
+			}
+			left := NewMetadata(keys, values)
+			for _, pattern := range []string{"reversed", "middle-swap", "last-swap"} {
+				t.Run(pattern, func(t *testing.T) {
+					right := left.clone()
+					if pattern == "reversed" {
+						slices.Reverse(right.Keys())
+						slices.Reverse(right.Values())
+					} else {
+						pos := size/2 - 1
+						if pattern == "last-swap" {
+							pos = size - 2
+						}
+						right.keys[pos], right.keys[pos+1] = right.keys[pos+1], right.keys[pos]
+						right.values[pos], right.values[pos+1] = right.values[pos+1], right.values[pos]
+					}
+					leftBefore, rightBefore := left.clone(), right.clone()
+					assert.True(t, left.Equal(right))
+					assert.True(t, right.Equal(left))
+					assert.Equal(t, leftBefore, left)
+					assert.Equal(t, rightBefore, right)
+					right.Values()[size/2] = "changed"
+					assert.False(t, left.Equal(right))
+					assert.False(t, right.Equal(left))
+				})
+			}
+		})
+	}
+}
+
+func TestMetadataEqualPermutations(t *testing.T) {
+	keys := []string{"z", "", "a", "键"}
+	values := []string{"value", "\x00", "", "值"}
+	var permutations [][]int
+	order := []int{0, 1, 2, 3}
+	var permute func(int)
+	permute = func(pos int) {
+		if pos == len(order) {
+			permutations = append(permutations, slices.Clone(order))
+			return
+		}
+		for i := pos; i < len(order); i++ {
+			order[pos], order[i] = order[i], order[pos]
+			permute(pos + 1)
+			order[pos], order[i] = order[i], order[pos]
+		}
+	}
+	permute(0)
+	makeMetadata := func(order []int) Metadata {
+		k, v := make([]string, len(order)), make([]string, len(order))
+		for i, idx := range order {
+			k[i], v[i] = keys[idx], values[idx]
+		}
+		return NewMetadata(k, v)
+	}
+	for _, leftOrder := range permutations {
+		for _, rightOrder := range permutations {
+			t.Run(fmt.Sprintf("%v/%v", leftOrder, rightOrder), func(t *testing.T) {
+				left, right := makeMetadata(leftOrder), makeMetadata(rightOrder)
+				assert.True(t, left.Equal(right))
+				assert.True(t, right.Equal(left))
+			})
+		}
+	}
+}
+
+func TestMetadataEqualSmallDoesNotAllocate(t *testing.T) {
+	left := NewMetadata(
+		[]string{"h", "a", "g", "b", "f", "c", "e", "d"},
+		[]string{"8", "1", "7", "2", "6", "3", "5", "4"},
+	)
+	right := NewMetadata(
+		[]string{"a", "b", "c", "d", "e", "f", "g", "h"},
+		[]string{"1", "2", "3", "4", "5", "6", "7", "8"},
+	)
+
+	allocs := testing.AllocsPerRun(100, func() {
+		if !left.Equal(right) {
+			panic("metadata should be equal")
+		}
+	})
+	assert.Zero(t, allocs)
+}
+
+func TestMetadataEqualAfterMutation(t *testing.T) {
+	left := NewMetadata([]string{"a", "b"}, []string{"1", "2"})
+	right := left.clone()
+	assert.True(t, left.Equal(right))
+	right.Values()[0] = "changed"
+	assert.False(t, left.Equal(right))
+	right.Values()[0] = "1"
+	right.Keys()[0] = "changed"
+	assert.False(t, left.Equal(right))
+	right.Keys()[0] = "a"
+	assert.True(t, left.Equal(right))
+	slices.Reverse(right.Keys())
+	slices.Reverse(right.Values())
+	assert.True(t, left.Equal(right))
+}
diff --git a/arrow/schema.go b/arrow/schema.go
index d3f52584..9172e0b5 100644
--- a/arrow/schema.go
+++ b/arrow/schema.go
@@ -123,6 +123,39 @@ func (md Metadata) clone() Metadata {
 	return o
 }
 
+const smallMetadataLimit = 12
+
+func sortSmallMetadataIndices(keys []string, idxes []int) {
+	for i := range idxes {
+		idxes[i] = i
+	}
+	for i := 1; i < len(idxes); i++ {
+		for j := i; j > 0 && keys[idxes[j]] < keys[idxes[j-1]]; j-- {
+			idxes[j], idxes[j-1] = idxes[j-1], idxes[j]
+		}
+	}
+}
+
+// Keep the small-sort arrays out of Equal's general frame.
+//
+//go:noinline
+func smallMetadataEqual(left, right Metadata) bool {
+	var leftStorage, rightStorage [smallMetadataLimit]int
+	leftIdxes := leftStorage[:len(left.keys)]
+	rightIdxes := rightStorage[:len(right.keys)]
+	sortSmallMetadataIndices(left.keys, leftIdxes)
+	sortSmallMetadataIndices(right.keys, rightIdxes)
+
+	for i := range leftIdxes {
+		j := leftIdxes[i]
+		k := rightIdxes[i]
+		if left.keys[j] != right.keys[k] || left.values[j] != right.values[k] {
+			return false
+		}
+	}
+	return true
+}
+
 func (md Metadata) sortedIndices() []int {
 	idxes := make([]int, len(md.keys))
 	for i := range idxes {
@@ -139,6 +172,9 @@ func (md Metadata) Equal(rhs Metadata) bool {
 	if md.Len() != rhs.Len() {
 		return false
 	}
+	if md.Len() <= smallMetadataLimit {
+		return smallMetadataEqual(md, rhs)
+	}
 
 	idxes := md.sortedIndices()
 	rhsIdxes := rhs.sortedIndices()