Merge #152399

152399: lookup join: add a cast when mapping references in computed columns r=mw5h a=mw5h

Previously, when using a lookup join on an index with a virtual column,
we would require the references inside that virtual column to be typed
identically on the left and right side. This was a change from our
previous behavior, which was to require they be equivalent, but not
identical. This tightening of the requirements was done to address issue
124732, which pertains to virtual columns using functions that are
senstive to the actual type of the input (e.g. 'pg_typeof').

By tightening the reference requirement, we excluded a set of queries
from using lookup joins that had previously been able to do so safely,
regressing their performance.

In this patch, we add a cast around the left hand side reference to the
right hand side's type, so that type sensitive functions return the
correct type. This allows us to go back to the old behavior of allowing
the types of these references to be equivalent rather than identical.

Fixes: #124732
Release note (performance improvement): Lookup joins can now be used on
tables with virtual columns even if the type of the search argument is
not identical to the column type referenced in the virtual column.

Co-authored-by: Matt White <[email protected]>

Author: craig[bot]

pkg/sql/logictest/testdata/logic_test/lookup_join

@@ -1708,8 +1708,10 @@ SELECT col1_6 FROM table_1_124732 INNER HASH JOIN table_3_124732 ON col3_0 = col
 ----
 -
 
-statement error pgcode XXUUU pq: could not produce a query plan conforming to the LOOKUP JOIN hint
+query T
 SELECT col1_6 FROM table_1_124732 INNER LOOKUP JOIN table_3_124732 ON col3_0 = col1_6;
+----
+-
 
 # Regression test for incorrectly remapping columns in a composite-sensitive
 # expression to produce a lookup join (#124732).

pkg/sql/opt/lookupjoin/constraint_builder.go

@@ -246,17 +246,6 @@ func (b *ConstraintBuilder) Build(
 		rightSideCols = append(rightSideCols, rightCol)
 	}
 
-	// rightEqIdenticalTypeCols is the set of columns in rightEq that have
-	// identical types to the corresponding columns in leftEq. This is used to
-	// determine if a computed column can be synthesized for a column in the
-	// index in order to allow a lookup join.
-	var rightEqIdenticalTypeCols opt.ColSet
-	for i := range rightEq {
-		if b.md.ColumnMeta(rightEq[i]).Type.Identical(b.md.ColumnMeta(leftEq[i]).Type) {
-			rightEqIdenticalTypeCols.Add(rightEq[i])
-		}
-	}
-
 	// All the lookup conditions must apply to the prefix of the index and so
 	// the projected columns created must be created in order.
 	for j := 0; j < numIndexKeyCols; j++ {
@@ -281,7 +270,7 @@ func (b *ConstraintBuilder) Build(
 		//
 		// NOTE: we must only consider equivalent columns with identical types,
 		// since column remapping is otherwise not valid.
-		if expr, ok := b.findComputedColJoinEquality(b.table, idxCol, rightEqIdenticalTypeCols); ok {
+		if expr, ok := b.findComputedColJoinEquality(b.table, idxCol, rightEq.ToSet()); ok {
 			colMeta := b.md.ColumnMeta(idxCol)
 			compEqCol := b.md.AddColumn(fmt.Sprintf("%s_eq", colMeta.Alias), colMeta.Type)
 
@@ -294,7 +283,22 @@ func (b *ConstraintBuilder) Build(
 
 			// Project the computed column expression, mapping all columns
 			// in rightEq to corresponding columns in leftEq.
-			projection := b.f.ConstructProjectionsItem(b.f.RemapCols(expr, b.eqColMap), compEqCol)
+			var replace norm.ReplaceFunc
+			replace = func(e opt.Expr) opt.Expr {
+				if v, ok := e.(*memo.VariableExpr); ok {
+					if col, ok := b.eqColMap.Get(int(v.Col)); ok {
+						// If the column is a computed column, we need to
+						// cast it to the type of the original column so that
+						// functions which are type sensitive still return the
+						// expected results.
+						return b.f.ConstructCast(b.f.ConstructVariable(opt.ColumnID(col)), b.md.ColumnMeta(v.Col).Type)
+					} else {
+						return e
+					}
+				}
+				return b.f.Replace(e, replace)
+			}
+			projection := b.f.ConstructProjectionsItem(b.f.Replace(expr, replace).(opt.ScalarExpr), compEqCol)
 			inputProjections = append(inputProjections, projection)
 			addEqualityColumns(compEqCol, idxCol)
 			derivedEquivCols.Add(compEqCol)

pkg/sql/opt/lookupjoin/testdata/computed

@@ -175,7 +175,11 @@ lookup expression:
 lookup-constraints left=(a regclass, b int) right=(x oid, v string not null as (x::string) stored) index=(v, x)
 x = a
 ----
-lookup join not possible
+key cols:
+  v = v_eq
+  x = a
+input projections:
+  v_eq = a::OID::STRING [type=STRING]
 
 # Computed columns cannot be remapped if the expression is composite-sensitive.
 lookup-constraints left=(a decimal, b int) right=(x decimal, v int not null as (x::int) stored) index=(v, x)