sql: prevent unsafe comparison constant unification

pkg/sql/logictest/testdata/logic_test/issue_172978

@@ -0,0 +1,56 @@
+# LogicTest: local
+# Regression tests for #172978.
+
+statement ok
+CREATE TABLE big (i INT8 PRIMARY KEY)
+
+statement ok
+INSERT INTO big VALUES (9007199254740993)
+
+query I
+SELECT i FROM big WHERE i = 9007199254740992.0::FLOAT8
+----
+9007199254740993
+
+query I
+SELECT i FROM big WHERE i = ANY(ARRAY[9007199254740992.0]::FLOAT[])
+----
+9007199254740993
+
+statement ok
+CREATE TABLE e (tz TIMESTAMPTZ PRIMARY KEY)
+
+statement ok
+INSERT INTO e VALUES ('2020-06-15 16:00:00+00')
+
+statement ok
+SET TIME ZONE 'UTC'
+
+statement ok
+PREPARE ps AS SELECT count(*) FROM e WHERE tz = '2020-06-15 12:00:00'::TIMESTAMP
+
+statement ok
+PREPARE pp AS SELECT count(*) FROM e WHERE tz = ANY(ARRAY['2020-06-15 12:00:00']::TIMESTAMP[])
+
+query I
+EXECUTE ps
+----
+0
+
+query I
+EXECUTE pp
+----
+0
+
+statement ok
+SET TIME ZONE 'America/New_York'
+
+query I
+EXECUTE ps
+----
+1
+
+query I
+EXECUTE pp
+----
+1

pkg/sql/opt/norm/scalar_funcs.go

@@ -11,6 +11,7 @@ import (
 	"github.com/cockroachdb/cockroach/pkg/sql/opt"
 	"github.com/cockroachdb/cockroach/pkg/sql/opt/memo"
 	"github.com/cockroachdb/cockroach/pkg/sql/opt/props"
+	"github.com/cockroachdb/cockroach/pkg/sql/sem/cast"
 	"github.com/cockroachdb/cockroach/pkg/sql/sem/eval"
 	"github.com/cockroachdb/cockroach/pkg/sql/sem/tree"
 	"github.com/cockroachdb/cockroach/pkg/sql/types"
@@ -117,8 +118,11 @@ func (c *CustomFuncs) IsConstValueEqual(const1, const2 opt.ScalarExpr) bool {
 }
 
 // UnifyComparison attempts to convert a constant expression to the type of the
-// variable expression, if that conversion can round-trip and is monotonic.
-// Otherwise it returns ok=false.
+// variable expression, if that conversion can round-trip and is monotonic. The
+// implicit conversion of the variable to the original type must also be
+// injective; otherwise distinct variable values can compare equal to the
+// constant before the rewrite but not afterwards. Otherwise it returns
+// ok=false.
 func (c *CustomFuncs) UnifyComparison(
 	v *memo.VariableExpr, cnst *memo.ConstExpr,
 ) (_ opt.ScalarExpr, ok bool) {
@@ -133,6 +137,18 @@ func (c *CustomFuncs) UnifyComparison(
 	if !isMonotonicConversion(originalType, desiredType) {
 		return nil, false
 	}
+	if !isInjectiveConversion(desiredType, originalType) {
+		return nil, false
+	}
+
+	// Do not fold stable casts into a reusable memo. In particular, converting a
+	// TIMESTAMP constant to TIMESTAMPTZ depends on the session time zone.
+	for _, conversion := range [][2]*types.T{{originalType, desiredType}, {desiredType, originalType}} {
+		volatility, ok := cast.LookupCastVolatility(conversion[0], conversion[1])
+		if !ok || !c.CanFoldOperator(volatility) {
+			return nil, false
+		}
+	}
 
 	// Check that the datum can round-trip between the types. If this is true, it
 	// means we don't lose any information needed to generate spans, and combined
@@ -157,6 +173,29 @@ func (c *CustomFuncs) UnifyComparison(
 	return c.f.ConstructConst(convertedDatum, desiredType), true
 }
 
+// isInjectiveConversion returns true if converting every value from FROM to TO
+// preserves its identity. This is required because comparison overloads cast
+// the variable to the constant's type before comparing it to the constant.
+//
+// Keep this list deliberately narrow. The round-trip test in UnifyComparison
+// proves only that the constant is representable in the variable's type; it
+// says nothing about other variable values. For example, INT8 to FLOAT8 is not
+// injective above 2^53, even when the FLOAT8 constant round-trips to INT8.
+func isInjectiveConversion(from, to *types.T) bool {
+	switch from.Family() {
+	case types.IntFamily:
+		return to.Family() == types.DecimalFamily
+	case types.DateFamily:
+		switch to.Family() {
+		case types.TimestampFamily, types.TimestampTZFamily:
+			return true
+		}
+	case types.TimestampFamily:
+		return to.Family() == types.TimestampTZFamily
+	}
+	return false
+}
+
 // SimplifyWhens removes known unreachable WHEN cases and constructs a new CASE
 // statement. Any known true condition is converted to the ELSE. If only the
 // ELSE remains, its expression is returned. condition must be a ConstValue.