Merge cockroachdb#107842 cockroachdb#108281

107842: opt: fix error with DISTINCT ON and ORDER BY ASC NULLS LAST r=rytaft a=rytaft

This commit fixes an error that could occur in the optbuilder when planning a query with `DISTINCT ON` and a non-standard nulls ordering.

The optimizer supports queries with a non-standard nulls ordering by projecting a column with the expression `(col IS NULL)` and adding it to the ordering. Since we require that `DISTINCT ON` columns must be the prefix of any ordering columns, we must account for the new ordering column when building `DISTINCT ON`. A previous bug fix for cockroachdb#90763 caused the new column to be simply ignored when building `DISTINCT ON`, but this was insufficient. We need to actually include the new column among the `DISTINCT ON` columns. This commit makes that change.

Fixes cockroachdb#107839

Release note (bug fix): Fixed a spurious error "no data source matches prefix" that could occur during planning for a query with `DISTINCT ON` and `ORDER BY ASC NULLS LAST` or `ORDER BY DESC NULLS FIRST`.

108281: batcheval: remove stray log line r=dt a=dt

Release note: none.
Epic: none.

Fixes cockroachdb#108275.

Co-authored-by: Rebecca Taft <[email protected]>
Co-authored-by: David Taylor <[email protected]>

Author: 3 people

pkg/kv/kvserver/batcheval/cmd_add_sstable.go

@@ -184,8 +184,6 @@ func EvalAddSSTable(
 		}, nil
 	}
 
-	log.Infof(ctx, "non-remote AddSSTable")
-
 	// Reject AddSSTable requests not writing at the request timestamp if requested.
 	if AddSSTableRequireAtRequestTimestamp.Get(&cArgs.EvalCtx.ClusterSettings().SV) &&
 		sstToReqTS.IsEmpty() {

pkg/sql/logictest/testdata/logic_test/distinct_on

@@ -405,3 +405,32 @@ Action
 Biography
 Crime
 NULL
+
+# Regression test for #107839. More fixes for NULLS FIRST/LAST with DISTINCT ON.
+statement ok
+CREATE TABLE t1 (id int, str text);
+CREATE TABLE t2 (id int, num int);
+INSERT INTO t1 VALUES (1, 'hello'), (2, NULL);
+INSERT INTO t2 VALUES (1, 1), (2, 2), (NULL, NULL)
+
+query II
+SELECT
+DISTINCT ON (t2.id)
+t2.*
+FROM t1, t2
+ORDER BY t2.id DESC NULLS FIRST
+----
+NULL  NULL
+2     2
+1     1
+
+query II
+SELECT
+DISTINCT ON (t2.id)
+t2.*
+FROM t1, t2
+ORDER BY t2.id ASC NULLS LAST
+----
+1     1
+2     2
+NULL  NULL

pkg/sql/opt/optbuilder/distinct.go

@@ -89,23 +89,29 @@ func (b *Builder) buildDistinctOn(
 	var seen opt.ColSet
 	for _, col := range inScope.ordering {
 		if !distinctOnCols.Contains(col.ID()) {
+			colIsValid := false
 			scopeCol := inScope.getColumn(col.ID())
 			if scopeCol != nil {
 				if isExpr, ok := scopeCol.scalar.(*memo.IsExpr); ok {
 					if _, ok := isExpr.Right.(*memo.NullExpr); ok {
 						if v, ok := isExpr.Left.(*memo.VariableExpr); ok {
 							if distinctOnCols.Contains(v.Col) {
 								// We have a col IS NULL expression (case 3 above).
-								continue
+								// Add the new column to distinctOnCols, since it doesn't change
+								// the semantics of the DISTINCT ON.
+								distinctOnCols.Add(col.ID())
+								colIsValid = true
 							}
 						}
 					}
 				}
 			}
-			panic(pgerror.Newf(
-				pgcode.InvalidColumnReference,
-				"SELECT DISTINCT ON expressions must match initial ORDER BY expressions",
-			))
+			if !colIsValid {
+				panic(pgerror.Newf(
+					pgcode.InvalidColumnReference,
+					"SELECT DISTINCT ON expressions must match initial ORDER BY expressions",
+				))
+			}
 		}
 		seen.Add(col.ID())
 		if seen.Equals(distinctOnCols) {

pkg/sql/opt/optbuilder/testdata/distinct_on

@@ -507,30 +507,87 @@ ORDER BY
   "genre" ASC NULLS LAST,
   "id" ASC NULLS LAST
 ----
+distinct-on
+ ├── columns: id:1!null name:2 genre:3  [hidden: nulls_ordering_genre:6!null]
+ ├── grouping columns: genre:3 nulls_ordering_genre:6!null
+ ├── internal-ordering: +7,+1 opt(3,6)
+ ├── ordering: +6,+3
+ ├── sort
+ │    ├── columns: id:1!null name:2 genre:3 nulls_ordering_genre:6!null nulls_ordering_id:7!null
+ │    ├── ordering: +6,+3,+7,+1
+ │    └── project
+ │         ├── columns: nulls_ordering_genre:6!null nulls_ordering_id:7!null id:1!null name:2 genre:3
+ │         ├── scan author
+ │         │    └── columns: id:1!null name:2 genre:3 crdb_internal_mvcc_timestamp:4 tableoid:5
+ │         └── projections
+ │              ├── genre:3 IS NULL [as=nulls_ordering_genre:6]
+ │              └── id:1 IS NULL [as=nulls_ordering_id:7]
+ └── aggregations
+      ├── first-agg [as=id:1]
+      │    └── id:1
+      └── first-agg [as=name:2]
+           └── name:2
+
+# Regression test for #107839. More fixes for NULLS FIRST/LAST with DISTINCT ON.
+
+exec-ddl
+CREATE TABLE t1 (id int, str text)
+----
+
+exec-ddl
+CREATE TABLE t2 (id int, num int)
+----
+
+build
+SELECT
+DISTINCT ON (t2.id)
+t1.id,
+t1.str
+FROM t1 JOIN t2 ON t1.id = t2.id
+ORDER BY t2.id DESC NULLS FIRST
+----
 sort
- ├── columns: id:1!null name:2 genre:3  [hidden: nulls_ordering_genre:8!null nulls_ordering_id:9!null]
- ├── ordering: +8,+3,+9,+1
- └── project
-      ├── columns: nulls_ordering_genre:8!null nulls_ordering_id:9!null id:1!null name:2 genre:3
-      ├── distinct-on
-      │    ├── columns: id:1!null name:2 genre:3
-      │    ├── grouping columns: genre:3
-      │    ├── internal-ordering: +6,+7,+1 opt(3)
-      │    ├── sort
-      │    │    ├── columns: id:1!null name:2 genre:3 nulls_ordering_genre:6!null nulls_ordering_id:7!null
-      │    │    ├── ordering: +7,+1 opt(3,6)
-      │    │    └── project
-      │    │         ├── columns: nulls_ordering_genre:6!null nulls_ordering_id:7!null id:1!null name:2 genre:3
-      │    │         ├── scan author
-      │    │         │    └── columns: id:1!null name:2 genre:3 crdb_internal_mvcc_timestamp:4 tableoid:5
-      │    │         └── projections
-      │    │              ├── genre:3 IS NULL [as=nulls_ordering_genre:6]
-      │    │              └── id:1 IS NULL [as=nulls_ordering_id:7]
-      │    └── aggregations
-      │         ├── first-agg [as=id:1]
-      │         │    └── id:1
-      │         └── first-agg [as=name:2]
-      │              └── name:2
-      └── projections
-           ├── genre:3 IS NULL [as=nulls_ordering_genre:8]
-           └── id:1 IS NULL [as=nulls_ordering_id:9]
+ ├── columns: id:1!null str:2  [hidden: t2.id:6!null nulls_ordering_id:11!null]
+ ├── ordering: -11,-6
+ └── distinct-on
+      ├── columns: t1.id:1!null str:2 t2.id:6!null nulls_ordering_id:11!null
+      ├── grouping columns: t2.id:6!null nulls_ordering_id:11!null
+      ├── project
+      │    ├── columns: nulls_ordering_id:11!null t1.id:1!null str:2 t2.id:6!null
+      │    ├── inner-join (hash)
+      │    │    ├── columns: t1.id:1!null str:2 t1.rowid:3!null t1.crdb_internal_mvcc_timestamp:4 t1.tableoid:5 t2.id:6!null num:7 t2.rowid:8!null t2.crdb_internal_mvcc_timestamp:9 t2.tableoid:10
+      │    │    ├── scan t1
+      │    │    │    └── columns: t1.id:1 str:2 t1.rowid:3!null t1.crdb_internal_mvcc_timestamp:4 t1.tableoid:5
+      │    │    ├── scan t2
+      │    │    │    └── columns: t2.id:6 num:7 t2.rowid:8!null t2.crdb_internal_mvcc_timestamp:9 t2.tableoid:10
+      │    │    └── filters
+      │    │         └── t1.id:1 = t2.id:6
+      │    └── projections
+      │         └── t2.id:6 IS NULL [as=nulls_ordering_id:11]
+      └── aggregations
+           ├── first-agg [as=t1.id:1]
+           │    └── t1.id:1
+           └── first-agg [as=str:2]
+                └── str:2
+
+# Because of the hack we use to allow ORDER BY NULLS FIRST / LAST with DISTINCT
+# ON, we also support this query. Postgres doesn't support it, but it doesn't
+# seem like a problem that we do.
+build
+SELECT DISTINCT ON (id) * FROM t1 ORDER BY id IS NULL, id
+----
+sort
+ ├── columns: id:1 str:2  [hidden: column6:6!null]
+ ├── ordering: +6,+1
+ └── distinct-on
+      ├── columns: id:1 str:2 column6:6!null
+      ├── grouping columns: id:1 column6:6!null
+      ├── project
+      │    ├── columns: column6:6!null id:1 str:2
+      │    ├── scan t1
+      │    │    └── columns: id:1 str:2 rowid:3!null crdb_internal_mvcc_timestamp:4 tableoid:5
+      │    └── projections
+      │         └── id:1 IS NULL [as=column6:6]
+      └── aggregations
+           └── first-agg [as=str:2]
+                └── str:2