sql: implement ConstructZigzagJoin in distsql_spec_exec_factory

Informs cockroachdb#47473

Release note: None

pkg/sql/distsql_physical_planner.go

@@ -37,6 +37,7 @@ import (
 	"github.com/cockroachdb/cockroach/pkg/sql/execinfra"
 	"github.com/cockroachdb/cockroach/pkg/sql/execinfrapb"
 	"github.com/cockroachdb/cockroach/pkg/sql/execstats"
+	"github.com/cockroachdb/cockroach/pkg/sql/opt/exec"
 	"github.com/cockroachdb/cockroach/pkg/sql/pgwire/pgcode"
 	"github.com/cockroachdb/cockroach/pkg/sql/pgwire/pgerror"
 	"github.com/cockroachdb/cockroach/pkg/sql/physicalplan"
@@ -2537,37 +2538,20 @@ func (dsp *DistSQLPlanner) createPlanForInvertedJoin(
 func (dsp *DistSQLPlanner) createPlanForZigzagJoin(
 	planCtx *PlanningCtx, n *zigzagJoinNode,
 ) (plan *PhysicalPlan, err error) {
-	plan = planCtx.NewPhysicalPlan()
 
-	tables := make([]descpb.TableDescriptor, len(n.sides))
-	indexOrdinals := make([]uint32, len(n.sides))
-	cols := make([]execinfrapb.Columns, len(n.sides))
-	numStreamCols := 0
+	pi := zigzagPlanningInfo{onCond: n.onCond, reqOrdering: exec.OutputOrdering(n.reqOrdering)}
+	pi.init(len(n.sides))
+
 	for i, side := range n.sides {
-		tables[i] = *side.scan.desc.TableDesc()
-		indexOrdinals[i], err = getIndexIdx(side.scan.index, side.scan.desc)
+		err = pi.initSide(i, side.scan.desc, side.scan.index, len(side.eqCols))
 		if err != nil {
 			return nil, err
 		}
-
-		cols[i].Columns = make([]uint32, len(side.eqCols))
 		for j, col := range side.eqCols {
-			cols[i].Columns[j] = uint32(col)
+			pi.eqCols[i].Columns[j] = uint32(col)
 		}
-
-		numStreamCols += len(side.scan.desc.PublicColumns())
 	}
 
-	// The zigzag join node only represents inner joins, so hardcode Type to
-	// InnerJoin.
-	zigzagJoinerSpec := execinfrapb.ZigzagJoinerSpec{
-		Tables:        tables,
-		IndexOrdinals: indexOrdinals,
-		EqColumns:     cols,
-		Type:          descpb.InnerJoin,
-	}
-	zigzagJoinerSpec.FixedValues = make([]*execinfrapb.ValuesCoreSpec, len(n.sides))
-
 	// The fixed values are represented as a Values node with one tuple.
 	for i := range n.sides {
 		fixedVals := n.sides[i].fixedVals
@@ -2576,7 +2560,7 @@ func (dsp *DistSQLPlanner) createPlanForZigzagJoin(
 		if err != nil {
 			return nil, err
 		}
-		zigzagJoinerSpec.FixedValues[i] = valuesSpec
+		pi.fixedValues[i] = valuesSpec
 	}
 
 	// The internal schema of the zigzag joiner is:
@@ -2587,12 +2571,7 @@ func (dsp *DistSQLPlanner) createPlanForZigzagJoin(
 	//    <side 1 index columns> ... <side 2 index columns> ...
 	// so the planToStreamColMap has to basically map index ordinals
 	// to table ordinals.
-	post := execinfrapb.PostProcessSpec{Projection: true}
-	numOutCols := len(n.columns)
-
-	post.OutputColumns = make([]uint32, numOutCols)
-	types := make([]*types.T, numOutCols)
-	planToStreamColMap := makePlanToStreamColMap(numOutCols)
+	pi.initOutput(len(n.columns))
 	colOffset := 0
 	i := 0
 
@@ -2605,30 +2584,91 @@ func (dsp *DistSQLPlanner) createPlanForZigzagJoin(
 		// opt/exec/execbuilder/relational_builder.go, similar to lookup joins.
 		for _, col := range side.scan.cols {
 			ord := tableOrdinal(side.scan.desc, col.GetID(), side.scan.colCfg.visibility)
-			post.OutputColumns[i] = uint32(colOffset + ord)
-			types[i] = col.GetType()
-			planToStreamColMap[i] = i
-
+			pi.outputColumns[i] = uint32(colOffset + ord)
+			pi.outputTypes[i] = col.GetType()
+			pi.planToStreamColMap[i] = i
 			i++
 		}
 
 		colOffset += len(side.scan.desc.PublicColumns())
 	}
 
+	return dsp.constructZigzagJoin(planCtx, pi)
+}
+
+type zigzagPlanningInfo struct {
+	tableDescriptors   []descpb.TableDescriptor
+	indexOrdinals      []uint32
+	fixedValues        []*execinfrapb.ValuesCoreSpec
+	eqCols             []execinfrapb.Columns
+	outputColumns      []uint32
+	outputTypes        []*types.T
+	planToStreamColMap []int
+	onCond             tree.TypedExpr
+	reqOrdering        exec.OutputOrdering
+}
+
+func (pi *zigzagPlanningInfo) init(nsides int) {
+	pi.tableDescriptors = make([]descpb.TableDescriptor, nsides)
+	pi.indexOrdinals = make([]uint32, nsides)
+	pi.fixedValues = make([]*execinfrapb.ValuesCoreSpec, nsides)
+	pi.eqCols = make([]execinfrapb.Columns, nsides)
+}
+
+func (pi *zigzagPlanningInfo) initSide(
+	i int, desc catalog.TableDescriptor, index catalog.Index, numEqCols int,
+) error {
+	pi.tableDescriptors[i] = *desc.TableDesc()
+	idxOrd, err := getIndexIdx(index, desc)
+	if err != nil {
+		return err
+	}
+	pi.indexOrdinals[i] = idxOrd
+	pi.eqCols[i].Columns = make([]uint32, numEqCols)
+
+	return nil
+}
+
+func (pi *zigzagPlanningInfo) initOutput(numCols int) {
+	pi.outputColumns = make([]uint32, numCols)
+	pi.outputTypes = make([]*types.T, numCols)
+	pi.planToStreamColMap = makePlanToStreamColMap(numCols)
+}
+
+func (dsp *DistSQLPlanner) constructZigzagJoin(
+	planCtx *PlanningCtx, pi zigzagPlanningInfo,
+) (*PhysicalPlan, error) {
+
+	plan := planCtx.NewPhysicalPlan()
+
+	post := execinfrapb.PostProcessSpec{Projection: true}
+	post.OutputColumns = pi.outputColumns
+
+	// The zigzag join node only represents inner joins, so hardcode Type to
+	// InnerJoin.
+	zigzagJoinerSpec := execinfrapb.ZigzagJoinerSpec{
+		Tables:        pi.tableDescriptors,
+		IndexOrdinals: pi.indexOrdinals,
+		EqColumns:     pi.eqCols,
+		Type:          descpb.InnerJoin,
+		FixedValues:   pi.fixedValues,
+	}
+
 	// Set the ON condition.
-	if n.onCond != nil {
+	if pi.onCond != nil {
 		// Note that the ON condition refers to the *internal* columns of the
 		// processor (before the OutputColumns projection).
-		indexVarMap := makePlanToStreamColMap(len(n.columns))
-		for i := range n.columns {
+		indexVarMap := makePlanToStreamColMap(len(pi.outputColumns))
+		for i := 0; i < len(pi.outputColumns); i++ {
 			indexVarMap[i] = int(post.OutputColumns[i])
 		}
-		zigzagJoinerSpec.OnExpr, err = physicalplan.MakeExpression(
-			n.onCond, planCtx, indexVarMap,
+		onExpr, err := physicalplan.MakeExpression(
+			pi.onCond, planCtx, indexVarMap,
 		)
 		if err != nil {
 			return nil, err
 		}
+		zigzagJoinerSpec.OnExpr = onExpr
 	}
 
 	// Figure out the node where this zigzag joiner goes.
@@ -2642,8 +2682,8 @@ func (dsp *DistSQLPlanner) createPlanForZigzagJoin(
 		Core:   execinfrapb.ProcessorCoreUnion{ZigzagJoiner: &zigzagJoinerSpec},
 	}}
 
-	plan.AddNoInputStage(corePlacement, post, types, execinfrapb.Ordering{})
-	plan.PlanToStreamColMap = planToStreamColMap
+	plan.AddNoInputStage(corePlacement, post, pi.outputTypes, execinfrapb.Ordering{})
+	plan.PlanToStreamColMap = pi.planToStreamColMap
 
 	return plan, nil
 }

pkg/sql/distsql_spec_exec_factory.go

@@ -14,6 +14,7 @@ import (
 	"github.com/cockroachdb/cockroach/pkg/roachpb"
 	"github.com/cockroachdb/cockroach/pkg/sql/catalog/colinfo"
 	"github.com/cockroachdb/cockroach/pkg/sql/catalog/descpb"
+	"github.com/cockroachdb/cockroach/pkg/sql/execinfra"
 	"github.com/cockroachdb/cockroach/pkg/sql/execinfrapb"
 	"github.com/cockroachdb/cockroach/pkg/sql/inverted"
 	"github.com/cockroachdb/cockroach/pkg/sql/opt"
@@ -695,7 +696,99 @@ func (e *distSQLSpecExecFactory) ConstructZigzagJoin(
 	onCond tree.TypedExpr,
 	reqOrdering exec.OutputOrdering,
 ) (exec.Node, error) {
-	return nil, unimplemented.NewWithIssue(47473, "experimental opt-driven distsql planning: zigzag join")
+	planCtx := e.getPlanCtx(cannotDistribute)
+
+	// Maybe someday N will be greater than 2, code it that way.
+	nsides := 2
+	pi := zigzagPlanningInfo{onCond: onCond, reqOrdering: reqOrdering}
+	pi.init(nsides)
+
+	tables := []cat.Table{leftTable, rightTable}
+	indexes := []cat.Index{leftIndex, rightIndex}
+	cols := []exec.TableColumnOrdinalSet{leftCols, rightCols}
+	eqCols := [][]exec.TableColumnOrdinal{leftEqCols, rightEqCols}
+	fixedVals := [][]tree.TypedExpr{leftFixedVals, rightFixedVals}
+	allCols := exec.TableColumnOrdinalSet{}
+
+	numOutputCols := 0
+
+	for s, table := range tables {
+		// RFC: does casting a cat.Table to an optTable ruin the ability to test this code from opttester?  opt_exec_factory
+		// does the same thing, just curious.
+		tdesc := table.(*optTable).desc
+		err := pi.initSide(s, tdesc, indexes[s].(*optIndex).idx, len(eqCols[s]))
+		if err != nil {
+			return nil, err
+		}
+		for j, col := range eqCols[s] {
+			pi.eqCols[s].Columns[j] = uint32(col)
+		}
+		numOutputCols += cols[s].Len()
+		allCols.UnionWith(cols[s])
+	}
+
+	// The fixed values are represented as a Values node with one tuple.
+	for s, vals := range fixedVals {
+		typs := make([]*types.T, len(vals))
+		for t := range typs {
+			col := indexes[s].Column(t)
+			typs[t] = col.DatumType()
+		}
+		valuesSpec, err := e.dsp.createValuesSpecFromTuples(planCtx, [][]tree.TypedExpr{vals}, typs)
+		if err != nil {
+			return nil, err
+		}
+		pi.fixedValues[s] = valuesSpec
+	}
+
+	// The internal schema of the zigzag joiner is:
+	//    <side 1 table columns> ... <side 2 table columns> ...
+	// with only the columns in the specified index populated.
+	//
+	// The schema of the zigzagJoinNode is:
+	//    <side 1 index columns> ... <side 2 index columns> ...
+	// so the planToStreamColMap has to basically map index ordinals
+	// to table ordinals.
+	pi.initOutput(numOutputCols)
+	colOffset := 0
+	i := 0
+
+	// Populate post.OutputColumns (the implicit projection), result types,
+	// and the planToStreamColMap for index columns from all sides.
+	for s, col := range cols {
+		tdesc := tables[s].(*optTable).desc
+		for c, ok := col.Next(0); ok; c, ok = col.Next(c + 1) {
+			column := tdesc.PublicColumns()[c]
+			// RFC: what's going on with visibility?  can't find where its set I think its just
+			// defaulting to the Go default so use that.
+			ord := tableOrdinal(tdesc, column.GetID(), execinfra.ScanVisibilityPublic)
+			pi.outputColumns[i] = uint32(colOffset + ord)
+			pi.outputTypes[i] = column.GetType()
+			pi.planToStreamColMap[i] = i
+			i++
+		}
+
+		colOffset += len(tdesc.PublicColumns())
+	}
+
+	p, err := e.dsp.constructZigzagJoin(planCtx, pi)
+	if err != nil {
+		return nil, err
+	}
+
+	// We don't support different tables currently.
+	table := tables[0]
+	colCfg := makeScanColumnsConfig(table, allCols)
+	// Note that initColsForScan and setting ResultColumns below are equivalent
+	// to what scan.initTable call does in execFactory.ConstructScan.
+	tabDesc := table.(*optTable).desc
+	catCols, err := initColsForScan(tabDesc, colCfg)
+	if err != nil {
+		return nil, err
+	}
+	p.ResultColumns = colinfo.ResultColumnsFromColumns(tabDesc.GetID(), catCols)
+
+	return makePlanMaybePhysical(p, nil /* planNodesToClose */), nil
 }
 
 func (e *distSQLSpecExecFactory) ConstructLimit(

pkg/sql/logictest/testdata/logic_test/experimental_distsql_planning

@@ -118,3 +118,31 @@ SELECT DISTINCT v FROM kv
 2
 NULL
 3
+
+# Check that zigzag join works in always mode.
+statement ok
+RESET experimental_distsql_planning;
+CREATE TABLE a (n INT PRIMARY KEY, a INT, b INT, c STRING, INDEX a_idx(a), INDEX b_idx(b), INDEX bc_idx(b,c));
+INSERT INTO a SELECT a,a,a%3,'foo' FROM generate_series(1,10) AS g(a) ;
+SET enable_zigzag_join = true;
+SET experimental_distsql_planning = always
+
+query III rowsort
+SELECT n,a,b FROM a WHERE a = 5 AND b = 2
+----
+5  5  2
+
+query T
+EXPLAIN SELECT n,a,b FROM a WHERE a = 5 AND b = 2
+----
+distribution: local
+vectorized: true
+·
+• zigzag join
+  pred: (a = 5) AND (b = 2)
+  left table: a@a_idx
+  left columns: (n, a)
+  left fixed values: 1 column
+  right table: a@b_idx
+  right columns: (b)
+  right fixed values: 1 column

pkg/sql/opt/exec/execbuilder/testdata/zigzag_join

@@ -1,4 +1,4 @@
-# LogicTest: local
+# LogicTest: local local-spec-planning
 
 # Make sure that the zigzag join is used in the regression tests for #71093.
 statement ok