Merge pull request #152893 from mw5h/backportrelease-24.1-152399

release-24.1: lookup join: add a cast when mapping references in computed columns

Author: mw5h

pkg/sql/logictest/testdata/logic_test/lookup_join

@@ -1601,8 +1601,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

@@ -260,17 +260,6 @@ func (b *ConstraintBuilder) Build(
 		keyCols = nil
 	}
 
-	// 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++ {
@@ -294,7 +283,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)
 
@@ -307,7 +296,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

@@ -126,7 +126,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)