perf(table): index schemas by ID

schema.go

@@ -346,8 +346,7 @@ func (s *Schema) MarshalJSON() ([]byte, error) {
 
 	type Alias Schema
 
-	aliasCopy := *(*Alias)(s)
-	aliasCopy.IdentifierFieldIDs = ids
+	aliasCopy := Alias{ID: s.ID, IdentifierFieldIDs: ids}
 
 	return json.Marshal(struct {
 		Type   string        `json:"type"`

schema_test.go

@@ -24,6 +24,7 @@ import (
 	"path/filepath"
 	"runtime"
 	"strings"
+	"sync"
 	"testing"
 
 	"github.com/apache/iceberg-go"
@@ -2293,3 +2294,48 @@ func TestVisitGeoSchemaWithSchemaVisitorPerPrimitiveType(t *testing.T) {
 	assert.Equal(t, 1, v.geometryCalls)
 	assert.Equal(t, 1, v.geographyCalls)
 }
+
+func TestSchemaMarshalJSONConcurrentLazyLookups(t *testing.T) {
+	for range 32 {
+		schema := iceberg.NewSchemaWithIdentifiers(17, nil,
+			iceberg.NestedField{ID: 1, Name: "id", Type: iceberg.PrimitiveTypes.Int64, Required: true},
+			iceberg.NestedField{ID: 2, Name: "data", Type: iceberg.PrimitiveTypes.String},
+		)
+		start := make(chan struct{})
+		var wg sync.WaitGroup
+		for range 8 {
+			wg.Go(func() {
+				<-start
+				for range 8 {
+					_, err := json.Marshal(schema)
+					assert.NoError(t, err)
+				}
+			})
+			wg.Go(func() {
+				<-start
+				_, found := schema.FindFieldByID(1)
+				assert.True(t, found)
+				_, found = schema.FindFieldByName("data")
+				assert.True(t, found)
+				_, found = schema.FindFieldByNameCaseInsensitive("DATA")
+				assert.True(t, found)
+				name, found := schema.FindColumnName(2)
+				assert.True(t, found)
+				assert.Equal(t, "data", name)
+			})
+		}
+		close(start)
+		wg.Wait()
+
+		data, err := json.Marshal(schema)
+		require.NoError(t, err)
+		assert.JSONEq(t, `{
+			"type": "struct", "schema-id": 17, "identifier-field-ids": [],
+			"fields": [
+				{"id": 1, "name": "id", "type": "long", "required": true},
+				{"id": 2, "name": "data", "type": "string", "required": false}
+			]
+		}`, string(data))
+		assert.Nil(t, schema.IdentifierFieldIDs)
+	}
+}

table/metadata.go

@@ -142,6 +142,81 @@ func snapshotIndexPosition(index *snapshotIndexData, snapshots []Snapshot, id in
 	return 0, false
 }
 
+type schemaIndexData struct {
+	schemas map[int]*iceberg.Schema
+	// sourceCount is kept separately because duplicate IDs make the map
+	// smaller than the schema slice in invalid or in-package fixture state.
+	sourceCount int
+	// firstSchemaSlot identifies the slice used to build the index.
+	firstSchemaSlot **iceberg.Schema
+	// shared means the map is owned by more than one builder or metadata value.
+	shared bool
+}
+
+func schemaListFirst(schemas []*iceberg.Schema) **iceberg.Schema {
+	if len(schemas) == 0 {
+		return nil
+	}
+
+	return &schemas[0]
+}
+
+func buildSchemaIndex(schemas []*iceberg.Schema) *schemaIndexData {
+	byID := make(map[int]*iceberg.Schema, len(schemas))
+	for _, schema := range schemas {
+		if schema == nil {
+			continue
+		}
+		if _, exists := byID[schema.ID]; !exists {
+			byID[schema.ID] = schema
+		}
+	}
+
+	return &schemaIndexData{
+		schemas:         byID,
+		sourceCount:     len(schemas),
+		firstSchemaSlot: schemaListFirst(schemas),
+	}
+}
+
+func cloneSchemaIndex(index *schemaIndexData) *schemaIndexData {
+	if index == nil {
+		return nil
+	}
+
+	return &schemaIndexData{
+		schemas:         maps.Clone(index.schemas),
+		sourceCount:     index.sourceCount,
+		firstSchemaSlot: index.firstSchemaSlot,
+	}
+}
+
+func schemaIndexNeedsRebuild(index *schemaIndexData, schemas []*iceberg.Schema) bool {
+	if index == nil || index.sourceCount != len(schemas) {
+		return true
+	}
+
+	return index.firstSchemaSlot != schemaListFirst(schemas)
+}
+
+func schemaIndexLookup(index *schemaIndexData, schemas []*iceberg.Schema, id int) (*iceberg.Schema, bool) {
+	if index != nil {
+		if schema, ok := index.schemas[id]; ok && schema != nil && schema.ID == id {
+			return schema, true
+		}
+	}
+
+	// Builder lookups return mutable schemas, so a changed ID may not have
+	// an entry in the index even when the slice itself is unchanged.
+	for _, schema := range schemas {
+		if schema != nil && schema.ID == id {
+			return schema, true
+		}
+	}
+
+	return nil, false
+}
+
 // Metadata for an iceberg table as specified in the Iceberg spec
 //
 // https://iceberg.apache.org/spec/#iceberg-table-spec
@@ -257,6 +332,7 @@ type MetadataBuilder struct {
 	lastUpdatedMS      int64
 	lastColumnId       int
 	schemaList         []*iceberg.Schema
+	schemaIndex        *schemaIndexData // Derived from schemaList; not serialized.
 	currentSchemaID    int
 	specs              []iceberg.PartitionSpec
 	defaultSpecID      int
@@ -293,6 +369,7 @@ func NewMetadataBuilder(formatVersion int) (*MetadataBuilder, error) {
 	return &MetadataBuilder{
 		updates:            make([]Update, 0),
 		schemaList:         make([]*iceberg.Schema, 0),
+		schemaIndex:        buildSchemaIndex(nil),
 		specs:              make([]iceberg.PartitionSpec, 0),
 		props:              make(iceberg.Properties),
 		snapshotList:       make([]Snapshot, 0),
@@ -405,6 +482,7 @@ func MetadataBuilderFromBase(metadata Metadata, currentFileLocation string) (*Me
 		b.partitionStatsList = slices.Collect(metadata.PartitionStatistics())
 		b.encryptionKeyList = slices.Collect(metadata.EncryptionKeys())
 	}
+	b.schemaIndex = buildSchemaIndex(b.schemaList)
 
 	if currentFileLocation != "" {
 		b.previousFileEntry = &MetadataLogEntry{
@@ -466,6 +544,15 @@ func (b *MetadataBuilder) clone() *MetadataBuilder {
 		lastAddedPartitionID: clonePtr(b.lastAddedPartitionID),
 		lastAddedSortOrderID: clonePtr(b.lastAddedSortOrderID),
 	}
+	if b.schemaIndex != nil {
+		cloned.schemaIndex = &schemaIndexData{
+			schemas:         b.schemaIndex.schemas,
+			sourceCount:     b.schemaIndex.sourceCount,
+			firstSchemaSlot: schemaListFirst(cloned.schemaList),
+			shared:          true,
+		}
+		b.schemaIndex.shared = true
+	}
 	if b.snapshotIndex != nil {
 		cloned.snapshotIndex = &snapshotIndexData{
 			positions:     b.snapshotIndex.positions,
@@ -521,6 +608,19 @@ func (b *MetadataBuilder) newSnapshotID() int64 {
 	}
 }
 
+func (b *MetadataBuilder) ensureSchemaIndex() {
+	if schemaIndexNeedsRebuild(b.schemaIndex, b.schemaList) {
+		b.schemaIndex = buildSchemaIndex(b.schemaList)
+	}
+}
+
+func (b *MetadataBuilder) ensureSchemaIndexMutable() {
+	b.ensureSchemaIndex()
+	if b.schemaIndex.shared {
+		b.schemaIndex = cloneSchemaIndex(b.schemaIndex)
+	}
+}
+
 func (b *MetadataBuilder) ensureSnapshotIndex() {
 	if snapshotIndexNeedsRebuild(b.snapshotIndex, b.snapshotList) {
 		b.snapshotIndex = buildSnapshotIndex(b.snapshotList)
@@ -590,7 +690,11 @@ func (b *MetadataBuilder) AddSchema(schema *iceberg.Schema) error {
 
 	schema.ID = newSchemaID
 
+	b.ensureSchemaIndexMutable()
 	b.schemaList = append(b.schemaList, schema)
+	b.schemaIndex.schemas[newSchemaID] = schema
+	b.schemaIndex.sourceCount = len(b.schemaList)
+	b.schemaIndex.firstSchemaSlot = schemaListFirst(b.schemaList)
 	b.updates = append(b.updates, NewAddSchemaUpdate(schema))
 	b.lastAddedSchemaID = &newSchemaID
 
@@ -1250,6 +1354,10 @@ func (b *MetadataBuilder) SetLastUpdatedMS() *MetadataBuilder {
 }
 
 func (b *MetadataBuilder) buildCommonMetadata() (*commonMetadata, error) {
+	b.ensureSchemaIndex()
+	if b.schemaIndex != nil {
+		b.schemaIndex.shared = true
+	}
 	b.ensureSnapshotIndex()
 	if b.snapshotIndex != nil {
 		b.snapshotIndex.shared = true
@@ -1283,6 +1391,7 @@ func (b *MetadataBuilder) buildCommonMetadata() (*commonMetadata, error) {
 		LastUpdatedMS:      b.lastUpdatedMS,
 		LastColumnId:       b.lastColumnId,
 		SchemaList:         b.schemaList,
+		schemaIndex:        b.schemaIndex,
 		CurrentSchemaID:    b.currentSchemaID,
 		Specs:              b.specs,
 		DefaultSpecID:      defaultSpecID,
@@ -1350,10 +1459,12 @@ func (b *MetadataBuilder) updateSnapshotLog() error {
 }
 
 func (b *MetadataBuilder) GetSchemaByID(id int) (*iceberg.Schema, error) {
-	for _, s := range b.schemaList {
-		if s.ID == id {
-			return s, nil
-		}
+	index := b.schemaIndex
+	if schemaIndexNeedsRebuild(index, b.schemaList) {
+		index = buildSchemaIndex(b.schemaList)
+	}
+	if schema, ok := schemaIndexLookup(index, b.schemaList, id); ok {
+		return schema, nil
 	}
 
 	return nil, fmt.Errorf("%w: schema with id %d not found", iceberg.ErrInvalidArgument, id)
@@ -1598,6 +1709,7 @@ func (b *MetadataBuilder) RemoveSchemas(ints []int) error {
 	})
 
 	if len(removed) != 0 {
+		b.schemaIndex = buildSchemaIndex(b.schemaList)
 		b.updates = append(b.updates, NewRemoveSchemasUpdate(removed))
 	}
 
@@ -1957,6 +2069,7 @@ type commonMetadata struct {
 	// V3+ fields
 	NextRowID *int64 `json:"next-row-id,omitempty"` // V3: Next available row ID
 
+	schemaIndex   *schemaIndexData
 	snapshotIndex *snapshotIndexData
 }
 
@@ -2121,10 +2234,12 @@ func (c *commonMetadata) CurrentSchema() *iceberg.Schema {
 // versions through commonMetadata and is intentionally not part of Metadata,
 // so custom Metadata implementations keep the public Schemas fallback.
 func (c *commonMetadata) schemaByID(id int) *iceberg.Schema {
-	for _, schema := range c.SchemaList {
-		if schema.ID == id {
-			return cloneSchema(schema)
-		}
+	index := c.schemaIndex
+	if schemaIndexNeedsRebuild(index, c.SchemaList) {
+		index = buildSchemaIndex(c.SchemaList)
+	}
+	if schema, ok := schemaIndexLookup(index, c.SchemaList, id); ok {
+		return cloneSchema(schema)
 	}
 
 	return nil
@@ -2520,6 +2635,7 @@ func (c *commonMetadata) preValidate() {
 		c.SnapshotLog = []SnapshotLogEntry{}
 	}
 
+	c.schemaIndex = buildSchemaIndex(c.SchemaList)
 	c.snapshotIndex = buildSnapshotIndex(c.SnapshotList)
 }
 

table/metadata_builder_internal_test.go

@@ -3446,6 +3446,7 @@ func TestSetFormatVersionV2ToV3FromDeserializedMetadata(t *testing.T) {
 // both the drift guard and its filler consult it.
 var sharedCloneFields = map[string]struct{}{
 	"base":          {}, // immutable snapshot, shared by design.
+	"schemaIndex":   {}, // immutable schema references, copied on schema mutation.
 	"snapshotIndex": {}, // immutable index positions, copied on snapshot mutation.
 }