feat: Cache column offsets for Resource.Get/Set

scheduler/scheduler.go

@@ -249,6 +249,12 @@ func (s *Scheduler) Sync(ctx context.Context, client schema.ClientMeta, tables s
 		return fmt.Errorf("max depth exceeded, max depth is %d", s.maxDepth)
 	}
 
+	// Tables are final here, so cache column offsets for Resource.Get/Set. They must not
+	// be shared with a concurrent Sync; Tables.FilterDfs already hands out copies.
+	for _, table := range tables {
+		table.BuildColumnIndex()
+	}
+
 	// send migrate messages first
 	for _, tableOriginal := range tables.FlattenTables() {
 		migrateMessage := &message.SyncMigrateTable{

schema/column_index_test.go

@@ -0,0 +1,92 @@
+package schema
+
+import (
+	"testing"
+
+	"github.com/apache/arrow-go/v18/arrow"
+	"github.com/stretchr/testify/require"
+)
+
+func stringColumns(names ...string) ColumnList {
+	cols := make(ColumnList, len(names))
+	for i, name := range names {
+		cols[i] = Column{Name: name, Type: arrow.BinaryTypes.String}
+	}
+	return cols
+}
+
+func TestColumnIndex(t *testing.T) {
+	table := &Table{Name: "test", Columns: stringColumns("a", "b", "c")}
+
+	require.Equal(t, 1, table.ColumnIndex("b"))
+	require.Equal(t, -1, table.ColumnIndex("missing"))
+
+	table.BuildColumnIndex()
+	require.Equal(t, 0, table.ColumnIndex("a"))
+	require.Equal(t, 1, table.ColumnIndex("b"))
+	require.Equal(t, 2, table.ColumnIndex("c"))
+	require.Equal(t, -1, table.ColumnIndex("missing"))
+}
+
+func TestColumnIndexRelations(t *testing.T) {
+	table := &Table{
+		Name:    "parent",
+		Columns: stringColumns("a", "b"),
+		Relations: Tables{
+			{Name: "child", Columns: stringColumns("c", "d")},
+		},
+	}
+	table.BuildColumnIndex()
+	require.Equal(t, 1, table.Relations[0].ColumnIndex("d"))
+}
+
+// Columns mutated behind the back of every mutator must not make lookups wrong.
+func TestColumnIndexStaleCache(t *testing.T) {
+	table := &Table{Name: "test", Columns: stringColumns("a", "b", "c")}
+	table.BuildColumnIndex()
+
+	table.Columns = append(stringColumns("z"), table.Columns...)
+	require.Equal(t, 0, table.ColumnIndex("z"))
+	require.Equal(t, 1, table.ColumnIndex("a"))
+	require.Equal(t, 2, table.ColumnIndex("b"))
+	require.Equal(t, 3, table.ColumnIndex("c"))
+
+	table.Columns = table.Columns[:1]
+	require.Equal(t, 0, table.ColumnIndex("z"))
+	require.Equal(t, -1, table.ColumnIndex("c"))
+}
+
+func TestColumnIndexInvalidatedByMutators(t *testing.T) {
+	t.Run("AddCqIDs", func(t *testing.T) {
+		table := &Table{Name: "test", Columns: stringColumns("a", "b")}
+		table.BuildColumnIndex()
+		AddCqIDs(table)
+		require.Nil(t, table.columnIndex)
+		require.Equal(t, 2, table.ColumnIndex("a"))
+	})
+
+	t.Run("AddCqClientID", func(t *testing.T) {
+		table := &Table{Name: "test", Columns: stringColumns("a", "b")}
+		table.BuildColumnIndex()
+		AddCqClientID(table)
+		require.Nil(t, table.columnIndex)
+		require.Equal(t, 1, table.ColumnIndex("a"))
+	})
+
+	t.Run("OverwriteOrAddColumn", func(t *testing.T) {
+		table := &Table{Name: "test", Columns: stringColumns("a", "b")}
+		table.BuildColumnIndex()
+		table.OverwriteOrAddColumn(&Column{Name: "z", Type: arrow.BinaryTypes.String})
+		require.Nil(t, table.columnIndex)
+		require.Equal(t, 0, table.ColumnIndex("z"))
+		require.Equal(t, 1, table.ColumnIndex("a"))
+	})
+
+	t.Run("Copy", func(t *testing.T) {
+		table := &Table{Name: "test", Columns: stringColumns("a", "b")}
+		table.BuildColumnIndex()
+		c := table.Copy(nil)
+		require.Nil(t, c.columnIndex)
+		require.Equal(t, 1, c.ColumnIndex("b"))
+	})
+}

schema/resource.go

@@ -40,7 +40,7 @@ func NewResourceData(t *Table, parent *Resource, item any) *Resource {
 }
 
 func (r *Resource) Get(columnName string) scalar.Scalar {
-	index := r.Table.Columns.Index(columnName)
+	index := r.Table.ColumnIndex(columnName)
 	if index == -1 {
 		// we panic because we want to distinguish between code error and api error
 		// this also saves additional checks in our testing code
@@ -53,12 +53,17 @@ func (r *Resource) Get(columnName string) scalar.Scalar {
 // one of concrete  it returns an error just for backward compatibility
 // and panics in case it fails
 func (r *Resource) Set(columnName string, value any) error {
-	index := r.Table.Columns.Index(columnName)
+	index := r.Table.ColumnIndex(columnName)
 	if index == -1 {
 		// we panic because we want to distinguish between code error and api error
 		// this also saves additional checks in our testing code
 		panic(columnName + " column not found")
 	}
+	return r.setAtIndex(index, columnName, value)
+}
+
+// setAtIndex is Set for callers that already hold the column's offset.
+func (r *Resource) setAtIndex(index int, columnName string, value any) error {
 	if err := r.data[index].Set(value); err != nil {
 		panic(fmt.Errorf("failed to set column %s: %w", columnName, err))
 	}
@@ -113,22 +118,24 @@ func calculateCqIDValue(r *Resource, cols []string) hash.Hash {
 func (r *Resource) storeCQID(value uuid.UUID) error {
 	// We skip if _cq_id is not present.
 	// Mostly the problem here is because the transformation step is baked into the resolving step
-	if r.Table.Columns.Get(CqIDColumn.Name) == nil {
+	index := r.Table.ColumnIndex(CqIDColumn.Name)
+	if index == -1 {
 		return nil
 	}
 	b, err := value.MarshalBinary()
 	if err != nil {
 		return err
 	}
-	return r.Set(CqIDColumn.Name, b)
+	return r.setAtIndex(index, CqIDColumn.Name, b)
 }
 
 func (r *Resource) StoreCQClientID(clientID string) error {
 	// We skip if _cq_client_id is not present.
-	if r.Table.Columns.Get(CqClientIDColumn.Name) == nil {
+	index := r.Table.ColumnIndex(CqClientIDColumn.Name)
+	if index == -1 {
 		return nil
 	}
-	return r.Set(CqClientIDColumn.Name, clientID)
+	return r.setAtIndex(index, CqClientIDColumn.Name, clientID)
 }
 
 type PKError struct {

schema/resource_test.go

@@ -1,6 +1,7 @@
 package schema
 
 import (
+	"strconv"
 	"testing"
 
 	"github.com/apache/arrow-go/v18/arrow"
@@ -71,3 +72,23 @@ func TestResource_Validate(t *testing.T) {
 		})
 	}
 }
+
+func TestResource_SetWithColumnIndex(t *testing.T) {
+	for _, buildIndex := range []bool{false, true} {
+		t.Run("buildIndex="+strconv.FormatBool(buildIndex), func(t *testing.T) {
+			table := &Table{Name: "test", Columns: stringColumns("a", "b", "c")}
+			if buildIndex {
+				table.BuildColumnIndex()
+			}
+			r := NewResourceData(table, nil, nil)
+			for _, name := range []string{"a", "b", "c"} {
+				require.NoError(t, r.Set(name, "value-"+name))
+			}
+			for i, name := range []string{"a", "b", "c"} {
+				require.Equal(t, "value-"+name, r.Get(name).String())
+				require.Equal(t, "value-"+name, r.data[i].String())
+			}
+			require.Panics(t, func() { _ = r.Set("missing", "x") })
+		})
+	}
+}

schema/table.go

@@ -118,6 +118,34 @@ type Table struct {
 
 	// IgnorePKComponentsMismatchValidation is a flag that indicates if the table should skip validating usage of both primary key components and primary keys
 	IgnorePKComponentsMismatchValidation bool `json:"ignore_pk_components_mismatch_validation"`
+
+	// columnIndex is an optional name -> Columns offset cache built by BuildColumnIndex.
+	// ColumnIndex validates every hit, so a stale cache costs a scan, never correctness.
+	columnIndex map[string]int
+}
+
+// BuildColumnIndex caches column offsets for this table and its relations, making
+// ColumnIndex (and so Resource.Get/Set) a map lookup rather than a scan over Columns.
+// Call it once, from a single goroutine, after the table tree is final. It is optional:
+// lookups stay correct without it, and the mutators that shift offsets drop the cache.
+func (t *Table) BuildColumnIndex() {
+	idx := make(map[string]int, len(t.Columns))
+	for i := range t.Columns {
+		idx[t.Columns[i].Name] = i
+	}
+	t.columnIndex = idx
+	for _, rel := range t.Relations {
+		rel.BuildColumnIndex()
+	}
+}
+
+// ColumnIndex returns the offset of name in Columns, or -1. It uses the BuildColumnIndex
+// cache while that cache still agrees with Columns, and scans otherwise.
+func (t *Table) ColumnIndex(name string) int {
+	if i, ok := t.columnIndex[name]; ok && i < len(t.Columns) && t.Columns[i].Name == name {
+		return i
+	}
+	return t.Columns.Index(name)
 }
 
 var (
@@ -139,6 +167,7 @@ func AddCqIDs(table *Table) {
 		},
 		table.Columns...,
 	)
+	table.columnIndex = nil
 	for _, rel := range table.Relations {
 		AddCqIDs(rel)
 	}
@@ -149,6 +178,7 @@ func AddCqIDs(table *Table) {
 func AddCqClientID(t *Table) {
 	if t.Columns.Get(CqClientIDColumn.Name) == nil {
 		t.Columns = append(ColumnList{CqClientIDColumn}, t.Columns...)
+		t.columnIndex = nil
 	}
 	for _, rel := range t.Relations {
 		AddCqClientID(rel)
@@ -707,6 +737,7 @@ func (t *Table) OverwriteOrAddColumn(column *Column) {
 		}
 	}
 	t.Columns = append([]Column{*column}, t.Columns...)
+	t.columnIndex = nil
 }
 
 func (t *Table) PrimaryKeys() []string {
@@ -754,6 +785,7 @@ func (t *Table) TableNames() []string {
 func (t *Table) Copy(parent *Table) *Table {
 	c := *t
 	c.Parent = parent
+	c.columnIndex = nil // don't alias the source table's cache
 	c.Columns = make([]Column, len(t.Columns))
 	copy(c.Columns, t.Columns)
 	c.Relations = make([]*Table, len(t.Relations))

schema/table_test.go

@@ -7,6 +7,7 @@ import (
 	"github.com/apache/arrow-go/v18/arrow"
 	"github.com/cloudquery/plugin-sdk/v4/types"
 	"github.com/google/go-cmp/cmp"
+	"github.com/google/go-cmp/cmp/cmpopts"
 	"github.com/stretchr/testify/require"
 )
 
@@ -787,7 +788,7 @@ func TestTablesToAndFromArrow(t *testing.T) {
 		if err != nil {
 			t.Fatal(err)
 		}
-		if diff := cmp.Diff(table, tableFromArrow); diff != "" {
+		if diff := cmp.Diff(table, tableFromArrow, cmpopts.IgnoreUnexported(Table{})); diff != "" {
 			t.Errorf("diff (+got, -want): %v", diff)
 		}
 	}