Repository: vitessio/vitess
Pull request #5845 (2020-02-22 15:19:12)

vrepl: refactor vdiff for readability and add comments
MergeSort as a function that required a VCursor that performed
scatter queries was hard to use from outside of the engine code.

In the new change, MergeSort is a Primitive that directly drives
other sub-Primitives. This greatly simplifies the vdiff code,
removing the need for a mergeSorter and a resultStreamer.


Patch:
diff --git a/go/vt/vtgate/engine/merge_sort.go b/go/vt/vtgate/engine/merge_sort.go
index e55f213b89..b95190cc7e 100644
--- a/go/vt/vtgate/engine/merge_sort.go
+++ b/go/vt/vtgate/engine/merge_sort.go
@@ -14,8 +14,6 @@ See the License for the specific language governing permissions and
 limitations under the License.
 */
 
-// This file has the logic for performing merge-sorts of scatter queries.
-
 package engine
 
 import (
@@ -27,24 +25,62 @@ import (
 	"vitess.io/vitess/go/sqltypes"
 
 	querypb "vitess.io/vitess/go/vt/proto/query"
-	"vitess.io/vitess/go/vt/srvtopo"
+	vtrpcpb "vitess.io/vitess/go/vt/proto/vtrpc"
+	"vitess.io/vitess/go/vt/vterrors"
 )
 
-// MergeSort performs a merge-sort of rows returned by a streaming scatter query.
-// Each shard of the scatter query is treated as a stream. One row from each stream
-// is added to the merge-sorter heap. Every time a value is pulled out of the heap,
+// StreamExecutor is a subset of Primitive that MergeSort
+// requires its inputs to satisfy.
+type StreamExecutor interface {
+	StreamExecute(vcursor VCursor, bindVars map[string]*querypb.BindVariable, wantields bool, callback func(*sqltypes.Result) error) error
+}
+
+var _ Primitive = (*MergeSort)(nil)
+
+// MergeSort performs a merge-sort of rows returned by each Input. This should
+// only be used for StreamExecute. One row from each stream is added to the
+// merge-sorter heap. Every time a value is pulled out of the heap,
 // a new value is added to it from the stream that was the source of the value that
 // was pulled out. Since the input streams are sorted the same way that the heap is
 // sorted, this guarantees that the merged stream will also be sorted the same way.
-func MergeSort(vcursor VCursor, query string, orderBy []OrderbyParams, rss []*srvtopo.ResolvedShard, bvs []map[string]*querypb.BindVariable, callback func(*sqltypes.Result) error) error {
+// MergeSort only supports the StreamExecute function of a Primitive. So, it cannot
+// be used like other Primitives in VTGate. However, it satisfies the Primitive API
+// so that vdiff can use it. In that situation, only StreamExecute is used.
+type MergeSort struct {
+	Primitives []StreamExecutor
+	OrderBy    []OrderbyParams
+	noInputs
+}
+
+// RouteType satisfies Primitive.
+func (ms *MergeSort) RouteType() string { return "MergeSort" }
+
+// GetKeyspaceName satisfies Primitive.
+func (ms *MergeSort) GetKeyspaceName() string { return "" }
+
+// GetTableName satisfies Primitive.
+func (ms *MergeSort) GetTableName() string { return "" }
+
+// Execute is not supported.
+func (ms *MergeSort) Execute(vcursor VCursor, bindVars map[string]*querypb.BindVariable, wantfields bool) (*sqltypes.Result, error) {
+	return nil, vterrors.Errorf(vtrpcpb.Code_UNIMPLEMENTED, "Execute is not supported")
+}
+
+// GetFields is not supported.
+func (ms *MergeSort) GetFields(vcursor VCursor, bindVars map[string]*querypb.BindVariable) (*sqltypes.Result, error) {
+	return nil, vterrors.Errorf(vtrpcpb.Code_UNIMPLEMENTED, "GetFields is not supported")
+}
+
+// StreamExecute performs a streaming exec.
+func (ms *MergeSort) StreamExecute(vcursor VCursor, bindVars map[string]*querypb.BindVariable, wantfields bool, callback func(*sqltypes.Result) error) error {
 	ctx, cancel := context.WithCancel(vcursor.Context())
 	defer cancel()
 
-	handles := make([]*streamHandle, len(rss))
-	id := 0
-	for i, rs := range rss {
-		handles[id] = runOneStream(ctx, vcursor, query, rs, bvs[i])
-		id++
+	handles := make([]*streamHandle, len(ms.Primitives))
+	for i, input := range ms.Primitives {
+		handles[i] = runOneStream(vcursor, input, bindVars, wantfields)
+		// Need fields only from first handle, if wantfields was true.
+		wantfields = false
 	}
 
 	// Fetch field info from just one stream.
@@ -59,7 +95,7 @@ func MergeSort(vcursor VCursor, query string, orderBy []OrderbyParams, rss []*sr
 
 	sh := &scatterHeap{
 		rows:    make([]streamRow, 0, len(handles)),
-		orderBy: orderBy,
+		orderBy: ms.OrderBy,
 	}
 
 	// Prime the heap. One element must be pulled from
@@ -133,20 +169,21 @@ type streamHandle struct {
 }
 
 // runOnestream starts a streaming query on one shard, and returns a streamHandle for it.
-func runOneStream(ctx context.Context, vcursor VCursor, query string, rs *srvtopo.ResolvedShard, vars map[string]*querypb.BindVariable) *streamHandle {
+func runOneStream(vcursor VCursor, input StreamExecutor, bindVars map[string]*querypb.BindVariable, wantfields bool) *streamHandle {
 	handle := &streamHandle{
 		fields: make(chan []*querypb.Field, 1),
 		row:    make(chan []sqltypes.Value, 10),
 	}
+	ctx := vcursor.Context()
 
 	go func() {
 		defer close(handle.fields)
 		defer close(handle.row)
 
-		handle.err = vcursor.StreamExecuteMulti(
-			query,
-			[]*srvtopo.ResolvedShard{rs},
-			[]map[string]*querypb.BindVariable{vars},
+		handle.err = input.StreamExecute(
+			vcursor,
+			bindVars,
+			wantfields,
 			func(qr *sqltypes.Result) error {
 				if len(qr.Fields) != 0 {
 					select {
diff --git a/go/vt/vtgate/engine/route.go b/go/vt/vtgate/engine/route.go
index 9d5b45471b..4a849293e3 100644
--- a/go/vt/vtgate/engine/route.go
+++ b/go/vt/vtgate/engine/route.go
@@ -334,7 +334,20 @@ func (route *Route) StreamExecute(vcursor VCursor, bindVars map[string]*querypb.
 		})
 	}
 
-	return MergeSort(vcursor, route.Query, route.OrderBy, rss, bvs, func(qr *sqltypes.Result) error {
+	// There is an order by. We have to merge-sort.
+	prims := make([]StreamExecutor, 0, len(rss))
+	for i, rs := range rss {
+		prims = append(prims, &shardRoute{
+			query: route.Query,
+			rs:    rs,
+			bv:    bvs[i],
+		})
+	}
+	ms := MergeSort{
+		Primitives: prims,
+		OrderBy:    route.OrderBy,
+	}
+	return ms.StreamExecute(vcursor, bindVars, wantfields, func(qr *sqltypes.Result) error {
 		return callback(qr.Truncate(route.TruncateColumnCount))
 	})
 }
diff --git a/go/vt/vtgate/engine/shard_route.go b/go/vt/vtgate/engine/shard_route.go
new file mode 100644
index 0000000000..1d695489c0
--- /dev/null
+++ b/go/vt/vtgate/engine/shard_route.go
@@ -0,0 +1,38 @@
+/*
+Copyright 2020 The Vitess Authors.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+*/
+
+package engine
+
+import (
+	"vitess.io/vitess/go/sqltypes"
+	querypb "vitess.io/vitess/go/vt/proto/query"
+	"vitess.io/vitess/go/vt/srvtopo"
+)
+
+var _ StreamExecutor = (*shardRoute)(nil)
+
+// shardRoute is an internal primitive used by Route
+// for performing merge sorts.
+type shardRoute struct {
+	query string
+	rs    *srvtopo.ResolvedShard
+	bv    map[string]*querypb.BindVariable
+}
+
+// StreamExecute performs a streaming exec.
+func (sr *shardRoute) StreamExecute(vcursor VCursor, bindVars map[string]*querypb.BindVariable, wantfields bool, callback func(*sqltypes.Result) error) error {
+	return vcursor.StreamExecuteMulti(sr.query, []*srvtopo.ResolvedShard{sr.rs}, []map[string]*querypb.BindVariable{sr.bv}, callback)
+}
diff --git a/go/vt/wrangler/vdiff.go b/go/vt/wrangler/vdiff.go
index 7b85ed2c65..a9f2a54a2b 100644
--- a/go/vt/wrangler/vdiff.go
+++ b/go/vt/wrangler/vdiff.go
@@ -35,7 +35,6 @@ import (
 	tabletmanagerdatapb "vitess.io/vitess/go/vt/proto/tabletmanagerdata"
 	topodatapb "vitess.io/vitess/go/vt/proto/topodata"
 	"vitess.io/vitess/go/vt/sqlparser"
-	"vitess.io/vitess/go/vt/srvtopo"
 	"vitess.io/vitess/go/vt/topo"
 	"vitess.io/vitess/go/vt/topo/topoproto"
 	"vitess.io/vitess/go/vt/vterrors"
@@ -53,27 +52,51 @@ type DiffReport struct {
 	ExtraRowsTarget int
 }
 
+// vdiff contains the metadata for performing vdiff for one workflow.
 type vdiff struct {
 	mi             *migrater
 	sourceCell     string
 	targetCell     string
 	tabletTypesStr string
-	differs        map[string]*tableDiffer
-	sources        map[string]*dfParams
-	targets        map[string]*dfParams
+
+	// differs uses the target table name for its key.
+	differs map[string]*tableDiffer
+
+	// The key for sources and targets is the shard name.
+	// The source and target keyspaces are pulled from mi.
+	sources map[string]*shardStreamer
+	targets map[string]*shardStreamer
 }
 
+// tableDiffer performs a diff for one table in the workflow.
 type tableDiffer struct {
-	targetTable      string
+	targetTable string
+	// sourceExpression and targetExpression are select queries.
 	sourceExpression string
 	targetExpression string
-	compareCols      []int
-	comparePKs       []int
-	sourcePrimitive  engine.Primitive
-	targetPrimitive  engine.Primitive
+
+	// compareCols is the list of non-pk columns to compare.
+	// If the value is -1, it's a pk column and should not be
+	// compared.
+	compareCols []int
+	// comparePKs is the list of pk columns to compare. The logic
+	// for comparing pk columns is different from compareCols
+	comparePKs []int
+
+	// source Primitive and targetPrimitive are used for streaming
+	// results from source and target.
+	sourcePrimitive engine.Primitive
+	targetPrimitive engine.Primitive
 }
 
-type dfParams struct {
+// shardStreamer streams rows from one shard. This works for
+// the source as well as the target.
+// shardStreamer satisfies engine.StreamExecutor, and can be
+// added to Primitives of engine.MergeSort.
+// shardStreamer is a member of vdiff, and gets reused by
+// every tableDiffer. A new result channel gets instantiated
+// for every tableDiffer iteration.
+type shardStreamer struct {
 	master           *topo.TabletInfo
 	tablet           *topodatapb.Tablet
 	position         mysql.Position
@@ -85,6 +108,7 @@ type dfParams struct {
 // VDiff reports differences between the sources and targets of a vreplication workflow.
 func (wr *Wrangler) VDiff(ctx context.Context, targetKeyspace, workflow, sourceCell, targetCell, tabletTypesStr string,
 	filteredReplicationWaitTime, healthcheckTopologyRefresh, healthcheckRetryDelay, healthcheckTimeout time.Duration) (map[string]*DiffReport, error) {
+	// Assign defaults to sourceCell and targetCell if not specified.
 	if sourceCell == "" && targetCell == "" {
 		cells, err := wr.ts.GetCellInfoNames(ctx)
 		if err != nil {
@@ -103,6 +127,8 @@ func (wr *Wrangler) VDiff(ctx context.Context, targetKeyspace, workflow, sourceC
 	if targetCell == "" {
 		targetCell = sourceCell
 	}
+
+	// Reuse migrater code to fetch and validate initial metadata about the workflow.
 	mi, err := wr.buildMigrater(ctx, targetKeyspace, workflow)
 	if err != nil {
 		wr.Logger().Errorf("buildMigrater: %v", err)
@@ -112,22 +138,24 @@ func (wr *Wrangler) VDiff(ctx context.Context, targetKeyspace, workflow, sourceC
 		mi.wr.Logger().Errorf("validate: %v", err)
 		return nil, err
 	}
+
+	// Initialize vdiff.
 	df := &vdiff{
 		mi:             mi,
 		sourceCell:     sourceCell,
 		targetCell:     targetCell,
 		tabletTypesStr: tabletTypesStr,
-		sources:        make(map[string]*dfParams),
-		targets:        make(map[string]*dfParams),
+		sources:        make(map[string]*shardStreamer),
+		targets:        make(map[string]*shardStreamer),
 	}
 	for shard, source := range mi.sources {
-		df.sources[shard] = &dfParams{
+		df.sources[shard] = &shardStreamer{
 			master: source.master,
 		}
 	}
 	var oneTarget *miTarget
 	for shard, target := range mi.targets {
-		df.targets[shard] = &dfParams{
+		df.targets[shard] = &shardStreamer{
 			master: target.master,
 		}
 		oneTarget = target
@@ -141,8 +169,7 @@ func (wr *Wrangler) VDiff(ctx context.Context, targetKeyspace, workflow, sourceC
 	if err != nil {
 		return nil, vterrors.Wrap(err, "GetSchema")
 	}
-	df.differs, err = buildVDiffPlan(ctx, oneFilter, schm)
-	if err != nil {
+	if err = df.buildVDiffPlan(ctx, oneFilter, schm); err != nil {
 		return nil, vterrors.Wrap(err, "buildVDiffPlan")
 	}
 	if err := df.selectTablets(ctx, healthcheckTopologyRefresh, healthcheckRetryDelay, healthcheckTimeout); err != nil {
@@ -153,29 +180,38 @@ func (wr *Wrangler) VDiff(ctx context.Context, targetKeyspace, workflow, sourceC
 			wr.Logger().Errorf("Could not restart workflow %s: %v, please restart it manually", workflow, err)
 		}
 	}(ctx)
+
+	// Perform the diffs.
+	// We need a cancelable context to abort all running streams
+	// if one stream returns an error.
 	ctx, cancel := context.WithCancel(ctx)
 	defer cancel()
+
 	// TODO(sougou): parallelize
 	diffReports := make(map[string]*DiffReport)
 	for table, td := range df.differs {
+		// Stop the targets and record their source positions.
 		if err := df.stopTargets(ctx); err != nil {
 			return nil, vterrors.Wrap(err, "stopTargets")
 		}
-		sourceReader, err := df.startQueryStreams(ctx, df.mi.sourceKeyspace, df.sources, td.sourceExpression, filteredReplicationWaitTime)
-		if err != nil {
+		// Make sure all sources are past the target's positions and start a query stream that records the current source positions.
+		if err := df.startQueryStreams(ctx, df.mi.sourceKeyspace, df.sources, td.sourceExpression, filteredReplicationWaitTime); err != nil {
 			return nil, vterrors.Wrap(err, "startQueryStreams(sources)")
 		}
+		// Fast forward the targets to the newly recorded source positions.
 		if err := df.syncTargets(ctx, filteredReplicationWaitTime); err != nil {
 			return nil, vterrors.Wrap(err, "syncTargets")
 		}
-		targetReader, err := df.startQueryStreams(ctx, df.mi.targetKeyspace, df.targets, td.targetExpression, filteredReplicationWaitTime)
-		if err != nil {
+		// Sources and targets are in sync. Start query streams on the targets.
+		if err := df.startQueryStreams(ctx, df.mi.targetKeyspace, df.targets, td.targetExpression, filteredReplicationWaitTime); err != nil {
 			return nil, vterrors.Wrap(err, "startQueryStreams(targets)")
 		}
+		// Now that queries are running, target vreplication streams can be restarted.
 		if err := df.restartTargets(ctx); err != nil {
 			return nil, vterrors.Wrap(err, "restartTargets")
 		}
-		dr, err := td.diff(ctx, df.mi.wr, sourceReader, targetReader)
+		// Perform the diff of source and target streams.
+		dr, err := td.diff(ctx, df.mi.wr)
 		if err != nil {
 			return nil, vterrors.Wrap(err, "diff")
 		}
@@ -185,12 +221,13 @@ func (wr *Wrangler) VDiff(ctx context.Context, targetKeyspace, workflow, sourceC
 	return diffReports, nil
 }
 
-func buildVDiffPlan(ctx context.Context, filter *binlogdatapb.Filter, schm *tabletmanagerdatapb.SchemaDefinition) (map[string]*tableDiffer, error) {
-	differs := make(map[string]*tableDiffer)
+// buildVDiffPlan builds all the differs.
+func (df *vdiff) buildVDiffPlan(ctx context.Context, filter *binlogdatapb.Filter, schm *tabletmanagerdatapb.SchemaDefinition) error {
+	df.differs = make(map[string]*tableDiffer)
 	for _, table := range schm.TableDefinitions {
 		rule, err := vreplication.MatchTable(table.Name, filter)
 		if err != nil {
-			return nil, err
+			return err
 		}
 		if rule == nil {
 			continue
@@ -201,15 +238,16 @@ func buildVDiffPlan(ctx context.Context, filter *binlogdatapb.Filter, schm *tabl
 			buf.Myprintf("select * from %v", sqlparser.NewTableIdent(table.Name))
 			query = buf.String()
 		}
-		differs[table.Name], err = buildDifferPlan(table, query)
+		df.differs[table.Name], err = df.buildTablePlan(table, query)
 		if err != nil {
-			return nil, err
+			return err
 		}
 	}
-	return differs, nil
+	return nil
 }
 
-func buildDifferPlan(table *tabletmanagerdatapb.TableDefinition, query string) (*tableDiffer, error) {
+// buildTablePlan builds one tableDiffer.
+func (df *vdiff) buildTablePlan(table *tabletmanagerdatapb.TableDefinition, query string) (*tableDiffer, error) {
 	statement, err := sqlparser.Parse(query)
 	if err != nil {
 		return nil, err
@@ -223,10 +261,12 @@ func buildDifferPlan(table *tabletmanagerdatapb.TableDefinition, query string) (
 	}
 	sourceSelect := &sqlparser.Select{}
 	targetSelect := &sqlparser.Select{}
+	// aggregates contains the list if Aggregate functions, if any.
 	var aggregates []engine.AggregateParams
 	for _, selExpr := range sel.SelectExprs {
 		switch selExpr := selExpr.(type) {
 		case *sqlparser.StarExpr:
+			// If it's a '*' expression, expand column list from the schema.
 			for _, fld := range table.Fields {
 				aliased := &sqlparser.AliasedExpr{Expr: &sqlparser.ColName{Name: sqlparser.NewColIdent(fld.Name)}}
 				sourceSelect.SelectExprs = append(sourceSelect.SelectExprs, aliased)
@@ -243,6 +283,7 @@ func buildDifferPlan(table *tabletmanagerdatapb.TableDefinition, query string) (
 					return nil, fmt.Errorf("expression needs an alias: %v", sqlparser.String(selExpr))
 				}
 			}
+			// If the input was "select a as b", then source will use "a" and target will use "b".
 			sourceSelect.SelectExprs = append(sourceSelect.SelectExprs, selExpr)
 			targetSelect.SelectExprs = append(targetSelect.SelectExprs, &sqlparser.AliasedExpr{Expr: targetCol})
 
@@ -265,6 +306,7 @@ func buildDifferPlan(table *tabletmanagerdatapb.TableDefinition, query string) (
 		fields[strings.ToLower(field.Name)] = field.Type
 	}
 
+	// Start with adding all columns for comparison.
 	td.compareCols = make([]int, len(sourceSelect.SelectExprs))
 	for i := range td.compareCols {
 		colname := targetSelect.SelectExprs[i].(*sqlparser.AliasedExpr).Expr.(*sqlparser.ColName).Name.Lowered()
@@ -274,13 +316,17 @@ func buildDifferPlan(table *tabletmanagerdatapb.TableDefinition, query string) (
 		}
 		td.compareCols[i] = i
 		if sqltypes.IsText(typ) {
+			// For text columns, we need to additionally pull their weight string values for lexical comparisons.
 			sourceSelect.SelectExprs = append(sourceSelect.SelectExprs, wrapWeightString(sourceSelect.SelectExprs[i]))
 			targetSelect.SelectExprs = append(targetSelect.SelectExprs, wrapWeightString(targetSelect.SelectExprs[i]))
+			// Update the column number to point at the weight_string column instead.
 			td.compareCols[i] = len(sourceSelect.SelectExprs) - 1
 		}
 	}
 
 	sourceSelect.From = sel.From
+	// The target table name should the one that matched the rule.
+	// It can be different from the source table.
 	targetSelect.From = sqlparser.TableExprs{
 		&sqlparser.AliasedTableExpr{
 			Expr: &sqlparser.TableName{
@@ -311,17 +357,22 @@ func buildDifferPlan(table *tabletmanagerdatapb.TableDefinition, query string) (
 			Direction: sqlparser.AscScr,
 		})
 	}
+	// Remove in_keyrange. It's not understood by mysql.
 	sourceSelect.Where = removeKeyrange(sel.Where)
+	// The source should also perform the group by.
 	sourceSelect.GroupBy = sel.GroupBy
 	sourceSelect.OrderBy = orderby
 
+	// The target should perform the order by, but not the group by.
 	targetSelect.OrderBy = orderby
 
 	td.sourceExpression = sqlparser.String(sourceSelect)
 	td.targetExpression = sqlparser.String(targetSelect)
 
-	td.sourcePrimitive = newMergeSorter(td.comparePKs)
-	td.targetPrimitive = newMergeSorter(td.comparePKs)
+	td.sourcePrimitive = newMergeSorter(df.sources, td.comparePKs)
+	td.targetPrimitive = newMergeSorter(df.targets, td.comparePKs)
+	// If there were aggregate expressions, we have to re-aggregate
+	// the results, which engine.OrderedAggregate can do.
 	if len(aggregates) != 0 {
 		td.sourcePrimitive = &engine.OrderedAggregate{
 			Aggregates: aggregates,
@@ -333,6 +384,23 @@ func buildDifferPlan(table *tabletmanagerdatapb.TableDefinition, query string) (
 	return td, nil
 }
 
+// newMergeSorter creates an engine.MergeSort based on the shard streamers and pk columns.
+func newMergeSorter(participants map[string]*shardStreamer, comparePKs []int) *engine.MergeSort {
+	prims := make([]engine.StreamExecutor, 0, len(participants))
+	for _, participant := range participants {
+		prims = append(prims, participant)
+	}
+	ob := make([]engine.OrderbyParams, 0, len(comparePKs))
+	for _, cpk := range comparePKs {
+		ob = append(ob, engine.OrderbyParams{Col: cpk})
+	}
+	return &engine.MergeSort{
+		Primitives: prims,
+		OrderBy:    ob,
+	}
+}
+
+// selectTablets selects the tablets that will be used for the diff.
 func (df *vdiff) selectTablets(ctx context.Context, healthcheckTopologyRefresh, healthcheckRetryDelay, healthcheckTimeout time.Duration) error {
 	var wg sync.WaitGroup
 	var err1, err2 error
@@ -341,7 +409,7 @@ func (df *vdiff) selectTablets(ctx context.Context, healthcheckTopologyRefresh,
 	wg.Add(1)
 	go func() {
 		defer wg.Done()
-		err1 = df.forAll(df.sources, func(shard string, source *dfParams) error {
+		err1 = df.forAll(df.sources, func(shard string, source *shardStreamer) error {
 			tp, err := discovery.NewTabletPicker(ctx, df.mi.wr.ts, df.sourceCell, df.mi.sourceKeyspace, shard, df.tabletTypesStr, healthcheckTopologyRefresh, healthcheckRetryDelay, healthcheckTimeout)
 			if err != nil {
 				return err
@@ -360,7 +428,7 @@ func (df *vdiff) selectTablets(ctx context.Context, healthcheckTopologyRefresh,
 	wg.Add(1)
 	go func() {
 		defer wg.Done()
-		err2 = df.forAll(df.targets, func(shard string, target *dfParams) error {
+		err2 = df.forAll(df.targets, func(shard string, target *shardStreamer) error {
 			tp, err := discovery.NewTabletPicker(ctx, df.mi.wr.ts, df.targetCell, df.mi.targetKeyspace, shard, df.tabletTypesStr, healthcheckTopologyRefresh, healthcheckRetryDelay, healthcheckTimeout)
 			if err != nil {
 				return err
@@ -383,10 +451,11 @@ func (df *vdiff) selectTablets(ctx context.Context, healthcheckTopologyRefresh,
 	return err2
 }
 
+// stopTargets stops all the targets and records their source positions.
 func (df *vdiff) stopTargets(ctx context.Context) error {
 	var mu sync.Mutex
 
-	err := df.forAll(df.targets, func(shard string, target *dfParams) error {
+	err := df.forAll(df.targets, func(shard string, target *shardStreamer) error {
 		query := fmt.Sprintf("update _vt.vreplication set state='Stopped', message='for vdiff' where db_name=%s and workflow=%s", encodeString(target.master.DbName()), encodeString(df.mi.workflow))
 		_, err := df.mi.wr.tmc.VReplicationExec(ctx, target.master.Tablet, query)
 		if err != nil {
@@ -431,10 +500,13 @@ func (df *vdiff) stopTargets(ctx context.Context) error {
 	return nil
 }
 
-func (df *vdiff) startQueryStreams(ctx context.Context, keyspace string, participants map[string]*dfParams, query string, filteredReplicationWaitTime time.Duration) (*resultReader, error) {
+// starQueryStreams makes sure the sources are past the target's positions, starts the query streams,
+// and records the snapshot position of the query. It creates a result channel which StreamExecute
+// will use to serve rows.
+func (df *vdiff) startQueryStreams(ctx context.Context, keyspace string, participants map[string]*shardStreamer, query string, filteredReplicationWaitTime time.Duration) error {
 	waitCtx, cancel := context.WithTimeout(ctx, filteredReplicationWaitTime)
 	defer cancel()
-	err := df.forAll(participants, func(shard string, participant *dfParams) error {
+	return df.forAll(participants, func(shard string, participant *shardStreamer) error {
 		// Iteration for each participant.
 		if err := df.mi.wr.tmc.WaitForPosition(waitCtx, participant.tablet, mysql.EncodePosition(participant.position)); err != nil {
 			return vterrors.Wrapf(err, "WaitForPosition for tablet %v", topoproto.TabletAliasString(participant.tablet.Alias))
@@ -455,10 +527,6 @@ func (df *vdiff) startQueryStreams(ctx context.Context, keyspace string, partici
 		participant.snapshotPosition = gtid
 		return nil
 	})
-	if err != nil {
-		return nil, err
-	}
-	return newResultReader(ctx, participants), nil
 }
 
 // streamOne is called as a goroutine, and communicates its results through channels.
@@ -466,7 +534,8 @@ func (df *vdiff) startQueryStreams(ctx context.Context, keyspace string, partici
 // Then it streams results to participant.result.
 // Before returning, it sets participant.err, and closes all channels.
 // If any channel is closed, then participant.err can be checked if there was an error.
-func (df *vdiff) streamOne(ctx context.Context, keyspace, shard string, participant *dfParams, query string, gtidch chan string) {
+// The shardStreamer's StreamExecute consumes the result channel.
+func (df *vdiff) streamOne(ctx context.Context, keyspace, shard string, participant *shardStreamer, query string, gtidch chan string) {
 	defer close(participant.result)
 	defer close(gtidch)
 
@@ -485,7 +554,7 @@ func (df *vdiff) streamOne(ctx context.Context, keyspace, shard string, particip
 			TabletType: participant.tablet.Type,
 		}
 		var fields []*querypb.Field
-		err = conn.VStreamResults(ctx, target, query, func(vrs *binlogdatapb.VStreamResultsResponse) error {
+		return conn.VStreamResults(ctx, target, query, func(vrs *binlogdatapb.VStreamResultsResponse) error {
 			if vrs.Fields != nil {
 				fields = vrs.Fields
 				gtidch <- vrs.Gtid
@@ -506,10 +575,11 @@ func (df *vdiff) streamOne(ctx context.Context, keyspace, shard string, particip
 			}
 			return nil
 		})
-		return err
 	}()
 }
 
+// syncTargets fast-forwards the vreplication to the source snapshot positons
+// and waits for the selected tablets to catch up to that point.
 func (df *vdiff) syncTargets(ctx context.Context, filteredReplicationWaitTime time.Duration) error {
 	waitCtx, cancel := context.WithTimeout(ctx, filteredReplicationWaitTime)
 	defer cancel()
@@ -529,7 +599,7 @@ func (df *vdiff) syncTargets(ctx context.Context, filteredReplicationWaitTime ti
 		return err
 	}
 
-	err = df.forAll(df.targets, func(shard string, target *dfParams) error {
+	err = df.forAll(df.targets, func(shard string, target *shardStreamer) error {
 		pos, err := df.mi.wr.tmc.MasterPosition(ctx, target.master.Tablet)
 		if err != nil {
 			return err
@@ -544,20 +614,21 @@ func (df *vdiff) syncTargets(ctx context.Context, filteredReplicationWaitTime ti
 	return err
 }
 
+// restartTargets restarts the stopped target vreplication streams.
 func (df *vdiff) restartTargets(ctx context.Context) error {
-	return df.forAll(df.targets, func(shard string, target *dfParams) error {
+	return df.forAll(df.targets, func(shard string, target *shardStreamer) error {
 		query := fmt.Sprintf("update _vt.vreplication set state='Running', message='', stop_pos='' where db_name=%s and workflow=%s", encodeString(target.master.DbName()), encodeString(df.mi.workflow))
 		_, err := df.mi.wr.tmc.VReplicationExec(ctx, target.master.Tablet, query)
 		return err
 	})
 }
 
-func (df *vdiff) forAll(participants map[string]*dfParams, f func(string, *dfParams) error) error {
+func (df *vdiff) forAll(participants map[string]*shardStreamer, f func(string, *shardStreamer) error) error {
 	var wg sync.WaitGroup
 	allErrors := &concurrency.AllErrorRecorder{}
 	for shard, participant := range participants {
 		wg.Add(1)
-		go func(shard string, participant *dfParams) {
+		go func(shard string, participant *shardStreamer) {
 			defer wg.Done()
 
 			if err := f(shard, participant); err != nil {
@@ -572,6 +643,8 @@ func (df *vdiff) forAll(participants map[string]*dfParams, f func(string, *dfPar
 //-----------------------------------------------------------------
 // primitiveExecutor
 
+// primitiveExecutor starts execution on the top level primitive
+// and provides convenience functions for row-by-row iteration.
 type primitiveExecutor struct {
 	prim     engine.Primitive
 	rows     [][]sqltypes.Value
@@ -579,11 +652,12 @@ type primitiveExecutor struct {
 	err      error
 }
 
-func newPrimitiveExecutor(ctx context.Context, vcursor engine.VCursor, prim engine.Primitive) *primitiveExecutor {
+func newPrimitiveExecutor(ctx context.Context, prim engine.Primitive) *primitiveExecutor {
 	pe := &primitiveExecutor{
 		prim:     prim,
 		resultch: make(chan *sqltypes.Result, 1),
 	}
+	vcursor := &contextVCursor{ctx: ctx}
 	go func() {
 		defer close(pe.resultch)
 		pe.err = pe.prim.StreamExecute(vcursor, make(map[string]*querypb.BindVariable), false, func(qr *sqltypes.Result) error {
@@ -627,80 +701,23 @@ func (pe *primitiveExecutor) drain(ctx context.Context) (int, error) {
 }
 
 //-----------------------------------------------------------------
-// mergeSorter
-
-var _ engine.Primitive = (*mergeSorter)(nil)
-
-// mergeSorter performs a merge-sorted read from the participants.
-type mergeSorter struct {
-	engine.Primitive
-	orderBy []engine.OrderbyParams
-}
-
-func newMergeSorter(comparePKs []int) *mergeSorter {
-	ob := make([]engine.OrderbyParams, 0, len(comparePKs))
-	for _, col := range comparePKs {
-		ob = append(ob, engine.OrderbyParams{Col: col})
-	}
-	return &mergeSorter{
-		orderBy: ob,
-	}
-}
-
-func (ms *mergeSorter) StreamExecute(vcursor engine.VCursor, bindVars map[string]*querypb.BindVariable, wantields bool, callback func(*sqltypes.Result) error) error {
-	rr, ok := vcursor.(*resultReader)
-	if !ok {
-		return fmt.Errorf("internal error: vcursor is not a resultReader: %T", vcursor)
-	}
-	rss := make([]*srvtopo.ResolvedShard, 0, len(rr.participants))
-	bvs := make([]map[string]*querypb.BindVariable, 0, len(rr.participants))
-	for shard := range rr.participants {
-		rss = append(rss, &srvtopo.ResolvedShard{
-			Target: &querypb.Target{
-				Shard: shard,
-			},
-		})
-		bvs = append(bvs, bindVars)
-	}
-	return engine.MergeSort(vcursor, "", ms.orderBy, rss, bvs, callback)
-}
-
-//-----------------------------------------------------------------
-// resultReader
-
-// resultReader acts as a VCursor for the wrapping primitives.
-type resultReader struct {
-	engine.VCursor
-	ctx          context.Context
-	participants map[string]*dfParams
-}
+// shardStreamer
 
-func newResultReader(ctx context.Context, participants map[string]*dfParams) *resultReader {
-	return &resultReader{
-		ctx:          ctx,
-		participants: participants,
-	}
-}
-
-func (rr *resultReader) Context() context.Context {
-	return rr.ctx
-}
-
-func (rr *resultReader) StreamExecuteMulti(query string, rss []*srvtopo.ResolvedShard, bindVars []map[string]*querypb.BindVariable, callback func(reply *sqltypes.Result) error) error {
-	for result := range rr.participants[rss[0].Target.Shard].result {
+func (sm *shardStreamer) StreamExecute(vcursor engine.VCursor, bindVars map[string]*querypb.BindVariable, wantfields bool, callback func(*sqltypes.Result) error) error {
+	for result := range sm.result {
 		if err := callback(result); err != nil {
 			return err
 		}
 	}
-	return rr.participants[rss[0].Target.Shard].err
+	return sm.err
 }
 
 //-----------------------------------------------------------------
 // tableDiffer
 
-func (td *tableDiffer) diff(ctx context.Context, wr *Wrangler, sourceReader, targetReader *resultReader) (*DiffReport, error) {
-	sourceExecutor := newPrimitiveExecutor(ctx, sourceReader, td.sourcePrimitive)
-	targetExecutor := newPrimitiveExecutor(ctx, targetReader, td.targetPrimitive)
+func (td *tableDiffer) diff(ctx context.Context, wr *Wrangler) (*DiffReport, error) {
+	sourceExecutor := newPrimitiveExecutor(ctx, td.sourcePrimitive)
+	targetExecutor := newPrimitiveExecutor(ctx, td.targetPrimitive)
 	dr := &DiffReport{}
 	var sourceRow, targetRow []sqltypes.Value
 	var err error
@@ -807,6 +824,23 @@ func (td *tableDiffer) compare(sourceRow, targetRow []sqltypes.Value, cols []int
 	return 0, nil
 }
 
+//-----------------------------------------------------------------
+// contextVCursor
+
+// contextVCursor satisfies VCursor, but only implements Context().
+// MergeSort only requires Context to be implemented.
+type contextVCursor struct {
+	engine.VCursor
+	ctx context.Context
+}
+
+func (vc *contextVCursor) Context() context.Context {
+	return vc.ctx
+}
+
+//-----------------------------------------------------------------
+// Utility functions
+
 func removeKeyrange(where *sqlparser.Where) *sqlparser.Where {
 	if where == nil {
 		return nil


Test patch:
diff --git a/go/vt/vtgate/engine/merge_sort_test.go b/go/vt/vtgate/engine/merge_sort_test.go
index f8e0e34903..0b35a1e1a6 100644
--- a/go/vt/vtgate/engine/merge_sort_test.go
+++ b/go/vt/vtgate/engine/merge_sort_test.go
@@ -25,48 +25,41 @@ import (
 	"vitess.io/vitess/go/sqltypes"
 
 	querypb "vitess.io/vitess/go/vt/proto/query"
-	"vitess.io/vitess/go/vt/srvtopo"
 )
 
 // TestMergeSortNormal tests the normal flow of a merge
 // sort where all shards return ascending rows.
 func TestMergeSortNormal(t *testing.T) {
 	idColFields := sqltypes.MakeTestFields("id|col", "int32|varchar")
-	vc := &streamVCursor{
-		shardResults: map[string]*shardResult{
-			"0": {results: sqltypes.MakeTestStreamingResults(idColFields,
-				"1|a",
-				"7|g",
-			)},
-			"1": {results: sqltypes.MakeTestStreamingResults(idColFields,
-				"2|b",
-				"---",
-				"3|c",
-			)},
-			"2": {results: sqltypes.MakeTestStreamingResults(idColFields,
-				"4|d",
-				"6|f",
-			)},
-			"3": {results: sqltypes.MakeTestStreamingResults(idColFields,
-				"4|d",
-				"---",
-				"8|h",
-			)},
-		},
-	}
+	shardResults := []*shardResult{{
+		results: sqltypes.MakeTestStreamingResults(idColFields,
+			"1|a",
+			"7|g",
+		),
+	}, {
+		results: sqltypes.MakeTestStreamingResults(idColFields,
+			"2|b",
+			"---",
+			"3|c",
+		),
+	}, {
+		results: sqltypes.MakeTestStreamingResults(idColFields,
+			"4|d",
+			"6|f",
+		),
+	}, {
+		results: sqltypes.MakeTestStreamingResults(idColFields,
+			"4|d",
+			"---",
+			"8|h",
+		),
+	}}
 	orderBy := []OrderbyParams{{
 		Col: 0,
 	}}
-	rss := []*srvtopo.ResolvedShard{
-		{Target: &querypb.Target{Shard: "0"}},
-		{Target: &querypb.Target{Shard: "1"}},
-		{Target: &querypb.Target{Shard: "2"}},
-		{Target: &querypb.Target{Shard: "3"}},
-	}
-	bvs := []map[string]*querypb.BindVariable{nil, nil, nil, nil}
 
 	var results []*sqltypes.Result
-	err := MergeSort(vc, "", orderBy, rss, bvs, func(qr *sqltypes.Result) error {
+	err := testMergeSort(shardResults, orderBy, func(qr *sqltypes.Result) error {
 		results = append(results, qr)
 		return nil
 	})
@@ -99,42 +92,36 @@ func TestMergeSortNormal(t *testing.T) {
 // sort where all shards return descending rows.
 func TestMergeSortDescending(t *testing.T) {
 	idColFields := sqltypes.MakeTestFields("id|col", "int32|varchar")
-	vc := &streamVCursor{
-		shardResults: map[string]*shardResult{
-			"0": {results: sqltypes.MakeTestStreamingResults(idColFields,
-				"7|g",
-				"1|a",
-			)},
-			"1": {results: sqltypes.MakeTestStreamingResults(idColFields,
-				"3|c",
-				"---",
-				"2|b",
-			)},
-			"2": {results: sqltypes.MakeTestStreamingResults(idColFields,
-				"6|f",
-				"4|d",
-			)},
-			"3": {results: sqltypes.MakeTestStreamingResults(idColFields,
-				"8|h",
-				"---",
-				"4|d",
-			)},
-		},
-	}
+	shardResults := []*shardResult{{
+		results: sqltypes.MakeTestStreamingResults(idColFields,
+			"7|g",
+			"1|a",
+		),
+	}, {
+		results: sqltypes.MakeTestStreamingResults(idColFields,
+			"3|c",
+			"---",
+			"2|b",
+		),
+	}, {
+		results: sqltypes.MakeTestStreamingResults(idColFields,
+			"6|f",
+			"4|d",
+		),
+	}, {
+		results: sqltypes.MakeTestStreamingResults(idColFields,
+			"8|h",
+			"---",
+			"4|d",
+		),
+	}}
 	orderBy := []OrderbyParams{{
 		Col:  0,
 		Desc: true,
 	}}
-	rss := []*srvtopo.ResolvedShard{
-		{Target: &querypb.Target{Shard: "0"}},
-		{Target: &querypb.Target{Shard: "1"}},
-		{Target: &querypb.Target{Shard: "2"}},
-		{Target: &querypb.Target{Shard: "3"}},
-	}
-	bvs := []map[string]*querypb.BindVariable{nil, nil, nil, nil}
 
 	var results []*sqltypes.Result
-	err := MergeSort(vc, "", orderBy, rss, bvs, func(qr *sqltypes.Result) error {
+	err := testMergeSort(shardResults, orderBy, func(qr *sqltypes.Result) error {
 		results = append(results, qr)
 		return nil
 	})
@@ -165,33 +152,27 @@ func TestMergeSortDescending(t *testing.T) {
 
 func TestMergeSortEmptyResults(t *testing.T) {
 	idColFields := sqltypes.MakeTestFields("id|col", "int32|varchar")
-	vc := &streamVCursor{
-		shardResults: map[string]*shardResult{
-			"0": {results: sqltypes.MakeTestStreamingResults(idColFields,
-				"1|a",
-				"7|g",
-			)},
-			"1": {results: sqltypes.MakeTestStreamingResults(idColFields)},
-			"2": {results: sqltypes.MakeTestStreamingResults(idColFields,
-				"4|d",
-				"6|f",
-			)},
-			"3": {results: sqltypes.MakeTestStreamingResults(idColFields)},
-		},
-	}
+	shardResults := []*shardResult{{
+		results: sqltypes.MakeTestStreamingResults(idColFields,
+			"1|a",
+			"7|g",
+		),
+	}, {
+		results: sqltypes.MakeTestStreamingResults(idColFields),
+	}, {
+		results: sqltypes.MakeTestStreamingResults(idColFields,
+			"4|d",
+			"6|f",
+		),
+	}, {
+		results: sqltypes.MakeTestStreamingResults(idColFields),
+	}}
 	orderBy := []OrderbyParams{{
 		Col: 0,
 	}}
-	rss := []*srvtopo.ResolvedShard{
-		{Target: &querypb.Target{Shard: "0"}},
-		{Target: &querypb.Target{Shard: "1"}},
-		{Target: &querypb.Target{Shard: "2"}},
-		{Target: &querypb.Target{Shard: "3"}},
-	}
-	bvs := []map[string]*querypb.BindVariable{nil, nil, nil, nil}
 
 	var results []*sqltypes.Result
-	err := MergeSort(vc, "", orderBy, rss, bvs, func(qr *sqltypes.Result) error {
+	err := testMergeSort(shardResults, orderBy, func(qr *sqltypes.Result) error {
 		results = append(results, qr)
 		return nil
 	})
@@ -215,22 +196,15 @@ func TestMergeSortEmptyResults(t *testing.T) {
 // TestMergeSortResultFailures tests failures at various
 // stages of result return.
 func TestMergeSortResultFailures(t *testing.T) {
-	vc := &streamVCursor{
-		shardResults: make(map[string]*shardResult),
-	}
 	orderBy := []OrderbyParams{{
 		Col: 0,
 	}}
-	rss := []*srvtopo.ResolvedShard{
-		{Target: &querypb.Target{Shard: "0"}},
-	}
-	bvs := []map[string]*querypb.BindVariable{nil}
 
 	// Test early error.
-	vc.shardResults["0"] = &shardResult{
+	shardResults := []*shardResult{{
 		sendErr: errors.New("early error"),
-	}
-	err := MergeSort(vc, "", orderBy, rss, bvs, func(qr *sqltypes.Result) error { return nil })
+	}}
+	err := testMergeSort(shardResults, orderBy, func(qr *sqltypes.Result) error { return nil })
 	want := "early error"
 	if err == nil || err.Error() != want {
 		t.Errorf("MergeSort(): %v, want %v", err, want)
@@ -238,22 +212,22 @@ func TestMergeSortResultFailures(t *testing.T) {
 
 	// Test fail after fields.
 	idFields := sqltypes.MakeTestFields("id", "int32")
-	vc.shardResults["0"] = &shardResult{
+	shardResults = []*shardResult{{
 		results: sqltypes.MakeTestStreamingResults(idFields),
 		sendErr: errors.New("fail after fields"),
-	}
-	err = MergeSort(vc, "", orderBy, rss, bvs, func(qr *sqltypes.Result) error { return nil })
+	}}
+	err = testMergeSort(shardResults, orderBy, func(qr *sqltypes.Result) error { return nil })
 	want = "fail after fields"
 	if err == nil || err.Error() != want {
 		t.Errorf("MergeSort(): %v, want %v", err, want)
 	}
 
 	// Test fail after first row.
-	vc.shardResults["0"] = &shardResult{
+	shardResults = []*shardResult{{
 		results: sqltypes.MakeTestStreamingResults(idFields, "1"),
 		sendErr: errors.New("fail after first row"),
-	}
-	err = MergeSort(vc, "", orderBy, rss, bvs, func(qr *sqltypes.Result) error { return nil })
+	}}
+	err = testMergeSort(shardResults, orderBy, func(qr *sqltypes.Result) error { return nil })
 	want = "fail after first row"
 	if err == nil || err.Error() != want {
 		t.Errorf("MergeSort(): %v, want %v", err, want)
@@ -264,26 +238,20 @@ func TestMergeSortDataFailures(t *testing.T) {
 	// The first row being bad fails in a different code path than
 	// the case of subsequent rows. So, test the two cases separately.
 	idColFields := sqltypes.MakeTestFields("id|col", "int32|varchar")
-	vc := &streamVCursor{
-		shardResults: map[string]*shardResult{
-			"0": {results: sqltypes.MakeTestStreamingResults(idColFields,
-				"1|a",
-			)},
-			"1": {results: sqltypes.MakeTestStreamingResults(idColFields,
-				"2.1|b",
-			)},
-		},
-	}
+	shardResults := []*shardResult{{
+		results: sqltypes.MakeTestStreamingResults(idColFields,
+			"1|a",
+		),
+	}, {
+		results: sqltypes.MakeTestStreamingResults(idColFields,
+			"2.1|b",
+		),
+	}}
 	orderBy := []OrderbyParams{{
 		Col: 0,
 	}}
-	rss := []*srvtopo.ResolvedShard{
-		{Target: &querypb.Target{Shard: "0"}},
-		{Target: &querypb.Target{Shard: "1"}},
-	}
-	bvs := []map[string]*querypb.BindVariable{nil, nil}
 
-	err := MergeSort(vc, "", orderBy, rss, bvs, func(qr *sqltypes.Result) error { return nil })
+	err := testMergeSort(shardResults, orderBy, func(qr *sqltypes.Result) error { return nil })
 	want := `strconv.ParseInt: parsing "2.1": invalid syntax`
 	if err == nil || err.Error() != want {
 		t.Errorf("MergeSort(): %v, want %v", err, want)
@@ -291,50 +259,48 @@ func TestMergeSortDataFailures(t *testing.T) {
 
 	// Create a new VCursor because the previous MergeSort will still
 	// have lingering goroutines that can cause data race.
-	vc = &streamVCursor{
-		shardResults: map[string]*shardResult{
-			"0": {results: sqltypes.MakeTestStreamingResults(idColFields,
-				"1|a",
-				"1.1|a",
-			)},
-			"1": {results: sqltypes.MakeTestStreamingResults(idColFields,
-				"2|b",
-			)},
-		},
-	}
-	err = MergeSort(vc, "", orderBy, rss, bvs, func(qr *sqltypes.Result) error { return nil })
+	shardResults = []*shardResult{{
+		results: sqltypes.MakeTestStreamingResults(idColFields,
+			"1|a",
+			"1.1|a",
+		),
+	}, {
+		results: sqltypes.MakeTestStreamingResults(idColFields,
+			"2|b",
+		),
+	}}
+	err = testMergeSort(shardResults, orderBy, func(qr *sqltypes.Result) error { return nil })
 	want = `strconv.ParseInt: parsing "1.1": invalid syntax`
 	if err == nil || err.Error() != want {
 		t.Errorf("MergeSort(): %v, want %v", err, want)
 	}
 }
 
-type shardResult struct {
-	results []*sqltypes.Result
-	// sendErr is sent at the end of the stream if it's set.
-	sendErr error
+func testMergeSort(shardResults []*shardResult, orderBy []OrderbyParams, callback func(qr *sqltypes.Result) error) error {
+	prims := make([]StreamExecutor, 0, len(shardResults))
+	for _, sr := range shardResults {
+		prims = append(prims, sr)
+	}
+	ms := MergeSort{
+		Primitives: prims,
+		OrderBy:    orderBy,
+	}
+	return ms.StreamExecute(noopVCursor{}, nil, true, callback)
 }
 
-// streamVCursor fakes a VCursor that supports
-// a single-shard streaming query through StreamExecuteMulti.
-type streamVCursor struct {
-	noopVCursor
+type shardResult struct {
+	// shardRoute helps us avoid redefining the Primitive functions.
+	shardRoute
 
-	shardResults map[string]*shardResult
+	results []*sqltypes.Result
+	sendErr error
 }
 
-// StreamExecuteMulti streams a result from the specified shard.
-// The shard is specified by the only entry in shardVars. At the
-// end of a stream, if sendErr is set, that error is returned.
-func (t *streamVCursor) StreamExecuteMulti(query string, rss []*srvtopo.ResolvedShard, bindVars []map[string]*querypb.BindVariable, callback func(reply *sqltypes.Result) error) error {
-	shard := rss[0].Target.Shard
-	for _, r := range t.shardResults[shard].results {
+func (sr *shardResult) StreamExecute(vcursor VCursor, bindVars map[string]*querypb.BindVariable, wantfields bool, callback func(*sqltypes.Result) error) error {
+	for _, r := range sr.results {
 		if err := callback(r); err != nil {
 			return err
 		}
 	}
-	if t.shardResults[shard].sendErr != nil {
-		return t.shardResults[shard].sendErr
-	}
-	return nil
+	return sr.sendErr
 }
diff --git a/go/vt/wrangler/vdiff_test.go b/go/vt/wrangler/vdiff_test.go
index b2f92131c3..6a04f00056 100644
--- a/go/vt/wrangler/vdiff_test.go
+++ b/go/vt/wrangler/vdiff_test.go
@@ -74,8 +74,8 @@ func TestVDiffPlanSuccess(t *testing.T) {
 			targetExpression: "select c1, c2 from t1 order by c1 asc",
 			compareCols:      []int{-1, 1},
 			comparePKs:       []int{0},
-			sourcePrimitive:  newMergeSorter([]int{0}),
-			targetPrimitive:  newMergeSorter([]int{0}),
+			sourcePrimitive:  newMergeSorter(nil, []int{0}),
+			targetPrimitive:  newMergeSorter(nil, []int{0}),
 		},
 	}, {
 		input: &binlogdatapb.Rule{
@@ -89,8 +89,8 @@ func TestVDiffPlanSuccess(t *testing.T) {
 			targetExpression: "select c1, c2 from t1 order by c1 asc",
 			compareCols:      []int{-1, 1},
 			comparePKs:       []int{0},
-			sourcePrimitive:  newMergeSorter([]int{0}),
-			targetPrimitive:  newMergeSorter([]int{0}),
+			sourcePrimitive:  newMergeSorter(nil, []int{0}),
+			targetPrimitive:  newMergeSorter(nil, []int{0}),
 		},
 	}, {
 		input: &binlogdatapb.Rule{
@@ -104,8 +104,8 @@ func TestVDiffPlanSuccess(t *testing.T) {
 			targetExpression: "select c1, c2 from t1 order by c1 asc",
 			compareCols:      []int{-1, 1},
 			comparePKs:       []int{0},
-			sourcePrimitive:  newMergeSorter([]int{0}),
-			targetPrimitive:  newMergeSorter([]int{0}),
+			sourcePrimitive:  newMergeSorter(nil, []int{0}),
+			targetPrimitive:  newMergeSorter(nil, []int{0}),
 		},
 	}, {
 		input: &binlogdatapb.Rule{
@@ -119,8 +119,8 @@ func TestVDiffPlanSuccess(t *testing.T) {
 			targetExpression: "select c2, c1 from t1 order by c1 asc",
 			compareCols:      []int{0, -1},
 			comparePKs:       []int{1},
-			sourcePrimitive:  newMergeSorter([]int{1}),
-			targetPrimitive:  newMergeSorter([]int{1}),
+			sourcePrimitive:  newMergeSorter(nil, []int{1}),
+			targetPrimitive:  newMergeSorter(nil, []int{1}),
 		},
 	}, {
 		input: &binlogdatapb.Rule{
@@ -134,8 +134,8 @@ func TestVDiffPlanSuccess(t *testing.T) {
 			targetExpression: "select c1, c2 from t1 order by c1 asc",
 			compareCols:      []int{-1, 1},
 			comparePKs:       []int{0},
-			sourcePrimitive:  newMergeSorter([]int{0}),
-			targetPrimitive:  newMergeSorter([]int{0}),
+			sourcePrimitive:  newMergeSorter(nil, []int{0}),
+			targetPrimitive:  newMergeSorter(nil, []int{0}),
 		},
 	}, {
 		// non-pk text column.
@@ -150,8 +150,8 @@ func TestVDiffPlanSuccess(t *testing.T) {
 			targetExpression: "select c1, textcol, weight_string(textcol) from nonpktext order by c1 asc",
 			compareCols:      []int{-1, 2},
 			comparePKs:       []int{0},
-			sourcePrimitive:  newMergeSorter([]int{0}),
-			targetPrimitive:  newMergeSorter([]int{0}),
+			sourcePrimitive:  newMergeSorter(nil, []int{0}),
+			targetPrimitive:  newMergeSorter(nil, []int{0}),
 		},
 	}, {
 		// non-pk text column, different order.
@@ -166,8 +166,8 @@ func TestVDiffPlanSuccess(t *testing.T) {
 			targetExpression: "select textcol, c1, weight_string(textcol) from nonpktext order by c1 asc",
 			compareCols:      []int{2, -1},
 			comparePKs:       []int{1},
-			sourcePrimitive:  newMergeSorter([]int{1}),
-			targetPrimitive:  newMergeSorter([]int{1}),
+			sourcePrimitive:  newMergeSorter(nil, []int{1}),
+			targetPrimitive:  newMergeSorter(nil, []int{1}),
 		},
 	}, {
 		// pk text column.
@@ -182,8 +182,8 @@ func TestVDiffPlanSuccess(t *testing.T) {
 			targetExpression: "select textcol, c2, weight_string(textcol) from pktext order by textcol asc",
 			compareCols:      []int{-1, 1},
 			comparePKs:       []int{2},
-			sourcePrimitive:  newMergeSorter([]int{2}),
-			targetPrimitive:  newMergeSorter([]int{2}),
+			sourcePrimitive:  newMergeSorter(nil, []int{2}),
+			targetPrimitive:  newMergeSorter(nil, []int{2}),
 		},
 	}, {
 		// pk text column, different order.
@@ -198,8 +198,8 @@ func TestVDiffPlanSuccess(t *testing.T) {
 			targetExpression: "select c2, textcol, weight_string(textcol) from pktext order by textcol asc",
 			compareCols:      []int{0, -1},
 			comparePKs:       []int{2},
-			sourcePrimitive:  newMergeSorter([]int{2}),
-			targetPrimitive:  newMergeSorter([]int{2}),
+			sourcePrimitive:  newMergeSorter(nil, []int{2}),
+			targetPrimitive:  newMergeSorter(nil, []int{2}),
 		},
 	}, {
 		// text column as expression.
@@ -214,8 +214,8 @@ func TestVDiffPlanSuccess(t *testing.T) {
 			targetExpression: "select c2, textcol, weight_string(textcol) from pktext order by textcol asc",
 			compareCols:      []int{0, -1},
 			comparePKs:       []int{2},
-			sourcePrimitive:  newMergeSorter([]int{2}),
-			targetPrimitive:  newMergeSorter([]int{2}),
+			sourcePrimitive:  newMergeSorter(nil, []int{2}),
+			targetPrimitive:  newMergeSorter(nil, []int{2}),
 		},
 	}, {
 		input: &binlogdatapb.Rule{
@@ -228,8 +228,8 @@ func TestVDiffPlanSuccess(t *testing.T) {
 			targetExpression: "select c1, c2 from multipk order by c1 asc, c2 asc",
 			compareCols:      []int{-1, -1},
 			comparePKs:       []int{0, 1},
-			sourcePrimitive:  newMergeSorter([]int{0, 1}),
-			targetPrimitive:  newMergeSorter([]int{0, 1}),
+			sourcePrimitive:  newMergeSorter(nil, []int{0, 1}),
+			targetPrimitive:  newMergeSorter(nil, []int{0, 1}),
 		},
 	}, {
 		// in_keyrange
@@ -244,8 +244,8 @@ func TestVDiffPlanSuccess(t *testing.T) {
 			targetExpression: "select c1, c2 from t1 order by c1 asc",
 			compareCols:      []int{-1, 1},
 			comparePKs:       []int{0},
-			sourcePrimitive:  newMergeSorter([]int{0}),
-			targetPrimitive:  newMergeSorter([]int{0}),
+			sourcePrimitive:  newMergeSorter(nil, []int{0}),
+			targetPrimitive:  newMergeSorter(nil, []int{0}),
 		},
 	}, {
 		// in_keyrange on RHS of AND.
@@ -261,8 +261,8 @@ func TestVDiffPlanSuccess(t *testing.T) {
 			targetExpression: "select c1, c2 from t1 order by c1 asc",
 			compareCols:      []int{-1, 1},
 			comparePKs:       []int{0},
-			sourcePrimitive:  newMergeSorter([]int{0}),
-			targetPrimitive:  newMergeSorter([]int{0}),
+			sourcePrimitive:  newMergeSorter(nil, []int{0}),
+			targetPrimitive:  newMergeSorter(nil, []int{0}),
 		},
 	}, {
 		// in_keyrange on LHS of AND.
@@ -278,8 +278,8 @@ func TestVDiffPlanSuccess(t *testing.T) {
 			targetExpression: "select c1, c2 from t1 order by c1 asc",
 			compareCols:      []int{-1, 1},
 			comparePKs:       []int{0},
-			sourcePrimitive:  newMergeSorter([]int{0}),
-			targetPrimitive:  newMergeSorter([]int{0}),
+			sourcePrimitive:  newMergeSorter(nil, []int{0}),
+			targetPrimitive:  newMergeSorter(nil, []int{0}),
 		},
 	}, {
 		// in_keyrange on cascaded AND expression
@@ -295,8 +295,8 @@ func TestVDiffPlanSuccess(t *testing.T) {
 			targetExpression: "select c1, c2 from t1 order by c1 asc",
 			compareCols:      []int{-1, 1},
 			comparePKs:       []int{0},
-			sourcePrimitive:  newMergeSorter([]int{0}),
-			targetPrimitive:  newMergeSorter([]int{0}),
+			sourcePrimitive:  newMergeSorter(nil, []int{0}),
+			targetPrimitive:  newMergeSorter(nil, []int{0}),
 		},
 	}, {
 		// in_keyrange parenthesized
@@ -312,8 +312,8 @@ func TestVDiffPlanSuccess(t *testing.T) {
 			targetExpression: "select c1, c2 from t1 order by c1 asc",
 			compareCols:      []int{-1, 1},
 			comparePKs:       []int{0},
-			sourcePrimitive:  newMergeSorter([]int{0}),
-			targetPrimitive:  newMergeSorter([]int{0}),
+			sourcePrimitive:  newMergeSorter(nil, []int{0}),
+			targetPrimitive:  newMergeSorter(nil, []int{0}),
 		},
 	}, {
 		// group by
@@ -328,8 +328,8 @@ func TestVDiffPlanSuccess(t *testing.T) {
 			targetExpression: "select c1, c2 from t1 order by c1 asc",
 			compareCols:      []int{-1, 1},
 			comparePKs:       []int{0},
-			sourcePrimitive:  newMergeSorter([]int{0}),
-			targetPrimitive:  newMergeSorter([]int{0}),
+			sourcePrimitive:  newMergeSorter(nil, []int{0}),
+			targetPrimitive:  newMergeSorter(nil, []int{0}),
 		},
 	}, {
 		// aggregations
@@ -353,17 +353,18 @@ func TestVDiffPlanSuccess(t *testing.T) {
 					Col:    3,
 				}},
 				Keys:  []int{0},
-				Input: newMergeSorter([]int{0}),
+				Input: newMergeSorter(nil, []int{0}),
 			},
-			targetPrimitive: newMergeSorter([]int{0}),
+			targetPrimitive: newMergeSorter(nil, []int{0}),
 		},
 	}}
 	for _, tcase := range testcases {
 		filter := &binlogdatapb.Filter{Rules: []*binlogdatapb.Rule{tcase.input}}
-		differs, err := buildVDiffPlan(context.Background(), filter, schm)
+		df := &vdiff{}
+		err := df.buildVDiffPlan(context.Background(), filter, schm)
 		require.NoError(t, err, tcase.input)
-		require.Equal(t, 1, len(differs), tcase.input)
-		assert.Equal(t, tcase.td, differs[tcase.table], tcase.input)
+		require.Equal(t, 1, len(df.differs), tcase.input)
+		assert.Equal(t, tcase.td, df.differs[tcase.table], tcase.input)
 	}
 }
 
@@ -413,7 +414,8 @@ func TestVDiffPlanFailure(t *testing.T) {
 	}}
 	for _, tcase := range testcases {
 		filter := &binlogdatapb.Filter{Rules: []*binlogdatapb.Rule{tcase.input}}
-		_, err := buildVDiffPlan(context.Background(), filter, schm)
+		df := &vdiff{}
+		err := df.buildVDiffPlan(context.Background(), filter, schm)
 		assert.EqualError(t, err, tcase.err, tcase.input)
 	}
 }