Merge #74288

74288: colbuilder: fix window functions planning r=yuzefovich a=yuzefovich The problem with integer widths bites us again - in the vectorized window functions planning we could mistakenly think that the window function returns INT8 when it actually returned INT2 or INT4. This would result in an internal error if there was another window function on top of the "broken" one with the same PARTITION BY clause. Fixes: #74087. Release note (bug fix): Previously, CockroachDB could encounter an internal error when executing queries with multiple window functions and when one of those functions returned INT2 or INT4 type. Co-authored-by: Yahor Yuzefovich <[email protected]>
cockroachdb · Dec 29, 2021 · c40e28f · c40e28f
2 parents 01c0c73 + b2b0089
commit c40e28f
Show file tree

Hide file tree

Showing 2 changed files with 43 additions and 4 deletions.
diff --git a/pkg/sql/colexec/colbuilder/execplan.go b/pkg/sql/colexec/colbuilder/execplan.go
@@ -1274,6 +1274,11 @@ func NewColOperator(
 					PeersColIdx:     peersColIdx,
 				}
 
+				// Some window functions always return an INT result, so we use
+				// it as default and will update whenever necessary. We cannot
+				// use execinfrapb.GetWindowFunctionInfo because that function
+				// doesn't distinguish integers with different widths.
+				returnType := types.Int
 				if wf.Func.WindowFunc != nil {
 					// This is a 'pure' window function (e.g. not an aggregation).
 					windowFn := *wf.Func.WindowFunc
@@ -1296,6 +1301,7 @@ func NewColOperator(
 						)
 						result.Root, err = colexecwindow.NewRelativeRankOperator(
 							windowArgs, windowFn, wf.Ordering.Columns)
+						returnType = types.Float
 					case execinfrapb.WindowerSpec_NTILE:
 						opName := opNamePrefix + "ntile"
 						result.finishBufferedWindowerArgs(
@@ -1311,6 +1317,7 @@ func NewColOperator(
 						)
 						result.Root, err = colexecwindow.NewLagOperator(
 							windowArgs, argIdxs[0], argIdxs[1], argIdxs[2])
+						returnType = typs[argIdxs[0]]
 					case execinfrapb.WindowerSpec_LEAD:
 						opName := opNamePrefix + "lead"
 						result.finishBufferedWindowerArgs(
@@ -1319,6 +1326,7 @@ func NewColOperator(
 						)
 						result.Root, err = colexecwindow.NewLeadOperator(
 							windowArgs, argIdxs[0], argIdxs[1], argIdxs[2])
+						returnType = typs[argIdxs[0]]
 					case execinfrapb.WindowerSpec_FIRST_VALUE:
 						opName := opNamePrefix + "first_value"
 						result.finishBufferedWindowerArgs(
@@ -1327,6 +1335,7 @@ func NewColOperator(
 						)
 						result.Root, err = colexecwindow.NewFirstValueOperator(
 							windowArgs, wf.Frame, &wf.Ordering, argIdxs)
+						returnType = typs[argIdxs[0]]
 					case execinfrapb.WindowerSpec_LAST_VALUE:
 						opName := opNamePrefix + "last_value"
 						result.finishBufferedWindowerArgs(
@@ -1335,6 +1344,7 @@ func NewColOperator(
 						)
 						result.Root, err = colexecwindow.NewLastValueOperator(
 							windowArgs, wf.Frame, &wf.Ordering, argIdxs)
+						returnType = typs[argIdxs[0]]
 					case execinfrapb.WindowerSpec_NTH_VALUE:
 						opName := opNamePrefix + "nth_value"
 						result.finishBufferedWindowerArgs(
@@ -1343,6 +1353,7 @@ func NewColOperator(
 						)
 						result.Root, err = colexecwindow.NewNthValueOperator(
 							windowArgs, wf.Frame, &wf.Ordering, argIdxs)
+						returnType = typs[argIdxs[0]]
 					default:
 						return r, errors.AssertionFailedf("window function %s is not supported", wf.String())
 					}
@@ -1393,6 +1404,7 @@ func NewColOperator(
 						result.Root = colexecwindow.NewWindowAggregatorOperator(
 							windowArgs, aggType, wf.Frame, &wf.Ordering, argIdxs,
 							aggArgs.OutputTypes[0], aggFnsAlloc, toClose)
+						returnType = aggArgs.OutputTypes[0]
 					}
 				} else {
 					colexecerror.InternalError(errors.AssertionFailedf("window function spec is nil"))
@@ -1413,10 +1425,6 @@ func NewColOperator(
 					result.Root = colexecbase.NewSimpleProjectOp(result.Root, int(wf.OutputColIdx+tempColOffset), projection)
 				}
 
-				_, returnType, err := execinfrapb.GetWindowFunctionInfo(wf.Func, argTypes...)
-				if err != nil {
-					return r, err
-				}
 				result.ColumnTypes = appendOneType(result.ColumnTypes, returnType)
 				input = result.Root
 			}

diff --git a/pkg/sql/logictest/testdata/logic_test/window b/pkg/sql/logictest/testdata/logic_test/window
@@ -4301,3 +4301,34 @@ SELECT x, y, first_value(y) OVER (PARTITION BY x ROWS BETWEEN CURRENT ROW AND CU
 2  NotNull  NotNull
 2  NotNull  NotNull
 2  NotNull  NotNull
+
+# Regression test for incorrect type schema setup in the vectorized engine with
+# multiple window functions (#74087).
+statement ok
+CREATE TABLE t74087 AS
+  SELECT
+    g::INT2 AS _int2, g::INT4 AS _int4
+  FROM
+    ROWS FROM (generate_series(1, 5)) AS g;
+
+query IIIIIIIIIRRI rowsort
+SELECT
+  lag(_int2, _int4, _int2) OVER w,
+  lead(_int4, _int2, _int4) OVER w,
+  first_value(_int2) OVER w,
+  last_value(_int4) OVER w,
+  nth_value(_int2, _int4) OVER w,
+  min(_int4) OVER w,
+  row_number() OVER w,
+  rank() OVER w,
+  dense_rank() OVER w,
+  percent_rank() OVER w,
+  cume_dist() OVER w,
+  ntile(_int2) OVER w
+FROM t74087 WINDOW w AS (ORDER BY _int4);
+----
+1  2  1  1  1  1  1  1  1  0     0.2  1
+2  4  1  2  2  1  2  2  2  0.25  0.4  1
+3  3  1  3  3  1  3  3  3  0.5   0.6  1
+4  4  1  4  4  1  4  4  4  0.75  0.8  1
+5  5  1  5  5  1  5  5  5  1     1    1