opt: fix panic due to incorrect type of ArrayFlatten

[rytaft]

This commit fixes a panic caused by incorrect typing of an ArrayFlatten
expression. If the input to an ArrayFlatten expression is sorted, there
may be more than one output column (although the columns used for sorting
are hidden). If one of these hidden columns is chosen to infer the type
of the expression, the type could be incorrect. This commit fixes the
problem so that only the requested column is chosen for type inference.

Fixes #38867

Release note (bug fix): Fixed a panic due to incorrect type inference
of some ARRAY(...) expressions.

[cockroach-teamcity]

This change is 

[RaduBerinde]

Thanks for the fix!

Reviewable status: complete! 0 of 0 LGTMs obtained (waiting on @justinj, @RaduBerinde, and @rytaft)

---

pkg/sql/opt/norm/custom_funcs.go, line 1630 at r1 (raw file):

// SubqueryRequestedCol returns the requested column from a SubqueryPrivate.
func (c *CustomFuncs) SubqueryRequestedCol(sub *memo.SubqueryPrivate) opt.ColumnID {
	return sub.RequestedCol

I found that this field is not set for SubqueryOp..

[RaduBerinde]

hoistAll seems to have the same problem:

cockroach/pkg/sql/opt/norm/decorrelate.go

Line 805 in fe76c37

colID, _ := subqueryProps.OutputCols.Next(0)

Reviewable status: complete! 0 of 0 LGTMs obtained (waiting on @justinj and @rytaft)

[rytaft]

TFTR!

I tried to set RequestedCol for all types of subqueries, but it turned out to be a fair amount of trouble, especially for ANY and EXISTS. Since ArrayFlatten is the only subquery type where it's necessary, I think we should just make it clear that RequestedCol is only used by ArrayFlatten. I added a comment to that effect in the definition of SubqueryPrivate. I also updated hoistAll to make sure that if the subquery has type ArrayFlatten, we use RequestedCol instead of the first output column.

Reviewable status: complete! 0 of 0 LGTMs obtained (waiting on @justinj and @RaduBerinde)

---

pkg/sql/opt/norm/custom_funcs.go, line 1630 at r1 (raw file):

Previously, RaduBerinde wrote…

I found that this field is not set for SubqueryOp..

Added a comment to say this function should only be used with ArrayFlatten

[RaduBerinde]

Consider adding a SubqueryColumn method to SubqueryExpr and others.

Reviewable status: complete! 0 of 0 LGTMs obtained (waiting on @justinj and @rytaft)

---

pkg/sql/opt/norm/custom_funcs.go, line 1630 at r1 (raw file):

Previously, rytaft (Rebecca Taft) wrote…

Added a comment to say this function should only be used with ArrayFlatten

Maybe we can change it to take an ArrayFlattenExpr.

---

pkg/sql/opt/norm/decorrelate.go, line 805 at r2 (raw file):

		// Replace the Subquery operator with a Variable operator referring to
		// the output column of the hoisted query.
		colID, _ := subqueryProps.OutputCols.Next(0)

Ideally we should assert there is a single column (except for ArrayFlatten).

[rytaft]

Consider adding a SubqueryColumn method to SubqueryExpr and others.

I tried creating a Subquery interface with a SubqueryColumn method implemented by all four expressions, but it felt like a lot of overhead for the single use case of the hoistAll function. Instead I implemented your suggestion from the issue of adding a SingleColumn method to ColSet which panics if there is more than one column present.

Also added your other suggestion from the issue to update the comment for ConstructApplyJoin.

Reviewable status: complete! 0 of 0 LGTMs obtained (waiting on @justinj and @RaduBerinde)

---

pkg/sql/opt/norm/custom_funcs.go, line 1630 at r1 (raw file):

Previously, RaduBerinde wrote…

Maybe we can change it to take an ArrayFlattenExpr.

The problem is the one rule it is used with has ArrayFlatten as the root expression, so we can't pass that in.

---

pkg/sql/opt/norm/decorrelate.go, line 805 at r2 (raw file):

Previously, RaduBerinde wrote…

Ideally we should assert there is a single column (except for ArrayFlatten).

Changed to use SingleColumn method

[RaduBerinde]

LGTM, thanks!

nit: maybe try doing a grep for other instances of Next(0) that can be replaced with the new method.

[rytaft]

nit: maybe try doing a grep for other instances of Next(0) that can be replaced with the new method.

Good idea - fixed a couple of cases

bors r+

Reviewable status: complete! 0 of 0 LGTMs obtained (waiting on @justinj)

[craig]

Build succeeded

• GitHub CI (Cockroach)