diff --git a/internal/utils.go b/internal/utils.go index 65b393fe6..647010af6 100644 --- a/internal/utils.go +++ b/internal/utils.go @@ -29,6 +29,10 @@ import ( "golang.org/x/exp/constraints" ) +// SchemaRef marks schema access that may return references to internal state. +// It is restricted to packages within this module by Go's internal package rules. +type SchemaRef struct{} + // FloorDiv performs floored integer division, rounding toward negative infinity. // This matches Java's Math.floorDiv behavior for negative dividends. func FloorDiv[T constraints.Integer](a, b T) T { diff --git a/schema.go b/schema.go index 9f51474ed..ad69339b3 100644 --- a/schema.go +++ b/schema.go @@ -29,6 +29,8 @@ import ( "unicode" "unicode/utf16" "unicode/utf8" + + "github.com/apache/iceberg-go/internal" ) // Schema is an Iceberg table schema, represented as a struct with @@ -181,7 +183,13 @@ func (s *Schema) FlatFields() (iter.Seq[NestedField], error) { return nil, err } - return maps.Values(fields), nil + return func(yield func(NestedField) bool) { + for field := range maps.Values(fields) { + if !yield(cloneField(field)) { + return + } + } + }, nil } func (s *Schema) lazyIDToField() (map[int]NestedField, error) { @@ -267,8 +275,8 @@ func (s *Schema) AsStruct() StructType { return StructType{FieldList: cloneField func (s *Schema) asStructRef() StructType { return StructType{FieldList: s.fields} } func (s *Schema) NumFields() int { return len(s.fields) } -func (s *Schema) Field(i int) NestedField { return s.fields[i] } -func (s *Schema) Fields() []NestedField { return slices.Clone(s.fields) } +func (s *Schema) Field(i int) NestedField { return cloneField(s.fields[i]) } +func (s *Schema) Fields() []NestedField { return cloneFields(s.fields) } func (s *Schema) FieldIDs() []int { idx, _ := s.lazyNameToID() @@ -325,13 +333,54 @@ func cloneFields(fields []NestedField) []NestedField { cloned := make([]NestedField, len(fields)) for i, field := range fields { - cloned[i] = field - cloned[i].Type = cloneType(field.Type) + cloned[i] = cloneField(field) } return cloned } +func cloneField(field NestedField) NestedField { + return NestedField{ + ID: field.ID, + Name: field.Name, + Type: cloneType(field.Type), + Required: field.Required, + Doc: field.Doc, + InitialDefault: cloneDefault(field.InitialDefault), + WriteDefault: cloneDefault(field.WriteDefault), + } +} + +func cloneDefault(value any) any { + switch value := value.(type) { + case []byte: + return slices.Clone(value) + case BinaryLiteral: + return BinaryLiteral(slices.Clone([]byte(value))) + case FixedLiteral: + return FixedLiteral(slices.Clone([]byte(value))) + case []any: + cloned := make([]any, len(value)) + for i, item := range value { + cloned[i] = cloneDefault(item) + } + + return cloned + case map[string]any: + cloned := make(map[string]any, len(value)) + for key, item := range value { + cloned[key] = cloneDefault(item) + } + + return cloned + default: + // Iceberg scalar defaults are value types (bool, numeric values, + // strings, UUIDs, and Decimals). Any mutable reference type must get + // an explicit deep-copy case above. + return value + } +} + func cloneType(t Type) Type { if t == nil { return nil @@ -407,6 +456,19 @@ func (s *Schema) FindFieldByNameCaseInsensitive(name string) (NestedField, bool) // FindFieldByID is like [*Schema.FindColumnName], but returns the whole // field rather than just the field name. func (s *Schema) FindFieldByID(id int) (NestedField, bool) { + f, ok := s.FindFieldByIDRef(id, internal.SchemaRef{}) + if ok { + f = cloneField(f) + } + + return f, ok +} + +// FindFieldByIDRef returns a schema-owned field without cloning it. The returned +// field is a value copy, but its Type and everything reachable from that Type +// are shared with the schema and must be treated as read-only. This method is +// limited to internal callers that need to avoid clone-on-read overhead. +func (s *Schema) FindFieldByIDRef(id int, _ internal.SchemaRef) (NestedField, bool) { idx, _ := s.lazyIDToField() f, ok := idx[id] diff --git a/schema_test.go b/schema_test.go index 34dbe7414..2dff0bd0c 100644 --- a/schema_test.go +++ b/schema_test.go @@ -19,12 +19,15 @@ package iceberg_test import ( "encoding/json" + "fmt" "os" "path/filepath" + "runtime" "strings" "testing" "github.com/apache/iceberg-go" + "github.com/apache/iceberg-go/internal" "github.com/stretchr/testify/assert" "github.com/stretchr/testify/require" @@ -318,6 +321,120 @@ func TestSchemaAsStructClonesNestedMapTypes(t *testing.T) { assert.Equal(t, "name", clonedValueType.FieldList[0].Name) } +func TestSchemaFieldGettersReturnDefensiveCopies(t *testing.T) { + schema := iceberg.NewSchema(0, + iceberg.NestedField{ + ID: 1, + Name: "person", + Type: &iceberg.StructType{FieldList: []iceberg.NestedField{ + {ID: 2, Name: "name", Type: iceberg.PrimitiveTypes.String}, + }}, + }, + iceberg.NestedField{ + ID: 3, + Name: "payload", + Type: iceberg.PrimitiveTypes.Binary, + InitialDefault: iceberg.BinaryLiteral{1, 2}, + WriteDefault: map[string]any{"nested": []any{[]byte{3, 4}}}, + }, + ) + + assertIndependent := func(t *testing.T, personField, payloadField iceberg.NestedField) { + t.Helper() + person, ok := personField.Type.(*iceberg.StructType) + require.True(t, ok) + person.FieldList[0].Name = "hijacked" + payloadField.InitialDefault.(iceberg.BinaryLiteral)[0] = 9 + payloadField.WriteDefault.(map[string]any)["nested"].([]any)[0].([]byte)[0] = 9 + + actual, ok := schema.FindFieldByID(2) + require.True(t, ok) + assert.Equal(t, "name", actual.Name) + payload, ok := schema.FindFieldByID(3) + require.True(t, ok) + assert.Equal(t, iceberg.BinaryLiteral{1, 2}, payload.InitialDefault) + assert.Equal(t, map[string]any{"nested": []any{[]byte{3, 4}}}, payload.WriteDefault) + } + + t.Run("Field", func(t *testing.T) { + assertIndependent(t, schema.Field(0), schema.Field(1)) + }) + t.Run("Fields", func(t *testing.T) { + fields := schema.Fields() + assertIndependent(t, fields[0], fields[1]) + }) + t.Run("FindFieldByID", func(t *testing.T) { + person, ok := schema.FindFieldByID(1) + require.True(t, ok) + payload, ok := schema.FindFieldByID(3) + require.True(t, ok) + assertIndependent(t, person, payload) + }) + t.Run("FindFieldByName", func(t *testing.T) { + person, ok := schema.FindFieldByName("person") + require.True(t, ok) + payload, ok := schema.FindFieldByName("payload") + require.True(t, ok) + assertIndependent(t, person, payload) + }) + t.Run("FindFieldByNameCaseInsensitive", func(t *testing.T) { + person, ok := schema.FindFieldByNameCaseInsensitive("PERSON") + require.True(t, ok) + payload, ok := schema.FindFieldByNameCaseInsensitive("PAYLOAD") + require.True(t, ok) + assertIndependent(t, person, payload) + }) + t.Run("FindTypeByID", func(t *testing.T) { + typ, ok := schema.FindTypeByID(1) + require.True(t, ok) + person, ok := typ.(*iceberg.StructType) + require.True(t, ok) + person.FieldList[0].Name = "hijacked" + + actual, ok := schema.FindFieldByID(2) + require.True(t, ok) + assert.Equal(t, "name", actual.Name) + }) + t.Run("FlatFields", func(t *testing.T) { + fields, err := schema.FlatFields() + require.NoError(t, err) + byID := make(map[int]iceberg.NestedField) + for field := range fields { + byID[field.ID] = field + } + assertIndependent(t, byID[1], byID[3]) + }) +} + +func TestSchemaFieldRefLookupDoesNotAllocate(t *testing.T) { + children := make([]iceberg.NestedField, 50) + for i := range children { + children[i] = iceberg.NestedField{ + ID: i + 2, Name: fmt.Sprintf("child_%d", i), Type: iceberg.PrimitiveTypes.String, + } + } + schema := iceberg.NewSchema(0, iceberg.NestedField{ + ID: 1, Name: "parent", Type: &iceberg.StructType{FieldList: children}, + }) + _, ok := schema.FindFieldByIDRef(1, internal.SchemaRef{}) + require.True(t, ok) + + var field iceberg.NestedField + assert.Zero(t, testing.AllocsPerRun(100, func() { + field, ok = schema.FindFieldByIDRef(1, internal.SchemaRef{}) + runtime.KeepAlive(field) + })) + require.True(t, ok) + assert.Equal(t, 1, field.ID) + + parent := field.Type.(*iceberg.StructType) + parent.FieldList[0].Name = "shared_child" + shared, ok := schema.FindFieldByIDRef(1, internal.SchemaRef{}) + require.True(t, ok) + sharedParent := shared.Type.(*iceberg.StructType) + assert.Equal(t, "shared_child", sharedParent.FieldList[0].Name) +} + func TestSchemaIndexByIDVisitor(t *testing.T) { index, err := iceberg.IndexByID(tableSchemaNested) require.NoError(t, err) diff --git a/table/arrow_utils.go b/table/arrow_utils.go index 86cbaa286..df8527fa9 100644 --- a/table/arrow_utils.go +++ b/table/arrow_utils.go @@ -870,7 +870,7 @@ func (a arrowAccessor) FieldPartner(partnerStruct arrow.Array, fieldID int, _ st return nil } - field, ok := a.fileSchema.FindFieldByID(fieldID) + field, ok := a.fileSchema.FindFieldByIDRef(fieldID, internal.SchemaRef{}) if !ok { return nil } @@ -1062,7 +1062,7 @@ func (a *arrowProjectionVisitor) typeToArrowType(t iceberg.Type) arrow.DataType } func (a *arrowProjectionVisitor) castIfNeeded(field iceberg.NestedField, vals arrow.Array) arrow.Array { - fileField, ok := a.fileSchema.FindFieldByID(field.ID) + fileField, ok := a.fileSchema.FindFieldByIDRef(field.ID, internal.SchemaRef{}) if !ok { panic(fmt.Errorf("could not find field id %d in schema", field.ID)) }