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 f939fddad perf(table): specialize binary partition key lookups (#2155)
f939fddad is described below
commit f939fddad954f6cc554f7f37a4aff35c104608a4
Author: Minh Vu <[email protected]>
AuthorDate: Tue Oct 13 13:31:52 2026 +0200
perf(table): specialize binary partition key lookups (#2155)
* perf(table): specialize binary partition key lookups
Signed-off-by: Minh Vu <[email protected]>
* test(table): cover binary partition key boundaries
Verify null, empty, arbitrary bytes, and separator-containing keys at both
levels of binary partition grouping with exact row membership.
Signed-off-by: Minh Vu <[email protected]>
---------
Signed-off-by: Minh Vu <[email protected]>
Signed-off-by: Codex <[email protected]>
---
table/partition_binary_key_bench_test.go | 92 ++++++++++++++++++++++++++
table/partitioned_fanout_writer.go | 61 +++++++++++++----
table/partitioned_fanout_writer_review_test.go | 74 +++++++++++++++++++++
3 files changed, 215 insertions(+), 12 deletions(-)
diff --git a/table/partition_binary_key_bench_test.go
b/table/partition_binary_key_bench_test.go
new file mode 100644
index 000000000..7ff01eb58
--- /dev/null
+++ b/table/partition_binary_key_bench_test.go
@@ -0,0 +1,92 @@
+// 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 (
+ "encoding/binary"
+ "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 BenchmarkBinaryPartitionGrouping(b *testing.B) {
+ const rows = 65_536
+
+ tests := []struct {
+ name string
+ keySize int
+ partitions int
+ }{
+ {name: "16B_16_partitions", keySize: 16, partitions: 16},
+ {name: "16B_1024_partitions", keySize: 16, partitions: 1_024},
+ {name: "64B_1024_partitions", keySize: 64, partitions: 1_024},
+ {name: "16B_32768_partitions", keySize: 16, partitions: 32_768},
+ }
+
+ for _, test := range tests {
+ b.Run(test.name, func(b *testing.B) {
+ arrowSchema := arrow.NewSchema([]arrow.Field{
+ {Name: "part", Type: arrow.BinaryTypes.Binary},
+ }, nil)
+ icebergSchema := iceberg.NewSchema(0,
+ iceberg.NestedField{ID: 1, Name: "part", Type:
iceberg.PrimitiveTypes.Binary},
+ )
+ spec := iceberg.NewPartitionSpec(iceberg.PartitionField{
+ SourceIDs: []int{1}, FieldID: 1000, Name:
"part", Transform: iceberg.IdentityTransform{},
+ })
+
+ values := make([][]byte, test.partitions)
+ for i := range values {
+ value := make([]byte, test.keySize)
+ binary.LittleEndian.PutUint64(value, uint64(i))
+ values[i] = value
+ }
+
+ builder :=
array.NewBinaryBuilder(memory.DefaultAllocator, arrow.BinaryTypes.Binary)
+ for row := range rows {
+ builder.Append(values[row%test.partitions])
+ }
+ column := builder.NewArray()
+ builder.Release()
+
+ record := array.NewRecordBatch(arrowSchema,
[]arrow.Array{column}, rows)
+ column.Release()
+ defer record.Release()
+
+ plan, err := newPartitionExtractionPlan(spec,
icebergSchema, record.Schema())
+ if err != nil {
+ b.Fatal(err)
+ }
+
+ b.ReportAllocs()
+ b.ResetTimer()
+ for b.Loop() {
+ partitions, err :=
plan.getRecordPartitions(record)
+ if err != nil {
+ b.Fatal(err)
+ }
+ if len(partitions) != test.partitions {
+ b.Fatalf("got %d partitions, want %d",
len(partitions), test.partitions)
+ }
+ }
+ })
+ }
+}
diff --git a/table/partitioned_fanout_writer.go
b/table/partitioned_fanout_writer.go
index 50c5f7f9a..8d5795a3a 100644
--- a/table/partitioned_fanout_writer.go
+++ b/table/partitioned_fanout_writer.go
@@ -658,38 +658,68 @@ func bindPartitionTransform(transform iceberg.Transform,
sourceType iceberg.Type
// is in the order of the partition spec.
// The value is either a *partitionMapNode or a *partitionInfo.
type partitionMapNode struct {
- children map[any]any
+ children map[any]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.
partitionCount int
}
func newPartitionMapNode() *partitionMapNode {
- return &partitionMapNode{
- children: make(map[any]any),
+ 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.
+func (n *partitionMapNode) child(part any) (any, bool) {
+ if bytes, ok := part.([]byte); ok {
+ child, found := n.binaryChildren[string(bytes)]
+
+ return child, found
}
+
+ child, found := n.children[comparablePartitionKey(part)]
+
+ return child, found
+}
+
+func (n *partitionMapNode) setChild(part, child any) {
+ if bytes, ok := part.([]byte); ok {
+ 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.children[comparablePartitionKey(part)] = child
}
// getOrCreate navigates the tree and returns the partitionInfo for the given
partition key,
// creating nodes along the way if they don't exist
func (n *partitionMapNode) getOrCreate(partitionRec partitionRecord, fieldInfo
[]partitionFieldInfo, numRows int64) *partitionInfo {
- // Navigate through all but the last partition field
+ // Navigate through all but the last partition field.
node := n
for _, part := range partitionRec[:len(partitionRec)-1] {
- key := comparablePartitionKey(part)
- val, ok := node.children[key]
+ val, ok := node.child(part)
if !ok {
newNode := newPartitionMapNode()
- node.children[key] = newNode
+ node.setChild(part, newNode)
node = newNode
} else {
node = val.(*partitionMapNode)
}
}
- // Last level stores the actual partitionInfo
- lastKey := comparablePartitionKey(partitionRec[len(partitionRec)-1])
- partVal, ok := node.children[lastKey].(*partitionInfo)
+ // Last level stores the actual partitionInfo.
+ lastPart := partitionRec[len(partitionRec)-1]
+ child, _ := node.child(lastPart)
+ partVal, ok := child.(*partitionInfo)
if ok {
return partVal
}
@@ -710,7 +740,7 @@ func (n *partitionMapNode) getOrCreate(partitionRec
partitionRecord, fieldInfo [
partitionValues: partitionValues,
partitionRec: partRecCopy,
}
- node.children[lastKey] = partVal
+ node.setChild(lastPart, partVal)
n.partitionCount++
return partVal
@@ -748,7 +778,7 @@ func (n *partitionMapNode) collectPartitions()
[]*partitionInfo {
}
func (n *partitionMapNode) appendPartitions(result []*partitionInfo)
[]*partitionInfo {
- for _, v := range n.children {
+ appendChild := func(v any) {
switch node := v.(type) {
case *partitionInfo:
result = append(result, node)
@@ -757,6 +787,13 @@ func (n *partitionMapNode) appendPartitions(result
[]*partitionInfo) []*partitio
}
}
+ for _, v := range n.children {
+ appendChild(v)
+ }
+ for _, v := range n.binaryChildren {
+ appendChild(v)
+ }
+
return result
}
diff --git a/table/partitioned_fanout_writer_review_test.go
b/table/partitioned_fanout_writer_review_test.go
index 14e30bb08..61c08a3d3 100644
--- a/table/partitioned_fanout_writer_review_test.go
+++ b/table/partitioned_fanout_writer_review_test.go
@@ -87,6 +87,80 @@ func (s *FanoutWriterTestSuite)
TestBinaryPartitionValuesDoNotAliasArrowStorage(
}
}
+func (s *FanoutWriterTestSuite) TestBinaryPartitionKeysWorkAtEveryLevel() {
+ arrowSchema := arrow.NewSchema([]arrow.Field{
+ {Name: "first", Type: arrow.BinaryTypes.Binary, Nullable: true},
+ {Name: "second", Type: arrow.BinaryTypes.Binary, Nullable:
true},
+ }, nil)
+ record := s.createCustomTestRecord(arrowSchema, [][]any{
+ {[]byte("a"), []byte("x")},
+ {[]byte("a"), []byte("x")},
+ {[]byte("a"), []byte("y")},
+ {[]byte("b"), []byte("x")},
+ {[]byte("b"), []byte("x")},
+ {nil, []byte("x")},
+ {[]byte{}, []byte("x")},
+ {nil, []byte("x")},
+ {[]byte{}, []byte("x")},
+ {[]byte("a"), nil},
+ {[]byte("a"), []byte{}},
+ {[]byte("a"), nil},
+ {[]byte("a"), []byte{}},
+ {nil, nil},
+ {[]byte{}, []byte{}},
+ {[]byte("a"), []byte{0, 0xff}},
+ {[]byte("a"), []byte{0, 0xff}},
+ {[]byte("a"), []byte("x/y")},
+ {[]byte("a/x"), []byte("y")},
+ })
+ defer record.Release()
+
+ icebergSchema := iceberg.NewSchema(1,
+ iceberg.NestedField{ID: 1, Name: "first", Type:
iceberg.PrimitiveTypes.Binary},
+ iceberg.NestedField{ID: 2, Name: "second", Type:
iceberg.PrimitiveTypes.Binary},
+ )
+ spec := iceberg.NewPartitionSpec(
+ iceberg.PartitionField{
+ SourceIDs: []int{1}, FieldID: 1000, Name: "first",
Transform: iceberg.IdentityTransform{},
+ },
+ iceberg.PartitionField{
+ SourceIDs: []int{2}, FieldID: 1001, Name: "second",
Transform: iceberg.IdentityTransform{},
+ },
+ )
+
+ partitions, err := getRecordPartitions(spec, icebergSchema, record)
+ s.Require().NoError(err)
+ s.Require().Len(partitions, 12)
+
+ // Keep NULL separate from an empty binary value at either tree level.
+ // A tuple key also preserves arbitrary bytes and embedded separators.
+ rowsByPartition := make(map[[2]any][]int64)
+ for _, partition := range partitions {
+ var key [2]any
+ for i, value := range partition.partitionRec {
+ if value != nil {
+ key[i] = string(value.([]byte))
+ }
+ }
+ rowsByPartition[key] = partition.rows
+ }
+
+ s.Equal(map[[2]any][]int64{
+ {"a", "x"}: {0, 1},
+ {"a", "y"}: {2},
+ {"b", "x"}: {3, 4},
+ {nil, "x"}: {5, 7},
+ {"", "x"}: {6, 8},
+ {"a", nil}: {9, 11},
+ {"a", ""}: {10, 12},
+ {nil, nil}: {13},
+ {"", ""}: {14},
+ {"a", "\x00\xff"}: {15, 16},
+ {"a", "x/y"}: {17},
+ {"a/x", "y"}: {18},
+ }, rowsByPartition)
+}
+
func (s *FanoutWriterTestSuite)
TestFixedSizeBinaryPartitionReportsWidthMismatch() {
arrowSchema := arrow.NewSchema([]arrow.Field{{
Name: "part",
