opt: generate constrained scans for indexed, boolean expressions

pkg/sql/opt/idxconstraint/index_constraints.go

@@ -706,6 +706,12 @@ func (c *indexConstraintCtx) makeSpansForExpr(
 		return c.makeSpansForExpr(offset, t.And, out)
 	}
 
+	// Support e as (c = TRUE) if c is an indexed, boolean, computed expression
+	// equivalent to e. This is similar to the VariableExpr case above.
+	if c.colType(offset).Family() == types.BoolFamily && c.isExpressionIndexColumn(e, offset) {
+		return c.makeSpansForSingleColumnDatum(offset, opt.EqOp, tree.DBoolTrue, out)
+	}
+
 	if e.ChildCount() < 2 {
 		c.unconstrained(offset, out)
 		return false
@@ -1280,6 +1286,15 @@ func (c *indexConstraintCtx) isIndexColumn(e opt.Expr, offset int) bool {
 	if v, ok := e.(*memo.VariableExpr); ok && v.Col == c.columns[offset].ID() {
 		return true
 	}
+	if c.isExpressionIndexColumn(e, offset) {
+		return true
+	}
+	return false
+}
+
+// isExpressionIndexColumn returns true if e is computed column expression that
+// corresponds to index column <offset>.
+func (c *indexConstraintCtx) isExpressionIndexColumn(e opt.Expr, offset int) bool {
 	if c.computedCols != nil && e == c.computedCols[c.columns[offset].ID()] {
 		return true
 	}

pkg/sql/opt/idxconstraint/testdata/computed

@@ -30,6 +30,17 @@ index-constraints vars=(j json, field int as ((j->'foo')::string::int) stored) i
 ----
 [/10 - /10]
 
+# Generate constraints for indexed, boolean computed expressions.
+index-constraints vars=(b bool, c bool as (coalesce(b, false)) stored) index=(c)
+coalesce(b, false)
+----
+[/true - /true]
+
+index-constraints vars=(b bool, c bool as (coalesce(b, false)) stored) index=(c)
+NOT coalesce(b, false)
+----
+[/false - /false]
+
 # ---------------------------------------------------
 # Unit tests for computed column predicate derivation
 # ---------------------------------------------------

pkg/sql/opt/xform/testdata/rules/computed

@@ -278,18 +278,12 @@ CREATE TABLE t83390 (
 opt
 SELECT * FROM t83390@idx WHERE a IS NULL
 ----
-select
+index-join t83390
  ├── columns: k:1!null a:2
  ├── key: (1)
  ├── fd: ()-->(2)
- ├── index-join t83390
- │    ├── columns: k:1!null a:2
- │    ├── key: (1)
- │    ├── fd: (1)-->(2)
- │    └── scan t83390@idx
- │         ├── columns: k:1!null
- │         ├── constraint: /5/1: [/true - /true]
- │         ├── flags: force-index=idx
- │         └── key: (1)
- └── filters
-      └── a:2 IS NULL [outer=(2), constraints=(/2: [/NULL - /NULL]; tight), fd=()-->(2)]
+ └── scan t83390@idx
+      ├── columns: k:1!null
+      ├── constraint: /5/1: [/true - /true]
+      ├── flags: force-index=idx
+      └── key: (1)

pkg/sql/opt/xform/testdata/rules/select

@@ -2613,6 +2613,43 @@ project
       └── filters
            └── z:3 = 5 [outer=(3), constraints=(/3: [/5 - /5]; tight), fd=()-->(3)]
 
+# Regression test for #114798. A constrained scan should be planned on boolean
+# virtual computed columns.
+
+exec-ddl
+CREATE TABLE t114798 (
+  k INT PRIMARY KEY,
+  b BOOL,
+  s BOOL AS (COALESCE(b, false)) STORED,
+  v BOOL AS (COALESCE(b, true)) VIRTUAL,
+  INDEX (s),
+  INDEX (v)
+)
+----
+
+opt expect=GenerateConstrainedScans
+SELECT k FROM t114798 WHERE COALESCE(b, false)
+----
+project
+ ├── columns: k:1!null
+ ├── key: (1)
+ └── scan t114798@t114798_s_idx
+      ├── columns: k:1!null
+      ├── constraint: /3/1: [/true - /true]
+      └── key: (1)
+
+opt expect=GenerateConstrainedScans
+SELECT k FROM t114798 WHERE COALESCE(b, true)
+----
+project
+ ├── columns: k:1!null
+ ├── key: (1)
+ └── scan t114798@t114798_v_idx
+      ├── columns: k:1!null
+      ├── constraint: /4/1: [/true - /true]
+      └── key: (1)
+
+
 # --------------------------------------------------
 # GenerateInvertedIndexScans
 # --------------------------------------------------