Merge #32500

32500: opt: support COLLATE expressions r=justinj a=justinj

This commit adds support for COLLATE expressions. In particular, this
allows doing comparisons against existing collated string indexes, which
was not possible before.

Release note (sql change): Queries involving COLLATE expressions are now
supported by the cost-based optimizer.

Co-authored-by: Justin Jaffray <justin@cockroachlabs.com>

pkg/sql/opt/exec/execbuilder/scalar_builder.go

@@ -56,6 +56,7 @@ func init() {
 		opt.AnyOp:             (*Builder).buildAny,
 		opt.AnyScalarOp:       (*Builder).buildAnyScalar,
 		opt.IndirectionOp:     (*Builder).buildIndirection,
+		opt.CollateOp:         (*Builder).buildCollate,
 		opt.UnsupportedExprOp: (*Builder).buildUnsupportedExpr,
 
 		// Item operators.
@@ -402,6 +403,15 @@ func (b *Builder) buildIndirection(
 	return tree.NewTypedIndirectionExpr(expr, index), nil
 }
 
+func (b *Builder) buildCollate(ctx *buildScalarCtx, scalar opt.ScalarExpr) (tree.TypedExpr, error) {
+	expr, err := b.buildScalar(ctx, scalar.Child(0).(opt.ScalarExpr))
+	if err != nil {
+		return nil, err
+	}
+
+	return tree.NewTypedCollateExpr(expr, scalar.(*memo.CollateExpr).Locale), nil
+}
+
 func (b *Builder) buildUnsupportedExpr(
 	ctx *buildScalarCtx, scalar opt.ScalarExpr,
 ) (tree.TypedExpr, error) {

pkg/sql/opt/memo/expr_format.go

@@ -433,7 +433,7 @@ func (f *ExprFmtCtx) FormatScalarProps(scalar opt.ScalarExpr) {
 func (f *ExprFmtCtx) formatScalarPrivate(scalar opt.ScalarExpr) {
 	var private interface{}
 	switch t := scalar.(type) {
-	case *NullExpr, *TupleExpr:
+	case *NullExpr, *TupleExpr, *CollateExpr:
 		// Private is redundant with logical type property.
 		private = nil
 

pkg/sql/opt/memo/typing.go

@@ -122,6 +122,7 @@ func init() {
 	typingFuncMap[opt.SubqueryOp] = typeSubquery
 	typingFuncMap[opt.ColumnAccessOp] = typeColumnAccess
 	typingFuncMap[opt.IndirectionOp] = typeIndirection
+	typingFuncMap[opt.CollateOp] = typeCollate
 
 	// Override default typeAsAggregate behavior for aggregate functions with
 	// a large number of possible overloads or where ReturnType depends on
@@ -177,6 +178,11 @@ func typeIndirection(e opt.ScalarExpr) types.T {
 	return types.UnwrapType(e.Child(0).(opt.ScalarExpr).DataType()).(types.TArray).Typ
 }
 
+// typeCollate returns the collated string typed with the given locale.
+func typeCollate(e opt.ScalarExpr) types.T {
+	return types.TCollatedString{Locale: e.(*CollateExpr).Locale}
+}
+
 // typeAsFirstArg returns the type of the expression's 0th argument.
 func typeAsFirstArg(e opt.ScalarExpr) types.T {
 	return e.Child(0).(opt.ScalarExpr).DataType()

pkg/sql/opt/norm/custom_funcs.go

@@ -1102,6 +1102,22 @@ func (c *CustomFuncs) ConvertConstArrayToTuple(scalar opt.ScalarExpr) opt.Scalar
 	return c.f.ConstructTuple(elems, types.TTuple{Types: ts})
 }
 
+// CastToCollatedString returns the given string or collated string as a
+// collated string constant with the given locale.
+func (c *CustomFuncs) CastToCollatedString(str opt.ScalarExpr, locale string) opt.ScalarExpr {
+	var value string
+	switch t := str.(*memo.ConstExpr).Value.(type) {
+	case *tree.DString:
+		value = string(*t)
+	case *tree.DCollatedString:
+		value = t.Contents
+	default:
+		panic(fmt.Sprintf("unexpected type for COLLATE: %T", str.(*memo.ConstExpr).Value))
+	}
+
+	return c.f.ConstructConst(tree.NewDCollatedString(value, locale, &c.f.evalCtx.CollationEnv))
+}
+
 // ----------------------------------------------------------------------
 //
 // Numeric Rules

pkg/sql/opt/norm/rules/scalar.opt

@@ -286,3 +286,13 @@
     (ConvertConstArrayToTuple $ary)
     $cmp
 )
+
+# FoldCollate converts a Collate expr over an uncollated string into a collated string
+# constant.
+[FoldCollate, Normalize]
+(Collate
+    $input:(Const)
+    $locale:*
+)
+=>
+(CastToCollatedString $input $locale)

pkg/sql/opt/norm/testdata/rules/scalar

@@ -1109,3 +1109,62 @@ select
  │    └── key: (1)
  └── filters
       └── k = ANY '{1,2,3}'::INT[] [type=bool, outer=(1)]
+
+# --------------------------------------------------
+# FoldCollate
+# --------------------------------------------------
+
+norm expect=FoldCollate
+SELECT 'hello' COLLATE en_u_ks_level1
+----
+project
+ ├── columns: "?column?":1(collatedstring{en_u_ks_level1}!null)
+ ├── cardinality: [1 - 1]
+ ├── key: ()
+ ├── fd: ()-->(1)
+ ├── values
+ │    ├── cardinality: [1 - 1]
+ │    ├── key: ()
+ │    └── tuple [type=tuple]
+ └── projections
+      └── const: 'hello' COLLATE en_u_ks_level1 [type=collatedstring{en_u_ks_level1}]
+
+norm expect=FoldCollate
+SELECT ('hello' COLLATE en_u_ks_level1) COLLATE en_u_ks_level1
+----
+project
+ ├── columns: "?column?":1(collatedstring{en_u_ks_level1}!null)
+ ├── cardinality: [1 - 1]
+ ├── key: ()
+ ├── fd: ()-->(1)
+ ├── values
+ │    ├── cardinality: [1 - 1]
+ │    ├── key: ()
+ │    └── tuple [type=tuple]
+ └── projections
+      └── const: 'hello' COLLATE en_u_ks_level1 [type=collatedstring{en_u_ks_level1}]
+
+norm expect=FoldCollate
+SELECT ('hello' COLLATE en) COLLATE en_u_ks_level1
+----
+project
+ ├── columns: "?column?":1(collatedstring{en_u_ks_level1}!null)
+ ├── cardinality: [1 - 1]
+ ├── key: ()
+ ├── fd: ()-->(1)
+ ├── values
+ │    ├── cardinality: [1 - 1]
+ │    ├── key: ()
+ │    └── tuple [type=tuple]
+ └── projections
+      └── const: 'hello' COLLATE en_u_ks_level1 [type=collatedstring{en_u_ks_level1}]
+
+norm expect-not=FoldCollate
+SELECT s COLLATE en_u_ks_level1 FROM a
+----
+project
+ ├── columns: s:7(collatedstring{en_u_ks_level1})
+ ├── scan a
+ │    └── columns: a.s:4(string)
+ └── projections
+      └── a.s COLLATE en_u_ks_level1 [type=collatedstring{en_u_ks_level1}, outer=(4)]

pkg/sql/opt/ops/scalar.opt

@@ -655,6 +655,18 @@ define FunctionPrivate {
 	Overload   FuncOverload
 }
 
+# Collate is an expression of the form
+#
+#     x COLLATE y
+#
+# Where x is a "string type" (meaning either a normal string or a collated string),
+# and y is a locale. It evaluates to the string collated to the given locale.
+[Scalar]
+define Collate {
+    Input  ScalarExpr
+    Locale string
+}
+
 [Scalar]
 define Coalesce {
     Args ScalarListExpr

pkg/sql/opt/optbuilder/scalar.go

@@ -106,6 +106,10 @@ func (b *Builder) buildScalar(
 		}
 		out = b.factory.ConstructArray(els, arrayType)
 
+	case *tree.CollateExpr:
+		in := b.buildScalar(t.Expr.(tree.TypedExpr), inScope, nil, nil, colRefs)
+		out = b.factory.ConstructCollate(in, t.Locale)
+
 	case *tree.ArrayFlatten:
 		if b.AllowUnsupportedExpr {
 			out = b.factory.ConstructUnsupportedExpr(t)

pkg/sql/opt/optbuilder/testdata/scalar

@@ -705,7 +705,8 @@ concat [type=jsonb]
 build-scalar allow-unsupported
 'hello' COLLATE en
 ----
-unsupported-expr: 'hello' COLLATE en [type=collatedstring{en}]
+collate [type=collatedstring{en}]
+ └── const: 'hello' [type=string]
 
 build-scalar
 random()
@@ -1055,3 +1056,9 @@ tuple [type=tuple{tuple{bool, bool, bool}, tuple{bool, bool, bool}, tuple{bool,
       ├── false [type=bool]
       ├── true [type=bool]
       └── false [type=bool]
+
+build-scalar vars=(string)
+@1 COLLATE en
+----
+collate [type=collatedstring{en}]
+ └── variable: @1 [type=string]

pkg/sql/opt/xform/testdata/rules/scan

@@ -512,3 +512,43 @@ memo (optimized, ~5KB, required=[presentation: i:2,k:1])
  ├── G6: (ge G7 G8)
  ├── G7: (variable s)
  └── G8: (const 'foo')
+
+# Collated strings are treated properly.
+exec-ddl
+CREATE TABLE x (s STRING COLLATE en_u_ks_level1 PRIMARY KEY)
+----
+TABLE x
+ ├── s collatedstring{en_u_ks_level1} not null
+ └── INDEX primary
+      └── s collatedstring{en_u_ks_level1} not null
+
+opt
+SELECT s FROM x WHERE s < 'hello' COLLATE en_u_ks_level1
+----
+scan x
+ ├── columns: s:1(collatedstring{en_u_ks_level1}!null)
+ ├── constraint: /1: [ - /'hello' COLLATE en_u_ks_level1)
+ └── key: (1)
+
+opt
+SELECT s FROM x WHERE s = 'hello' COLLATE en_u_ks_level1
+----
+scan x
+ ├── columns: s:1(collatedstring{en_u_ks_level1}!null)
+ ├── constraint: /1: [/'hello' COLLATE en_u_ks_level1 - /'hello' COLLATE en_u_ks_level1]
+ ├── cardinality: [0 - 1]
+ ├── key: ()
+ └── fd: ()-->(1)
+
+# Can't generate spans for other collations.
+opt
+SELECT s FROM x WHERE s COLLATE en = 'hello' COLLATE en
+----
+select
+ ├── columns: s:1(collatedstring{en_u_ks_level1}!null)
+ ├── key: (1)
+ ├── scan x
+ │    ├── columns: s:1(collatedstring{en_u_ks_level1}!null)
+ │    └── key: (1)
+ └── filters
+      └── s COLLATE en = 'hello' COLLATE en [type=bool, outer=(1)]

pkg/sql/sem/tree/eval.go

@@ -2491,7 +2491,7 @@ type EvalContext struct {
 	// EXPLAIN(TYPES[, NORMALIZE]).
 	SkipNormalize bool
 
-	collationEnv CollationEnvironment
+	CollationEnv CollationEnvironment
 
 	TestingKnobs EvalContextTestingKnobs
 
@@ -3124,7 +3124,7 @@ func PerformCast(ctx *EvalContext, d Datum, t coltypes.CastTargetType) (Datum, e
 			if c.N > 0 && c.N < uint(len(s)) {
 				s = s[:c.N]
 			}
-			return NewDCollatedString(s, c.Locale, &ctx.collationEnv), nil
+			return NewDCollatedString(s, c.Locale, &ctx.CollationEnv), nil
 		case *coltypes.TName:
 			return NewDName(s), nil
 		}
@@ -3460,9 +3460,9 @@ func (expr *CollateExpr) Eval(ctx *EvalContext) (Datum, error) {
 	}
 	switch d := unwrapped.(type) {
 	case *DString:
-		return NewDCollatedString(string(*d), expr.Locale, &ctx.collationEnv), nil
+		return NewDCollatedString(string(*d), expr.Locale, &ctx.CollationEnv), nil
 	case *DCollatedString:
-		return NewDCollatedString(d.Contents, expr.Locale, &ctx.collationEnv), nil
+		return NewDCollatedString(d.Contents, expr.Locale, &ctx.CollationEnv), nil
 	default:
 		return nil, pgerror.NewErrorf(pgerror.CodeDatatypeMismatchError, "incompatible type for COLLATE: %s", d)
 	}

pkg/sql/sem/tree/expr.go

@@ -449,6 +449,16 @@ func NewTypedIndirectionExpr(expr, index TypedExpr) *IndirectionExpr {
 	return node
 }
 
+// NewTypedCollateExpr returns a new CollateExpr that is verified to be well-typed.
+func NewTypedCollateExpr(expr TypedExpr, locale string) *CollateExpr {
+	node := &CollateExpr{
+		Expr:   expr,
+		Locale: locale,
+	}
+	node.typ = types.TCollatedString{Locale: locale}
+	return node
+}
+
 func (node *ComparisonExpr) memoizeFn() {
 	fOp, fLeft, fRight, _, _ := foldComparisonExpr(node.Operator, node.Left, node.Right)
 	leftRet, rightRet := fLeft.(TypedExpr).ResolvedType(), fRight.(TypedExpr).ResolvedType()

pkg/sql/sem/tree/parse_string.go

@@ -42,7 +42,7 @@ func ParseStringAs(t types.T, s string, evalCtx *EvalContext) (Datum, error) {
 				return nil, err
 			}
 		case types.TCollatedString:
-			d = NewDCollatedString(s, t.Locale, &evalCtx.collationEnv)
+			d = NewDCollatedString(s, t.Locale, &evalCtx.CollationEnv)
 		default:
 			d, err = parseStringAs(t, s, evalCtx)
 			if d == nil && err == nil {