diff --git a/scheduler/scheduler.go b/scheduler/scheduler.go index 13913dbc80..c3924bc161 100644 --- a/scheduler/scheduler.go +++ b/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{ diff --git a/schema/column_index_test.go b/schema/column_index_test.go new file mode 100644 index 0000000000..c2aa352613 --- /dev/null +++ b/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")) + }) +} diff --git a/schema/resource.go b/schema/resource.go index 0ac7eb7ccc..a19474c000 100644 --- a/schema/resource.go +++ 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 { diff --git a/schema/resource_test.go b/schema/resource_test.go index 2d64b2cb83..deb041d559 100644 --- a/schema/resource_test.go +++ b/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") }) + }) + } +} diff --git a/schema/table.go b/schema/table.go index 48712ac35b..599f8b638a 100644 --- a/schema/table.go +++ b/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)) diff --git a/schema/table_test.go b/schema/table_test.go index 2d491a245a..7e5baf0599 100644 --- a/schema/table_test.go +++ b/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) } }