opt: optbuilder support for scanning virtual columns

This change adds support for scanning tables with virtual columns. The virtual columns are projected on top of the Scan. Release note: None
cockroachdb · Dec 7, 2020 · bee0662 · bee0662
1 parent 094f595
commit bee0662
Show file tree

Hide file tree

Showing 2 changed files with 126 additions and 67 deletions.
diff --git a/pkg/sql/opt/optbuilder/select.go b/pkg/sql/opt/optbuilder/select.go
@@ -117,7 +117,7 @@ func (b *Builder) buildDataSource(
 					includeMutations:       false,
 					includeSystem:          true,
 					includeVirtualInverted: false,
-					includeVirtualComputed: false,
+					includeVirtualComputed: true,
 				}),
 				indexFlags, locking, inScope,
 			)
@@ -401,7 +401,7 @@ func (b *Builder) buildScanFromTableRef(
 			includeMutations:       false,
 			includeSystem:          true,
 			includeVirtualInverted: false,
-			includeVirtualComputed: false,
+			includeVirtualComputed: true,
 		})
 	}
 
@@ -453,20 +453,26 @@ func (b *Builder) buildScan(
 
 	outScope = inScope.push()
 
-	var tabColIDs opt.ColSet
+	// We collect VirtualComputed columns separately; these cannot be scanned,
+	// they can only be projected afterward.
+	var tabColIDs, virtualColIDs opt.ColSet
 	outScope.cols = make([]scopeColumn, len(ordinals))
 	for i, ord := range ordinals {
 		col := tab.Column(ord)
 		colID := tabID.ColumnID(ord)
-		tabColIDs.Add(colID)
 		name := col.ColName()
 		kind := col.Kind()
+		if kind != cat.VirtualComputed {
+			tabColIDs.Add(colID)
+		} else {
+			virtualColIDs.Add(colID)
+		}
 		outScope.cols[i] = scopeColumn{
 			id:           colID,
 			name:         name,
 			table:        tabMeta.Alias,
 			typ:          col.DatumType(),
-			hidden:       col.IsHidden() || kind != cat.Ordinary,
+			hidden:       col.IsHidden() || (kind != cat.Ordinary && kind != cat.VirtualComputed),
 			kind:         kind,
 			mutation:     kind == cat.WriteOnly || kind == cat.DeleteOnly,
 			tableOrdinal: ord,
@@ -485,80 +491,96 @@ func (b *Builder) buildScan(
 		private := memo.ScanPrivate{Table: tabID, Cols: tabColIDs}
 		outScope.expr = b.factory.ConstructScan(&private)
 
-		// Virtual tables should not be collected as view dependencies.
-	} else {
-		private := memo.ScanPrivate{Table: tabID, Cols: tabColIDs}
-		if indexFlags != nil {
-			private.Flags.NoIndexJoin = indexFlags.NoIndexJoin
-			if indexFlags.Index != "" || indexFlags.IndexID != 0 {
-				idx := -1
-				for i := 0; i < tab.IndexCount(); i++ {
-					if tab.Index(i).Name() == tree.Name(indexFlags.Index) ||
-						tab.Index(i).ID() == cat.StableID(indexFlags.IndexID) {
-						idx = i
-						break
-					}
+		// Note: virtual tables should not be collected as view dependencies.
+		return outScope
+	}
+
+	private := memo.ScanPrivate{Table: tabID, Cols: tabColIDs}
+	if indexFlags != nil {
+		private.Flags.NoIndexJoin = indexFlags.NoIndexJoin
+		if indexFlags.Index != "" || indexFlags.IndexID != 0 {
+			idx := -1
+			for i := 0; i < tab.IndexCount(); i++ {
+				if tab.Index(i).Name() == tree.Name(indexFlags.Index) ||
+					tab.Index(i).ID() == cat.StableID(indexFlags.IndexID) {
+					idx = i
+					break
 				}
-				if idx == -1 {
-					var err error
-					if indexFlags.Index != "" {
-						err = errors.Errorf("index %q not found", tree.ErrString(&indexFlags.Index))
-					} else {
-						err = errors.Errorf("index [%d] not found", indexFlags.IndexID)
-					}
-					panic(err)
+			}
+			if idx == -1 {
+				var err error
+				if indexFlags.Index != "" {
+					err = errors.Errorf("index %q not found", tree.ErrString(&indexFlags.Index))
+				} else {
+					err = errors.Errorf("index [%d] not found", indexFlags.IndexID)
 				}
-				private.Flags.ForceIndex = true
-				private.Flags.Index = idx
-				private.Flags.Direction = indexFlags.Direction
+				panic(err)
 			}
+			private.Flags.ForceIndex = true
+			private.Flags.Index = idx
+			private.Flags.Direction = indexFlags.Direction
 		}
-		if locking.isSet() {
-			private.Locking = locking.get()
-		}
+	}
+	if locking.isSet() {
+		private.Locking = locking.get()
+	}
 
-		b.addCheckConstraintsForTable(tabMeta)
-		b.addComputedColsForTable(tabMeta)
+	b.addCheckConstraintsForTable(tabMeta)
+	b.addComputedColsForTable(tabMeta)
 
-		outScope.expr = b.factory.ConstructScan(&private)
+	outScope.expr = b.factory.ConstructScan(&private)
 
-		// Add the partial indexes after constructing the scan so we can use the
-		// logical properties of the scan to fully normalize the index
-		// predicates. Partial index predicates are only added if the outScope
-		// contains all the table's ordinary columns. If it does not, partial
-		// index predicates cannot be built because they may reference columns
-		// not in outScope. In the most common case, the outScope has the same
-		// number of columns as the table and we can skip checking that each
-		// ordinary column exists in outScope.
-		containsAllOrdinaryTableColumns := true
-		if len(outScope.cols) != tab.ColumnCount() {
-			for i := 0; i < tab.ColumnCount(); i++ {
-				col := tab.Column(i)
-				if col.Kind() == cat.Ordinary && !outScope.colSet().Contains(tabID.ColumnID(col.Ordinal())) {
-					containsAllOrdinaryTableColumns = false
-					break
-				}
+	if !virtualColIDs.Empty() {
+		// Project the expressions for the virtual columns (and pass through all
+		// scanned columns).
+		proj := make(memo.ProjectionsExpr, 0, virtualColIDs.Len())
+		virtualColIDs.ForEach(func(col opt.ColumnID) {
+			item := memo.ProjectionsItem{
+				Element: tabMeta.ComputedCols[col],
+				Col:     col,
+				Typ:     tab.Column(tabID.ColumnOrdinal(col)).DatumType(),
 			}
-		}
-		if containsAllOrdinaryTableColumns {
-			b.addPartialIndexPredicatesForTable(tabMeta, outScope)
-		}
+			proj = append(proj, item)
+		})
+		outScope.expr = b.factory.ConstructProject(outScope.expr, proj, tabColIDs)
+	}
 
-		if b.trackViewDeps {
-			dep := opt.ViewDep{DataSource: tab}
-			dep.ColumnIDToOrd = make(map[opt.ColumnID]int)
-			// We will track the ColumnID to Ord mapping so Ords can be added
-			// when a column is referenced.
-			for i, col := range outScope.cols {
-				dep.ColumnIDToOrd[col.id] = ordinals[i]
-			}
-			if private.Flags.ForceIndex {
-				dep.SpecificIndex = true
-				dep.Index = private.Flags.Index
+	// Add the partial indexes after constructing the scan so we can use the
+	// logical properties of the scan to fully normalize the index
+	// predicates. Partial index predicates are only added if the outScope
+	// contains all the table's ordinary columns. If it does not, partial
+	// index predicates cannot be built because they may reference columns
+	// not in outScope. In the most common case, the outScope has the same
+	// number of columns as the table and we can skip checking that each
+	// ordinary column exists in outScope.
+	containsAllOrdinaryTableColumns := true
+	if len(outScope.cols) != tab.ColumnCount() {
+		for i := 0; i < tab.ColumnCount(); i++ {
+			col := tab.Column(i)
+			if col.Kind() == cat.Ordinary && !outScope.colSet().Contains(tabID.ColumnID(col.Ordinal())) {
+				containsAllOrdinaryTableColumns = false
+				break
 			}
-			b.viewDeps = append(b.viewDeps, dep)
 		}
 	}
+	if containsAllOrdinaryTableColumns {
+		b.addPartialIndexPredicatesForTable(tabMeta, outScope)
+	}
+
+	if b.trackViewDeps {
+		dep := opt.ViewDep{DataSource: tab}
+		dep.ColumnIDToOrd = make(map[opt.ColumnID]int)
+		// We will track the ColumnID to Ord mapping so Ords can be added
+		// when a column is referenced.
+		for i, col := range outScope.cols {
+			dep.ColumnIDToOrd[col.id] = ordinals[i]
+		}
+		if private.Flags.ForceIndex {
+			dep.SpecificIndex = true
+			dep.Index = private.Flags.Index
+		}
+		b.viewDeps = append(b.viewDeps, dep)
+	}
 	return outScope
 }
 

diff --git a/pkg/sql/opt/optbuilder/testdata/virtual-columns b/pkg/sql/opt/optbuilder/testdata/virtual-columns
@@ -0,0 +1,37 @@
+exec-ddl
+CREATE TABLE t (
+  a INT PRIMARY KEY,
+  b INT,
+  c INT AS (a+b) VIRTUAL
+)
+----
+
+build
+SELECT * FROM t
+----
+project
+ ├── columns: a:1!null b:2 c:3
+ └── project
+      ├── columns: c:3 a:1!null b:2 crdb_internal_mvcc_timestamp:4
+      ├── scan t
+      │    ├── columns: a:1!null b:2 crdb_internal_mvcc_timestamp:4
+      │    └── computed column expressions
+      │         └── c:3
+      │              └── a:1 + b:2
+      └── projections
+           └── a:1 + b:2 [as=c:3]
+
+build
+SELECT c FROM t
+----
+project
+ ├── columns: c:3
+ └── project
+      ├── columns: c:3 a:1!null b:2 crdb_internal_mvcc_timestamp:4
+      ├── scan t
+      │    ├── columns: a:1!null b:2 crdb_internal_mvcc_timestamp:4
+      │    └── computed column expressions
+      │         └── c:3
+      │              └── a:1 + b:2
+      └── projections
+           └── a:1 + b:2 [as=c:3]