fix(plan): separate shuffle admission from filter cardinality

pkg/sql/plan/shuffle.go

@@ -588,6 +588,35 @@ func determineShuffleForJoin(node *plan.Node, builder *QueryBuilder) {
 	determineShuffleForJoinWithColRefMode(node, builder, false)
 }
 
+// shuffleJoinBuildSizeForAdmission keeps the cardinality estimate used by join
+// ordering separate from the memory-risk estimate used to admit shuffle. A
+// residual predicate that compares columns from different relations cannot be
+// estimated from marginal column statistics: its output can be as large as its
+// input. In that case use the filter input only for the shuffle threshold.
+func shuffleJoinBuildSizeForAdmission(node *plan.Node, builder *QueryBuilder, afterRemap bool) float64 {
+	buildSize := node.Stats.HashmapStats.HashmapSize
+	if afterRemap || node.IsRightJoin || len(node.Children) != 2 {
+		return buildSize
+	}
+
+	build := builder.qry.Nodes[node.Children[1]]
+	if build.NodeType != plan.Node_FILTER || len(build.Children) != 1 {
+		return buildSize
+	}
+	for _, filter := range build.FilterList {
+		tags := make(map[int32]int)
+		increaseTagCnt(filter, 1, tags)
+		if len(tags) > 1 {
+			input := builder.qry.Nodes[build.Children[0]]
+			if input.Stats != nil && input.Stats.Outcnt >= threshHoldForHashShuffle &&
+				input.Stats.Outcnt > buildSize {
+				return input.Stats.Outcnt
+			}
+		}
+	}
+	return buildSize
+}
+
 func isSupportedShuffleJoinKeyType(typ int32) bool {
 	switch types.T(typ) {
 	case types.T_int64, types.T_int32, types.T_int16,
@@ -686,7 +715,7 @@ func selectShuffleJoinCondition(
 	leftTags, rightTags map[int32]bool,
 	afterRemap bool,
 	previousHashmapStats *plan.HashMapStats,
-) (int, plan.HashMapStats) {
+) (int, plan.HashMapStats, bool) {
 	firstSupportedIdx := -1
 	var firstSupportedStats plan.HashMapStats
 
@@ -721,11 +750,11 @@ func selectShuffleJoinCondition(
 			firstSupportedStats = candidateStats
 		}
 		if eligible {
-			return i, candidateStats
+			return i, candidateStats, true
 		}
 	}
 
-	return firstSupportedIdx, firstSupportedStats
+	return firstSupportedIdx, firstSupportedStats, false
 }
 
 // determineShuffleForJoinWithColRefMode plans join shuffle either before or
@@ -797,7 +826,7 @@ func determineShuffleForJoinWithColRefMode(node *plan.Node, builder *QueryBuilde
 	if node.JoinType == plan.Node_MARK && !markJoinSupportsShuffle(node, builder, leftTags, rightTags, afterRemap) {
 		return
 	}
-	idx, candidateHashmapStats := selectShuffleJoinCondition(
+	idx, candidateHashmapStats, candidateEligible := selectShuffleJoinCondition(
 		node, builder, node.OnList, leftTags, rightTags, afterRemap,
 		previousHashmapStats,
 	)
@@ -812,7 +841,11 @@ func determineShuffleForJoinWithColRefMode(node *plan.Node, builder *QueryBuilde
 		leftchild := builder.qry.Nodes[node.Children[0]]
 		rightchild := builder.qry.Nodes[node.Children[1]]
 		factor := math.Pow((leftchild.Stats.Outcnt / rightchild.Stats.Outcnt), 0.4)
-		if node.Stats.HashmapStats.HashmapSize < threshHoldForShuffleJoin*factor {
+		buildSize := node.Stats.HashmapStats.HashmapSize
+		if candidateEligible {
+			buildSize = shuffleJoinBuildSizeForAdmission(node, builder, afterRemap)
+		}
+		if buildSize < threshHoldForShuffleJoin*factor {
 			return
 		}
 	}

pkg/sql/plan/shuffle_test.go

@@ -934,6 +934,57 @@ func TestDetermineShuffleForJoinNormalizesReversedConditionAfterRemap(t *testing
 	require.Equal(t, int32(1), condition.GetF().Args[1].GetCol().RelPos)
 }
 
+func TestDetermineShuffleForJoinUsesCrossRelationFilterInputForAdmission(t *testing.T) {
+	tests := []struct {
+		name              string
+		secondTag         int32
+		inputRows         float64
+		wantAdmissionRows float64
+		wantShuffle       bool
+	}{
+		{name: "single relation uses point estimate", secondTag: 2, inputRows: 60_000_000, wantAdmissionRows: 3_000_000, wantShuffle: false},
+		{name: "small cross relation filter stays resident", secondTag: 3, inputRows: 1_000_000, wantAdmissionRows: 50_000, wantShuffle: false},
+		{name: "large cross relation filter uses input risk", secondTag: 3, inputRows: 60_000_000, wantAdmissionRows: 60_000_000, wantShuffle: true},
+	}
+
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			predicate, err := BindFuncExprImplByPlanExpr(context.Background(), "!=", []*plan.Expr{
+				{Typ: plan.Type{Id: int32(types.T_int64)}, Expr: &plan.Expr_Col{Col: &plan.ColRef{RelPos: 2, ColPos: 1}}},
+				{Typ: plan.Type{Id: int32(types.T_int64)}, Expr: &plan.Expr_Col{Col: &plan.ColRef{RelPos: tt.secondTag, ColPos: 1}}},
+			})
+			require.NoError(t, err)
+
+			probe := makeShuffleJoinTestChild(1, 10_000_000_000)
+			buildInput := makeShuffleJoinTestChild(2, tt.inputRows)
+			buildInput.BindingTags = []int32{2, 3}
+			filter := &plan.Node{
+				NodeType:   plan.Node_FILTER,
+				Children:   []int32{1},
+				FilterList: []*plan.Expr{predicate},
+				Stats:      DefaultStats(),
+			}
+			join := &plan.Node{
+				NodeType: plan.Node_JOIN,
+				JoinType: plan.Node_INNER,
+				Children: []int32{0, 2},
+				OnList:   []*plan.Expr{makeShuffleJoinEquality(t, types.T_int64, 100_000, 1, 2, 0)},
+				Stats:    DefaultStats(),
+			}
+			builder := &QueryBuilder{qry: &plan.Query{Nodes: []*plan.Node{probe, buildInput, filter, join}}}
+
+			ReCalcNodeStats(2, builder, false, false, false)
+			ReCalcNodeStats(3, builder, false, false, false)
+			require.Equal(t, 0.05, filter.Stats.Selectivity)
+			require.Equal(t, tt.inputRows*0.05, join.Stats.HashmapStats.HashmapSize)
+			require.Equal(t, tt.wantAdmissionRows, shuffleJoinBuildSizeForAdmission(join, builder, false))
+
+			determineShuffleForJoin(join, builder)
+			require.Equal(t, tt.wantShuffle, join.Stats.HashmapStats.Shuffle)
+		})
+	}
+}
+
 func TestDetermineShuffleForJoinSkipsCandidateRejectedByFinalRecheck(t *testing.T) {
 	tests := []struct {
 		name            string
@@ -1072,9 +1123,10 @@ func TestSelectShuffleJoinConditionAdversarialPermutations(t *testing.T) {
 				}},
 			}
 
-			idx, _ := selectShuffleJoinCondition(node, builder, conditions, leftTags, rightTags, false, nil)
+			idx, _, eligible := selectShuffleJoinCondition(node, builder, conditions, leftTags, rightTags, false, nil)
 
 			require.NotEqual(t, -1, idx)
+			require.True(t, eligible)
 			require.Same(t, reusable, conditions[idx])
 			return
 		}