Skip to content

Commit

Permalink
sqlsmith: add support for computed columns
Browse files Browse the repository at this point in the history 
This changes the random table generator to also create computed
columns (either STORED or VIRTUAL). Some example of definitions:
 - `col1_14 STRING NOT NULL AS (lower(CAST(col1_8 AS STRING))) VIRTUAL`
 - `col1_10 FLOAT8 AS (col1_5 + (-0.16607712222551002):::FLOAT8) STORED`
 - `col1_13 INT4 AS (col1_0 + col1_10) STORED`

Release justification: non-production code change.

Release note: None
  • Loading branch information
@RaduBerinde
RaduBerinde committed Mar 5, 2021
1 parent c0d8340 commit e8d1a90
Show file tree
Hide file tree
Showing 2 changed files with 149 additions and 19 deletions.
10 changes: 10 additions & 0 deletions pkg/sql/mutations/mutations.go
View file
Original file line number Diff line number Diff line change
Expand Up @@ -313,6 +313,8 @@ func statisticsMutator(
return stmts, changed
}

// foreignKeyMutator is a MultiStatementMutation implementation which adds
// foreign key references between existing columns.
func foreignKeyMutator(
rng *rand.Rand, stmts []tree.Statement,
) (mutated []tree.Statement, changed bool) {
Expand Down Expand Up @@ -369,6 +371,10 @@ func foreignKeyMutator(
// Choose a random column subset.
var fkCols []*tree.ColumnTableDef
for _, c := range cols[table.Table] {
if c.Computed.Computed {
// We don't support FK references from computed columns (#46672).
continue
}
if usedCols[table.Table][c.Name] {
continue
}
Expand Down Expand Up @@ -429,6 +435,10 @@ func foreignKeyMutator(
fkCol := fkCols[len(usingCols)]
found := false
for refI, refCol := range availCols {
if refCol.Computed.Virtual {
// We don't support FK references to virtual columns (#51296).
continue
}
fkColType := tree.MustBeStaticallyKnownType(fkCol.Type)
refColType := tree.MustBeStaticallyKnownType(refCol.Type)
if fkColType.Equivalent(refColType) && colinfo.ColumnTypeIsIndexable(refColType) {
Expand Down
158 changes: 139 additions & 19 deletions pkg/sql/rowenc/testutils.go
View file
Original file line number Diff line number Diff line change
Expand Up @@ -1146,6 +1146,15 @@ func RandCreateTableWithInterleave(
}
}
}

// colIdx generates numbers that are incorporated into column names.
colIdx := func(ordinal int) int {
if generateColumnIndexNumber != nil {
return int(generateColumnIndexNumber())
}
return ordinal
}

var interleaveDef *tree.InterleaveDef
if interleaveIntoPK != nil && len(interleaveIntoPK.Columns) > 0 {
// Make the interleave prefix, which has to be exactly the columns in the
Expand All @@ -1165,11 +1174,7 @@ func RandCreateTableWithInterleave(
// Loop until we generate an indexable column type.
var extraCol *tree.ColumnTableDef
for {
colIdx := i + prefixLength
if generateColumnIndexNumber != nil {
colIdx = int(generateColumnIndexNumber())
}
extraCol = randColumnTableDef(rng, tableIdx, colIdx)
extraCol = randColumnTableDef(rng, tableIdx, colIdx(i+prefixLength))
extraColType := tree.MustBeStaticallyKnownType(extraCol.Type)
if colinfo.ColumnTypeIsIndexable(extraColType) {
break
Expand Down Expand Up @@ -1205,12 +1210,10 @@ func RandCreateTableWithInterleave(
}
} else {
// Make new defs from scratch.
for i := 0; i < nColumns; i++ {
colIdx := i
if generateColumnIndexNumber != nil {
colIdx = int(generateColumnIndexNumber())
}
columnDef := randColumnTableDef(rng, tableIdx, colIdx)
nComputedColumns := randutil.RandIntInRange(rng, 0, (nColumns+1)/2)
nNormalColumns := nColumns - nComputedColumns
for i := 0; i < nNormalColumns; i++ {
columnDef := randColumnTableDef(rng, tableIdx, colIdx(i))
columnDefs = append(columnDefs, columnDef)
defs = append(defs, columnDef)
}
Expand All @@ -1236,6 +1239,14 @@ func RandCreateTableWithInterleave(
}
}
}

// Make defs for computed columns.
normalColDefs := columnDefs
for i := nNormalColumns; i < nColumns; i++ {
columnDef := randComputedColumnTableDef(rng, normalColDefs, tableIdx, colIdx(i))
columnDefs = append(columnDefs, columnDef)
defs = append(defs, columnDef)
}
}

// Make indexes.
Expand Down Expand Up @@ -1374,11 +1385,15 @@ func IndexStoringMutator(rng *rand.Rand, stmts []tree.Statement) ([]tree.Stateme
}
return colMap
}
generateStoringCols := func(rng *rand.Rand, tableCols []tree.Name, indexCols map[tree.Name]struct{}) []tree.Name {
generateStoringCols := func(rng *rand.Rand, tableInfo tableInfo, indexCols map[tree.Name]struct{}) []tree.Name {
var storingCols []tree.Name
for _, col := range tableCols {
_, ok := indexCols[col]
if ok {
for colOrdinal, col := range tableInfo.columnNames {
if _, ok := indexCols[col]; ok {
// Skip PK columns and columns already in the index.
continue
}
if tableInfo.columnsTableDefs[colOrdinal].Computed.Virtual {
// Virtual columns can't be stored.
continue
}
if rng.Intn(2) == 0 {
Expand All @@ -1403,7 +1418,7 @@ func IndexStoringMutator(rng *rand.Rand, stmts []tree.Statement) ([]tree.Stateme
for _, elem := range ast.Columns {
indexCols[elem.Column] = struct{}{}
}
ast.Storing = generateStoringCols(rng, tableInfo.columnNames, indexCols)
ast.Storing = generateStoringCols(rng, tableInfo, indexCols)
changed = true
}
case *tree.CreateTable:
Expand All @@ -1430,7 +1445,7 @@ func IndexStoringMutator(rng *rand.Rand, stmts []tree.Statement) ([]tree.Stateme
for _, elem := range idx.Columns {
indexCols[elem.Column] = struct{}{}
}
idx.Storing = generateStoringCols(rng, tableInfo.columnNames, indexCols)
idx.Storing = generateStoringCols(rng, tableInfo, indexCols)
changed = true
}
}
Expand Down Expand Up @@ -1493,8 +1508,8 @@ func PartialIndexMutator(rng *rand.Rand, stmts []tree.Statement) ([]tree.Stateme
return stmts, changed
}

// randColumnTableDef produces a random ColumnTableDef, with a random type and
// nullability.
// randColumnTableDef produces a random ColumnTableDef for a non-computed
// column, with a random type and nullability.
func randColumnTableDef(rand *rand.Rand, tableIdx int, colIdx int) *tree.ColumnTableDef {
columnDef := &tree.ColumnTableDef{
// We make a unique name for all columns by prefixing them with the table
Expand All @@ -1506,6 +1521,111 @@ func randColumnTableDef(rand *rand.Rand, tableIdx int, colIdx int) *tree.ColumnT
return columnDef
}

// randComputedColumnTableDef produces a random ColumnTableDef for a computed
// column (either STORED or VIRTUAL). The computed expressions refer to columns
// in normalColDefs.
func randComputedColumnTableDef(
rng *rand.Rand, normalColDefs []*tree.ColumnTableDef, tableIdx int, colIdx int,
) *tree.ColumnTableDef {
newDef := randColumnTableDef(rng, tableIdx, colIdx)
newDef.Computed.Computed = true
newDef.Computed.Virtual = (rng.Intn(2) == 0)

if rng.Intn(2) == 0 {
// Try to find a set of numeric columns with the same type; the computed
// expression will be of the form "a+b+c".
var cols []*tree.ColumnTableDef
var fam types.Family
for _, idx := range rng.Perm(len(normalColDefs)) {
x := normalColDefs[idx]
xFam := x.Type.(*types.T).Family()

if len(cols) == 0 {
switch xFam {
case types.IntFamily, types.FloatFamily, types.DecimalFamily:
fam = xFam
cols = append(cols, x)
}
} else if fam == xFam {
cols = append(cols, x)
if len(cols) > 1 && rng.Intn(2) == 0 {
break
}
}
}
if len(cols) > 1 {
var expr tree.Expr
expr = tree.NewUnresolvedName(string(cols[0].Name))
for _, x := range cols[1:] {
expr = &tree.BinaryExpr{
Operator: tree.Plus,
Left: expr,
Right: tree.NewUnresolvedName(string(x.Name)),
}
}
newDef.Type = cols[0].Type
newDef.Computed.Expr = expr
return newDef
}
}

// Pick a single column and create a computed column that depends on it.
// The expression is as follows:
// - for numeric types (int, float, decimal), the expression is "x+1";
// - for string type, the expression is "lower(x)";
// - for types that can be cast to string in computed columns, the expression
// is "lower(x::string)";
// - otherwise, the expression is "IF(x IS NULL, 'foo', 'bar')".
x := normalColDefs[randutil.RandIntInRange(rng, 0, len(normalColDefs))]
xTyp := x.Type.(*types.T)

switch xTyp.Family() {
case types.IntFamily, types.FloatFamily, types.DecimalFamily:
newDef.Type = xTyp
newDef.Computed.Expr = &tree.BinaryExpr{
Operator: tree.Plus,
Left: tree.NewUnresolvedName(string(x.Name)),
Right: RandDatum(rng, xTyp, false /* nullOk */),
}

case types.StringFamily:
newDef.Type = types.String
newDef.Computed.Expr = &tree.FuncExpr{
Func: tree.WrapFunction("lower"),
Exprs: tree.Exprs{tree.NewUnresolvedName(string(x.Name))},
}

default:
volatility, ok := tree.LookupCastVolatility(xTyp, types.String)
if ok && volatility <= tree.VolatilityImmutable {
// We can cast to string; use lower(x::string)
newDef.Type = types.String
newDef.Computed.Expr = &tree.FuncExpr{
Func: tree.WrapFunction("lower"),
Exprs: tree.Exprs{
&tree.CastExpr{
Expr: tree.NewUnresolvedName(string(x.Name)),
Type: types.String,
},
},
}
} else {
// We cannot cast this type to string in a computed column expression.
// Use IF(x IS NULL, 'foo', 'bar').
newDef.Type = types.String
newDef.Computed.Expr = &tree.IfExpr{
Cond: &tree.IsNullExpr{
Expr: tree.NewUnresolvedName(string(x.Name)),
},
True: RandDatum(rng, types.String, true /* nullOK */),
Else: RandDatum(rng, types.String, true /* nullOK */),
}
}
}

return newDef
}

// randIndexTableDefFromCols creates an IndexTableDef with a random subset of
// the given columns and a random direction.
func randIndexTableDefFromCols(
Expand Down

0 comments on commit e8d1a90

Please sign in to comment.