You are viewing a plain text version of this content. The canonical link for it is here.
Posted to issues@iceberg.apache.org by "fallintoplace (via GitHub)" <gi...@apache.org> on 2026/08/11 20:59:05 UTC
[PR] perf(table): bind partition transforms once [iceberg-go]
fallintoplace opened a new pull request, #1768:
URL: https://github.com/apache/iceberg-go/pull/1768
Depends on #1765. This is stacked on top of it, so only the second commit is new here.
## What changed
- Bind built-in partition transform functions when the extraction plan is created.
- Feed native Iceberg values directly into identity, bucket, truncate, and time transforms.
- Keep the generic literal path for custom or unsupported transforms.
- Keep unknown and void transforms producing null partition values.
## Why
The row loop currently converts each Arrow value into a literal and calls `Transform.Apply`. Some transforms do more setup there. For example, `TruncateTransform.Apply` rebuilds its transformer for every value.
The extraction plan has the source type already, so it can bind the transform once and reuse the function for every batch. Java's `StructTransform` uses the same bind-once shape.
## Benchmark
Apple M1 Pro, 65,536 rows, median of 3 runs. The baseline is #1765.
```
go test ./table -run '^$' -bench '^BenchmarkPartitionTransforms$' -benchmem -benchtime=1s -count=3 -cpu=1
```
| Transform | #1765 | this PR |
| --- | ---: | ---: |
| identity int64 | 2,648,326 ns/op | 2,531,346 ns/op |
| bucket string | 5,862,667 ns/op | 5,710,295 ns/op |
| truncate string | 15,767,652 ns/op, 7,271,656 B/op, 394,124 allocs/op | 7,205,658 ns/op, 3,077,352 B/op, 131,980 allocs/op |
| day timestamp | 3,898,547 ns/op | 3,468,050 ns/op |
| hour timestamp_ns | 3,812,458 ns/op | 3,621,065 ns/op |
The largest change is `truncate[string]`, which is about 2.2x faster with 67% fewer allocations.
## Tests
- `go test ./...`
- `go test -race ./table -run '^TestFanoutWriter$' -count=1`
- `go vet ./...`
- `git diff --check`
--
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: issues-unsubscribe@iceberg.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: issues-unsubscribe@iceberg.apache.org
For additional commands, e-mail: issues-help@iceberg.apache.org
Re: [PR] perf(table): bind partition transforms once [iceberg-go]
Posted by "laskoviymishka (via GitHub)" <gi...@apache.org>.
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: issues-unsubscribe@iceberg.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: issues-unsubscribe@iceberg.apache.org
For additional commands, e-mail: issues-help@iceberg.apache.org
Re: [PR] perf(table): bind partition transforms once [iceberg-go]
Posted by "laskoviymishka (via GitHub)" <gi...@apache.org>.
laskoviymishka commented on PR #1768:
URL: https://github.com/apache/iceberg-go/pull/1768#issuecomment-5438167008
@badalprasadsingh need to resolve confllicts, once done - ready to merge
--
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: issues-unsubscribe@iceberg.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: issues-unsubscribe@iceberg.apache.org
For additional commands, e-mail: issues-help@iceberg.apache.org
Re: [PR] perf(table): bind partition transforms once [iceberg-go]
Posted by "badalprasadsingh (via GitHub)" <gi...@apache.org>.
badalprasadsingh commented on PR #1768:
URL: https://github.com/apache/iceberg-go/pull/1768#issuecomment-5438232491
> @badalprasadsingh need to resolve confllicts, once done - ready to merge
@fallintoplace :))
--
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: issues-unsubscribe@iceberg.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: issues-unsubscribe@iceberg.apache.org
For additional commands, e-mail: issues-help@iceberg.apache.org
Re: [PR] perf(table): bind partition transforms once [iceberg-go]
Posted by "laskoviymishka (via GitHub)" <gi...@apache.org>.
laskoviymishka merged PR #1768:
URL: https://github.com/apache/iceberg-go/pull/1768
--
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: issues-unsubscribe@iceberg.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: issues-unsubscribe@iceberg.apache.org
For additional commands, e-mail: issues-help@iceberg.apache.org
Re: [PR] perf(table): bind partition transforms once [iceberg-go]
Posted by "fallintoplace (via GitHub)" <gi...@apache.org>.
fallintoplace closed pull request #1768: perf(table): bind partition transforms once
URL: https://github.com/apache/iceberg-go/pull/1768
--
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: issues-unsubscribe@iceberg.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: issues-unsubscribe@iceberg.apache.org
For additional commands, e-mail: issues-help@iceberg.apache.org
Re: [PR] perf(table): bind partition transforms once [iceberg-go]
Posted by "laskoviymishka (via GitHub)" <gi...@apache.org>.
laskoviymishka commented on PR #1768:
URL: https://github.com/apache/iceberg-go/pull/1768#issuecomment-5438285438
Ups!
--
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: issues-unsubscribe@iceberg.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: issues-unsubscribe@iceberg.apache.org
For additional commands, e-mail: issues-help@iceberg.apache.org