diff --git a/table/arrow_utils.go b/table/arrow_utils.go index 7fab00d6b..6d8ea0380 100644 --- a/table/arrow_utils.go +++ b/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 diff --git a/table/arrow_utils_bench_test.go b/table/arrow_utils_bench_test.go new file mode 100644 index 000000000..5f17737e1 --- /dev/null +++ b/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...) +} diff --git a/table/arrow_utils_internal_test.go b/table/arrow_utils_internal_test.go index 2c9b92f86..c6ae656b2 100644 --- a/table/arrow_utils_internal_test.go +++ b/table/arrow_utils_internal_test.go @@ -19,10 +19,8 @@ package table import ( "bytes" - "cmp" "encoding/json" "math" - "slices" "testing" "github.com/apache/arrow-go/v18/arrow/array" @@ -428,34 +426,157 @@ var tableSchemaNested = iceberg.NewSchemaWithIdentifiers(1, ) func TestStatsTypes(t *testing.T) { - statsCols, err := iceberg.PreOrderVisit(tableSchemaNested, - &arrowStatsCollector{schema: tableSchemaNested, defaultMode: "full"}) - + plan, err := computeStatsPlan(tableSchemaNested, iceberg.Properties{ + DefaultWriteMetricsModeKey: string(internal.MetricModeFull), + }) require.NoError(t, err) - // field ids should be sorted - assert.True(t, slices.IsSortedFunc(statsCols, func(a, b internal.StatisticsCollector) int { - return cmp.Compare(a.FieldID, b.FieldID) - })) + full := internal.MetricsMode{Typ: internal.MetricModeFull} + counts := internal.MetricsMode{Typ: internal.MetricModeCounts} + assert.Equal(t, map[int]internal.StatisticsCollector{ + 1: {FieldID: 1, IcebergTyp: iceberg.PrimitiveTypes.String, ColName: "foo", Mode: full}, + 2: {FieldID: 2, IcebergTyp: iceberg.PrimitiveTypes.Int32, ColName: "bar", Mode: full}, + 3: {FieldID: 3, IcebergTyp: iceberg.PrimitiveTypes.Bool, ColName: "baz", Mode: full}, + 5: {FieldID: 5, IcebergTyp: iceberg.PrimitiveTypes.String, ColName: "qux.element", Mode: counts}, + 7: {FieldID: 7, IcebergTyp: iceberg.PrimitiveTypes.String, ColName: "quux.key", Mode: counts}, + 9: {FieldID: 9, IcebergTyp: iceberg.PrimitiveTypes.String, ColName: "quux.value.key", Mode: counts}, + 10: {FieldID: 10, IcebergTyp: iceberg.PrimitiveTypes.Int32, ColName: "quux.value.value", Mode: counts}, + 13: {FieldID: 13, IcebergTyp: iceberg.PrimitiveTypes.Float32, ColName: "location.element.latitude", Mode: counts}, + 14: {FieldID: 14, IcebergTyp: iceberg.PrimitiveTypes.Float32, ColName: "location.element.longitude", Mode: counts}, + 16: {FieldID: 16, IcebergTyp: iceberg.PrimitiveTypes.String, ColName: "person.name", Mode: counts}, + 17: {FieldID: 17, IcebergTyp: iceberg.PrimitiveTypes.Int32, ColName: "person.age", Mode: counts}, + }, plan) +} + +func TestComputeStatsPlanMetricsModes(t *testing.T) { + schema := iceberg.NewSchema(0, + iceberg.NestedField{ID: 1, Name: "text", Type: iceberg.PrimitiveTypes.String}, + iceberg.NestedField{ID: 2, Name: "binary", Type: iceberg.PrimitiveTypes.Binary}, + iceberg.NestedField{ID: 3, Name: "number", Type: iceberg.PrimitiveTypes.Int32}, + iceberg.NestedField{ID: 4, Name: "variant", Type: iceberg.VariantType{}}, + iceberg.NestedField{ID: 5, Name: "dotted.name", Type: iceberg.PrimitiveTypes.String}, + iceberg.NestedField{ID: 6, Name: "nested", Type: &iceberg.StructType{FieldList: []iceberg.NestedField{ + {ID: 7, Name: "text", Type: iceberg.PrimitiveTypes.String}, + {ID: 8, Name: "variant", Type: iceberg.VariantType{}}, + }}}, + ) + full := internal.MetricsMode{Typ: internal.MetricModeFull} + counts := internal.MetricsMode{Typ: internal.MetricModeCounts} + none := internal.MetricsMode{Typ: internal.MetricModeNone} + truncate := internal.MetricsMode{Typ: internal.MetricModeTruncate, Len: 16} + truncateTwo := internal.MetricsMode{Typ: internal.MetricModeTruncate, Len: 2} + + for _, tt := range []struct { + name string + props iceberg.Properties + expected map[int]internal.MetricsMode + }{ + { + name: "default", + expected: map[int]internal.MetricsMode{1: truncate, 2: truncate, 3: full, 4: truncate, 5: counts, 7: counts, 8: truncate}, + }, + { + name: "full", + props: iceberg.Properties{DefaultWriteMetricsModeKey: "full"}, + expected: map[int]internal.MetricsMode{1: full, 2: full, 3: full, 4: full, 5: counts, 7: counts, 8: full}, + }, + { + name: "counts", + props: iceberg.Properties{DefaultWriteMetricsModeKey: "counts"}, + expected: map[int]internal.MetricsMode{1: counts, 2: counts, 3: counts, 4: counts, 5: counts, 7: counts, 8: counts}, + }, + { + name: "none", + props: iceberg.Properties{DefaultWriteMetricsModeKey: "none"}, + expected: map[int]internal.MetricsMode{1: none, 2: none, 3: none, 4: none, 5: none, 7: none, 8: none}, + }, + { + name: "custom truncation", + props: iceberg.Properties{DefaultWriteMetricsModeKey: " TRUNCATE(2) "}, + expected: map[int]internal.MetricsMode{1: truncateTwo, 2: truncateTwo, 3: full, 4: truncateTwo, 5: counts, 7: counts, 8: truncateTwo}, + }, + { + name: "column overrides", + props: iceberg.Properties{ + DefaultWriteMetricsModeKey: "none", + MetricsModeColumnConfPrefix + ".text": "truncate(2)", + MetricsModeColumnConfPrefix + ".binary": "full", + MetricsModeColumnConfPrefix + ".number": "truncate(2)", + MetricsModeColumnConfPrefix + ".variant": "counts", + MetricsModeColumnConfPrefix + ".dotted.name": "full", + MetricsModeColumnConfPrefix + ".nested.text": "full", + MetricsModeColumnConfPrefix + ".nested.variant": "truncate(2)", + MetricsModeColumnConfPrefix + ".not_present": "truncate(0)", + "unrelated.property": "truncate(0)", + }, + expected: map[int]internal.MetricsMode{1: truncateTwo, 2: full, 3: full, 4: counts, 5: counts, 7: counts, 8: truncateTwo}, + }, + { + name: "nested none overrides", + props: iceberg.Properties{ + DefaultWriteMetricsModeKey: "full", + MetricsModeColumnConfPrefix + ".nested.text": "none", + MetricsModeColumnConfPrefix + ".nested.variant": "none", + }, + expected: map[int]internal.MetricsMode{1: full, 2: full, 3: full, 4: full, 5: counts, 7: none, 8: none}, + }, + } { + t.Run(tt.name, func(t *testing.T) { + plan, err := computeStatsPlan(schema, tt.props) + require.NoError(t, err) + require.Len(t, plan, len(tt.expected)) + for fieldID, expected := range tt.expected { + require.Contains(t, plan, fieldID) + assert.Equal(t, expected, plan[fieldID].Mode, "field %d", fieldID) + } + assert.Nil(t, plan[4].IcebergTyp) + assert.Nil(t, plan[8].IcebergTyp) + assert.Equal(t, "variant", plan[4].ColName) + assert.Equal(t, "nested.variant", plan[8].ColName) + }) + } +} - actual := make([]iceberg.Type, len(statsCols)) - for i, col := range statsCols { - actual[i] = col.IcebergTyp +func TestComputeStatsPlanInvalidModes(t *testing.T) { + schema := iceberg.NewSchema(0, + iceberg.NestedField{ID: 1, Name: "text", Type: iceberg.PrimitiveTypes.String}, + iceberg.NestedField{ID: 2, Name: "variant", Type: iceberg.VariantType{}}, + ) + for _, tt := range []struct { + name string + props iceberg.Properties + }{ + {name: "invalid default", props: iceberg.Properties{DefaultWriteMetricsModeKey: "truncate(0)"}}, + {name: "invalid primitive override", props: iceberg.Properties{MetricsModeColumnConfPrefix + ".text": "truncate(0)"}}, + {name: "invalid variant override", props: iceberg.Properties{MetricsModeColumnConfPrefix + ".variant": "truncate(0)"}}, + { + name: "invalid default with valid overrides", + props: iceberg.Properties{ + DefaultWriteMetricsModeKey: "truncate(0)", + MetricsModeColumnConfPrefix + ".text": "full", + MetricsModeColumnConfPrefix + ".variant": "counts", + }, + }, + } { + t.Run(tt.name, func(t *testing.T) { + plan, err := computeStatsPlan(schema, tt.props) + require.ErrorContains(t, err, "invalid truncate length: 0") + assert.Nil(t, plan) + }) } +} + +func TestComputeStatsPlanIgnoresInvalidUnusedColumnMode(t *testing.T) { + schema := iceberg.NewSchema(0, iceberg.NestedField{ + ID: 1, Name: "id", Type: iceberg.PrimitiveTypes.Int32, Required: true, + }) - assert.Equal(t, []iceberg.Type{ - iceberg.PrimitiveTypes.String, - iceberg.PrimitiveTypes.Int32, - iceberg.PrimitiveTypes.Bool, - iceberg.PrimitiveTypes.String, - iceberg.PrimitiveTypes.String, - iceberg.PrimitiveTypes.String, - iceberg.PrimitiveTypes.Int32, - iceberg.PrimitiveTypes.Float32, - iceberg.PrimitiveTypes.Float32, - iceberg.PrimitiveTypes.String, - iceberg.PrimitiveTypes.Int32, - }, actual) + _, err := computeStatsPlan(schema, iceberg.Properties{ + DefaultWriteMetricsModeKey: "counts", + MetricsModeColumnConfPrefix + ".not_present": "truncate(0)", + }) + + require.NoError(t, err) } func TestIcebergCRSToGeoArrowMetadata(t *testing.T) {