laskoviymishka commented on code in PR #1912: URL: https://github.com/apache/iceberg-go/pull/1912#discussion_r3873479868
########## table/writer_schema_projection_test.go: ########## @@ -0,0 +1,369 @@ +// 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 ( + "context" + "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" + iceio "github.com/apache/iceberg-go/io" + tblutils "github.com/apache/iceberg-go/table/internal" + "github.com/google/uuid" + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" +) + +func writerProjectionOptions() SchemaOptions { + return SchemaOptions{ + DowncastTimestamp: true, + IncludeFieldIDs: true, + UseWriteDefault: true, + } +} + +func writerFieldIDMeta(id string) arrow.Metadata { + return arrow.MetadataFrom(map[string]string{ArrowParquetFieldIDKey: id}) +} + +func projectionTestRecord(t *testing.T, mem memory.Allocator, schema *arrow.Schema, json []byte) arrow.RecordBatch { + t.Helper() + + builder := array.NewRecordBuilder(mem, schema) + defer builder.Release() + require.NoError(t, builder.UnmarshalJSON(json)) + + return builder.NewRecordBatch() +} + +func projectionTestSchema(t *testing.T, schema *iceberg.Schema) *arrow.Schema { + t.Helper() + + arrowSchema, err := SchemaToArrowSchemaWithOptions(schema, ArrowSchemaOptions{IncludeFieldIDs: true}) + require.NoError(t, err) + + return arrowSchema +} + +func assertProjectedBatch(t *testing.T, requested, provided *iceberg.Schema, batch arrow.RecordBatch, target *arrow.Schema, wantReuse bool) arrow.RecordBatch { + t.Helper() + + got, err := toRequestedSchema(context.Background(), requested, provided, batch, writerProjectionOptions(), target) + require.NoError(t, err) + if wantReuse { + assert.Same(t, batch, got) + } else { + assert.NotSame(t, batch, got) + } + assert.True(t, arrowSchemaEqual(got.Schema(), target), "expected schema: %s\ngot: %s", target, got.Schema()) + + return got +} + +func TestToRequestedSchemaWriteFastPath(t *testing.T) { + tests := []struct { + name string + requested *iceberg.Schema + provided *iceberg.Schema + batch func(*testing.T, memory.Allocator) arrow.RecordBatch + wantReuse bool + check func(*testing.T, arrow.RecordBatch) + }{ + { + name: "exact Arrow schema", + requested: iceberg.NewSchema(0, + iceberg.NestedField{ID: 1, Name: "id", Type: iceberg.PrimitiveTypes.Int64, Required: true}, + ), + provided: iceberg.NewSchema(0, + iceberg.NestedField{ID: 1, Name: "id", Type: iceberg.PrimitiveTypes.Int64, Required: true}, + ), + batch: func(t *testing.T, mem memory.Allocator) arrow.RecordBatch { + schema := projectionTestSchema(t, iceberg.NewSchema(0, + iceberg.NestedField{ID: 1, Name: "id", Type: iceberg.PrimitiveTypes.Int64, Required: true}, + )) + + return projectionTestRecord(t, mem, schema, []byte(`{"id": 1}`)) + }, + wantReuse: true, Review Comment: This single-column Int64 case is the only one that actually takes the fast path; every other case is a rejection. If `arrowSchemaEqual` had a latent bug on multi-field or nested schemas we'd never see it here. I'd add a multi-column exact-match and one nested (struct or list) exact-match to the `wantReuse: true` set. And a negative case worth having explicitly: a `Required: true` field where the incoming batch carries `Nullable: true`, which has to reject the fast path so full projection emits `Nullable: false` (required, not optional, in the written Parquet). Right now that rejection is untested. ########## table/arrow_utils.go: ########## @@ -1457,6 +1457,19 @@ type SchemaOptions struct { // ToRequestedSchema will construct a new record batch matching the requested iceberg schema // casting columns if necessary as appropriate. func ToRequestedSchema(ctx context.Context, requested, fileSchema *iceberg.Schema, batch arrow.RecordBatch, opts SchemaOptions) (arrow.RecordBatch, error) { + return toRequestedSchema(ctx, requested, fileSchema, batch, opts, nil) +} + +// toRequestedSchema is the internal write-path variant of ToRequestedSchema. +// When targetArrowSchema is provided, an exact Arrow schema match can reuse the +// input batch without walking the Iceberg schema or rebuilding any arrays. +func toRequestedSchema(ctx context.Context, requested, fileSchema *iceberg.Schema, batch arrow.RecordBatch, opts SchemaOptions, targetArrowSchema *arrow.Schema) (arrow.RecordBatch, error) { + if targetArrowSchema != nil && arrowSchemaEqual(batch.Schema(), targetArrowSchema) { Review Comment: The optimization is sound as long as one invariant holds: the precomputed `targetArrowSchema` from `SchemaToArrowSchemaWithOptions` has to be byte-for-byte what `arrowProjectionVisitor` would have produced for a batch that needs no transformation. If those two ever drift (say `SchemaToArrowSchemaWithOptions` starts emitting a looser schema than the visitor), a batch matching the looser schema takes the fast path and skips a transformation it actually needed. I traced both paths and they agree today, so this isn't a live bug. But nothing enforces the coupling, and the existing tests can't catch drift because they build the input batch from `SchemaToArrowSchemaWithOptions` too, so they're comparing that function against itself. I'd add one cross-check test that breaks the circularity: push a batch through the slow path (`ToRequestedSchema` with a `nil` target), then assert the returned schema equals `SchemaToArrowSchemaWithOptions(requested, {IncludeFieldIDs: true})` for a non-trivial shape (a struct or list). That locks the equality-check contract to the real projection output. wdyt? ########## table/writer_invariants_bench_test.go: ########## @@ -90,3 +91,93 @@ func BenchmarkDefaultDataFileWriter(b *testing.B) { } b.ReportMetric(float64(rows), "rows/op") } + +func BenchmarkToRequestedSchemaWriteFastPath(b *testing.B) { + requested := iceberg.NewSchema(0, iceberg.NestedField{ + ID: 1, Name: "id", Type: iceberg.PrimitiveTypes.Int64, + }) + requestedArrowSchema, err := SchemaToArrowSchemaWithOptions(requested, ArrowSchemaOptions{IncludeFieldIDs: true}) + if err != nil { + b.Fatal(err) + } + + provided := iceberg.NewSchema(0, iceberg.NestedField{ + ID: 1, Name: "id", Type: iceberg.PrimitiveTypes.Int32, + }) + providedArrowSchema, err := SchemaToArrowSchemaWithOptions(provided, ArrowSchemaOptions{IncludeFieldIDs: true}) + if err != nil { + b.Fatal(err) + } + + opts := SchemaOptions{ + DowncastTimestamp: true, + IncludeFieldIDs: true, + UseWriteDefault: true, + } + for _, rows := range []int{0, 1, 16, 1024, 65536} { + b.Run(fmt.Sprintf("rows=%d", rows), func(b *testing.B) { + exact := benchmarkIntRecord(b, requestedArrowSchema, rows, true) + defer exact.Release() + conversion := benchmarkIntRecord(b, providedArrowSchema, rows, false) + defer conversion.Release() + + b.Run("projection", func(b *testing.B) { + b.ReportAllocs() + for b.Loop() { + out, err := ToRequestedSchema(b.Context(), requested, requested, exact, opts) + if err != nil { + b.Fatal(err) + } + out.Release() + } + b.ReportMetric(float64(rows), "rows/op") + }) + + b.Run("exact_fast_path", func(b *testing.B) { + b.ReportAllocs() + for b.Loop() { + out, err := toRequestedSchema(b.Context(), requested, requested, exact, opts, requestedArrowSchema) + if err != nil { + b.Fatal(err) + } + out.Release() + } + b.ReportMetric(float64(rows), "rows/op") + }) + + b.Run("conversion", func(b *testing.B) { + b.ReportAllocs() + for b.Loop() { + out, err := toRequestedSchema(b.Context(), requested, provided, conversion, opts, requestedArrowSchema) + if err != nil { + b.Fatal(err) + } + out.Release() + } + b.ReportMetric(float64(rows), "rows/op") + }) + Review Comment: gofumpt is failing on this one blank line; it doesn't allow a trailing blank at the end of a closure body. Deleting it turns CI green. ########## table/arrow_utils.go: ########## @@ -1481,6 +1494,12 @@ func ToRequestedSchema(ctx context.Context, requested, fileSchema *iceberg.Schem return out, nil } +func arrowSchemaEqual(left, right *arrow.Schema) bool { + // Schema.Equal intentionally ignores top-level metadata, but projection Review Comment: The comment has the direction backwards: it reads as if projection is what lacks the metadata, but the point is the other way. The input batch might carry top-level metadata that full projection would strip (`RecordFromStructArray` carries none), so we guard here to stop the fast path silently preserving it. Something like: ```go // Schema.Equal ignores top-level metadata. Full projection strips it, so // guard it here; otherwise the fast path would preserve metadata that // projection would have dropped. ``` While we're here, both params get dereferenced immediately, so a one-line nil precondition (or guard) wouldn't hurt. It's unexported and every call site passes non-nil today, so this one's minor. ########## table/writer_schema_projection_test.go: ########## @@ -0,0 +1,369 @@ +// 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 ( + "context" + "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" + iceio "github.com/apache/iceberg-go/io" + tblutils "github.com/apache/iceberg-go/table/internal" + "github.com/google/uuid" + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" +) + +func writerProjectionOptions() SchemaOptions { + return SchemaOptions{ + DowncastTimestamp: true, + IncludeFieldIDs: true, + UseWriteDefault: true, + } +} + +func writerFieldIDMeta(id string) arrow.Metadata { + return arrow.MetadataFrom(map[string]string{ArrowParquetFieldIDKey: id}) +} + +func projectionTestRecord(t *testing.T, mem memory.Allocator, schema *arrow.Schema, json []byte) arrow.RecordBatch { + t.Helper() + + builder := array.NewRecordBuilder(mem, schema) + defer builder.Release() + require.NoError(t, builder.UnmarshalJSON(json)) + + return builder.NewRecordBatch() +} + +func projectionTestSchema(t *testing.T, schema *iceberg.Schema) *arrow.Schema { + t.Helper() + + arrowSchema, err := SchemaToArrowSchemaWithOptions(schema, ArrowSchemaOptions{IncludeFieldIDs: true}) + require.NoError(t, err) + + return arrowSchema +} + +func assertProjectedBatch(t *testing.T, requested, provided *iceberg.Schema, batch arrow.RecordBatch, target *arrow.Schema, wantReuse bool) arrow.RecordBatch { + t.Helper() + + got, err := toRequestedSchema(context.Background(), requested, provided, batch, writerProjectionOptions(), target) + require.NoError(t, err) + if wantReuse { + assert.Same(t, batch, got) + } else { + assert.NotSame(t, batch, got) + } + assert.True(t, arrowSchemaEqual(got.Schema(), target), "expected schema: %s\ngot: %s", target, got.Schema()) + + return got +} + +func TestToRequestedSchemaWriteFastPath(t *testing.T) { + tests := []struct { + name string + requested *iceberg.Schema + provided *iceberg.Schema + batch func(*testing.T, memory.Allocator) arrow.RecordBatch + wantReuse bool + check func(*testing.T, arrow.RecordBatch) + }{ + { + name: "exact Arrow schema", + requested: iceberg.NewSchema(0, + iceberg.NestedField{ID: 1, Name: "id", Type: iceberg.PrimitiveTypes.Int64, Required: true}, + ), + provided: iceberg.NewSchema(0, + iceberg.NestedField{ID: 1, Name: "id", Type: iceberg.PrimitiveTypes.Int64, Required: true}, + ), + batch: func(t *testing.T, mem memory.Allocator) arrow.RecordBatch { + schema := projectionTestSchema(t, iceberg.NewSchema(0, + iceberg.NestedField{ID: 1, Name: "id", Type: iceberg.PrimitiveTypes.Int64, Required: true}, + )) + + return projectionTestRecord(t, mem, schema, []byte(`{"id": 1}`)) + }, + wantReuse: true, + }, + { + name: "reordered columns", + requested: iceberg.NewSchema(0, + iceberg.NestedField{ID: 1, Name: "id", Type: iceberg.PrimitiveTypes.Int64}, + iceberg.NestedField{ID: 2, Name: "name", Type: iceberg.PrimitiveTypes.String}, + ), + provided: iceberg.NewSchema(0, + iceberg.NestedField{ID: 2, Name: "name", Type: iceberg.PrimitiveTypes.String}, + iceberg.NestedField{ID: 1, Name: "id", Type: iceberg.PrimitiveTypes.Int64}, + ), + batch: func(t *testing.T, mem memory.Allocator) arrow.RecordBatch { + schema := arrow.NewSchema([]arrow.Field{ + {Name: "name", Type: arrow.BinaryTypes.String, Nullable: true, Metadata: writerFieldIDMeta("2")}, + {Name: "id", Type: arrow.PrimitiveTypes.Int64, Nullable: true, Metadata: writerFieldIDMeta("1")}, + }, nil) + + return projectionTestRecord(t, mem, schema, []byte(`{"id": 1, "name": "one"}`)) + }, + }, + { + name: "missing optional field", + requested: iceberg.NewSchema(0, + iceberg.NestedField{ID: 1, Name: "id", Type: iceberg.PrimitiveTypes.Int64}, + iceberg.NestedField{ID: 2, Name: "name", Type: iceberg.PrimitiveTypes.String}, + ), + provided: iceberg.NewSchema(0, + iceberg.NestedField{ID: 1, Name: "id", Type: iceberg.PrimitiveTypes.Int64}, + ), + batch: func(t *testing.T, mem memory.Allocator) arrow.RecordBatch { + schema := arrow.NewSchema([]arrow.Field{{ + Name: "id", Type: arrow.PrimitiveTypes.Int64, Nullable: true, Metadata: writerFieldIDMeta("1"), + }}, nil) + + return projectionTestRecord(t, mem, schema, []byte(`{"id": 1}`)) + }, + check: func(t *testing.T, got arrow.RecordBatch) { + assert.True(t, got.Column(1).IsNull(0)) + }, + }, + { + name: "missing write-default field", + requested: iceberg.NewSchema(0, + iceberg.NestedField{ID: 1, Name: "id", Type: iceberg.PrimitiveTypes.Int64}, + iceberg.NestedField{ID: 2, Name: "name", Type: iceberg.PrimitiveTypes.String, Required: true, WriteDefault: "default"}, + ), + provided: iceberg.NewSchema(0, + iceberg.NestedField{ID: 1, Name: "id", Type: iceberg.PrimitiveTypes.Int64}, + ), + batch: func(t *testing.T, mem memory.Allocator) arrow.RecordBatch { + schema := arrow.NewSchema([]arrow.Field{{ + Name: "id", Type: arrow.PrimitiveTypes.Int64, Nullable: true, Metadata: writerFieldIDMeta("1"), + }}, nil) + + return projectionTestRecord(t, mem, schema, []byte(`{"id": 1}`)) + }, + check: func(t *testing.T, got arrow.RecordBatch) { + assert.Equal(t, "default", got.Column(1).(*array.String).Value(0)) + }, + }, + { + name: "timestamp nanoseconds require downcast", + requested: iceberg.NewSchema(0, + iceberg.NestedField{ID: 1, Name: "ts", Type: iceberg.PrimitiveTypes.Timestamp}, + ), + provided: iceberg.NewSchema(0, + iceberg.NestedField{ID: 1, Name: "ts", Type: iceberg.PrimitiveTypes.TimestampNs}, + ), + batch: func(t *testing.T, mem memory.Allocator) arrow.RecordBatch { + schema := arrow.NewSchema([]arrow.Field{{ + Name: "ts", Type: &arrow.TimestampType{Unit: arrow.Nanosecond}, Nullable: true, Metadata: writerFieldIDMeta("1"), + }}, nil) + builder := array.NewRecordBuilder(mem, schema) + builder.Field(0).(*array.TimestampBuilder).Append(-1_500) + batch := builder.NewRecordBatch() + builder.Release() + + return batch + }, + check: func(t *testing.T, got arrow.RecordBatch) { + assert.Equal(t, arrow.Timestamp(-2), got.Column(0).(*array.Timestamp).Value(0)) + }, + }, + { + name: "large list requires offset conversion", + requested: iceberg.NewSchema(0, + iceberg.NestedField{ID: 1, Name: "items", Type: &iceberg.ListType{ + ElementID: 2, Element: iceberg.PrimitiveTypes.Int32, ElementRequired: true, + }}, + ), + provided: iceberg.NewSchema(0, + iceberg.NestedField{ID: 1, Name: "items", Type: &iceberg.ListType{ + ElementID: 2, Element: iceberg.PrimitiveTypes.Int32, ElementRequired: true, + }}, + ), + batch: func(t *testing.T, mem memory.Allocator) arrow.RecordBatch { + target := projectionTestSchema(t, iceberg.NewSchema(0, + iceberg.NestedField{ID: 1, Name: "items", Type: &iceberg.ListType{ + ElementID: 2, Element: iceberg.PrimitiveTypes.Int32, ElementRequired: true, + }}, + )) + listType := target.Field(0).Type.(*arrow.ListType) + schema := arrow.NewSchema([]arrow.Field{{ + Name: "items", Type: arrow.LargeListOfField(listType.ElemField()), Nullable: true, + Metadata: writerFieldIDMeta("1"), + }}, nil) + + return projectionTestRecord(t, mem, schema, []byte(`{"items": [1, 2]}`)) + }, + }, + { + name: "field ID metadata mismatch", + requested: iceberg.NewSchema(0, + iceberg.NestedField{ID: 1, Name: "id", Type: iceberg.PrimitiveTypes.Int64}, + ), + provided: iceberg.NewSchema(0, + iceberg.NestedField{ID: 1, Name: "id", Type: iceberg.PrimitiveTypes.Int64}, + ), + batch: func(t *testing.T, mem memory.Allocator) arrow.RecordBatch { + schema := arrow.NewSchema([]arrow.Field{{ + Name: "id", Type: arrow.PrimitiveTypes.Int64, Nullable: true, + }}, nil) + + return projectionTestRecord(t, mem, schema, []byte(`{"id": 1}`)) + }, + }, + { + name: "top-level metadata mismatch", + requested: iceberg.NewSchema(0, + iceberg.NestedField{ID: 1, Name: "id", Type: iceberg.PrimitiveTypes.Int64}, + ), + provided: iceberg.NewSchema(0, + iceberg.NestedField{ID: 1, Name: "id", Type: iceberg.PrimitiveTypes.Int64}, + ), + batch: func(t *testing.T, mem memory.Allocator) arrow.RecordBatch { + field := arrow.Field{ + Name: "id", Type: arrow.PrimitiveTypes.Int64, Nullable: true, Metadata: writerFieldIDMeta("1"), + } + metadata := arrow.MetadataFrom(map[string]string{"source": "input"}) + schema := arrow.NewSchema([]arrow.Field{field}, &metadata) + + return projectionTestRecord(t, mem, schema, []byte(`{"id": 1}`)) + }, + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + mem := memory.NewCheckedAllocator(memory.NewGoAllocator()) + defer mem.AssertSize(t, 0) + + target := projectionTestSchema(t, tt.requested) + batch := tt.batch(t, mem) + got := assertProjectedBatch(t, tt.requested, tt.provided, batch, target, tt.wantReuse) + if tt.check != nil { + tt.check(t, got) + } + got.Release() + batch.Release() + }) + } +} + +type captureWriteDataFileFormat struct { + tblutils.FileFormat + batches []arrow.RecordBatch + writer *captureFileWriter +} + +func (f *captureWriteDataFileFormat) WriteDataFile(_ context.Context, _ iceio.WriteFileIO, _ map[int]any, _ tblutils.WriteFileInfo, batches []arrow.RecordBatch) (iceberg.DataFile, error) { + f.batches = batches + + return nil, nil +} + +func (f *captureWriteDataFileFormat) NewFileWriter(_ context.Context, _ iceio.WriteFileIO, _ map[int]any, _ tblutils.WriteFileInfo, _ *arrow.Schema) (tblutils.FileWriter, error) { + f.writer = &captureFileWriter{} + + return f.writer, nil +} + +type captureFileWriter struct { + batch arrow.RecordBatch +} + +func (w *captureFileWriter) Write(batch arrow.RecordBatch) error { + w.batch = batch + + return nil +} + +func (w *captureFileWriter) BytesWritten() int64 { return 0 } +func (w *captureFileWriter) Close() (iceberg.DataFile, error) { return nil, nil } +func (w *captureFileWriter) Abort() error { return nil } + +func TestDefaultDataFileWriterReusesExactBatch(t *testing.T) { + mem := memory.NewCheckedAllocator(memory.NewGoAllocator()) + defer mem.AssertSize(t, 0) + + schema := iceberg.NewSchema(0, + iceberg.NestedField{ID: 1, Name: "id", Type: iceberg.PrimitiveTypes.Int64, Required: true}, + ) + spec := iceberg.NewPartitionSpec() + metadata, err := NewMetadata(schema, &spec, UnsortedSortOrder, t.TempDir(), iceberg.Properties{}) + require.NoError(t, err) + metaBuilder, err := MetadataBuilderFromBase(metadata, "") + require.NoError(t, err) + + format := &captureWriteDataFileFormat{FileFormat: tblutils.GetFileFormat(iceberg.ParquetFile)} + writer, err := newDataFileWriter(t.TempDir(), iceio.LocalFS{}, metaBuilder, iceberg.Properties{}, withFormat(format)) + require.NoError(t, err) + + arrowSchema := projectionTestSchema(t, schema) + record := projectionTestRecord(t, mem, arrowSchema, []byte(`{"id": 1}`)) + defer record.Release() + record.Retain() + + _, err = writer.writeFile(t.Context(), nil, WriteTask{ + Uuid: uuid.New(), ID: 1, FileCount: 1, Schema: schema, + Batches: []arrow.RecordBatch{record}, + }) + require.NoError(t, err) + require.Len(t, format.batches, 1) + assert.Same(t, record, format.batches[0]) +} + +func TestRollingDataWriterReusesExactBatch(t *testing.T) { + schema := iceberg.NewSchema(0, + iceberg.NestedField{ID: 1, Name: "id", Type: iceberg.PrimitiveTypes.Int64, Required: true}, + ) + spec := iceberg.NewPartitionSpec() + metadata, err := NewMetadata(schema, &spec, UnsortedSortOrder, t.TempDir(), iceberg.Properties{}) + require.NoError(t, err) + metaBuilder, err := MetadataBuilderFromBase(metadata, "") + require.NoError(t, err) + arrowSchema := projectionTestSchema(t, schema) + writeUUID := uuid.New() + factory, err := newWriterFactory(t.TempDir(), recordWritingArgs{ + sc: arrowSchema, + fs: iceio.LocalFS{}, + writeUUID: &writeUUID, + counter: func(yield func(int) bool) { + for i := 0; ; i++ { + if !yield(i) { + return + } + } + }, + }, metaBuilder, schema, 1024*1024) + require.NoError(t, err) + format := &captureWriteDataFileFormat{FileFormat: tblutils.GetFileFormat(iceberg.ParquetFile)} + factory.format = format + + output := make(chan iceberg.DataFile, 1) + writer := factory.newRollingDataWriter(t.Context(), "", nil, output) + record := projectionTestRecord(t, memory.DefaultAllocator, arrowSchema, []byte(`{"id": 1}`)) Review Comment: This is the one test exercising the async stream path, the trickiest refcount route, where the fast-path `Retain` has to balance against the stream release, `writeConverted`, and the abort/cleanup defer. And it's the only one on `memory.DefaultAllocator`, which silently swallows an imbalance. `TestDefaultDataFileWriterReusesExactBatch` already uses a checked allocator; I'd do the same here (`memory.NewCheckedAllocator(memory.NewGoAllocator())` with `defer mem.AssertSize(t, 0)`) so the test actually proves the balance it's meant to. ########## table/writer_schema_projection_test.go: ########## @@ -0,0 +1,369 @@ +// 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 ( + "context" + "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" + iceio "github.com/apache/iceberg-go/io" + tblutils "github.com/apache/iceberg-go/table/internal" + "github.com/google/uuid" + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" +) + +func writerProjectionOptions() SchemaOptions { + return SchemaOptions{ + DowncastTimestamp: true, + IncludeFieldIDs: true, + UseWriteDefault: true, + } +} + +func writerFieldIDMeta(id string) arrow.Metadata { + return arrow.MetadataFrom(map[string]string{ArrowParquetFieldIDKey: id}) +} + +func projectionTestRecord(t *testing.T, mem memory.Allocator, schema *arrow.Schema, json []byte) arrow.RecordBatch { + t.Helper() + + builder := array.NewRecordBuilder(mem, schema) + defer builder.Release() + require.NoError(t, builder.UnmarshalJSON(json)) + + return builder.NewRecordBatch() +} + +func projectionTestSchema(t *testing.T, schema *iceberg.Schema) *arrow.Schema { + t.Helper() + + arrowSchema, err := SchemaToArrowSchemaWithOptions(schema, ArrowSchemaOptions{IncludeFieldIDs: true}) + require.NoError(t, err) + + return arrowSchema +} + +func assertProjectedBatch(t *testing.T, requested, provided *iceberg.Schema, batch arrow.RecordBatch, target *arrow.Schema, wantReuse bool) arrow.RecordBatch { + t.Helper() + + got, err := toRequestedSchema(context.Background(), requested, provided, batch, writerProjectionOptions(), target) + require.NoError(t, err) + if wantReuse { + assert.Same(t, batch, got) + } else { + assert.NotSame(t, batch, got) + } + assert.True(t, arrowSchemaEqual(got.Schema(), target), "expected schema: %s\ngot: %s", target, got.Schema()) + + return got +} + +func TestToRequestedSchemaWriteFastPath(t *testing.T) { + tests := []struct { + name string + requested *iceberg.Schema + provided *iceberg.Schema + batch func(*testing.T, memory.Allocator) arrow.RecordBatch + wantReuse bool + check func(*testing.T, arrow.RecordBatch) + }{ + { + name: "exact Arrow schema", + requested: iceberg.NewSchema(0, + iceberg.NestedField{ID: 1, Name: "id", Type: iceberg.PrimitiveTypes.Int64, Required: true}, + ), + provided: iceberg.NewSchema(0, + iceberg.NestedField{ID: 1, Name: "id", Type: iceberg.PrimitiveTypes.Int64, Required: true}, + ), + batch: func(t *testing.T, mem memory.Allocator) arrow.RecordBatch { + schema := projectionTestSchema(t, iceberg.NewSchema(0, + iceberg.NestedField{ID: 1, Name: "id", Type: iceberg.PrimitiveTypes.Int64, Required: true}, + )) + + return projectionTestRecord(t, mem, schema, []byte(`{"id": 1}`)) + }, + wantReuse: true, + }, + { + name: "reordered columns", + requested: iceberg.NewSchema(0, + iceberg.NestedField{ID: 1, Name: "id", Type: iceberg.PrimitiveTypes.Int64}, + iceberg.NestedField{ID: 2, Name: "name", Type: iceberg.PrimitiveTypes.String}, + ), + provided: iceberg.NewSchema(0, + iceberg.NestedField{ID: 2, Name: "name", Type: iceberg.PrimitiveTypes.String}, + iceberg.NestedField{ID: 1, Name: "id", Type: iceberg.PrimitiveTypes.Int64}, + ), + batch: func(t *testing.T, mem memory.Allocator) arrow.RecordBatch { + schema := arrow.NewSchema([]arrow.Field{ + {Name: "name", Type: arrow.BinaryTypes.String, Nullable: true, Metadata: writerFieldIDMeta("2")}, + {Name: "id", Type: arrow.PrimitiveTypes.Int64, Nullable: true, Metadata: writerFieldIDMeta("1")}, + }, nil) + + return projectionTestRecord(t, mem, schema, []byte(`{"id": 1, "name": "one"}`)) + }, + }, + { + name: "missing optional field", + requested: iceberg.NewSchema(0, + iceberg.NestedField{ID: 1, Name: "id", Type: iceberg.PrimitiveTypes.Int64}, + iceberg.NestedField{ID: 2, Name: "name", Type: iceberg.PrimitiveTypes.String}, + ), + provided: iceberg.NewSchema(0, + iceberg.NestedField{ID: 1, Name: "id", Type: iceberg.PrimitiveTypes.Int64}, + ), + batch: func(t *testing.T, mem memory.Allocator) arrow.RecordBatch { + schema := arrow.NewSchema([]arrow.Field{{ + Name: "id", Type: arrow.PrimitiveTypes.Int64, Nullable: true, Metadata: writerFieldIDMeta("1"), + }}, nil) + + return projectionTestRecord(t, mem, schema, []byte(`{"id": 1}`)) + }, + check: func(t *testing.T, got arrow.RecordBatch) { + assert.True(t, got.Column(1).IsNull(0)) + }, + }, + { + name: "missing write-default field", + requested: iceberg.NewSchema(0, + iceberg.NestedField{ID: 1, Name: "id", Type: iceberg.PrimitiveTypes.Int64}, + iceberg.NestedField{ID: 2, Name: "name", Type: iceberg.PrimitiveTypes.String, Required: true, WriteDefault: "default"}, + ), + provided: iceberg.NewSchema(0, + iceberg.NestedField{ID: 1, Name: "id", Type: iceberg.PrimitiveTypes.Int64}, + ), + batch: func(t *testing.T, mem memory.Allocator) arrow.RecordBatch { + schema := arrow.NewSchema([]arrow.Field{{ + Name: "id", Type: arrow.PrimitiveTypes.Int64, Nullable: true, Metadata: writerFieldIDMeta("1"), + }}, nil) + + return projectionTestRecord(t, mem, schema, []byte(`{"id": 1}`)) + }, + check: func(t *testing.T, got arrow.RecordBatch) { + assert.Equal(t, "default", got.Column(1).(*array.String).Value(0)) + }, + }, + { + name: "timestamp nanoseconds require downcast", + requested: iceberg.NewSchema(0, + iceberg.NestedField{ID: 1, Name: "ts", Type: iceberg.PrimitiveTypes.Timestamp}, + ), + provided: iceberg.NewSchema(0, + iceberg.NestedField{ID: 1, Name: "ts", Type: iceberg.PrimitiveTypes.TimestampNs}, + ), + batch: func(t *testing.T, mem memory.Allocator) arrow.RecordBatch { + schema := arrow.NewSchema([]arrow.Field{{ + Name: "ts", Type: &arrow.TimestampType{Unit: arrow.Nanosecond}, Nullable: true, Metadata: writerFieldIDMeta("1"), + }}, nil) + builder := array.NewRecordBuilder(mem, schema) + builder.Field(0).(*array.TimestampBuilder).Append(-1_500) + batch := builder.NewRecordBatch() + builder.Release() + + return batch + }, + check: func(t *testing.T, got arrow.RecordBatch) { + assert.Equal(t, arrow.Timestamp(-2), got.Column(0).(*array.Timestamp).Value(0)) + }, + }, + { + name: "large list requires offset conversion", + requested: iceberg.NewSchema(0, + iceberg.NestedField{ID: 1, Name: "items", Type: &iceberg.ListType{ + ElementID: 2, Element: iceberg.PrimitiveTypes.Int32, ElementRequired: true, + }}, + ), + provided: iceberg.NewSchema(0, + iceberg.NestedField{ID: 1, Name: "items", Type: &iceberg.ListType{ + ElementID: 2, Element: iceberg.PrimitiveTypes.Int32, ElementRequired: true, + }}, + ), + batch: func(t *testing.T, mem memory.Allocator) arrow.RecordBatch { + target := projectionTestSchema(t, iceberg.NewSchema(0, + iceberg.NestedField{ID: 1, Name: "items", Type: &iceberg.ListType{ + ElementID: 2, Element: iceberg.PrimitiveTypes.Int32, ElementRequired: true, + }}, + )) + listType := target.Field(0).Type.(*arrow.ListType) + schema := arrow.NewSchema([]arrow.Field{{ + Name: "items", Type: arrow.LargeListOfField(listType.ElemField()), Nullable: true, + Metadata: writerFieldIDMeta("1"), + }}, nil) + + return projectionTestRecord(t, mem, schema, []byte(`{"items": [1, 2]}`)) + }, + }, + { + name: "field ID metadata mismatch", + requested: iceberg.NewSchema(0, + iceberg.NestedField{ID: 1, Name: "id", Type: iceberg.PrimitiveTypes.Int64}, + ), + provided: iceberg.NewSchema(0, + iceberg.NestedField{ID: 1, Name: "id", Type: iceberg.PrimitiveTypes.Int64}, + ), + batch: func(t *testing.T, mem memory.Allocator) arrow.RecordBatch { + schema := arrow.NewSchema([]arrow.Field{{ + Name: "id", Type: arrow.PrimitiveTypes.Int64, Nullable: true, + }}, nil) + + return projectionTestRecord(t, mem, schema, []byte(`{"id": 1}`)) + }, + }, + { + name: "top-level metadata mismatch", + requested: iceberg.NewSchema(0, + iceberg.NestedField{ID: 1, Name: "id", Type: iceberg.PrimitiveTypes.Int64}, + ), + provided: iceberg.NewSchema(0, + iceberg.NestedField{ID: 1, Name: "id", Type: iceberg.PrimitiveTypes.Int64}, + ), + batch: func(t *testing.T, mem memory.Allocator) arrow.RecordBatch { + field := arrow.Field{ + Name: "id", Type: arrow.PrimitiveTypes.Int64, Nullable: true, Metadata: writerFieldIDMeta("1"), + } + metadata := arrow.MetadataFrom(map[string]string{"source": "input"}) + schema := arrow.NewSchema([]arrow.Field{field}, &metadata) + + return projectionTestRecord(t, mem, schema, []byte(`{"id": 1}`)) + }, + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + mem := memory.NewCheckedAllocator(memory.NewGoAllocator()) + defer mem.AssertSize(t, 0) + + target := projectionTestSchema(t, tt.requested) + batch := tt.batch(t, mem) + got := assertProjectedBatch(t, tt.requested, tt.provided, batch, target, tt.wantReuse) + if tt.check != nil { + tt.check(t, got) + } + got.Release() + batch.Release() + }) + } +} + +type captureWriteDataFileFormat struct { + tblutils.FileFormat + batches []arrow.RecordBatch + writer *captureFileWriter +} + +func (f *captureWriteDataFileFormat) WriteDataFile(_ context.Context, _ iceio.WriteFileIO, _ map[int]any, _ tblutils.WriteFileInfo, batches []arrow.RecordBatch) (iceberg.DataFile, error) { + f.batches = batches + + return nil, nil +} + +func (f *captureWriteDataFileFormat) NewFileWriter(_ context.Context, _ iceio.WriteFileIO, _ map[int]any, _ tblutils.WriteFileInfo, _ *arrow.Schema) (tblutils.FileWriter, error) { + f.writer = &captureFileWriter{} + + return f.writer, nil +} + +type captureFileWriter struct { + batch arrow.RecordBatch +} + +func (w *captureFileWriter) Write(batch arrow.RecordBatch) error { + w.batch = batch + + return nil +} + +func (w *captureFileWriter) BytesWritten() int64 { return 0 } +func (w *captureFileWriter) Close() (iceberg.DataFile, error) { return nil, nil } +func (w *captureFileWriter) Abort() error { return nil } + +func TestDefaultDataFileWriterReusesExactBatch(t *testing.T) { + mem := memory.NewCheckedAllocator(memory.NewGoAllocator()) + defer mem.AssertSize(t, 0) + + schema := iceberg.NewSchema(0, + iceberg.NestedField{ID: 1, Name: "id", Type: iceberg.PrimitiveTypes.Int64, Required: true}, + ) + spec := iceberg.NewPartitionSpec() + metadata, err := NewMetadata(schema, &spec, UnsortedSortOrder, t.TempDir(), iceberg.Properties{}) + require.NoError(t, err) + metaBuilder, err := MetadataBuilderFromBase(metadata, "") + require.NoError(t, err) + + format := &captureWriteDataFileFormat{FileFormat: tblutils.GetFileFormat(iceberg.ParquetFile)} + writer, err := newDataFileWriter(t.TempDir(), iceio.LocalFS{}, metaBuilder, iceberg.Properties{}, withFormat(format)) + require.NoError(t, err) + + arrowSchema := projectionTestSchema(t, schema) + record := projectionTestRecord(t, mem, arrowSchema, []byte(`{"id": 1}`)) + defer record.Release() + record.Retain() Review Comment: A one-line comment here would help. This `Retain` counterbalances the releases inside `writeFile` (the per-batch `rec.Release()` plus the `task.Batches` cleanup defer), and on the fast path `rec` is this same `record`. Without a note, someone tidying the test could delete it and get a double-release panic. -- 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]
