fix panic for virtual columns in dolt_diff_<tbl> (#8673)

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 {

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",

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
 		}

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{