Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
32 changes: 32 additions & 0 deletions pkg/container/hashtable/common.go
Original file line number Diff line number Diff line change
Expand Up @@ -37,6 +37,38 @@ func maxElemCnt(cellCnt, cellSize uint64) uint64 {
return cellCnt * 4 / 5
}

// EstimateInt64HashMapSize returns the cell allocation required by an
// Int64HashMap for the given number of distinct keys. Keep this calculation
// in the hashtable package so planner memory guards follow the same growth and
// load-factor rules as the runtime implementation.
func EstimateInt64HashMapSize(cardinality uint64) uint64 {
return estimateHashMapSize(cardinality, intCellSize)
}

// EstimateStringHashMapSize returns the cell allocation required by a
// StringHashMap for the given number of distinct keys.
func EstimateStringHashMapSize(cardinality uint64) uint64 {
return estimateHashMapSize(cardinality, strCellSize)
}

func estimateHashMapSize(cardinality, cellSize uint64) uint64 {
const maxUint64 = ^uint64(0)

cellCnt := uint64(kInitialCellCnt)
for {
if cellCnt > maxUint64/cellSize {
return maxUint64
}
if maxElemCnt(cellCnt, cellSize) >= cardinality {
return cellCnt * cellSize
}
if cellCnt > maxUint64/2 {
return maxUint64
}
cellCnt <<= 1
}
}

func toUnsafePointer[T any](p *T) unsafe.Pointer {
return unsafe.Pointer(p)
}
34 changes: 34 additions & 0 deletions pkg/container/hashtable/common_test.go
Original file line number Diff line number Diff line change
@@ -0,0 +1,34 @@
// Copyright 2026 Matrix Origin
//
// 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 hashtable

import (
"testing"

"github.com/stretchr/testify/require"
)

func TestEstimateHashMapSizeFollowsRuntimeGrowth(t *testing.T) {
require.Equal(t, uint64(16*1024), EstimateInt64HashMapSize(0))
require.Equal(t, uint64(16*1024), EstimateInt64HashMapSize(512))
require.Equal(t, uint64(32*1024), EstimateInt64HashMapSize(513))

require.Equal(t, uint64(32*1024), EstimateStringHashMapSize(0))
require.Equal(t, uint64(32*1024), EstimateStringHashMapSize(512))
require.Equal(t, uint64(64*1024), EstimateStringHashMapSize(513))

require.Equal(t, ^uint64(0), EstimateInt64HashMapSize(^uint64(0)))
require.Equal(t, ^uint64(0), EstimateStringHashMapSize(^uint64(0)))
}
2 changes: 1 addition & 1 deletion pkg/sql/colexec/dedupjoin/join.go
Original file line number Diff line number Diff line change
Expand Up @@ -363,7 +363,7 @@ func (ctr *container) finalize(ap *DedupJoin, proc *process.Process) error {
if ctr.matched == nil {
return nil
}
if ap.NumCPU > 1 {
if ap.needsFinalizeMerge() {
if !ap.IsMerger {
msg := &WorkerJoinMsg{matched: ctr.matched}
if len(ap.OldColCapturePlaceholderIdxList) > 0 {
Expand Down
91 changes: 91 additions & 0 deletions pkg/sql/colexec/dedupjoin/join_test.go
Original file line number Diff line number Diff line change
Expand Up @@ -151,6 +151,97 @@ func TestDedupResetClearsBucketState(t *testing.T) {
proc.Free()
}

func TestDedupShuffleWorkersFinalizeTheirOwnPartitions(t *testing.T) {
tests := []struct {
name string
isShuffle bool
wantOutput bool
wantMessage bool
}{
{
name: "broadcast worker defers to merger",
wantMessage: true,
},
{
name: "shuffle worker emits local partition",
isShuffle: true,
wantOutput: true,
},
}

for _, test := range tests {
t.Run(test.name, func(t *testing.T) {
proc := testutil.NewProcessWithMPool(t, "", mpool.MustNewZero())
baseline := proc.Mp().CurrNB()
typ := types.T_int32.ToType()
bat := batch.NewOffHeapWithSize(1)
bat.Vecs[0] = vector.NewOffHeapVecWithType(typ)
require.NoError(t, vector.AppendFixed(bat.Vecs[0], int32(42), false, proc.Mp()))
bat.SetRowCount(1)

jm := message.NewJoinMap(message.GroupSels{}, nil, nil, nil, []*batch.Batch{bat}, proc.Mp())
jm.SetRowCount(1)
jm.IncRef(1)
matched := &bitmap.Bitmap{}
matched.InitWithSize(1)
ch := make(chan *WorkerJoinMsg, 1)
arg := &DedupJoin{
RightTypes: []types.Type{typ},
Result: []colexec.ResultPos{{Rel: 1, Pos: 0}},
OnDuplicateAction: plan.Node_FAIL,
NumCPU: 2,
IsMerger: false,
IsShuffle: test.isShuffle,
Channel: ch,
}
arg.ctr.mp = jm
arg.ctr.batches = jm.GetBatches()
arg.ctr.batchRowCount = jm.GetRowCount()
arg.ctr.matched = matched

require.NoError(t, arg.ctr.finalize(arg, proc))
if test.wantOutput {
require.Len(t, arg.ctr.buf, 1)
require.Equal(t, []int32{42}, vector.MustFixedColNoTypeCheck[int32](arg.ctr.buf[0].Vecs[0]))
} else {
require.Nil(t, arg.ctr.buf)
}
require.Equal(t, test.wantMessage, len(ch) == 1)

arg.Free(proc, false, nil)
require.Equal(t, baseline, proc.Mp().CurrNB())
proc.Free()
})
}
}

func TestDedupResetNotifiesOnlySharedBuildMerger(t *testing.T) {
for _, test := range []struct {
name string
isShuffle bool
wantMessage bool
}{
{name: "broadcast worker notifies merger", wantMessage: true},
{name: "shuffle worker owns its partition", isShuffle: true},
} {
t.Run(test.name, func(t *testing.T) {
proc := testutil.NewProcessWithMPool(t, "", mpool.MustNewZero())
ch := make(chan *WorkerJoinMsg, 1)
arg := &DedupJoin{
NumCPU: 2,
IsMerger: false,
IsShuffle: test.isShuffle,
Channel: ch,
}

arg.Reset(proc, false, nil)

require.Equal(t, test.wantMessage, len(ch) == 1)
proc.Free()
})
}
}

func TestDedupPrepareFailureCanRetry(t *testing.T) {
proc := testutil.NewProcessWithMPool(t, "", mpool.MustNewZero())
typ := types.T_int32.ToType()
Expand Down
10 changes: 9 additions & 1 deletion pkg/sql/colexec/dedupjoin/types.go
Original file line number Diff line number Diff line change
Expand Up @@ -184,9 +184,17 @@ func (dedupJoin *DedupJoin) Release() {
}
}

// needsFinalizeMerge reports whether parallel workers share one build map and
// therefore must merge their matched state before a single worker emits the
// build rows. Shuffle workers own disjoint build partitions, so each worker
// must finalize and emit its partition independently.
func (dedupJoin *DedupJoin) needsFinalizeMerge() bool {
return dedupJoin.NumCPU > 1 && !dedupJoin.IsShuffle
}

func (dedupJoin *DedupJoin) Reset(proc *process.Process, pipelineFailed bool, err error) {
ctr := &dedupJoin.ctr
if !ctr.handledLast && dedupJoin.NumCPU > 1 && !dedupJoin.IsMerger {
if !ctr.handledLast && dedupJoin.needsFinalizeMerge() && !dedupJoin.IsMerger {
dedupJoin.Channel <- nil
}
if dedupJoin.OpAnalyzer != nil {
Expand Down
1 change: 1 addition & 0 deletions pkg/sql/plan/query_builder.go
Original file line number Diff line number Diff line change
Expand Up @@ -2749,6 +2749,7 @@ func (builder *QueryBuilder) createQuery() (*Query, error) {
builder.optimizeDistinctAgg(rootID)
ReCalcNodeStats(rootID, builder, true, false, true)
builder.determineBuildAndProbeSide(rootID, true)
builder.disableMemoryUnsafeRightDedup(rootID)

builder.qry.Steps[i] = rootID

Expand Down
126 changes: 126 additions & 0 deletions pkg/sql/plan/right_dedup_memory_test.go
Original file line number Diff line number Diff line change
@@ -0,0 +1,126 @@
// Copyright 2026 Matrix Origin
//
// 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 plan

import (
"math"
"testing"

"github.com/matrixorigin/matrixone/pkg/container/types"
planpb "github.com/matrixorigin/matrixone/pkg/pb/plan"
"github.com/stretchr/testify/require"
)

func TestDisableMemoryUnsafeRightDedupUsesCombinedMapSize(t *testing.T) {
const combinedMapBytes = 64*1024 + 128*1024

t.Run("combined maps fit", func(t *testing.T) {
builder, joins := makeChainedRightDedupBuilder(combinedMapBytes)

builder.disableMemoryUnsafeRightDedup(4)

require.True(t, joins[0].IsRightJoin)
require.True(t, joins[1].IsRightJoin)
})

t.Run("combined maps exceed budget", func(t *testing.T) {
builder, joins := makeChainedRightDedupBuilder(combinedMapBytes - 1)

builder.disableMemoryUnsafeRightDedup(4)

require.False(t, joins[0].IsRightJoin)
require.False(t, joins[1].IsRightJoin)
})
}

func TestDisableMemoryUnsafeRightDedupHonorsRowThreshold(t *testing.T) {
t.Run("combined keys fit", func(t *testing.T) {
builder, joins := makeChainedRightDedupBuilder(2201)

builder.disableMemoryUnsafeRightDedup(4)

require.True(t, joins[0].IsRightJoin)
require.True(t, joins[1].IsRightJoin)
})

t.Run("combined keys reach threshold", func(t *testing.T) {
builder, joins := makeChainedRightDedupBuilder(2200)

builder.disableMemoryUnsafeRightDedup(4)

require.False(t, joins[0].IsRightJoin)
require.False(t, joins[1].IsRightJoin)
})
}

func TestDisableMemoryUnsafeRightDedupRejectsUnknownCardinality(t *testing.T) {
tests := []struct {
name string
stats *planpb.Stats
}{
{name: "non-finite", stats: &planpb.Stats{Outcnt: math.NaN()}},
{name: "default sentinel", stats: DefaultStats()},
}
for _, test := range tests {
t.Run(test.name, func(t *testing.T) {
builder, joins := makeChainedRightDedupBuilder(1 << 30)
builder.qry.Nodes[0].Stats = test.stats

builder.disableMemoryUnsafeRightDedup(4)

require.False(t, joins[0].IsRightJoin)
require.False(t, joins[1].IsRightJoin)
})
}
}

func makeChainedRightDedupBuilder(joinSpillMem int64) (*QueryBuilder, []*planpb.Node) {
source := &planpb.Node{NodeType: planpb.Node_VALUE_SCAN, Stats: &planpb.Stats{Outcnt: 1000}}
targetPK := &planpb.Node{NodeType: planpb.Node_TABLE_SCAN, Stats: &planpb.Stats{Outcnt: 100}}
pkDedup := &planpb.Node{
NodeType: planpb.Node_JOIN,
JoinType: planpb.Node_DEDUP,
Children: []int32{1, 0},
OnList: []*planpb.Expr{makeRightDedupEquality(types.T_int64)},
OnDuplicateAction: planpb.Node_FAIL,
IsRightJoin: true,
Stats: &planpb.Stats{Outcnt: 100}, // stale pre-swap RIGHT DEDUP estimate
}
targetUnique := &planpb.Node{NodeType: planpb.Node_TABLE_SCAN, Stats: &planpb.Stats{Outcnt: 100}}
uniqueDedup := &planpb.Node{
NodeType: planpb.Node_JOIN,
JoinType: planpb.Node_DEDUP,
Children: []int32{3, 2},
OnList: []*planpb.Expr{makeRightDedupEquality(types.T_varchar)},
OnDuplicateAction: planpb.Node_FAIL,
IsRightJoin: true,
Stats: &planpb.Stats{Outcnt: 100},
}
return &QueryBuilder{
qry: &planpb.Query{Nodes: []*planpb.Node{source, targetPK, pkDedup, targetUnique, uniqueDedup}},
joinSpillMem: joinSpillMem,
}, []*planpb.Node{pkDedup, uniqueDedup}
}

func makeRightDedupEquality(typ types.T) *planpb.Expr {
planType := planpb.Type{Id: int32(typ)}
return &planpb.Expr{
Typ: planpb.Type{Id: int32(types.T_bool)},
Expr: &planpb.Expr_F{F: &planpb.Function{Args: []*planpb.Expr{
{Typ: planType, Expr: &planpb.Expr_Col{Col: &planpb.ColRef{RelPos: 0, ColPos: 0}}},
{Typ: planType, Expr: &planpb.Expr_Col{Col: &planpb.ColRef{RelPos: 1, ColPos: 0}}},
}}},
}
}
17 changes: 15 additions & 2 deletions pkg/sql/plan/shuffle.go
Original file line number Diff line number Diff line change
Expand Up @@ -502,8 +502,9 @@ func determineShuffleForJoin(node *plan.Node, builder *QueryBuilder) {
}
} else {
rightChild := builder.qry.Nodes[node.Children[1]]
if rightChild.Stats.Outcnt > 320000 {
//dedup join always go hash shuffle, optimize this in the future
if rightChild.Stats.Outcnt > 320000 && !dedupJoinUsesUnsupportedFloatShuffle(node) {
// Large DEDUP joins normally use hash shuffle. FLOAT hash shuffle is
// not supported yet, so those joins stay single-CN and spill locally.
node.Stats.HashmapStats.Shuffle = true
node.Stats.HashmapStats.ShuffleColIdx = 0
node.Stats.HashmapStats.ShuffleType = plan.ShuffleType_Hash
Expand Down Expand Up @@ -626,6 +627,18 @@ func determineShuffleForJoin(node *plan.Node, builder *QueryBuilder) {
}
}

func dedupJoinUsesUnsupportedFloatShuffle(node *plan.Node) bool {
if len(node.OnList) == 0 {
return false
}
condition := node.OnList[0].GetF()
if condition == nil || len(condition.Args) == 0 {
return false
}
keyType := types.T(condition.Args[0].Typ.Id)
return keyType == types.T_float32 || keyType == types.T_float64
}

// find mergegroup or mergegroup->filter node
func dontShuffle(node *plan.Node, builder *QueryBuilder) bool {
if node.NodeType == plan.Node_AGG && !node.Stats.HashmapStats.Shuffle {
Expand Down
Loading
Loading