sql: WIP: add cast to massage computed col type to desired

cockroachdb · Jun 25, 2021 · 031c5b1 · 031c5b1
1 parent f661758
commit 031c5b1
Show file tree

Hide file tree

Showing 4 changed files with 58 additions and 4 deletions.
diff --git a/pkg/sql/opt/optbuilder/select.go b/pkg/sql/opt/optbuilder/select.go
@@ -705,6 +705,8 @@ func (b *Builder) addComputedColsForTable(tabMeta *opt.TableMeta) {
 			colID := tabMeta.MetaID.ColumnID(i)
 			var scalar opt.ScalarExpr
 			b.factory.FoldingControl().TemporarilyDisallowStableFolds(func() {
+				// TBD: should resolveAndRequireType do this?
+				texpr := tree.ReType(texpr, tabCol.DatumType())
 				scalar = b.buildScalar(texpr, tableScope, nil, nil, nil)
 			})
 			// Check if the expression contains non-immutable operators.

diff --git a/pkg/sql/opt/optbuilder/testdata/update b/pkg/sql/opt/optbuilder/testdata/update
@@ -1631,7 +1631,7 @@ update decimals
       │    │    │    │    │    ├── columns: a:7!null b:8 c:9 d:10 crdb_internal_mvcc_timestamp:11 tableoid:12
       │    │    │    │    │    └── computed column expressions
       │    │    │    │    │         └── d:10
-      │    │    │    │    │              └── a:7::DECIMAL + c:9::DECIMAL
+      │    │    │    │    │              └── (a:7::DECIMAL + c:9::DECIMAL)::DECIMAL(10,1)
       │    │    │    │    └── projections
       │    │    │    │         ├── 1.1 [as=a_new:13]
       │    │    │    │         └── ARRAY[0.95,NULL,15] [as=b_new:14]

diff --git a/pkg/sql/opt/optbuilder/testdata/update-col-cast-bug b/pkg/sql/opt/optbuilder/testdata/update-col-cast-bug
@@ -0,0 +1,52 @@
+exec-ddl
+create table t (a int2, b int2, c int2 as (a + b) virtual)
+----
+
+build format=show-types
+update t set a = (with cte as (select 1:::int8) select t.c from cte limit 1)
+----
+with &1 (cte)
+ ├── project
+ │    ├── columns: int8:13(int!null)
+ │    ├── values
+ │    │    └── () [type=tuple]
+ │    └── projections
+ │         └── 1 [as=int8:13, type=int]
+ └── update t
+      ├── columns: <none>
+      ├── fetch columns: a:7(int2) b:8(int2) t.c:9(int2) rowid:10(int)
+      ├── update-mapping:
+      │    ├── a_new:16 => a:1
+      │    └── c_comp:17 => t.c:3
+      └── project
+           ├── columns: c_comp:17(int) a:7(int2) b:8(int2) t.c:9(int2) rowid:10(int!null) crdb_internal_mvcc_timestamp:11(decimal) tableoid:12(oid) a_new:16(int2)
+           ├── project
+           │    ├── columns: a_new:16(int2) a:7(int2) b:8(int2) t.c:9(int2) rowid:10(int!null) crdb_internal_mvcc_timestamp:11(decimal) tableoid:12(oid)
+           │    ├── project
+           │    │    ├── columns: t.c:9(int2) a:7(int2) b:8(int2) rowid:10(int!null) crdb_internal_mvcc_timestamp:11(decimal) tableoid:12(oid)
+           │    │    ├── scan t
+           │    │    │    ├── columns: a:7(int2) b:8(int2) rowid:10(int!null) crdb_internal_mvcc_timestamp:11(decimal) tableoid:12(oid)
+           │    │    │    └── computed column expressions
+           │    │    │         └── t.c:9
+           │    │    │              └── (a:7::INT8 + b:8::INT8)::INT2 [type=int2]
+           │    │    └── projections
+           │    │         └── (a:7::INT8 + b:8::INT8)::INT2 [as=t.c:9, type=int2]
+           │    └── projections
+           │         └── subquery [as=a_new:16, type=int2]
+           │              └── max1-row
+           │                   ├── columns: c:15(int2)
+           │                   └── limit
+           │                        ├── columns: c:15(int2)
+           │                        ├── project
+           │                        │    ├── columns: c:15(int2)
+           │                        │    ├── limit hint: 1.00
+           │                        │    ├── with-scan &1 (cte)
+           │                        │    │    ├── columns: int8:14(int!null)
+           │                        │    │    ├── mapping:
+           │                        │    │    │    └──  int8:13(int) => int8:14(int)
+           │                        │    │    └── limit hint: 1.00
+           │                        │    └── projections
+           │                        │         └── t.c:9 [as=c:15, type=int2]
+           │                        └── 1 [type=int]
+           └── projections
+                └── a_new:16::INT8 + b:8::INT8 [as=c_comp:17, type=int]
diff --git a/pkg/sql/opt/optbuilder/testdata/upsert b/pkg/sql/opt/optbuilder/testdata/upsert
@@ -1742,7 +1742,7 @@ upsert decimals
       │    │    │    │    ├── columns: decimals.a:15!null decimals.b:16 c:17 d:18 crdb_internal_mvcc_timestamp:19 tableoid:20
       │    │    │    │    └── computed column expressions
       │    │    │    │         └── d:18
-      │    │    │    │              └── decimals.a:15::DECIMAL + c:17::DECIMAL
+      │    │    │    │              └── (decimals.a:15::DECIMAL + c:17::DECIMAL)::DECIMAL(10,1)
       │    │    │    └── filters
       │    │    │         └── a:10 = decimals.a:15
       │    │    └── projections
@@ -1814,7 +1814,7 @@ upsert decimals
       │    │    │    │    ├── columns: decimals.a:15!null b:16 c:17 d:18 crdb_internal_mvcc_timestamp:19 tableoid:20
       │    │    │    │    └── computed column expressions
       │    │    │    │         └── d:18
-      │    │    │    │              └── decimals.a:15::DECIMAL + c:17::DECIMAL
+      │    │    │    │              └── (decimals.a:15::DECIMAL + c:17::DECIMAL)::DECIMAL(10,1)
       │    │    │    └── filters
       │    │    │         └── a:10 = decimals.a:15
       │    │    └── projections
@@ -1894,7 +1894,7 @@ upsert decimals
       │    │    │    │    │    │    ├── columns: decimals.a:15!null decimals.b:16 c:17 d:18 crdb_internal_mvcc_timestamp:19 tableoid:20
       │    │    │    │    │    │    └── computed column expressions
       │    │    │    │    │    │         └── d:18
-      │    │    │    │    │    │              └── decimals.a:15::DECIMAL + c:17::DECIMAL
+      │    │    │    │    │    │              └── (decimals.a:15::DECIMAL + c:17::DECIMAL)::DECIMAL(10,1)
       │    │    │    │    │    └── filters
       │    │    │    │    │         └── a:10 = decimals.a:15
       │    │    │    │    └── projections