sql: support adding virtual columns

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)