diff --git a/table/position_delta_writer.go b/table/position_delta_writer.go index e7c97541f..14dc00a2e 100644 --- a/table/position_delta_writer.go +++ b/table/position_delta_writer.go @@ -114,13 +114,16 @@ func (w *PositionDeltaWriter) Reinsert(batch arrow.RecordBatch) error { return fmt.Errorf("%w: writer is already closed", ErrInvalidOperation) } - indices := batch.Schema().FieldIndices(iceberg.RowIDColumnName) - if len(indices) == 0 { + rowIDIndex, found, err := rowIDColumnIndex(batch.Schema()) + if err != nil { + return err + } + if !found { return fmt.Errorf("%w: Reinsert batch must contain %s column", iceberg.ErrInvalidArgument, iceberg.RowIDColumnName) } - col := batch.Column(indices[0]) + col := batch.Column(rowIDIndex) if col.NullN() > 0 { return fmt.Errorf("%w: Reinsert batch %s column must not contain null values", iceberg.ErrInvalidArgument, iceberg.RowIDColumnName) @@ -140,8 +143,12 @@ func (w *PositionDeltaWriter) Insert(batch arrow.RecordBatch) error { return fmt.Errorf("%w: writer is already closed", ErrInvalidOperation) } - if indices := batch.Schema().FieldIndices(iceberg.RowIDColumnName); len(indices) > 0 { - if batch.Column(indices[0]).NullN() != int(batch.NumRows()) { + rowIDIndex, found, err := rowIDColumnIndex(batch.Schema()) + if err != nil { + return err + } + if found { + if batch.Column(rowIDIndex).NullN() != int(batch.NumRows()) { return fmt.Errorf("%w: Insert batch %s column must be all null (use Reinsert for preserved IDs)", iceberg.ErrInvalidArgument, iceberg.RowIDColumnName) } @@ -245,7 +252,11 @@ func (w *PositionDeltaWriter) buildUnifiedIterator() iter.Seq2[arrow.RecordBatch // already has _row_id (validated to be all-null by Insert), it is returned // retained as-is. func appendNullRowIDColumn(alloc memory.Allocator, batch arrow.RecordBatch) (arrow.RecordBatch, error) { - if indices := batch.Schema().FieldIndices(iceberg.RowIDColumnName); len(indices) > 0 { + _, found, err := rowIDColumnIndex(batch.Schema()) + if err != nil { + return nil, err + } + if found { batch.Retain() return batch, nil @@ -278,3 +289,16 @@ func appendNullRowIDColumn(alloc memory.Allocator, batch arrow.RecordBatch) (arr return array.NewRecordBatch(newSchema, cols, nrows), nil } + +func rowIDColumnIndex(schema *arrow.Schema) (int, bool, error) { + indices := schema.FieldIndices(iceberg.RowIDColumnName) + switch len(indices) { + case 0: + return 0, false, nil + case 1: + return indices[0], true, nil + default: + return 0, false, fmt.Errorf("%w: batch contains multiple %s columns", + iceberg.ErrInvalidArgument, iceberg.RowIDColumnName) + } +} diff --git a/table/position_delta_writer_test.go b/table/position_delta_writer_test.go index 55eb524a0..629d3931a 100644 --- a/table/position_delta_writer_test.go +++ b/table/position_delta_writer_test.go @@ -280,3 +280,49 @@ func TestPositionDeltaWriter_ReinsertRejectsMissingColumn(t *testing.T) { require.Error(t, err) assert.Contains(t, err.Error(), "must contain _row_id") } + +func TestPositionDeltaWriter_RejectsDuplicateRowIDColumns(t *testing.T) { + for _, layout := range []struct { + name string + appendRows func(first, second *array.Int64Builder) + }{ + {"both null", func(first, second *array.Int64Builder) { + first.AppendNull() + second.AppendNull() + }}, + {"both non-null", func(first, second *array.Int64Builder) { + first.Append(10) + second.Append(11) + }}, + } { + t.Run(layout.name, func(t *testing.T) { + schema := arrow.NewSchema([]arrow.Field{ + {Name: "id", Type: arrow.PrimitiveTypes.Int64}, + {Name: iceberg.RowIDColumnName, Type: arrow.PrimitiveTypes.Int64, Nullable: true}, + {Name: iceberg.RowIDColumnName, Type: arrow.PrimitiveTypes.Int64, Nullable: true}, + }, nil) + builder := array.NewRecordBuilder(memory.DefaultAllocator, schema) + defer builder.Release() + builder.Field(0).(*array.Int64Builder).Append(1) + layout.appendRows(builder.Field(1).(*array.Int64Builder), builder.Field(2).(*array.Int64Builder)) + batch := builder.NewRecordBatch() + defer batch.Release() + + for _, operation := range []struct { + name string + call func(*table.PositionDeltaWriter) error + }{ + {"Reinsert", func(w *table.PositionDeltaWriter) error { return w.Reinsert(batch) }}, + {"Insert", func(w *table.PositionDeltaWriter) error { return w.Insert(batch) }}, + } { + t.Run(operation.name, func(t *testing.T) { + writer, err := table.NewPositionDeltaWriter(newV3RowLineageTestTable(t)) + require.NoError(t, err) + err = operation.call(writer) + require.ErrorIs(t, err, iceberg.ErrInvalidArgument) + require.ErrorContains(t, err, "multiple _row_id columns") + }) + } + }) + } +}