sql: support adding virtual columns

Support adding and removing virtual columns (when no indexes on those
columns are in play).

Note that because these columns don't need backfill in existing
indexes, they don't need special WriteOnly / DeleteOnly logic (other
than having the right visibility).

Release note: None

pkg/sql/add_column.go

@@ -17,6 +17,7 @@ import (
 	"github.com/cockroachdb/cockroach/pkg/sql/pgwire/pgcode"
 	"github.com/cockroachdb/cockroach/pkg/sql/pgwire/pgerror"
 	"github.com/cockroachdb/cockroach/pkg/sql/sem/tree"
+	"github.com/cockroachdb/cockroach/pkg/sql/sessiondata"
 	"github.com/cockroachdb/cockroach/pkg/sql/sqlerrors"
 	"github.com/cockroachdb/cockroach/pkg/util/errorutil/unimplemented"
 	"github.com/cockroachdb/errors"
@@ -29,6 +30,7 @@ func (p *planner) addColumnImpl(
 	tn *tree.TableName,
 	desc *tabledesc.Mutable,
 	t *tree.AlterTableAddColumn,
+	sessionData *sessiondata.SessionData,
 ) error {
 	d := t.ColumnDef
 	version := params.ExecCfg().Settings.Version.ActiveVersionOrEmpty(params.ctx)
@@ -143,7 +145,7 @@ func (p *planner) addColumnImpl(
 	}
 
 	if d.IsComputed() {
-		if d.IsVirtual() {
+		if d.IsVirtual() && !sessionData.VirtualColumnsEnabled {
 			return unimplemented.NewWithIssue(57608, "virtual computed columns")
 		}
 		computedColValidator := schemaexpr.MakeComputedColumnValidator(

pkg/sql/alter_table.go

@@ -169,7 +169,7 @@ func (n *alterTableNode) startExec(params runParams) error {
 			}
 			var err error
 			params.p.runWithOptions(resolveFlags{contextDatabaseID: n.tableDesc.ParentID}, func() {
-				err = params.p.addColumnImpl(params, n, tn, n.tableDesc, t)
+				err = params.p.addColumnImpl(params, n, tn, n.tableDesc, t, params.SessionData())
 			})
 			if err != nil {
 				return err

pkg/sql/backfill.go

@@ -215,7 +215,7 @@ func (sc *SchemaChanger) runBackfill(ctx context.Context) error {
 		case descpb.DescriptorMutation_ADD:
 			switch t := m.Descriptor_.(type) {
 			case *descpb.DescriptorMutation_Column:
-				if tabledesc.ColumnNeedsBackfill(m.GetColumn()) {
+				if tabledesc.ColumnNeedsBackfill(m.Direction, m.GetColumn()) {
 					needColumnBackfill = true
 				}
 			case *descpb.DescriptorMutation_Index:
@@ -249,7 +249,9 @@ func (sc *SchemaChanger) runBackfill(ctx context.Context) error {
 		case descpb.DescriptorMutation_DROP:
 			switch t := m.Descriptor_.(type) {
 			case *descpb.DescriptorMutation_Column:
-				needColumnBackfill = true
+				if tabledesc.ColumnNeedsBackfill(m.Direction, m.GetColumn()) {
+					needColumnBackfill = true
+				}
 			case *descpb.DescriptorMutation_Index:
 				if !canClearRangeForDrop(t.Index) {
 					droppedIndexDescs = append(droppedIndexDescs, *t.Index)
@@ -1803,7 +1805,7 @@ func runSchemaChangesInTxn(
 			case *descpb.DescriptorMutation_ComputedColumnSwap:
 				return AlterColTypeInTxnNotSupportedErr
 			case *descpb.DescriptorMutation_Column:
-				if doneColumnBackfill || !tabledesc.ColumnNeedsBackfill(m.GetColumn()) {
+				if doneColumnBackfill || !tabledesc.ColumnNeedsBackfill(m.Direction, m.GetColumn()) {
 					break
 				}
 				if err := columnBackfillInTxn(ctx, planner.Txn(), planner.EvalContext(), planner.SemaCtx(), immutDesc, traceKV); err != nil {
@@ -1830,7 +1832,7 @@ func runSchemaChangesInTxn(
 			// Drop the name and drop the associated data later.
 			switch t := m.Descriptor_.(type) {
 			case *descpb.DescriptorMutation_Column:
-				if doneColumnBackfill {
+				if doneColumnBackfill || !tabledesc.ColumnNeedsBackfill(m.Direction, m.GetColumn()) {
 					break
 				}
 				if err := columnBackfillInTxn(

pkg/sql/catalog/tabledesc/structured.go

@@ -3777,9 +3777,25 @@ func (desc *wrapper) MakeFirstMutationPublic(includeConstraints bool) (*Mutable,
 	return table, nil
 }
 
-// ColumnNeedsBackfill returns true if adding the given column requires a
-// backfill (dropping a column always requires a backfill).
-func ColumnNeedsBackfill(desc *descpb.ColumnDescriptor) bool {
+// ColumnNeedsBackfill returns true if adding or dropping (according to
+// the direction) the given column requires backfill.
+func ColumnNeedsBackfill(
+	direction descpb.DescriptorMutation_Direction, desc *descpb.ColumnDescriptor,
+) bool {
+	if desc.Virtual {
+		// Virtual columns are not stored in the primary index, so they do not need
+		// backfill.
+		return false
+	}
+	if direction == descpb.DescriptorMutation_DROP {
+		// In all other cases, DROP requires backfill.
+		return true
+	}
+	// ADD requires backfill for:
+	//  - columns with non-NULL default value
+	//  - computed columns
+	//  - non-nullable columns (note: if a non-nullable column doesn't have a
+	//    default value, the backfill will fail unless the table is empty).
 	if desc.HasNullDefault() {
 		return false
 	}
@@ -3800,10 +3816,7 @@ func (desc *wrapper) HasColumnBackfillMutation() bool {
 			// Index backfills don't affect changefeeds.
 			continue
 		}
-		// It's unfortunate that there's no one method we can call to check if a
-		// mutation will be a backfill or not, but this logic was extracted from
-		// backfill.go.
-		if m.Direction == descpb.DescriptorMutation_DROP || ColumnNeedsBackfill(col) {
+		if ColumnNeedsBackfill(m.Direction, col) {
 			return true
 		}
 	}

pkg/sql/catalog/tabledesc/structured_test.go

@@ -1823,32 +1823,66 @@ func TestColumnNeedsBackfill(t *testing.T) {
 	// Define variable strings here such that we can pass their address below.
 	null := "NULL"
 	four := "4:::INT8"
+
 	// Create Column Descriptors that reflect the definition of a column with a
 	// default value of NULL that was set implicitly, one that was set explicitly,
 	// and one that has an INT default value, respectively.
-	implicitNull := &descpb.ColumnDescriptor{
-		Name: "im", ID: 2, Type: types.Int, DefaultExpr: nil, Nullable: true, ComputeExpr: nil,
-	}
-	explicitNull := &descpb.ColumnDescriptor{
-		Name: "ex", ID: 3, Type: types.Int, DefaultExpr: &null, Nullable: true, ComputeExpr: nil,
-	}
-	defaultNotNull := &descpb.ColumnDescriptor{
-		Name: "four", ID: 4, Type: types.Int, DefaultExpr: &four, Nullable: true, ComputeExpr: nil,
-	}
-	// Verify that a backfill doesn't occur according to the ColumnNeedsBackfill
-	// function for the default NULL values, and that it does occur for an INT
-	// default value.
-	if ColumnNeedsBackfill(implicitNull) != false {
-		t.Fatal("Expected implicit SET DEFAULT NULL to not require a backfill," +
-			" ColumnNeedsBackfill states that it does.")
-	}
-	if ColumnNeedsBackfill(explicitNull) != false {
-		t.Fatal("Expected explicit SET DEFAULT NULL to not require a backfill," +
-			" ColumnNeedsBackfill states that it does.")
+	testCases := []struct {
+		info string
+		desc descpb.ColumnDescriptor
+		// add is true of we expect backfill when adding this column.
+		add bool
+		// drop is true of we expect backfill when adding this column.
+		drop bool
+	}{
+		{
+			info: "implicit SET DEFAULT NULL",
+			desc: descpb.ColumnDescriptor{
+				Name: "am", ID: 2, Type: types.Int, DefaultExpr: nil, Nullable: true, ComputeExpr: nil,
+			},
+			add:  false,
+			drop: true,
+		}, {
+			info: "explicit SET DEFAULT NULL",
+			desc: descpb.ColumnDescriptor{
+				Name: "ex", ID: 3, Type: types.Int, DefaultExpr: &null, Nullable: true, ComputeExpr: nil,
+			},
+			add:  false,
+			drop: true,
+		},
+		{
+			info: "explicit SET DEFAULT non-NULL",
+			desc: descpb.ColumnDescriptor{
+				Name: "four", ID: 4, Type: types.Int, DefaultExpr: &four, Nullable: true, ComputeExpr: nil,
+			},
+			add:  true,
+			drop: true,
+		},
+		{
+			info: "computed stored",
+			desc: descpb.ColumnDescriptor{
+				Name: "stored", ID: 5, Type: types.Int, DefaultExpr: nil, ComputeExpr: &four,
+			},
+			add:  true,
+			drop: true,
+		},
+		{
+			info: "computed virtual",
+			desc: descpb.ColumnDescriptor{
+				Name: "virtual", ID: 6, Type: types.Int, DefaultExpr: nil, ComputeExpr: &four, Virtual: true,
+			},
+			add:  false,
+			drop: false,
+		},
 	}
-	if ColumnNeedsBackfill(defaultNotNull) != true {
-		t.Fatal("Expected explicit SET DEFAULT NULL to require a backfill," +
-			" ColumnNeedsBackfill states that it does not.")
+
+	for _, tc := range testCases {
+		if ColumnNeedsBackfill(descpb.DescriptorMutation_ADD, &tc.desc) != tc.add {
+			t.Errorf("expected ColumnNeedsBackfill to be %v for adding %s", tc.add, tc.info)
+		}
+		if ColumnNeedsBackfill(descpb.DescriptorMutation_DROP, &tc.desc) != tc.drop {
+			t.Errorf("expected ColumnNeedsBackfill to be %v for dropping %s", tc.drop, tc.info)
+		}
 	}
 }
 

pkg/sql/logictest/testdata/logic_test/virtual_columns

@@ -247,3 +247,79 @@ UPDATE t_idx SET b=b+1 RETURNING w
 ----
 3
 7
+
+# Test schema changes with virtual columns.
+
+statement ok
+CREATE TABLE sc (a INT PRIMARY KEY, b INT)
+
+statement ok
+INSERT INTO sc VALUES (1, 10), (2, 20), (3, 30);
+
+statement ok
+ALTER TABLE sc ADD COLUMN v INT AS (a+b) VIRTUAL
+
+query III rowsort,colnames
+SELECT * FROM sc
+----
+a  b   v
+1  10  11
+2  20  22
+3  30  33
+
+statement ok
+ALTER TABLE sc ADD COLUMN x INT AS (a+1) VIRTUAL, ADD COLUMN y INT AS (b+1) VIRTUAL, ADD COLUMN z INT AS (a+b) VIRTUAL
+
+query IIIIII rowsort,colnames
+SELECT * FROM sc
+----
+a  b   v   x  y   z
+1  10  11  2  11  11
+2  20  22  3  21  22
+3  30  33  4  31  33
+
+statement error computed columns cannot reference other computed columns
+ALTER TABLE sc ADD COLUMN u INT AS (a+v) VIRTUAL
+
+statement ok
+ALTER TABLE sc DROP COLUMN z
+
+query IIIII rowsort,colnames
+SELECT * FROM sc
+----
+a  b   v   x  y
+1  10  11  2  11
+2  20  22  3  21
+3  30  33  4  31
+
+statement ok
+ALTER TABLE sc DROP COLUMN x, DROP COLUMN y
+
+query III rowsort,colnames
+SELECT * FROM sc
+----
+a  b   v
+1  10  11
+2  20  22
+3  30  33
+
+# Add virtual columns inside an explicit transactions.
+statement ok
+BEGIN
+
+statement ok
+ALTER TABLE sc ADD COLUMN w1 INT AS (a*b) VIRTUAL
+
+statement ok
+ALTER TABLE sc ADD COLUMN w2 INT AS (b*2) VIRTUAL
+
+statement ok
+COMMIT
+
+query IIIII rowsort,colnames
+SELECT * FROM sc
+----
+a  b   v   w1  w2
+1  10  11  10  20
+2  20  22  40  40
+3  30  33  90  60

pkg/sql/opt_catalog.go

@@ -670,9 +670,9 @@ func newOptTable(
 				desc.ComputeExpr,
 			)
 		} else {
-			if kind != cat.Ordinary {
-				return nil, errors.AssertionFailedf("virtual mutation column")
-			}
+			// Note: a WriteOnly or DeleteOnly mutation column doesn't require any
+			// special treatment inside the optimizer, other than having the correct
+			// visibility.
 			ot.columns[ordinal].InitVirtualComputed(
 				ordinal,
 				cat.StableID(desc.ID),

pkg/sql/schema_changer_test.go

@@ -4571,6 +4571,44 @@ func TestAddComputedColumn(t *testing.T) {
 	sqlDB.CheckQueryResults(t, `SELECT * FROM t.test ORDER BY a`, [][]string{{"2", "7"}, {"10", "15"}})
 }
 
+// TestNoBackfillForVirtualColumn verifies that adding or dropping a virtual
+// column does not involve a backfill.
+func TestNoBackfillForVirtualColumn(t *testing.T) {
+	defer leaktest.AfterTest(t)()
+	defer log.Scope(t).Close(t)
+
+	sawBackfill := false
+	params, _ := tests.CreateTestServerParams()
+	params.Knobs = base.TestingKnobs{
+		DistSQL: &execinfra.TestingKnobs{
+			RunBeforeBackfillChunk: func(sp roachpb.Span) error {
+				sawBackfill = true
+				return nil
+			},
+		},
+	}
+	tc := serverutils.StartNewTestCluster(t, 1, base.TestClusterArgs{ServerArgs: params})
+	defer tc.Stopper().Stop(context.Background())
+	db := tc.ServerConn(0)
+	sqlDB := sqlutils.MakeSQLRunner(db)
+
+	sqlDB.Exec(t, `CREATE DATABASE t`)
+	sqlDB.Exec(t, `CREATE TABLE t.test (a INT PRIMARY KEY)`)
+	sqlDB.Exec(t, `INSERT INTO t.test VALUES (1), (2), (3)`)
+	sqlDB.Exec(t, `SET experimental_enable_virtual_columns = true`)
+
+	sawBackfill = false
+	sqlDB.Exec(t, `ALTER TABLE t.test ADD COLUMN b INT AS (a + 5) VIRTUAL`)
+	if sawBackfill {
+		t.Fatal("saw backfill when adding virtual column")
+	}
+
+	sqlDB.Exec(t, `ALTER TABLE t.test DROP COLUMN b`)
+	if sawBackfill {
+		t.Fatal("saw backfill when dropping virtual column")
+	}
+}
+
 func TestSchemaChangeAfterCreateInTxn(t *testing.T) {
 	defer leaktest.AfterTest(t)()
 	defer log.Scope(t).Close(t)