Merge pull request #34168 from RaduBerinde/backport2.1-34123

release-2.1: opt: add rule to simplify DistinctOn with no grouping columns

pkg/sql/logictest/testdata/logic_test/distinct_on

@@ -361,3 +361,14 @@ SELECT DISTINCT ON (x, y, z) pk1 FROM (SELECT * FROM xyz WHERE x >= 2) ORDER BY
 5
 6
 7
+
+# Regression tests for #34112: distinct on constant column.
+query II
+SELECT DISTINCT ON (x) x, y FROM xyz WHERE x = 1 ORDER BY x, y
+----
+1 1
+
+query I
+SELECT count(*) FROM (SELECT DISTINCT ON (x) x, y FROM xyz WHERE x = 1 ORDER BY x, y)
+----
+1

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

@@ -597,12 +597,17 @@ func (b *Builder) buildGroupBy(ev memo.ExprView) (execPlan, error) {
 }
 
 func (b *Builder) buildDistinct(ev memo.ExprView) (execPlan, error) {
+	def := ev.Private().(*memo.GroupByDef)
+	if def.GroupingCols.Empty() {
+		// A DistinctOn with no grouping columns should have been converted to a
+		// LIMIT 1 by normalization rules.
+		return execPlan{}, fmt.Errorf("cannot execute distinct on no columns")
+	}
 	input, err := b.buildGroupByInput(ev)
 	if err != nil {
 		return execPlan{}, err
 	}
 
-	def := ev.Private().(*memo.GroupByDef)
 	distinctCols := input.getColumnOrdinalSet(def.GroupingCols)
 	orderedCols := input.getColumnOrdinalSet(def.StreamingAggCols(&ev.Physical().Ordering))
 	node, err := b.factory.ConstructDistinct(input.root, distinctCols, orderedCols)

pkg/sql/opt/exec/execbuilder/testdata/distinct_on

@@ -279,9 +279,10 @@ render          ·            ·            (min)      ·
 query TTTTT
 EXPLAIN (VERBOSE) SELECT DISTINCT ON(min(x)) min(x) FROM xyz GROUP BY y HAVING min(x) = 1
 ----
-distinct                  ·            ·            (min)      weak-key(min)
- └── render               ·            ·            (min)      ·
-      │                   render 0     agg0         ·          ·
+render                    ·            ·            (min)      ·
+ │                        render 0     agg0         ·          ·
+ └── limit                ·            ·            (y, agg0)  ·
+      │                   count        1            ·          ·
       └── filter          ·            ·            (y, agg0)  ·
            │              filter       agg0 = 1     ·          ·
            └── group      ·            ·            (y, agg0)  ·
@@ -397,3 +398,32 @@ render               ·            ·            (pk1)           ·
            └── scan  ·            ·            (x, y, z, pk1)  ·
 ·                    table        xyz@primary  ·               ·
 ·                    spans        ALL          ·               ·
+
+# Regression tests for #34112: distinct on constant column.
+query TTTTT
+EXPLAIN (VERBOSE) SELECT DISTINCT ON (x) x, y FROM xyz WHERE x = 1 ORDER BY x, y
+----
+limit           ·       ·            (x, y)  +y
+ │              count   1            ·       ·
+ └── sort       ·       ·            (x, y)  +y
+      │         order   +y           ·       ·
+      └── scan  ·       ·            (x, y)  ·
+·               table   xyz@primary  ·       ·
+·               spans   ALL          ·       ·
+·               filter  x = 1        ·       ·
+
+query TTTTT
+EXPLAIN (VERBOSE) SELECT count(*) FROM (SELECT DISTINCT ON (x) x, y FROM xyz WHERE x = 1 ORDER BY x, y)
+----
+group                     ·            ·             (count)  ·
+ │                        aggregate 0  count_rows()  ·        ·
+ │                        scalar       ·             ·        ·
+ └── render               ·            ·             ()       ·
+      └── limit           ·            ·             (x, y)   +y
+           │              count        1             ·        ·
+           └── sort       ·            ·             (x, y)   +y
+                │         order        +y            ·        ·
+                └── scan  ·            ·             (x, y)   ·
+·                         table        xyz@primary   ·        ·
+·                         spans        ALL           ·        ·
+·                         filter       x = 1         ·        ·

pkg/sql/opt/norm/groupby.go

@@ -223,3 +223,41 @@ func (c *CustomFuncs) hasRemovableAggDistinct(
 	}
 	return false, 0
 }
+
+// HasNoGroupingCols returns true if the GroupingCols in the private are empty.
+func (c *CustomFuncs) HasNoGroupingCols(def memo.PrivateID) bool {
+	return c.f.mem.LookupPrivate(def).(*memo.GroupByDef).GroupingCols.Empty()
+}
+
+// GroupingInputOrdering returns the Ordering in the private.
+func (c *CustomFuncs) GroupingInputOrdering(def memo.PrivateID) memo.PrivateID {
+	groupByDef := c.f.mem.LookupPrivate(def).(*memo.GroupByDef)
+	return c.f.mem.InternOrderingChoice(&groupByDef.Ordering)
+}
+
+// ConstructProjectionFromDistinctOn converts a DistinctOn to a projection; this
+// is correct when the input has at most one row. Note that DistinctOn can only
+// have aggregations of type FirstAgg or ConstAgg.
+func (c *CustomFuncs) ConstructProjectionFromDistinctOn(input, aggs memo.GroupID) memo.GroupID {
+	aggsExpr := c.f.mem.NormExpr(aggs).AsAggregations()
+	aggsElems := c.f.mem.LookupList(aggsExpr.Aggs())
+	aggsColList := c.ExtractColList(aggsExpr.Cols())
+
+	var def memo.ProjectionsOpDef
+	var projections []memo.GroupID
+	for i, outputCol := range aggsColList {
+		aggExpr := memo.MakeNormExprView(c.mem, aggsElems[i])
+		varExpr := memo.ExtractVarFromAggInput(aggExpr.Child(0))
+		inputCol := varExpr.Private().(opt.ColumnID)
+		if inputCol == outputCol {
+			def.PassthroughCols.Add(int(inputCol))
+		} else {
+			def.SynthesizedCols = append(def.SynthesizedCols, outputCol)
+			projections = append(projections, varExpr.Group())
+		}
+	}
+	return c.f.ConstructProject(
+		input,
+		c.f.ConstructProjections(c.f.InternList(projections), c.f.InternProjectionsOpDef(&def)),
+	)
+}

pkg/sql/opt/norm/rules/groupby.opt

@@ -2,7 +2,7 @@
 # groupby.opt contains normalization rules for the GroupBy operator.
 # =============================================================================
 
-# ConvertGroupByToDsitinct converts a GroupBy operator that has no aggregations
+# ConvertGroupByToDistinct converts a GroupBy operator that has no aggregations
 # to an equivalent DistinctOn operator.
 [ConvertGroupByToDistinct, Normalize]
 (GroupBy
@@ -76,3 +76,20 @@ $input
     (RemoveAggDistinctForKeys $aggregations $def $input)
     $def
 )
+
+# EliminateDistinctOnNoColumns eliminates a DistinctOn with no grouping columns,
+# replacing it with a projection and a LIMIT 1. For example:
+#   SELECT DISTINCT ON (a) a, b FROM ab WHERE a=1
+# is equivalent to:
+#   SELECT a, b FROM ab WHERE a=1 LIMIT 1
+[EliminateDistinctOnNoColumns, Normalize]
+(DistinctOn
+    $input:*
+    $aggregations:*
+    $groupingPrivate:* & (HasNoGroupingCols $groupingPrivate)
+)
+=>
+(ConstructProjectionFromDistinctOn
+    (Limit $input (Const 1) (GroupingInputOrdering $groupingPrivate))
+    $aggregations
+)

pkg/sql/opt/norm/testdata/rules/decorrelate

@@ -2896,17 +2896,17 @@ FROM a
 WHERE k=10 AND (SELECT y FROM xy WHERE y=k LIMIT 1) = i AND (SELECT x FROM xy LIMIT 1) = 100
 ----
 project
- ├── columns: k:1(int) i:2(int!null) f:3(float) s:4(string) j:5(jsonb)
+ ├── columns: k:1(int!null) i:2(int!null) f:3(float) s:4(string) j:5(jsonb)
  ├── cardinality: [0 - 1]
  ├── key: ()
  ├── fd: ()-->(1-5)
  └── select
-      ├── columns: k:1(int) i:2(int!null) f:3(float) s:4(string) j:5(jsonb) xy.y:7(int!null)
+      ├── columns: k:1(int!null) i:2(int!null) f:3(float) s:4(string) j:5(jsonb) xy.y:7(int!null)
       ├── cardinality: [0 - 1]
       ├── key: ()
       ├── fd: ()-->(1-5,7)
-      ├── distinct-on
-      │    ├── columns: k:1(int) i:2(int) f:3(float) s:4(string) j:5(jsonb) xy.y:7(int)
+      ├── limit
+      │    ├── columns: k:1(int!null) i:2(int) f:3(float) s:4(string) j:5(jsonb) xy.y:7(int)
       │    ├── cardinality: [0 - 1]
       │    ├── key: ()
       │    ├── fd: ()-->(1-5,7)
@@ -2952,19 +2952,7 @@ project
       │    │         ├── left ordering: +1
       │    │         ├── right ordering: +7
       │    │         └── true [type=bool]
-      │    └── aggregations [outer=(1-5,7)]
-      │         ├── const-agg [type=int, outer=(2)]
-      │         │    └── variable: i [type=int, outer=(2)]
-      │         ├── const-agg [type=float, outer=(3)]
-      │         │    └── variable: f [type=float, outer=(3)]
-      │         ├── const-agg [type=string, outer=(4)]
-      │         │    └── variable: s [type=string, outer=(4)]
-      │         ├── const-agg [type=jsonb, outer=(5)]
-      │         │    └── variable: j [type=jsonb, outer=(5)]
-      │         ├── first-agg [type=int, outer=(7)]
-      │         │    └── variable: xy.y [type=int, outer=(7)]
-      │         └── const-agg [type=int, outer=(1)]
-      │              └── variable: k [type=int, outer=(1)]
+      │    └── const: 1 [type=int]
       └── filters [type=bool, outer=(2,7), constraints=(/2: (/NULL - ]; /7: (/NULL - ]), fd=(2)==(7), (7)==(2)]
            └── i = xy.y [type=bool, outer=(2,7), constraints=(/2: (/NULL - ]; /7: (/NULL - ])]
 

pkg/sql/opt/norm/testdata/rules/groupby

@@ -526,3 +526,39 @@ project
            └── avg [type=decimal, outer=(4)]
                 └── agg-distinct [type=int, outer=(4)]
                      └── variable: z [type=int, outer=(4)]
+
+# --------------------------------------------------
+# EliminateDistinctOnNoColumns
+# --------------------------------------------------
+
+opt expect=EliminateDistinctOnNoColumns
+SELECT DISTINCT ON (a) a, b FROM abc WHERE a = 1
+----
+scan abc
+ ├── columns: a:1(int!null) b:2(int!null)
+ ├── constraint: /1/2/3: [/1 - /1]
+ ├── limit: 1
+ ├── key: ()
+ └── fd: ()-->(1,2)
+
+opt expect=EliminateDistinctOnNoColumns
+SELECT DISTINCT ON (b) b, c FROM abc WHERE b = 1 ORDER BY b, c
+----
+limit
+ ├── columns: b:2(int!null) c:3(int!null)
+ ├── internal-ordering: +3 opt(2)
+ ├── cardinality: [0 - 1]
+ ├── key: ()
+ ├── fd: ()-->(2,3)
+ ├── sort
+ │    ├── columns: b:2(int!null) c:3(int!null)
+ │    ├── fd: ()-->(2)
+ │    ├── ordering: +3 opt(2)
+ │    └── select
+ │         ├── columns: b:2(int!null) c:3(int!null)
+ │         ├── fd: ()-->(2)
+ │         ├── scan abc
+ │         │    └── columns: b:2(int!null) c:3(int!null)
+ │         └── filters [type=bool, outer=(2), constraints=(/2: [/1 - /1]; tight), fd=()-->(2)]
+ │              └── b = 1 [type=bool, outer=(2), constraints=(/2: [/1 - /1]; tight)]
+ └── const: 1 [type=int]

pkg/sql/opt/norm/testdata/rules/select

@@ -834,8 +834,8 @@ group-by
 opt expect=PushSelectIntoGroupBy
 SELECT * FROM (SELECT i FROM a GROUP BY i) a WHERE i=1
 ----
-distinct-on
- ├── columns: i:2(int)
+limit
+ ├── columns: i:2(int!null)
  ├── cardinality: [0 - 1]
  ├── key: ()
  ├── fd: ()-->(2)
@@ -846,9 +846,7 @@ distinct-on
  │    │    └── columns: i:2(int)
  │    └── filters [type=bool, outer=(2), constraints=(/2: [/1 - /1]; tight), fd=()-->(2)]
  │         └── i = 1 [type=bool, outer=(2), constraints=(/2: [/1 - /1]; tight)]
- └── aggregations [outer=(2)]
-      └── const-agg [type=int, outer=(2)]
-           └── variable: i [type=int, outer=(2)]
+ └── const: 1 [type=int]
 
 # Push down only conditions that do not depend on aggregations.
 opt expect=PushSelectIntoGroupBy