laskoviymishka commented on code in PR #1768: URL: https://github.com/apache/iceberg-go/pull/1768#discussion_r3850569225
########## table/partition_extraction_bench_test.go: ########## @@ -0,0 +1,197 @@ +// 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/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 BenchmarkPartitionExtraction(b *testing.B) { + const partitionFields = 4 + + arrowFields := make([]arrow.Field, partitionFields) + icebergFields := make([]iceberg.NestedField, partitionFields) + specFields := make([]iceberg.PartitionField, partitionFields) + for i := range partitionFields { + name := fmt.Sprintf("part_%d", i) + arrowFields[i] = arrow.Field{Name: name, Type: arrow.PrimitiveTypes.Int32} + icebergFields[i] = iceberg.NestedField{ID: i + 1, Name: name, Type: iceberg.PrimitiveTypes.Int32} + specFields[i] = iceberg.PartitionField{ + SourceIDs: []int{i + 1}, + FieldID: 1000 + i, + Name: name, + Transform: iceberg.IdentityTransform{}, + } + } + + arrowSchema := arrow.NewSchema(arrowFields, nil) + icebergSchema := iceberg.NewSchema(0, icebergFields...) + spec := iceberg.NewPartitionSpec(specFields...) + + for _, rows := range []int{0, 1, 16, 1024} { + b.Run(fmt.Sprintf("rows_%d", rows), func(b *testing.B) { + columns := make([]arrow.Array, partitionFields) + for i := range columns { + builder := array.NewInt32Builder(memory.DefaultAllocator) + for row := range rows { + builder.Append(int32(row % 8)) + } + columns[i] = builder.NewArray() + builder.Release() + } + + record := array.NewRecordBatch(arrowSchema, columns, int64(rows)) + for _, column := range columns { + column.Release() + } + defer record.Release() + writer := newPartitionedFanoutWriter(spec, icebergSchema, nil, nil) Review Comment: This folds the one-time plan build into the measured loop: the first `b.Loop()` iteration runs `planOnce.Do` and the whole plan construction. `BenchmarkPartitionTransforms` already warms up with a `getPartitions` call before `ResetTimer`; I'd do the same here so the `rows_0`/`rows_1` numbers aren't dominated by plan setup. ########## table/partitioned_fanout_writer.go: ########## @@ -320,60 +349,149 @@ func getRecordPartitions(spec iceberg.PartitionSpec, schema *iceberg.Schema, rec if !ok { continue } - colIndices := record.Schema().FieldIndices(colName) + colIndices := recordSchema.FieldIndices(colName) if len(colIndices) == 0 { return nil, fmt.Errorf("failed to find source column %q in record schema", colName) } sourceType, ok := schema.FindTypeByID(sourceField.SourceID()) if !ok { return nil, fmt.Errorf("failed to find type for source field ID %d in schema", sourceField.SourceID()) } - partitionColumns[i] = record.Column(colIndices[0]) partitionFieldsInfo[i] = partitionFieldInfo{ sourceField: sourceField, sourceName: colName, fieldID: sourceField.FieldID, sourceType: sourceType, + columnIndex: colIndices[0], + valueAt: bindPartitionValue(sourceField.Transform, sourceType), + } + } + + return &partitionExtractionPlan{ + spec: spec, + schema: schema, + recordSchema: recordSchema, + fields: partitionFieldsInfo, + }, nil +} + +func (p *partitionExtractionPlan) getRecordPartitions(record arrow.RecordBatch) ([]*partitionInfo, error) { + // Preserve support for iterators whose batch schema changes. The usual path compares + // schema pointers; equivalent independently-built schemas also reuse the plan. + if record.Schema() != p.recordSchema && !record.Schema().Equal(p.recordSchema) { + plan, err := newPartitionExtractionPlan(p.spec, p.schema, record.Schema()) + if err != nil { + return nil, err + } + + return plan.getRecordPartitions(record) + } + + partitionMap := newPartitionMapNode() + partitionRec := make(partitionRecord, len(p.fields)) + partitionColumns := make([]arrow.Array, len(p.fields)) + for i, fieldInfo := range p.fields { + if fieldInfo.columnIndex >= 0 { + partitionColumns[i] = record.Column(fieldInfo.columnIndex) } } for row := range record.NumRows() { - for i := range partitionFields { + for i, fieldInfo := range p.fields { col := partitionColumns[i] if col != nil && !col.IsNull(int(row)) { - fieldInfo := partitionFieldsInfo[i] - sourceField := fieldInfo.sourceField - val, err := getArrowValueAsIcebergLiteral(col, int(row), fieldInfo.sourceType) + value, err := fieldInfo.valueAt(col, int(row)) if err != nil { return nil, fmt.Errorf( "failed to convert source column %q (field ID %d) from Arrow type %s to Iceberg type %s: %w", fieldInfo.sourceName, - sourceField.SourceID(), + fieldInfo.sourceField.SourceID(), col.DataType(), fieldInfo.sourceType, err, ) } - transformedLiteral := sourceField.Transform.Apply(iceberg.Optional[iceberg.Literal]{Valid: true, Val: val}) - if transformedLiteral.Valid { - partitionRec[i] = transformedLiteral.Val.Any() - } else { - partitionRec[i] = nil - } + partitionRec[i] = value } else { partitionRec[i] = nil } } // Get or create partition info for this partition key - partVal := partitionMap.getOrCreate(partitionRec, partitionFieldsInfo) + partVal := partitionMap.getOrCreate(partitionRec, p.fields) partVal.rows = append(partVal.rows, row) } return partitionMap.collectPartitions(), nil } +func bindPartitionValue(transform iceberg.Transform, sourceType iceberg.Type) func(arrow.Array, int) (any, error) { + bound, ok := bindPartitionTransform(transform, sourceType) + if !ok { + return func(column arrow.Array, row int) (any, error) { + value, err := getArrowValueAsIcebergLiteral(column, row, sourceType) + if err != nil { + return nil, err + } + + transformed := transform.Apply(iceberg.Optional[iceberg.Literal]{Valid: true, Val: value}) + if !transformed.Valid { + return nil, nil + } + + return transformed.Val.Any(), nil + } + } + + return func(column arrow.Array, row int) (any, error) { + value, err := getArrowValueAsIcebergValue(column, row, sourceType) + if err != nil { + return nil, err + } + + return bound(value), nil + } +} + +func bindPartitionTransform(transform iceberg.Transform, sourceType iceberg.Type) (func(any) any, bool) { + optionalInt := func(transformer func(any) iceberg.Optional[int32]) func(any) any { + return func(value any) any { + transformed := transformer(value) + if !transformed.Valid { + return nil + } + + return transformed.Val + } + } + + switch typed := transform.(type) { + case iceberg.IdentityTransform: + return func(value any) any { return value }, true + case iceberg.VoidTransform: + return func(any) any { return nil }, true + case iceberg.BucketTransform: + return optionalInt(typed.Transformer(sourceType)), true Review Comment: Truncate guards on the `Transformer` error and falls back to `Apply`; Bucket calls `Transformer(sourceType)` unconditionally. Can `BucketTransform.Transformer` hand back a transformer that fails on an unsupported source type here? If so we'd panic in the row loop instead of falling back the way Truncate does. wdyt? ########## table/partitioned_fanout_writer_test.go: ########## @@ -683,6 +683,168 @@ func (s *FanoutWriterTestSuite) TestGetRecordPartitionsWithDroppedLeadingSourceC s.Equal("foo=null/bar=7/baz=true", spec.PartitionToPath(partitions[0].partitionRec, icebergSchema)) } +func (s *FanoutWriterTestSuite) TestPartitionedWriterReusesExtractionPlan() { + icebergSchema := iceberg.NewSchema(0, + iceberg.NestedField{ID: 1, Name: "part", Type: iceberg.PrimitiveTypes.Int32}, + iceberg.NestedField{ID: 2, Name: "value", Type: iceberg.PrimitiveTypes.String}, + ) + spec := iceberg.NewPartitionSpec(iceberg.PartitionField{ + SourceIDs: []int{1}, FieldID: 1000, Transform: iceberg.IdentityTransform{}, Name: "part", + }) + + arrowSchema := arrow.NewSchema([]arrow.Field{ + {Name: "part", Type: arrow.PrimitiveTypes.Int32}, + {Name: "value", Type: arrow.BinaryTypes.String}, + }, nil) + firstRecord := s.createCustomTestRecord(arrowSchema, [][]any{{int32(7), "a"}}) + defer firstRecord.Release() + + writer := newPartitionedFanoutWriter(spec, icebergSchema, nil, nil) + partitions, err := writer.getPartitions(firstRecord) + s.Require().NoError(err) + s.Require().Len(partitions, 1) + s.Equal(int32(7), partitions[0].partitionRec.Get(0)) + firstPlan := writer.plan + s.Require().NotNil(firstPlan) + + equivalentSchema := arrow.NewSchema(arrowSchema.Fields(), nil) + secondRecord := s.createCustomTestRecord(equivalentSchema, [][]any{{int32(8), "b"}}) + defer secondRecord.Release() + + partitions, err = writer.getPartitions(secondRecord) + s.Require().NoError(err) + s.Require().Len(partitions, 1) + s.Equal(int32(8), partitions[0].partitionRec.Get(0)) + s.Same(firstPlan, writer.plan) +} + +func (s *FanoutWriterTestSuite) TestPartitionExtractionPlanHandlesReorderedRecordSchema() { + icebergSchema := iceberg.NewSchema(0, + iceberg.NestedField{ID: 1, Name: "part", Type: iceberg.PrimitiveTypes.Int32}, + iceberg.NestedField{ID: 2, Name: "value", Type: iceberg.PrimitiveTypes.String}, + ) + spec := iceberg.NewPartitionSpec(iceberg.PartitionField{ + SourceIDs: []int{1}, FieldID: 1000, Transform: iceberg.IdentityTransform{}, Name: "part", + }) + + originalSchema := arrow.NewSchema([]arrow.Field{ + {Name: "part", Type: arrow.PrimitiveTypes.Int32}, + {Name: "value", Type: arrow.BinaryTypes.String}, + }, nil) + plan, err := newPartitionExtractionPlan(spec, icebergSchema, originalSchema) + s.Require().NoError(err) + s.Equal(0, plan.fields[0].columnIndex) + + reorderedSchema := arrow.NewSchema([]arrow.Field{ + {Name: "value", Type: arrow.BinaryTypes.String}, + {Name: "part", Type: arrow.PrimitiveTypes.Int32}, + }, nil) + record := s.createCustomTestRecord(reorderedSchema, [][]any{{"a", int32(7)}, {"b", int32(8)}}) + defer record.Release() + + partitions, err := plan.getRecordPartitions(record) + s.Require().NoError(err) + s.Require().Len(partitions, 2) + values := []int32{ + partitions[0].partitionRec.Get(0).(int32), + partitions[1].partitionRec.Get(0).(int32), + } + s.ElementsMatch([]int32{7, 8}, values) +} + +func (s *FanoutWriterTestSuite) TestBoundPartitionTransformsMatchGenericApply() { + unknown, err := iceberg.ParseTransform("custom-transform") + s.Require().NoError(err) + + tests := []struct { + name string + transform iceberg.Transform + sourceType iceberg.Type + value iceberg.Literal + fallback bool + }{ + { + name: "identity", transform: iceberg.IdentityTransform{}, + sourceType: iceberg.PrimitiveTypes.Int64, value: iceberg.Int64Literal(34), + }, + { + name: "void", transform: iceberg.VoidTransform{}, + sourceType: iceberg.PrimitiveTypes.String, value: iceberg.StringLiteral("abc"), + }, + { + name: "bucket", transform: iceberg.BucketTransform{NumBuckets: 16}, + sourceType: iceberg.PrimitiveTypes.String, value: iceberg.StringLiteral("abc"), + }, + { + name: "truncate", transform: iceberg.TruncateTransform{Width: 3}, + sourceType: iceberg.PrimitiveTypes.String, value: iceberg.StringLiteral("abcdef"), + }, + { + name: "year", transform: iceberg.YearTransform{}, + sourceType: iceberg.PrimitiveTypes.Date, value: iceberg.DateLiteral(19_358), + }, + { + name: "month", transform: iceberg.MonthTransform{}, + sourceType: iceberg.PrimitiveTypes.Timestamp, value: iceberg.TimestampLiteral(1_672_531_200_000_000), + }, + { + name: "day nanoseconds", transform: iceberg.DayTransform{}, + sourceType: iceberg.PrimitiveTypes.TimestampNs, value: iceberg.TimestampNsLiteral(1_672_531_200_000_000_000), + }, + { + name: "hour", transform: iceberg.HourTransform{}, + sourceType: iceberg.PrimitiveTypes.Timestamp, value: iceberg.TimestampLiteral(1_672_531_200_000_000), + }, + { + name: "invalid truncate fallback", transform: iceberg.TruncateTransform{Width: 0}, + sourceType: iceberg.PrimitiveTypes.String, value: iceberg.StringLiteral("abc"), fallback: true, + }, + { + name: "unknown fallback", transform: unknown, + sourceType: iceberg.PrimitiveTypes.String, value: iceberg.StringLiteral("abc"), + }, + } + + for _, test := range tests { + s.Run(test.name, func() { + expected := test.transform.Apply(iceberg.Optional[iceberg.Literal]{Valid: true, Val: test.value}) + var expectedValue any + if expected.Valid { + expectedValue = expected.Val.Any() + } + + bound, ok := bindPartitionTransform(test.transform, test.sourceType) + if test.fallback { + s.False(ok) + + return + } + + s.Require().True(ok) + s.Equal(expectedValue, bound(test.value.Any())) Review Comment: This proves parity at the literal level: `bound(test.value.Any())` feeds a literal's native value, so it never touches `getArrowValueAsIcebergValue`, which is the actual production input. That leaves the cases where equivalence only holds incidentally untested: bucket on int32/date binds `hashHelperInt` at a narrower width than `Apply`'s int64 path (equal only because two's-complement conversion preserves the value mod 2^64), and bucket/truncate on decimal/uuid/fixed lean on `v.(Decimal)` / `uuid.UUID` assertions in the extractor that nothing here hits. Every case here is String or timestamp. I'd add bucket cases for int32, decimal, and uuid that build a real Arrow array, run it through `getArrowValueAsIcebergValue`, and compare against `Apply(getArrowValueAsIcebergLiteral(...))`. That pins the equivalence we're currently relying on rather than trusting it. ########## table/partitioned_fanout_writer.go: ########## @@ -320,60 +349,149 @@ func getRecordPartitions(spec iceberg.PartitionSpec, schema *iceberg.Schema, rec if !ok { continue } - colIndices := record.Schema().FieldIndices(colName) + colIndices := recordSchema.FieldIndices(colName) if len(colIndices) == 0 { return nil, fmt.Errorf("failed to find source column %q in record schema", colName) } sourceType, ok := schema.FindTypeByID(sourceField.SourceID()) if !ok { return nil, fmt.Errorf("failed to find type for source field ID %d in schema", sourceField.SourceID()) } - partitionColumns[i] = record.Column(colIndices[0]) partitionFieldsInfo[i] = partitionFieldInfo{ sourceField: sourceField, sourceName: colName, fieldID: sourceField.FieldID, sourceType: sourceType, + columnIndex: colIndices[0], + valueAt: bindPartitionValue(sourceField.Transform, sourceType), + } + } + + return &partitionExtractionPlan{ + spec: spec, + schema: schema, + recordSchema: recordSchema, + fields: partitionFieldsInfo, + }, nil +} + +func (p *partitionExtractionPlan) getRecordPartitions(record arrow.RecordBatch) ([]*partitionInfo, error) { + // Preserve support for iterators whose batch schema changes. The usual path compares + // schema pointers; equivalent independently-built schemas also reuse the plan. + if record.Schema() != p.recordSchema && !record.Schema().Equal(p.recordSchema) { Review Comment: The new tests cover the pointer-equal and equivalent-schema reuse paths, but nothing sends a genuinely divergent schema through the fanout writer to exercise this rebuild branch. I'd add a two-batch case with different-but-both-valid schemas and check both produce correct partition assignments. ########## table/partitioned_fanout_writer.go: ########## @@ -320,60 +349,149 @@ func getRecordPartitions(spec iceberg.PartitionSpec, schema *iceberg.Schema, rec if !ok { continue } - colIndices := record.Schema().FieldIndices(colName) + colIndices := recordSchema.FieldIndices(colName) if len(colIndices) == 0 { return nil, fmt.Errorf("failed to find source column %q in record schema", colName) } sourceType, ok := schema.FindTypeByID(sourceField.SourceID()) if !ok { return nil, fmt.Errorf("failed to find type for source field ID %d in schema", sourceField.SourceID()) } - partitionColumns[i] = record.Column(colIndices[0]) partitionFieldsInfo[i] = partitionFieldInfo{ sourceField: sourceField, sourceName: colName, fieldID: sourceField.FieldID, sourceType: sourceType, + columnIndex: colIndices[0], + valueAt: bindPartitionValue(sourceField.Transform, sourceType), + } + } + + return &partitionExtractionPlan{ + spec: spec, + schema: schema, + recordSchema: recordSchema, + fields: partitionFieldsInfo, + }, nil +} + +func (p *partitionExtractionPlan) getRecordPartitions(record arrow.RecordBatch) ([]*partitionInfo, error) { + // Preserve support for iterators whose batch schema changes. The usual path compares + // schema pointers; equivalent independently-built schemas also reuse the plan. + if record.Schema() != p.recordSchema && !record.Schema().Equal(p.recordSchema) { + plan, err := newPartitionExtractionPlan(p.spec, p.schema, record.Schema()) + if err != nil { + return nil, err + } + + return plan.getRecordPartitions(record) + } + + partitionMap := newPartitionMapNode() + partitionRec := make(partitionRecord, len(p.fields)) + partitionColumns := make([]arrow.Array, len(p.fields)) + for i, fieldInfo := range p.fields { + if fieldInfo.columnIndex >= 0 { + partitionColumns[i] = record.Column(fieldInfo.columnIndex) } } for row := range record.NumRows() { - for i := range partitionFields { + for i, fieldInfo := range p.fields { col := partitionColumns[i] if col != nil && !col.IsNull(int(row)) { - fieldInfo := partitionFieldsInfo[i] - sourceField := fieldInfo.sourceField - val, err := getArrowValueAsIcebergLiteral(col, int(row), fieldInfo.sourceType) + value, err := fieldInfo.valueAt(col, int(row)) if err != nil { return nil, fmt.Errorf( "failed to convert source column %q (field ID %d) from Arrow type %s to Iceberg type %s: %w", fieldInfo.sourceName, - sourceField.SourceID(), + fieldInfo.sourceField.SourceID(), col.DataType(), fieldInfo.sourceType, err, ) } - transformedLiteral := sourceField.Transform.Apply(iceberg.Optional[iceberg.Literal]{Valid: true, Val: val}) - if transformedLiteral.Valid { - partitionRec[i] = transformedLiteral.Val.Any() - } else { - partitionRec[i] = nil - } + partitionRec[i] = value } else { partitionRec[i] = nil } } // Get or create partition info for this partition key - partVal := partitionMap.getOrCreate(partitionRec, partitionFieldsInfo) + partVal := partitionMap.getOrCreate(partitionRec, p.fields) partVal.rows = append(partVal.rows, row) } return partitionMap.collectPartitions(), nil } +func bindPartitionValue(transform iceberg.Transform, sourceType iceberg.Type) func(arrow.Array, int) (any, error) { + bound, ok := bindPartitionTransform(transform, sourceType) + if !ok { + return func(column arrow.Array, row int) (any, error) { + value, err := getArrowValueAsIcebergLiteral(column, row, sourceType) + if err != nil { + return nil, err + } + + transformed := transform.Apply(iceberg.Optional[iceberg.Literal]{Valid: true, Val: value}) + if !transformed.Valid { + return nil, nil + } + + return transformed.Val.Any(), nil + } + } + + return func(column arrow.Array, row int) (any, error) { + value, err := getArrowValueAsIcebergValue(column, row, sourceType) + if err != nil { + return nil, err + } + + return bound(value), nil Review Comment: `getArrowValueAsIcebergValue` returns `(nil, nil)` on a null cell, and this closure hands that straight to `bound`: for bucket/truncate that's a type assertion on `nil`, i.e. a panic. It's safe today only because the sole caller guards with `col != nil && !col.IsNull(...)` before dispatching, while the fallback closure routes through `Apply` and is fine. I'd give this closure the same nil check the literal path has so it's self-contained: ```go value, err := getArrowValueAsIcebergValue(column, row, sourceType) if err != nil { return nil, err } if value == nil { return nil, nil } return bound(value), nil ``` ########## table/partitioned_fanout_writer_test.go: ########## @@ -683,6 +683,168 @@ func (s *FanoutWriterTestSuite) TestGetRecordPartitionsWithDroppedLeadingSourceC s.Equal("foo=null/bar=7/baz=true", spec.PartitionToPath(partitions[0].partitionRec, icebergSchema)) } +func (s *FanoutWriterTestSuite) TestPartitionedWriterReusesExtractionPlan() { + icebergSchema := iceberg.NewSchema(0, + iceberg.NestedField{ID: 1, Name: "part", Type: iceberg.PrimitiveTypes.Int32}, + iceberg.NestedField{ID: 2, Name: "value", Type: iceberg.PrimitiveTypes.String}, + ) + spec := iceberg.NewPartitionSpec(iceberg.PartitionField{ + SourceIDs: []int{1}, FieldID: 1000, Transform: iceberg.IdentityTransform{}, Name: "part", + }) + + arrowSchema := arrow.NewSchema([]arrow.Field{ + {Name: "part", Type: arrow.PrimitiveTypes.Int32}, + {Name: "value", Type: arrow.BinaryTypes.String}, + }, nil) + firstRecord := s.createCustomTestRecord(arrowSchema, [][]any{{int32(7), "a"}}) + defer firstRecord.Release() + + writer := newPartitionedFanoutWriter(spec, icebergSchema, nil, nil) + partitions, err := writer.getPartitions(firstRecord) + s.Require().NoError(err) + s.Require().Len(partitions, 1) + s.Equal(int32(7), partitions[0].partitionRec.Get(0)) + firstPlan := writer.plan + s.Require().NotNil(firstPlan) + + equivalentSchema := arrow.NewSchema(arrowSchema.Fields(), nil) + secondRecord := s.createCustomTestRecord(equivalentSchema, [][]any{{int32(8), "b"}}) + defer secondRecord.Release() + + partitions, err = writer.getPartitions(secondRecord) + s.Require().NoError(err) + s.Require().Len(partitions, 1) + s.Equal(int32(8), partitions[0].partitionRec.Get(0)) + s.Same(firstPlan, writer.plan) +} + +func (s *FanoutWriterTestSuite) TestPartitionExtractionPlanHandlesReorderedRecordSchema() { + icebergSchema := iceberg.NewSchema(0, + iceberg.NestedField{ID: 1, Name: "part", Type: iceberg.PrimitiveTypes.Int32}, + iceberg.NestedField{ID: 2, Name: "value", Type: iceberg.PrimitiveTypes.String}, + ) + spec := iceberg.NewPartitionSpec(iceberg.PartitionField{ + SourceIDs: []int{1}, FieldID: 1000, Transform: iceberg.IdentityTransform{}, Name: "part", + }) + + originalSchema := arrow.NewSchema([]arrow.Field{ + {Name: "part", Type: arrow.PrimitiveTypes.Int32}, + {Name: "value", Type: arrow.BinaryTypes.String}, + }, nil) + plan, err := newPartitionExtractionPlan(spec, icebergSchema, originalSchema) + s.Require().NoError(err) + s.Equal(0, plan.fields[0].columnIndex) + + reorderedSchema := arrow.NewSchema([]arrow.Field{ + {Name: "value", Type: arrow.BinaryTypes.String}, + {Name: "part", Type: arrow.PrimitiveTypes.Int32}, + }, nil) + record := s.createCustomTestRecord(reorderedSchema, [][]any{{"a", int32(7)}, {"b", int32(8)}}) + defer record.Release() + + partitions, err := plan.getRecordPartitions(record) + s.Require().NoError(err) + s.Require().Len(partitions, 2) + values := []int32{ + partitions[0].partitionRec.Get(0).(int32), + partitions[1].partitionRec.Get(0).(int32), + } + s.ElementsMatch([]int32{7, 8}, values) +} + +func (s *FanoutWriterTestSuite) TestBoundPartitionTransformsMatchGenericApply() { + unknown, err := iceberg.ParseTransform("custom-transform") + s.Require().NoError(err) + + tests := []struct { + name string + transform iceberg.Transform + sourceType iceberg.Type + value iceberg.Literal + fallback bool + }{ + { + name: "identity", transform: iceberg.IdentityTransform{}, + sourceType: iceberg.PrimitiveTypes.Int64, value: iceberg.Int64Literal(34), + }, + { + name: "void", transform: iceberg.VoidTransform{}, + sourceType: iceberg.PrimitiveTypes.String, value: iceberg.StringLiteral("abc"), + }, + { + name: "bucket", transform: iceberg.BucketTransform{NumBuckets: 16}, + sourceType: iceberg.PrimitiveTypes.String, value: iceberg.StringLiteral("abc"), + }, + { + name: "truncate", transform: iceberg.TruncateTransform{Width: 3}, + sourceType: iceberg.PrimitiveTypes.String, value: iceberg.StringLiteral("abcdef"), + }, + { + name: "year", transform: iceberg.YearTransform{}, + sourceType: iceberg.PrimitiveTypes.Date, value: iceberg.DateLiteral(19_358), + }, + { + name: "month", transform: iceberg.MonthTransform{}, + sourceType: iceberg.PrimitiveTypes.Timestamp, value: iceberg.TimestampLiteral(1_672_531_200_000_000), + }, + { + name: "day nanoseconds", transform: iceberg.DayTransform{}, + sourceType: iceberg.PrimitiveTypes.TimestampNs, value: iceberg.TimestampNsLiteral(1_672_531_200_000_000_000), + }, + { + name: "hour", transform: iceberg.HourTransform{}, + sourceType: iceberg.PrimitiveTypes.Timestamp, value: iceberg.TimestampLiteral(1_672_531_200_000_000), + }, + { + name: "invalid truncate fallback", transform: iceberg.TruncateTransform{Width: 0}, + sourceType: iceberg.PrimitiveTypes.String, value: iceberg.StringLiteral("abc"), fallback: true, + }, + { + name: "unknown fallback", transform: unknown, Review Comment: This case has `fallback: false`, so it asserts `ok == true`: it's checking the `UnknownTransform` arm returns a binding, which is the opposite of a fallback. The name reads just like the `"invalid truncate fallback"` case above it; I'd rename to something like `"unknown_transform"` to avoid the mixup. ########## table/partitioned_fanout_writer.go: ########## @@ -320,60 +349,149 @@ func getRecordPartitions(spec iceberg.PartitionSpec, schema *iceberg.Schema, rec if !ok { continue } - colIndices := record.Schema().FieldIndices(colName) + colIndices := recordSchema.FieldIndices(colName) if len(colIndices) == 0 { return nil, fmt.Errorf("failed to find source column %q in record schema", colName) } sourceType, ok := schema.FindTypeByID(sourceField.SourceID()) if !ok { return nil, fmt.Errorf("failed to find type for source field ID %d in schema", sourceField.SourceID()) } - partitionColumns[i] = record.Column(colIndices[0]) partitionFieldsInfo[i] = partitionFieldInfo{ sourceField: sourceField, sourceName: colName, fieldID: sourceField.FieldID, sourceType: sourceType, + columnIndex: colIndices[0], + valueAt: bindPartitionValue(sourceField.Transform, sourceType), + } + } + + return &partitionExtractionPlan{ + spec: spec, + schema: schema, + recordSchema: recordSchema, + fields: partitionFieldsInfo, + }, nil +} + +func (p *partitionExtractionPlan) getRecordPartitions(record arrow.RecordBatch) ([]*partitionInfo, error) { + // Preserve support for iterators whose batch schema changes. The usual path compares + // schema pointers; equivalent independently-built schemas also reuse the plan. + if record.Schema() != p.recordSchema && !record.Schema().Equal(p.recordSchema) { + plan, err := newPartitionExtractionPlan(p.spec, p.schema, record.Schema()) + if err != nil { + return nil, err + } + + return plan.getRecordPartitions(record) + } + + partitionMap := newPartitionMapNode() + partitionRec := make(partitionRecord, len(p.fields)) + partitionColumns := make([]arrow.Array, len(p.fields)) + for i, fieldInfo := range p.fields { + if fieldInfo.columnIndex >= 0 { + partitionColumns[i] = record.Column(fieldInfo.columnIndex) } } for row := range record.NumRows() { - for i := range partitionFields { + for i, fieldInfo := range p.fields { col := partitionColumns[i] if col != nil && !col.IsNull(int(row)) { - fieldInfo := partitionFieldsInfo[i] - sourceField := fieldInfo.sourceField - val, err := getArrowValueAsIcebergLiteral(col, int(row), fieldInfo.sourceType) + value, err := fieldInfo.valueAt(col, int(row)) if err != nil { return nil, fmt.Errorf( "failed to convert source column %q (field ID %d) from Arrow type %s to Iceberg type %s: %w", fieldInfo.sourceName, - sourceField.SourceID(), + fieldInfo.sourceField.SourceID(), col.DataType(), fieldInfo.sourceType, err, ) } - transformedLiteral := sourceField.Transform.Apply(iceberg.Optional[iceberg.Literal]{Valid: true, Val: val}) - if transformedLiteral.Valid { - partitionRec[i] = transformedLiteral.Val.Any() - } else { - partitionRec[i] = nil - } + partitionRec[i] = value } else { partitionRec[i] = nil } } // Get or create partition info for this partition key - partVal := partitionMap.getOrCreate(partitionRec, partitionFieldsInfo) + partVal := partitionMap.getOrCreate(partitionRec, p.fields) partVal.rows = append(partVal.rows, row) } return partitionMap.collectPartitions(), nil } +func bindPartitionValue(transform iceberg.Transform, sourceType iceberg.Type) func(arrow.Array, int) (any, error) { + bound, ok := bindPartitionTransform(transform, sourceType) + if !ok { + return func(column arrow.Array, row int) (any, error) { + value, err := getArrowValueAsIcebergLiteral(column, row, sourceType) + if err != nil { + return nil, err + } + + transformed := transform.Apply(iceberg.Optional[iceberg.Literal]{Valid: true, Val: value}) + if !transformed.Valid { + return nil, nil + } + + return transformed.Val.Any(), nil + } + } + + return func(column arrow.Array, row int) (any, error) { + value, err := getArrowValueAsIcebergValue(column, row, sourceType) + if err != nil { + return nil, err + } + + return bound(value), nil + } +} + +func bindPartitionTransform(transform iceberg.Transform, sourceType iceberg.Type) (func(any) any, bool) { + optionalInt := func(transformer func(any) iceberg.Optional[int32]) func(any) any { + return func(value any) any { + transformed := transformer(value) + if !transformed.Valid { + return nil + } + + return transformed.Val + } + } + + switch typed := transform.(type) { + case iceberg.IdentityTransform: + return func(value any) any { return value }, true + case iceberg.VoidTransform: + return func(any) any { return nil }, true + case iceberg.BucketTransform: + return optionalInt(typed.Transformer(sourceType)), true + case iceberg.TruncateTransform: + transformer, err := typed.Transformer(sourceType) + if err == nil { + return transformer, true + } + case iceberg.UnknownTransform: + return func(any) any { return nil }, true + } + + if typed, ok := transform.(iceberg.TimeTransform); ok { Review Comment: Year/Month/Day/Hour aren't in the switch above; they land here via `TimeTransform`. Took me a second to convince myself the switch was complete; a one-line comment noting the four time transforms are handled by this interface check would save the next reader that detour. -- This is an automated message from the Apache Git Service. To respond to the message, please log on to GitHub and use the URL above to go to the specific comment. To unsubscribe, e-mail: [email protected] For queries about this service, please contact Infrastructure at: [email protected] --------------------------------------------------------------------- To unsubscribe, e-mail: [email protected] For additional commands, e-mail: [email protected]
