sql: rewrite setNeededColumns() and complete the feature

This patch rewrites and largely extends the feature introduced in
1993af0: marking unused planNode
columns as "omitted" to save I/O operations.

This is done by:

- removing `setNeededColumns()` from the planNode interface, and
  replacing its implementations by planNodes by a single recursive
  function `setNeededColumns()`. The new function is written so that
  it is guaranteed to traverse every planNode in the tree.

- adding some missing implementations of `setNeededColumns()`, for
  `unionNode`, `distinctNode`, `selectNode`, `sortNode`,
  `createTableNode`. This yields some previously unattainable
  optimizations. (`windowNode` and `groupNode` are still conspicuously
  missing.)

- rewriting the initialization logic for `indexJoinNode` and ensuring
  its `setNeededColumns()` code can be reused multiple times.

- adding a new `initNeededColumns()` method on `planner` that is
  called for every plan constructed, which calls both
  `setNeededColumns()` on the plan and `initNeededColumns()` on all
  the sub-query plans. This ensures that all plans undergo this
  optimization.

- adding a new EXPLAIN option `NOOPTIMIZE` that shows the planNode
  when `initNeededColumns()` has not been called.

pkg/sql/alter_table.go

@@ -402,7 +402,6 @@ func (n *alterTableNode) Ordering() orderingInfo       { return orderingInfo{} }
 func (n *alterTableNode) Values() parser.DTuple        { return parser.DTuple{} }
 func (n *alterTableNode) DebugValues() debugValues     { return debugValues{} }
 func (n *alterTableNode) SetLimitHint(_ int64, _ bool) {}
-func (n *alterTableNode) setNeededColumns(_ []bool)    {}
 func (n *alterTableNode) MarkDebug(mode explainMode)   {}
 
 func applyColumnMutation(

pkg/sql/copy.go

@@ -53,7 +53,6 @@ func (n *copyNode) Columns() ResultColumns     { return n.resultColumns }
 func (*copyNode) Ordering() orderingInfo       { return orderingInfo{} }
 func (*copyNode) Values() parser.DTuple        { return nil }
 func (*copyNode) SetLimitHint(_ int64, _ bool) {}
-func (*copyNode) setNeededColumns(_ []bool)    {}
 func (*copyNode) MarkDebug(_ explainMode)      {}
 func (*copyNode) expandPlan() error            { return nil }
 func (*copyNode) Next() (bool, error)          { return false, nil }

pkg/sql/create.go

@@ -143,7 +143,6 @@ func (n *createDatabaseNode) Ordering() orderingInfo       { return orderingInfo
 func (n *createDatabaseNode) Values() parser.DTuple        { return parser.DTuple{} }
 func (n *createDatabaseNode) DebugValues() debugValues     { return debugValues{} }
 func (n *createDatabaseNode) SetLimitHint(_ int64, _ bool) {}
-func (n *createDatabaseNode) setNeededColumns(_ []bool)    {}
 func (n *createDatabaseNode) MarkDebug(mode explainMode)   {}
 
 type createIndexNode struct {
@@ -259,7 +258,6 @@ func (n *createIndexNode) Ordering() orderingInfo       { return orderingInfo{}
 func (n *createIndexNode) Values() parser.DTuple        { return parser.DTuple{} }
 func (n *createIndexNode) DebugValues() debugValues     { return debugValues{} }
 func (n *createIndexNode) SetLimitHint(_ int64, _ bool) {}
-func (n *createIndexNode) setNeededColumns(_ []bool)    {}
 func (n *createIndexNode) MarkDebug(mode explainMode)   {}
 
 type createUserNode struct {
@@ -348,7 +346,6 @@ func (n *createUserNode) Ordering() orderingInfo       { return orderingInfo{} }
 func (n *createUserNode) Values() parser.DTuple        { return parser.DTuple{} }
 func (n *createUserNode) DebugValues() debugValues     { return debugValues{} }
 func (n *createUserNode) SetLimitHint(_ int64, _ bool) {}
-func (n *createUserNode) setNeededColumns(_ []bool)    {}
 func (n *createUserNode) MarkDebug(mode explainMode)   {}
 
 type createViewNode struct {
@@ -498,7 +495,6 @@ func (n *createViewNode) Columns() ResultColumns       { return make(ResultColum
 func (n *createViewNode) Ordering() orderingInfo       { return orderingInfo{} }
 func (n *createViewNode) Values() parser.DTuple        { return parser.DTuple{} }
 func (n *createViewNode) DebugValues() debugValues     { return debugValues{} }
-func (n *createViewNode) setNeededColumns(_ []bool)    {}
 func (n *createViewNode) SetLimitHint(_ int64, _ bool) {}
 func (n *createViewNode) MarkDebug(mode explainMode)   {}
 
@@ -729,7 +725,6 @@ func (n *createTableNode) Ordering() orderingInfo       { return orderingInfo{}
 func (n *createTableNode) Values() parser.DTuple        { return parser.DTuple{} }
 func (n *createTableNode) DebugValues() debugValues     { return debugValues{} }
 func (n *createTableNode) SetLimitHint(_ int64, _ bool) {}
-func (n *createTableNode) setNeededColumns(_ []bool)    {}
 func (n *createTableNode) MarkDebug(mode explainMode)   {}
 
 type indexMatch bool

pkg/sql/delayed.go

@@ -32,7 +32,6 @@ type delayedNode struct {
 type nodeConstructor func(p *planner) (planNode, error)
 
 func (d *delayedNode) SetLimitHint(_ int64, _ bool) {}
-func (d *delayedNode) setNeededColumns(_ []bool)    {}
 
 func (d *delayedNode) expandPlan() error {
 	v, err := d.constructor(d.p)

pkg/sql/delete.go

@@ -241,4 +241,3 @@ func (d *deleteNode) DebugValues() debugValues {
 func (d *deleteNode) Ordering() orderingInfo { return orderingInfo{} }
 
 func (d *deleteNode) SetLimitHint(numRows int64, soft bool) {}
-func (d *deleteNode) setNeededColumns(_ []bool)             {}

pkg/sql/distinct.go

@@ -244,8 +244,6 @@ func (n *distinctNode) SetLimitHint(numRows int64, soft bool) {
 	n.plan.SetLimitHint(numRows, true)
 }
 
-func (*distinctNode) setNeededColumns(_ []bool) {}
-
 func (n *distinctNode) Close() {
 	n.plan.Close()
 	n.prefixSeen = nil

pkg/sql/drop.go

@@ -217,7 +217,6 @@ func (n *dropDatabaseNode) Ordering() orderingInfo       { return orderingInfo{}
 func (n *dropDatabaseNode) Values() parser.DTuple        { return parser.DTuple{} }
 func (n *dropDatabaseNode) DebugValues() debugValues     { return debugValues{} }
 func (n *dropDatabaseNode) SetLimitHint(_ int64, _ bool) {}
-func (n *dropDatabaseNode) setNeededColumns(_ []bool)    {}
 func (n *dropDatabaseNode) MarkDebug(mode explainMode)   {}
 
 type dropIndexNode struct {
@@ -394,7 +393,6 @@ func (n *dropIndexNode) Ordering() orderingInfo       { return orderingInfo{} }
 func (n *dropIndexNode) Values() parser.DTuple        { return parser.DTuple{} }
 func (n *dropIndexNode) DebugValues() debugValues     { return debugValues{} }
 func (n *dropIndexNode) SetLimitHint(_ int64, _ bool) {}
-func (n *dropIndexNode) setNeededColumns(_ []bool)    {}
 func (n *dropIndexNode) MarkDebug(mode explainMode)   {}
 
 type dropViewNode struct {
@@ -507,7 +505,6 @@ func (n *dropViewNode) Ordering() orderingInfo       { return orderingInfo{} }
 func (n *dropViewNode) Values() parser.DTuple        { return parser.DTuple{} }
 func (n *dropViewNode) DebugValues() debugValues     { return debugValues{} }
 func (n *dropViewNode) SetLimitHint(_ int64, _ bool) {}
-func (n *dropViewNode) setNeededColumns(_ []bool)    {}
 func (n *dropViewNode) MarkDebug(mode explainMode)   {}
 
 type dropTableNode struct {
@@ -732,7 +729,6 @@ func (n *dropTableNode) Ordering() orderingInfo       { return orderingInfo{} }
 func (n *dropTableNode) Values() parser.DTuple        { return parser.DTuple{} }
 func (n *dropTableNode) DebugValues() debugValues     { return debugValues{} }
 func (n *dropTableNode) SetLimitHint(_ int64, _ bool) {}
-func (n *dropTableNode) setNeededColumns(_ []bool)    {}
 func (n *dropTableNode) MarkDebug(mode explainMode)   {}
 
 // dropTableOrViewPrepare/dropTableImpl is used to drop a single table by

pkg/sql/empty.go

@@ -32,7 +32,6 @@ func (*emptyNode) Ordering() orderingInfo       { return orderingInfo{} }
 func (*emptyNode) Values() parser.DTuple        { return nil }
 func (*emptyNode) Start() error                 { return nil }
 func (*emptyNode) SetLimitHint(_ int64, _ bool) {}
-func (*emptyNode) setNeededColumns(_ []bool)    {}
 func (*emptyNode) MarkDebug(_ explainMode)      {}
 func (*emptyNode) expandPlan() error            { return nil }
 func (*emptyNode) Close()                       {}

pkg/sql/explain.go

@@ -44,6 +44,7 @@ var explainStrings = []string{"", "debug", "plan", "trace", "types"}
 func (p *planner) Explain(n *parser.Explain, autoCommit bool) (planNode, error) {
 	mode := explainNone
 
+	optimized := true
 	expanded := true
 	normalizeExprs := true
 	explainer := explainer{
@@ -92,6 +93,8 @@ func (p *planner) Explain(n *parser.Explain, autoCommit bool) (planNode, error)
 			expanded = false
 		} else if strings.EqualFold(opt, "NONORMALIZE") {
 			normalizeExprs = false
+		} else if strings.EqualFold(opt, "NOOPTIMIZE") {
+			optimized = false
 		} else {
 			return nil, fmt.Errorf("unsupported EXPLAIN option: %s", opt)
 		}
@@ -130,7 +133,7 @@ func (p *planner) Explain(n *parser.Explain, autoCommit bool) (planNode, error)
 		// We may want to show placeholder types, so ensure no values
 		// are missing.
 		p.semaCtx.Placeholders.FillUnassigned()
-		return p.makeExplainPlanNode(explainer, expanded, plan), nil
+		return p.makeExplainPlanNode(explainer, expanded, optimized, plan), nil
 
 	case explainTrace:
 		return p.makeTraceNode(plan, p.txn), nil
@@ -258,4 +261,3 @@ func (n *explainDebugNode) Values() parser.DTuple {
 func (*explainDebugNode) MarkDebug(_ explainMode)      {}
 func (*explainDebugNode) DebugValues() debugValues     { return debugValues{} }
 func (*explainDebugNode) SetLimitHint(_ int64, _ bool) {}
-func (*explainDebugNode) setNeededColumns(_ []bool)    {}

pkg/sql/explain_plan.go

@@ -68,7 +68,9 @@ type explainer struct {
 }
 
 // newExplainPlanNode instantiates a planNode that runs an EXPLAIN query.
-func (p *planner) makeExplainPlanNode(explainer explainer, expanded bool, plan planNode) planNode {
+func (p *planner) makeExplainPlanNode(
+	explainer explainer, expanded, optimized bool, plan planNode,
+) planNode {
 	columns := ResultColumns{
 		// Level is the depth of the node in the tree.
 		{Name: "Level", Typ: parser.TypeInt},
@@ -95,6 +97,7 @@ func (p *planner) makeExplainPlanNode(explainer explainer, expanded bool, plan p
 		p:         p,
 		explainer: explainer,
 		expanded:  expanded,
+		optimized: optimized,
 		plan:      plan,
 		results:   p.newContainerValuesNode(columns, 0),
 	}
@@ -258,9 +261,18 @@ func formatColumns(cols ResultColumns, printTypes bool) string {
 type explainPlanNode struct {
 	p         *planner
 	explainer explainer
-	expanded  bool
-	plan      planNode
-	results   *valuesNode
+
+	// plan is the sub-node being explained.
+	plan planNode
+
+	// results is the container for EXPLAIN's output.
+	results *valuesNode
+
+	// expanded indicates whether to invoke expandPlan() on the sub-node.
+	expanded bool
+
+	// optimized indicates whether to invoke initNeededColumns() on the sub-node.
+	optimized bool
 }
 
 func (e *explainPlanNode) Next() (bool, error)          { return e.results.Next() }
@@ -269,10 +281,14 @@ func (e *explainPlanNode) Ordering() orderingInfo       { return e.results.Order
 func (e *explainPlanNode) Values() parser.DTuple        { return e.results.Values() }
 func (e *explainPlanNode) DebugValues() debugValues     { return debugValues{} }
 func (e *explainPlanNode) SetLimitHint(n int64, s bool) { e.results.SetLimitHint(n, s) }
-func (e *explainPlanNode) setNeededColumns(_ []bool)    {}
 func (e *explainPlanNode) MarkDebug(mode explainMode)   {}
 func (e *explainPlanNode) expandPlan() error {
+
 	if e.expanded {
+		if e.optimized {
+			e.p.initNeededColumns(e.plan, allColumns(e.plan), false)
+		}
+
 		if err := e.plan.expandPlan(); err != nil {
 			return err
 		}
@@ -282,6 +298,11 @@ func (e *explainPlanNode) expandPlan() error {
 		// interested in.
 		e.plan.SetLimitHint(math.MaxInt64, true)
 	}
+
+	if e.optimized {
+		e.p.initNeededColumns(e.plan, allColumns(e.plan), true)
+	}
+
 	return nil
 }
 

pkg/sql/expr_filter.go

@@ -108,6 +108,9 @@ func exprCheckVars(expr parser.Expr, conv varConvertFunc) bool {
 // Convert the variables in the given expression; the expression must only contain
 // variables known to the conversion function (exprCheckVars should be used first).
 func exprConvertVars(expr parser.TypedExpr, conv varConvertFunc) parser.TypedExpr {
+	if expr == nil {
+		return expr
+	}
 	v := varConvertVisitor{justCheck: false, conv: conv}
 	ret, _ := parser.WalkExpr(&v, expr)
 	return ret.(parser.TypedExpr)

pkg/sql/generator.go

@@ -136,4 +136,3 @@ func (n *valueGenerator) Values() parser.DTuple        { return n.gen.Values() }
 func (n *valueGenerator) MarkDebug(_ explainMode)      {}
 func (n *valueGenerator) Columns() ResultColumns       { return n.columns }
 func (n *valueGenerator) SetLimitHint(_ int64, _ bool) {}
-func (n *valueGenerator) setNeededColumns(_ []bool)    {}

pkg/sql/group.go

@@ -414,7 +414,6 @@ func (n *groupNode) computeAggregates() error {
 }
 
 func (*groupNode) SetLimitHint(_ int64, _ bool) {}
-func (*groupNode) setNeededColumns(_ []bool)    {}
 
 func (n *groupNode) Close() {
 	n.plan.Close()

pkg/sql/index_join.go

@@ -37,6 +37,7 @@ type indexJoinNode struct {
 	table            *scanNode
 	primaryKeyPrefix roachpb.Key
 	colIDtoRowIndex  map[sqlbase.ColumnID]int
+	valNeededIndex   []bool
 	explain          explainMode
 	debugVals        debugValues
 }
@@ -77,49 +78,48 @@ func (p *planner) makeIndexJoin(
 		colIDtoRowIndex[colID] = idx
 	}
 
-	// Transfer needed columns set to the table node.
-	table.setNeededColumns(origScan.valNeededForCol)
-
 	// For the index node, we need values for columns that are part of the index.
 	// TODO(radu): we could reduce this further - we only need the PK columns plus
 	// whatever filters may be used by the filter below.
 	valNeededIndex := make([]bool, len(origScan.valNeededForCol))
 	for _, idx := range colIDtoRowIndex {
 		valNeededIndex[idx] = true
 	}
-	indexScan.setNeededColumns(valNeededIndex)
 
 	if origScan.filter != nil {
-		// Transfer the filter to the table node. We must first convert the
-		// IndexedVars associated with indexNode.
-		convFunc := func(expr parser.VariableExpr) (ok bool, newExpr parser.VariableExpr) {
-			iv := expr.(*parser.IndexedVar)
-			return true, table.filterVars.IndexedVar(iv.Idx)
-		}
-		table.filter = exprConvertVars(origScan.filter, convFunc)
+		// Now we split the filter by extracting the part that can be
+		// evaluated using just the index columns.
 
-		// Now we split the filter by extracting the part that can be evaluated using just the index
-		// columns.
+		// Since we are re-populating the IndexedVars, reset the helper
+		// first so that the new vars are properly accounted for.
+		indexScan.filterVars.Reset()
 		splitFunc := func(expr parser.VariableExpr) (ok bool, newExpr parser.VariableExpr) {
 			colIdx := expr.(*parser.IndexedVar).Idx
-			if !indexScan.valNeededForCol[colIdx] {
+			if !valNeededIndex[colIdx] {
 				return false, nil
 			}
 			return true, indexScan.filterVars.IndexedVar(colIdx)
 		}
-		indexScan.filter, table.filter = splitFilter(table.filter, splitFunc)
+		indexScan.filter, table.filter = splitFilter(origScan.filter, splitFunc)
 	}
 
+	// Ensure that the indexed vars are transferred to the scanNodes fully.
+	table.filterVars.Reset()
+	table.filter = table.filterVars.Rebind(table.filter)
+
 	indexScan.initOrdering(exactPrefix)
 
 	primaryKeyPrefix := roachpb.Key(sqlbase.MakeIndexKeyPrefix(&table.desc, table.index.ID))
 
-	return &indexJoinNode{
+	node := &indexJoinNode{
 		index:            indexScan,
 		table:            table,
 		primaryKeyPrefix: primaryKeyPrefix,
 		colIDtoRowIndex:  colIDtoRowIndex,
-	}, indexScan
+		valNeededIndex:   valNeededIndex,
+	}
+
+	return node, indexScan
 }
 
 func (n *indexJoinNode) Columns() ResultColumns {
@@ -243,8 +243,6 @@ func (n *indexJoinNode) SetLimitHint(numRows int64, soft bool) {
 	n.index.SetLimitHint(numRows, soft)
 }
 
-func (*indexJoinNode) setNeededColumns(_ []bool) {}
-
 func (n *indexJoinNode) Close() {
 	n.index.Close()
 	n.table.Close()

pkg/sql/insert.go

@@ -578,4 +578,3 @@ func (n *insertNode) DebugValues() debugValues {
 func (n *insertNode) Ordering() orderingInfo { return orderingInfo{} }
 
 func (n *insertNode) SetLimitHint(numRows int64, soft bool) {}
-func (n *insertNode) setNeededColumns(_ []bool)             {}

pkg/sql/join.go

@@ -273,16 +273,6 @@ func (p *planner) makeJoin(
 // SetLimitHint implements the planNode interface.
 func (n *joinNode) SetLimitHint(numRows int64, soft bool) {}
 
-// setNeededColumns implements the planNode interface.
-func (n *joinNode) setNeededColumns(needed []bool) {
-	leftNeeded, rightNeeded := n.pred.getNeededColumns(needed)
-	n.left.plan.setNeededColumns(leftNeeded)
-	n.right.plan.setNeededColumns(rightNeeded)
-	for i, v := range needed {
-		n.columns[i].omitted = !v
-	}
-}
-
 // expandPlan implements the planNode interface.
 func (n *joinNode) expandPlan() error {
 	if err := n.planner.expandSubqueryPlans(n.pred.filter); err != nil {

pkg/sql/join_predicate.go

@@ -399,8 +399,15 @@ func (p *joinPredicate) eval(
 	return true, nil
 }
 
-// getNeededColumns figures out the columns needed for the two sources.
+// getNeededColumns figures out the columns needed for the two
+// sources.  This takes into account both the equality columns and the
+// predicate expression.
 func (p *joinPredicate) getNeededColumns(neededJoined []bool) ([]bool, []bool) {
+	// Reset the helper and rebind the variable to detect which columns
+	// are effectively needed.
+	p.iVarHelper.Reset()
+	p.filter = p.iVarHelper.Rebind(p.filter)
+
 	// The columns that are part of the expression are always needed.
 	neededJoined = append([]bool(nil), neededJoined...)
 	for i := range neededJoined {

pkg/sql/limit.go

@@ -289,5 +289,3 @@ func (n *limitNode) SetLimitHint(count int64, soft bool) {
 	}
 	n.plan.SetLimitHint(getLimit(hintCount, n.offset), soft)
 }
-
-func (*limitNode) setNeededColumns(_ []bool) {}