From 61461b59d9bb06e56a274e83504469c3212e33e2 Mon Sep 17 00:00:00 2001 From: James Cor Date: Thu, 12 Dec 2024 13:01:30 -0800 Subject: [PATCH] fix panic for virtual columns in `dolt_diff_` (#8673) --- go/libraries/doltcore/schema/schema.go | 19 ++++--- go/libraries/doltcore/schema/schema_test.go | 13 +++-- .../doltcore/sqle/dtables/prolly_row_conv.go | 16 ++++-- .../sqle/enginetest/dolt_queries_diff.go | 56 +++++++++++++++++++ 4 files changed, 86 insertions(+), 18 deletions(-) diff --git a/go/libraries/doltcore/schema/schema.go b/go/libraries/doltcore/schema/schema.go index f15b775f6a1..bbdbdfd6e14 100644 --- a/go/libraries/doltcore/schema/schema.go +++ b/go/libraries/doltcore/schema/schema.go @@ -307,13 +307,12 @@ func ArePrimaryKeySetsDiffable(format *types.NomsBinFormat, fromSch, toSch Schem // use to map key, value val.Tuple's of schema |inSch| to |outSch|. The first // ordinal map is for keys, and the second is for values. If a column of |inSch| // is missing in |outSch| then that column's index in the ordinal map holds -1. -func MapSchemaBasedOnTagAndName(inSch, outSch Schema) ([]int, []int, error) { +func MapSchemaBasedOnTagAndName(inSch, outSch Schema) (val.OrdinalMapping, val.OrdinalMapping, error) { keyMapping := make([]int, inSch.GetPKCols().Size()) - valMapping := make([]int, inSch.GetNonPKCols().Size()) // if inSch or outSch is empty schema. This can be from added or dropped table. if len(inSch.GetAllCols().cols) == 0 || len(outSch.GetAllCols().cols) == 0 { - return keyMapping, valMapping, nil + return keyMapping, make([]int, inSch.GetNonPKCols().Size()), nil } err := inSch.GetPKCols().Iter(func(tag uint64, col Column) (stop bool, err error) { @@ -330,14 +329,20 @@ func MapSchemaBasedOnTagAndName(inSch, outSch Schema) ([]int, []int, error) { return nil, nil, err } - err = inSch.GetNonPKCols().Iter(func(tag uint64, col Column) (stop bool, err error) { - i := inSch.GetNonPKCols().TagToIdx[col.Tag] - if col, ok := outSch.GetNonPKCols().GetByName(col.Name); ok { - j := outSch.GetNonPKCols().TagToIdx[col.Tag] + inNonPKCols := inSch.GetNonPKCols() + outNonPKCols := outSch.GetNonPKCols() + valMapping := make([]int, inSch.GetNonPKCols().Size()) + err = inNonPKCols.Iter(func(tag uint64, col Column) (stop bool, err error) { + i := inNonPKCols.TagToIdx[col.Tag] + if col.Virtual { + valMapping[i] = -1 + } else if col, ok := outNonPKCols.GetByName(col.Name); ok { + j := outNonPKCols.TagToIdx[col.Tag] valMapping[i] = j } else { valMapping[i] = -1 } + return false, nil }) if err != nil { diff --git a/go/libraries/doltcore/schema/schema_test.go b/go/libraries/doltcore/schema/schema_test.go index 2182fd8efae..5705e4ca624 100644 --- a/go/libraries/doltcore/schema/schema_test.go +++ b/go/libraries/doltcore/schema/schema_test.go @@ -25,6 +25,7 @@ import ( "github.com/dolthub/dolt/go/libraries/doltcore/schema/typeinfo" "github.com/dolthub/dolt/go/store/types" + "github.com/dolthub/dolt/go/store/val" ) const ( @@ -231,7 +232,7 @@ func TestArePrimaryKeySetsDiffable(t *testing.T) { From Schema To Schema Diffable bool - KeyMap []int + KeyMap val.OrdinalMapping }{ { Name: "Basic", @@ -240,7 +241,7 @@ func TestArePrimaryKeySetsDiffable(t *testing.T) { To: MustSchemaFromCols(NewColCollection( NewColumn("pk", 0, types.IntKind, true))), Diffable: true, - KeyMap: []int{0}, + KeyMap: val.OrdinalMapping{0}, }, { Name: "PK-Column renames", @@ -249,7 +250,7 @@ func TestArePrimaryKeySetsDiffable(t *testing.T) { To: MustSchemaFromCols(NewColCollection( NewColumn("pk2", 1, types.IntKind, true))), Diffable: true, - KeyMap: []int{0}, + KeyMap: val.OrdinalMapping{0}, }, { Name: "Only pk ordering should matter for diffability", @@ -259,7 +260,7 @@ func TestArePrimaryKeySetsDiffable(t *testing.T) { To: MustSchemaFromCols(NewColCollection( NewColumn("pk", 1, types.IntKind, true))), Diffable: true, - KeyMap: []int{0}, + KeyMap: val.OrdinalMapping{0}, }, { Name: "Only pk ordering should matter for diffability - inverse", @@ -269,7 +270,7 @@ func TestArePrimaryKeySetsDiffable(t *testing.T) { NewColumn("col1", 2, types.IntKind, false), NewColumn("pk", 1, types.IntKind, true))), Diffable: true, - KeyMap: []int{0}, + KeyMap: val.OrdinalMapping{0}, }, { Name: "Only pk ordering should matter for diffability - compound", @@ -281,7 +282,7 @@ func TestArePrimaryKeySetsDiffable(t *testing.T) { NewColumn("pk1", 0, types.IntKind, true), NewColumn("pk2", 2, types.IntKind, true))), Diffable: true, - KeyMap: []int{0, 1}, + KeyMap: val.OrdinalMapping{0, 1}, }, { Name: "Tag mismatches", diff --git a/go/libraries/doltcore/sqle/dtables/prolly_row_conv.go b/go/libraries/doltcore/sqle/dtables/prolly_row_conv.go index d0f5714ccb2..f842494ef08 100644 --- a/go/libraries/doltcore/sqle/dtables/prolly_row_conv.go +++ b/go/libraries/doltcore/sqle/dtables/prolly_row_conv.go @@ -102,21 +102,27 @@ func NewProllyRowConverter(inSch, outSch schema.Schema, warnFn rowconv.WarnFunct // PutConverted converts the |key| and |value| val.Tuple from |inSchema| to |outSchema| // and places the converted row in |dstRow|. -func (c ProllyRowConverter) PutConverted(ctx context.Context, key, value val.Tuple, dstRow []interface{}) error { - err := c.putFields(ctx, key, c.keyProj, c.keyDesc, c.pkTargetTypes, dstRow) +func (c ProllyRowConverter) PutConverted(ctx context.Context, key, value val.Tuple, dstRow sql.Row) error { + err := c.putFields(ctx, key, c.keyProj, c.keyDesc, c.pkTargetTypes, dstRow, true) if err != nil { return err } - return c.putFields(ctx, value, c.valProj, c.valDesc, c.nonPkTargetTypes, dstRow) + return c.putFields(ctx, value, c.valProj, c.valDesc, c.nonPkTargetTypes, dstRow, false) } -func (c ProllyRowConverter) putFields(ctx context.Context, tup val.Tuple, proj val.OrdinalMapping, desc val.TupleDesc, targetTypes []sql.Type, dstRow []interface{}) error { +func (c ProllyRowConverter) putFields(ctx context.Context, tup val.Tuple, proj val.OrdinalMapping, desc val.TupleDesc, targetTypes []sql.Type, dstRow sql.Row, isPk bool) error { + virtualOffset := 0 for i, j := range proj { if j == -1 { + // Skip over virtual columns in non-pk cols as they are not stored + if !isPk && c.inSchema.GetNonPKCols().GetByIndex(i).Virtual { + virtualOffset++ + } continue } - f, err := tree.GetField(ctx, desc, i, tup, c.ns) + + f, err := tree.GetField(ctx, desc, i-virtualOffset, tup, c.ns) if err != nil { return err } diff --git a/go/libraries/doltcore/sqle/enginetest/dolt_queries_diff.go b/go/libraries/doltcore/sqle/enginetest/dolt_queries_diff.go index e38dd4eefcf..9824ae67d99 100644 --- a/go/libraries/doltcore/sqle/enginetest/dolt_queries_diff.go +++ b/go/libraries/doltcore/sqle/enginetest/dolt_queries_diff.go @@ -1496,6 +1496,62 @@ on a.to_pk = b.to_pk;`, }, }, }, + { + Name: "diff table function works with virtual generated columns", + SetUpScript: []string{ + "create table t (i int primary key, j int, k int, jk int generated always as (10 * j + k), l int);", + "call dolt_commit('-Am', 'created table')", + "insert into t(i, j, k, l) values (1, 2, 3, 4);", + }, + Assertions: []queries.ScriptTestAssertion{ + { + Query: "select * from t;", + Expected: []sql.Row{ + {1, 2, 3, 23, 4}, + }, + }, + { + Query: "select to_i, to_jk, from_i, from_jk from dolt_diff_t;", + Expected: []sql.Row{ + {1, nil, nil, nil}, + }, + }, + { + Query: "select to_i, to_jk, from_i, from_jk from dolt_diff('HEAD', 'WORKING', 't');", + Expected: []sql.Row{ + {1, nil, nil, nil}, + }, + }, + }, + }, + { + Name: "diff table function works with stored generated columns", + SetUpScript: []string{ + "create table t (i int primary key, j int, k int, jk int generated always as (10 * j + k) stored, l int);", + "call dolt_commit('-Am', 'created table')", + "insert into t(i, j, k, l) values (1, 2, 3, 4);", + }, + Assertions: []queries.ScriptTestAssertion{ + { + Query: "select * from t;", + Expected: []sql.Row{ + {1, 2, 3, 23, 4}, + }, + }, + { + Query: "select to_i, to_jk, from_i, from_jk from dolt_diff_t;", + Expected: []sql.Row{ + {1, 23, nil, nil}, + }, + }, + { + Query: "select to_i, to_jk, from_i, from_jk from dolt_diff('HEAD', 'WORKING', 't');", + Expected: []sql.Row{ + {1, 23, nil, nil}, + }, + }, + }, + }, } var DiffStatTableFunctionScriptTests = []queries.ScriptTest{