perf(table): compile metrics modes once (#1991)

* perf(table): compile metrics modes once

* fix(table): address metrics plan review feedback

Signed-off-by: Minh Vu <vuhoangminh97@gmail.com>

---------

Signed-off-by: Minh Vu <vuhoangminh97@gmail.com>

Author: fallintoplace

table/arrow_utils.go

@@ -1664,66 +1664,33 @@ func must[T any](v T, err error) T {
 }
 
 type arrowStatsCollector struct {
-	fieldID     int
-	schema      *iceberg.Schema
-	props       iceberg.Properties
-	defaultMode string
+	fieldID          int
+	schema           *iceberg.Schema
+	defaultMode      tblutils.MetricsMode
+	defaultModeError error
+	columnModes      map[string]tblutils.MetricsMode
+	columnModeErrors map[string]error
 }
 
-func (a *arrowStatsCollector) Schema(_ *iceberg.Schema, results func() []tblutils.StatisticsCollector) []tblutils.StatisticsCollector {
-	return results()
-}
-
-func (a *arrowStatsCollector) Struct(_ iceberg.StructType, results []func() []tblutils.StatisticsCollector) []tblutils.StatisticsCollector {
-	result := make([]tblutils.StatisticsCollector, 0, len(results))
-	for _, res := range results {
-		result = append(result, res()...)
+func (a *arrowStatsCollector) resolveColumnMetricsMode(colName string) tblutils.MetricsMode {
+	if a.defaultModeError != nil {
+		panic(a.defaultModeError)
 	}
 
-	return result
-}
-
-func (a *arrowStatsCollector) Field(field iceberg.NestedField, fieldRes func() []tblutils.StatisticsCollector) []tblutils.StatisticsCollector {
-	a.fieldID = field.ID
-
-	return fieldRes()
-}
-
-func (a *arrowStatsCollector) List(list iceberg.ListType, elemResult func() []tblutils.StatisticsCollector) []tblutils.StatisticsCollector {
-	a.fieldID = list.ElementID
-
-	return elemResult()
-}
-
-func (a *arrowStatsCollector) Map(m iceberg.MapType, keyResult, valResult func() []tblutils.StatisticsCollector) []tblutils.StatisticsCollector {
-	a.fieldID = m.KeyID
-	keyRes := keyResult()
-
-	a.fieldID = m.ValueID
-	valRes := valResult()
-
-	return append(keyRes, valRes...)
-}
-
-func (a *arrowStatsCollector) resolveColumnMetricsMode(colName string) tblutils.MetricsMode {
-	metMode, err := tblutils.MatchMetricsMode(a.defaultMode)
-	if err != nil {
-		panic(err)
+	if metMode, ok := a.columnModes[colName]; ok {
+		return metMode
 	}
-	if colMode, ok := a.props[MetricsModeColumnConfPrefix+"."+colName]; ok {
-		metMode, err = tblutils.MatchMetricsMode(colMode)
-		if err != nil {
-			panic(err)
-		}
+	if err, ok := a.columnModeErrors[colName]; ok {
+		panic(err)
 	}
 
-	return metMode
+	return a.defaultMode
 }
 
-func (a *arrowStatsCollector) Primitive(dt iceberg.PrimitiveType) []tblutils.StatisticsCollector {
+func (a *arrowStatsCollector) primitiveCollector(dt iceberg.PrimitiveType, isNested bool) (tblutils.StatisticsCollector, bool) {
 	colName, ok := a.schema.FindColumnName(a.fieldID)
 	if !ok {
-		return []tblutils.StatisticsCollector{}
+		return tblutils.StatisticsCollector{}, false
 	}
 
 	metMode := a.resolveColumnMetricsMode(colName)
@@ -1737,47 +1704,141 @@ func (a *arrowStatsCollector) Primitive(dt iceberg.PrimitiveType) []tblutils.Sta
 		}
 	}
 
-	isNested := strings.Contains(colName, ".")
 	if isNested && (metMode.Typ == tblutils.MetricModeTruncate || metMode.Typ == tblutils.MetricModeFull) {
 		metMode = tblutils.MetricsMode{Typ: tblutils.MetricModeCounts}
 	}
 
-	return []tblutils.StatisticsCollector{{
+	return tblutils.StatisticsCollector{
 		FieldID:    a.fieldID,
 		IcebergTyp: dt,
 		ColName:    colName,
 		Mode:       metMode,
-	}}
+	}, true
 }
 
-func (a *arrowStatsCollector) Variant(_ iceberg.VariantType) []tblutils.StatisticsCollector {
+func (a *arrowStatsCollector) variantCollector() (tblutils.StatisticsCollector, bool) {
 	colName, ok := a.schema.FindColumnName(a.fieldID)
 	if !ok {
-		return []tblutils.StatisticsCollector{}
+		return tblutils.StatisticsCollector{}, false
 	}
 
-	return []tblutils.StatisticsCollector{{
+	return tblutils.StatisticsCollector{
 		FieldID: a.fieldID,
 		ColName: colName,
 		Mode:    a.resolveColumnMetricsMode(colName),
-	}}
+	}, true
+}
+
+func statsPlanFieldCount(field iceberg.NestedField) int {
+	switch typ := field.Type.(type) {
+	case *iceberg.StructType:
+		count := 0
+		for _, nestedField := range typ.FieldList {
+			count += statsPlanFieldCount(nestedField)
+		}
+
+		return count
+	case *iceberg.ListType:
+		return statsPlanFieldCount(typ.ElementField())
+	case *iceberg.MapType:
+		return statsPlanFieldCount(typ.KeyField()) + statsPlanFieldCount(typ.ValueField())
+	default:
+		return 1
+	}
+}
+
+func collectStatsPlanField(visitor *arrowStatsCollector, result map[int]tblutils.StatisticsCollector, field iceberg.NestedField, isNested bool) {
+	switch typ := field.Type.(type) {
+	case *iceberg.StructType:
+		for _, nestedField := range typ.FieldList {
+			collectStatsPlanField(visitor, result, nestedField, true)
+		}
+	case *iceberg.ListType:
+		collectStatsPlanField(visitor, result, typ.ElementField(), true)
+	case *iceberg.MapType:
+		collectStatsPlanField(visitor, result, typ.KeyField(), true)
+		collectStatsPlanField(visitor, result, typ.ValueField(), true)
+	case iceberg.VariantType:
+		visitor.fieldID = field.ID
+		if collector, ok := visitor.variantCollector(); ok {
+			result[collector.FieldID] = collector
+		}
+	default:
+		visitor.fieldID = field.ID
+		collector, ok := visitor.primitiveCollector(field.Type.(iceberg.PrimitiveType), isNested)
+		if ok {
+			result[collector.FieldID] = collector
+		}
+	}
 }
 
-func computeStatsPlan(sc *iceberg.Schema, props iceberg.Properties) (map[int]tblutils.StatisticsCollector, error) {
-	result := make(map[int]tblutils.StatisticsCollector)
+func computeStatsPlan(sc *iceberg.Schema, props iceberg.Properties) (result map[int]tblutils.StatisticsCollector, err error) {
+	defer func() {
+		if r := recover(); r != nil {
+			result = nil
+			switch e := r.(type) {
+			case string:
+				err = fmt.Errorf("%w: %s", iceberg.ErrInvalidSchema, e)
+			case error:
+				err = fmt.Errorf("error encountered during schema visitor: %w", e)
+			}
+		}
+	}()
+
+	if sc == nil {
+		return nil, fmt.Errorf("%w: cannot visit nil schema", iceberg.ErrInvalidArgument)
+	}
+
+	defaultMode, defaultModeErr := tblutils.MatchMetricsMode(
+		props.Get(DefaultWriteMetricsModeKey, DefaultWriteMetricsModeDefault))
+	overrideCount := 0
+	for key := range props {
+		if strings.HasPrefix(key, MetricsModeColumnConfPrefix+".") {
+			overrideCount++
+		}
+	}
+
+	var columnModes map[string]tblutils.MetricsMode
+	var columnModeErrors map[string]error
+	for key, rawMode := range props {
+		colName, ok := strings.CutPrefix(key, MetricsModeColumnConfPrefix+".")
+		if !ok {
+			continue
+		}
+
+		mode, err := tblutils.MatchMetricsMode(rawMode)
+		if err != nil {
+			if columnModeErrors == nil {
+				columnModeErrors = make(map[string]error, overrideCount)
+			}
+
+			columnModeErrors[colName] = err
+
+			continue
+		}
+		if columnModes == nil {
+			columnModes = make(map[string]tblutils.MetricsMode, overrideCount)
+		}
+		columnModes[colName] = mode
+	}
 
 	visitor := &arrowStatsCollector{
-		schema: sc, props: props,
-		defaultMode: props.Get(DefaultWriteMetricsModeKey, DefaultWriteMetricsModeDefault),
+		schema:           sc,
+		defaultMode:      defaultMode,
+		defaultModeError: defaultModeErr,
+		columnModes:      columnModes,
+		columnModeErrors: columnModeErrors,
 	}
 
-	collectors, err := iceberg.PreOrderVisit(sc, visitor)
-	if err != nil {
-		return nil, err
+	fields := sc.FieldsRef(internal.SchemaRef{})
+	resultCount := 0
+	for _, field := range fields {
+		resultCount += statsPlanFieldCount(field)
 	}
 
-	for _, entry := range collectors {
-		result[entry.FieldID] = entry
+	result = make(map[int]tblutils.StatisticsCollector, resultCount)
+	for _, field := range fields {
+		collectStatsPlanField(visitor, result, field, strings.Contains(field.Name, "."))
 	}
 
 	return result, nil

table/arrow_utils_bench_test.go

@@ -0,0 +1,77 @@
+// 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"
+	"testing"
+
+	"github.com/apache/iceberg-go"
+)
+
+func BenchmarkComputeStatsPlan(b *testing.B) {
+	for _, fieldCount := range []int{100, 1000, 10000} {
+		for _, benchmarkCase := range []struct {
+			name                string
+			defaultMode         string
+			overrideStride      int
+			unrelatedProperties int
+		}{
+			{name: "default", defaultMode: "truncate(16)"},
+			{name: "one_percent_overrides", defaultMode: "truncate(16)", overrideStride: 100},
+			{name: "one_percent_overrides_many_properties", defaultMode: "truncate(16)", overrideStride: 100, unrelatedProperties: 1000},
+		} {
+			b.Run(fmt.Sprintf("fields=%d/%s", fieldCount, benchmarkCase.name), func(b *testing.B) {
+				schema := benchmarkMetricsSchema(fieldCount)
+				props := iceberg.Properties{DefaultWriteMetricsModeKey: benchmarkCase.defaultMode}
+				for i := 0; benchmarkCase.overrideStride > 0 && i < fieldCount; i += benchmarkCase.overrideStride {
+					props[MetricsModeColumnConfPrefix+fmt.Sprintf(".field_%d", i)] = "counts"
+				}
+				for i := range benchmarkCase.unrelatedProperties {
+					props[fmt.Sprintf("unrelated.property.%d", i)] = "value"
+				}
+
+				b.ReportAllocs()
+				b.ResetTimer()
+				for range b.N {
+					plan, err := computeStatsPlan(schema, props)
+					if err != nil {
+						b.Fatal(err)
+					}
+					if len(plan) != fieldCount {
+						b.Fatalf("expected %d stats columns, got %d", fieldCount, len(plan))
+					}
+				}
+			})
+		}
+	}
+}
+
+func benchmarkMetricsSchema(fieldCount int) *iceberg.Schema {
+	fields := make([]iceberg.NestedField, fieldCount)
+	for i := range fields {
+		fields[i] = iceberg.NestedField{
+			ID:       i + 1,
+			Name:     fmt.Sprintf("field_%d", i),
+			Type:     iceberg.PrimitiveTypes.String,
+			Required: true,
+		}
+	}
+
+	return iceberg.NewSchema(0, fields...)
+}