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/iceberg-go.git
The following commit(s) were added to refs/heads/main by this push:
new 67592a526 perf(table): use typed maps for primitive fanout partition
keys (#2254)
67592a526 is described below
commit 67592a52608197136f1c24c40d432d24a4d8ceb2
Author: Minh Vu <[email protected]>
AuthorDate: Wed Oct 21 22:52:46 2026 +0200
perf(table): use typed maps for primitive fanout partition keys (#2254)
Signed-off-by: Minh Vu <[email protected]>
---
table/fanout_typed_partition_keys_bench_test.go | 165 ++++++++++++++++++++++++
table/partitioned_fanout_writer.go | 75 ++++++++---
2 files changed, 223 insertions(+), 17 deletions(-)
diff --git a/table/fanout_typed_partition_keys_bench_test.go
b/table/fanout_typed_partition_keys_bench_test.go
new file mode 100644
index 000000000..fdf85a256
--- /dev/null
+++ b/table/fanout_typed_partition_keys_bench_test.go
@@ -0,0 +1,165 @@
+// 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 table
+
+import (
+ "fmt"
+ "math"
+ "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/apache/iceberg-go"
+)
+
+func TestPartitionMapTypedKeyIsolation(t *testing.T) {
+ fieldInfo := []partitionFieldInfo{{fieldID: 1000}}
+ keys := []any{int32(1), int64(1), "1", []byte("1"), float32(1),
float64(1), nil}
+ root := newPartitionMapNode()
+ parts := make([]*partitionInfo, len(keys))
+ for i, key := range keys {
+ parts[i] = root.getOrCreate(partitionRecord{key}, fieldInfo,
int64(len(keys)))
+ parts[i].rows = append(parts[i].rows, int64(i))
+ }
+ if got := len(root.collectPartitions()); got != len(keys) {
+ t.Fatalf("got %d distinct typed partitions, want %d", got,
len(keys))
+ }
+ for i, key := range keys {
+ if got := root.getOrCreate(partitionRecord{key}, fieldInfo,
int64(len(keys))); got != parts[i] {
+ t.Fatalf("key %T %v changed partition", key, key)
+ }
+ }
+ // The fallback must still canonicalize NaNs.
+ nan1 := math.Float64frombits(0x7ff8000000000001)
+ nan2 := math.Float64frombits(0x7ff8000000000002)
+ if root.getOrCreate(partitionRecord{nan1}, fieldInfo, 8) !=
root.getOrCreate(partitionRecord{nan2}, fieldInfo, 8) {
+ t.Fatal("NaN payloads must share one partition")
+ }
+ // Mixed primitive levels must survive the result-tree traversal.
+ nested := newPartitionMapNode()
+ nestedFields := []partitionFieldInfo{{fieldID: 1000}, {fieldID: 1001}}
+ for _, key := range []partitionRecord{
+ {int64(2), "a"}, {int64(2), "b"}, {int64(3), "a"},
+ } {
+ nested.getOrCreate(key, nestedFields, 3)
+ }
+ if got := len(nested.collectPartitions()); got != 3 {
+ t.Fatalf("collected %d nested partitions, want 3", got)
+ }
+}
+
+// BenchmarkFanoutTypedPartitionKeys measures the real Arrow-to-Iceberg
+// extraction and fanout grouping path. Construct fixtures before timing.
+func BenchmarkFanoutTypedPartitionKeys(b *testing.B) {
+ const rows = 1 << 16
+ cases := []struct {
+ name string
+ kind string
+ cardinality int
+ clustered bool
+ }{
+ {"Int64_64", "int64", 64, false},
+ {"Int64_16384", "int64", 16384, false},
+ {"Int32_64", "int32", 64, false},
+ {"String_1024", "string", 1024, false},
+ {"Int64_Clustered", "int64", 64, true},
+ {"Float64_Fallback", "float64", 64, false},
+ }
+ for _, tc := range cases {
+ b.Run(tc.name, func(b *testing.B) {
+ b.StopTimer()
+ mem := memory.NewGoAllocator()
+ var column arrow.Array
+ var typ iceberg.Type
+ var arrowType arrow.DataType
+ asString := make([]string, tc.cardinality)
+ for i := range asString {
+ asString[i] = fmt.Sprintf("part-%04d", i)
+ }
+ partitionAt := func(i int) int {
+ if tc.clustered {
+ return (i * tc.cardinality) / rows
+ }
+
+ return i % tc.cardinality
+ }
+ switch tc.kind {
+ case "int32":
+ builder := array.NewInt32Builder(mem)
+ for i := range rows {
+ builder.Append(int32(partitionAt(i)))
+ }
+ column = builder.NewArray()
+ builder.Release()
+ typ, arrowType = iceberg.PrimitiveTypes.Int32,
arrow.PrimitiveTypes.Int32
+ case "int64":
+ builder := array.NewInt64Builder(mem)
+ for i := range rows {
+ builder.Append(int64(partitionAt(i)))
+ }
+ column = builder.NewArray()
+ builder.Release()
+ typ, arrowType = iceberg.PrimitiveTypes.Int64,
arrow.PrimitiveTypes.Int64
+ case "string":
+ builder := array.NewStringBuilder(mem)
+ for i := range rows {
+ builder.Append(asString[partitionAt(i)])
+ }
+ column = builder.NewArray()
+ builder.Release()
+ typ, arrowType = iceberg.PrimitiveTypes.String,
arrow.BinaryTypes.String
+ case "float64":
+ builder := array.NewFloat64Builder(mem)
+ for i := range rows {
+ builder.Append(float64(partitionAt(i)))
+ }
+ column = builder.NewArray()
+ builder.Release()
+ typ, arrowType =
iceberg.PrimitiveTypes.Float64, arrow.PrimitiveTypes.Float64
+ }
+ arrowSchema := arrow.NewSchema([]arrow.Field{{Name:
"part", Type: arrowType}}, nil)
+ record := array.NewRecordBatch(arrowSchema,
[]arrow.Array{column}, rows)
+ column.Release()
+ defer record.Release()
+ schema := iceberg.NewSchema(0, iceberg.NestedField{ID:
1, Name: "part", Type: typ})
+ spec := iceberg.NewPartitionSpec(iceberg.PartitionField{
+ SourceIDs: []int{1}, FieldID: 1000, Name:
"partition",
+ Transform: iceberg.IdentityTransform{},
+ })
+ plan, err := newPartitionExtractionPlan(spec, schema,
record.Schema())
+ if err != nil {
+ b.Fatal(err)
+ }
+ b.ReportAllocs()
+ b.StartTimer()
+ for range b.N {
+ parts, err := plan.getRecordPartitions(record)
+ if err != nil {
+ b.Fatal(err)
+ }
+ if len(parts) != tc.cardinality {
+ b.Fatalf("got %d partitions, want %d",
len(parts), tc.cardinality)
+ }
+ typedPartitionKeyBenchmarkSink = len(parts)
+ }
+ })
+ }
+}
+
+var typedPartitionKeyBenchmarkSink int
diff --git a/table/partitioned_fanout_writer.go
b/table/partitioned_fanout_writer.go
index 825601520..6583758a0 100644
--- a/table/partitioned_fanout_writer.go
+++ b/table/partitioned_fanout_writer.go
@@ -719,6 +719,9 @@ func bindPartitionTransform(transform iceberg.Transform,
sourceType iceberg.Type
// The value is either a *partitionMapNode or a *partitionInfo.
type partitionMapNode struct {
children map[any]any
+ int32Children map[int32]any
+ int64Children map[int64]any
+ stringChildren map[string]any
binaryChildren map[string]any
// partitionCount is maintained by the root node for the current batch
and
// sizes the single result slice returned by collectPartitions.
@@ -729,35 +732,64 @@ func newPartitionMapNode() *partitionMapNode {
return &partitionMapNode{}
}
-// child looks up binary partition values through a typed string map.
Converting a
-// byte slice to string for a map lookup does not allocate, while storing a new
-// partition still copies the key into map-owned storage.
+// child looks up common primitive keys through typed maps. Binary keys use a
+// separate string map: converting []byte to string for lookup does not
allocate,
+// while insertion copies the key into map-owned storage.
func (n *partitionMapNode) child(part any) (any, bool) {
- if bytes, ok := part.([]byte); ok {
- child, found := n.binaryChildren[string(bytes)]
+ // The common primitive partition keys can use typed maps instead of
+ // hashing interface keys for every row.
+ switch value := part.(type) {
+ case int32:
+ child, found := n.int32Children[value]
return child, found
- }
+ case int64:
+ child, found := n.int64Children[value]
+
+ return child, found
+ case string:
+ child, found := n.stringChildren[value]
+
+ return child, found
+ case []byte:
+ child, found := n.binaryChildren[string(value)]
- child, found := n.children[comparablePartitionKey(part)]
+ return child, found
+ default:
+ child, found := n.children[comparablePartitionKey(part)]
- return child, found
+ return child, found
+ }
}
func (n *partitionMapNode) setChild(part, child any) {
- if bytes, ok := part.([]byte); ok {
+ switch value := part.(type) {
+ case int32:
+ if n.int32Children == nil {
+ n.int32Children = make(map[int32]any)
+ }
+ n.int32Children[value] = child
+ case int64:
+ if n.int64Children == nil {
+ n.int64Children = make(map[int64]any)
+ }
+ n.int64Children[value] = child
+ case string:
+ if n.stringChildren == nil {
+ n.stringChildren = make(map[string]any)
+ }
+ n.stringChildren[value] = child
+ case []byte:
if n.binaryChildren == nil {
n.binaryChildren = make(map[string]any)
}
- n.binaryChildren[string(bytes)] = child
-
- return
- }
-
- if n.children == nil {
- n.children = make(map[any]any)
+ n.binaryChildren[string(value)] = child
+ default:
+ if n.children == nil {
+ n.children = make(map[any]any)
+ }
+ n.children[comparablePartitionKey(part)] = child
}
- n.children[comparablePartitionKey(part)] = child
}
// getOrCreate navigates the tree and returns the partitionInfo for the given
partition key,
@@ -850,6 +882,15 @@ func (n *partitionMapNode) appendPartitions(result
[]*partitionInfo) []*partitio
for _, v := range n.children {
appendChild(v)
}
+ for _, v := range n.int32Children {
+ appendChild(v)
+ }
+ for _, v := range n.int64Children {
+ appendChild(v)
+ }
+ for _, v := range n.stringChildren {
+ appendChild(v)
+ }
for _, v := range n.binaryChildren {
appendChild(v)
}
