Repository: pingcap/tidb
Pull request #48251 (2023-11-02 14:51:50)

ddl: make TestHalfwayCancelOperations faster
<!-- 

Thank you for contributing to TiDB!

PR Title Format:
1. pkg [, pkg2, pkg3]: what's changed
2. *: what's changed

-->

### What problem does this PR solve?
<!--

Please create an issue first to describe the problem.

There MUST be one line starting with "Issue Number:  " and 
linking the relevant issues via the "close" or "ref".

For more info, check https://pingcap.github.io/tidb-dev-guide/contribute-to-tidb/contribute-code.html#referring-to-an-issue.

-->

Issue Number: close https://github.com/pingcap/tidb/issues/46840

Problem Summary:
When we occur `exchangePartitionErr`, we will retry 512 times. 

### What is changed and how it works?
Set `DDLErrorCountLimit` to 3.

### Check List

Tests <!-- At least one of them must be included. -->

- [ ] Unit test
- [ ] Integration test
- [x] Manual test (add detailed scripts or steps below)

**Before this PR**
without race
```
PASS
ok  	github.com/pingcap/tidb/pkg/ddl/tests/fail	2.762s
```
with race
```
PASS
ok  	github.com/pingcap/tidb/pkg/ddl/tests/fail	10.506s
```
**After this PR**
without race
```
PASS
ok  	github.com/pingcap/tidb/pkg/ddl/tests/fail	2.068s
```
with race
```
PASS
ok  	github.com/pingcap/tidb/pkg/ddl/tests/fail	5.072s
```
- [ ] No need to test
  > - [ ] I checked and no code files have been changed.
  > <!-- Or your custom  "No need to test" reasons -->  

Side effects

- [ ] Performance regression: Consumes more CPU
- [ ] Performance regression: Consumes more Memory
- [ ] Breaking backward compatibility

Documentation

- [ ] Affects user behaviors
- [ ] Contains syntax changes
- [ ] Contains variable changes
- [ ] Contains experimental features
- [ ] Changes MySQL compatibility

### Release note

<!-- compatibility change, improvement, bugfix, and new feature need a release note -->

Please refer to [Release Notes Language Style Guide](https://pingcap.github.io/tidb-dev-guide/contribute-to-tidb/release-notes-style-guide.html) to write a quality release note.

```release-note
None
```


ti-chi-bot[bot]: [APPROVALNOTIFIER] This PR is **NOT APPROVED**

This pull-request has been approved by:
**Once this PR has been reviewed and has the lgtm label**, please ask for approval from [zimulala](https://github.com/zimulala). For more information see [the Kubernetes Code Review Process](https://book.prow.tidb.net/#/workflows/pr).

The full list of commands accepted by this bot can be found [here](https://prow.tidb.net/command-help?repo=pingcap%2Ftidb).

<details open>
Needs approval from an approver in each of these files:

- **[pkg/ddl/OWNERS](https://github.com/pingcap/tidb/blob/master/pkg/ddl/OWNERS)**

Approvers can indicate their approval by writing `/approve` in a comment
Approvers can cancel approval by writing `/approve cancel` in a comment
</details>
<!-- META={"approvers":["zimulala"]} -->
tiprow[bot]: Hi @zimulala. Thanks for your PR.

PRs from untrusted users cannot be marked as trusted with `/ok-to-test` in this repo meaning untrusted PR authors can never trigger tests themselves. Collaborators can still trigger tests on the PR using `/test all`.

I understand the commands that are listed [here](https://go.k8s.io/bot-commands?repo=pingcap%2Ftidb).

<details>

Instructions for interacting with me using PR comments are available [here](https://git.k8s.io/community/contributors/guide/pull-requests.md).  If you have questions or suggestions related to my behavior, please file an issue against the [kubernetes/test-infra](https://github.com/kubernetes/test-infra/issues/new?title=Prow%20issue:) repository.
</details>

codecov[bot]: ## [Codecov](https://app.codecov.io/gh/pingcap/tidb/pull/48251?src=pr&el=h1&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=pingcap) Report
> Merging [#48251](https://app.codecov.io/gh/pingcap/tidb/pull/48251?src=pr&el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=pingcap) (529f933) into [master](https://app.codecov.io/gh/pingcap/tidb/commit/eed2a7605064a155dfb1930a9b086f014b1ebc70?el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=pingcap) (eed2a76) will **increase** coverage by `0.0110%`.
> Report is 1 commits behind head on master.
> The diff coverage is `n/a`.

<details><summary>Additional details and impacted files</summary>


```diff
@@               Coverage Diff                @@
##             master     #48251        +/-   ##
================================================
+ Coverage   71.4958%   71.5069%   +0.0110%     
================================================
  Files          1402       1405         +3     
  Lines        406166     406843       +677     
================================================
+ Hits         290392     290921       +529     
- Misses        95949      96062       +113     
- Partials      19825      19860        +35     
```

| [Flag](https://app.codecov.io/gh/pingcap/tidb/pull/48251/flags?src=pr&el=flags&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=pingcap) | Coverage Δ | |
|---|---|---|
| [integration](https://app.codecov.io/gh/pingcap/tidb/pull/48251/flags?src=pr&el=flag&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=pingcap) | `78.1388% <ø> (?)` | |

Flags with carried forward coverage won't be shown. [Click here](https://docs.codecov.io/docs/carryforward-flags?utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=pingcap#carryforward-flags-in-the-pull-request-comment) to find out more.

| [Components](https://app.codecov.io/gh/pingcap/tidb/pull/48251/components?src=pr&el=components&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=pingcap) | Coverage Δ | |
|---|---|---|
| [dumpling](https://app.codecov.io/gh/pingcap/tidb/pull/48251/components?src=pr&el=component&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=pingcap) | `53.9874% <ø> (ø)` | |
| [parser](https://app.codecov.io/gh/pingcap/tidb/pull/48251/components?src=pr&el=component&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=pingcap) | `∅ <ø> (∅)` | |
| [br](https://app.codecov.io/gh/pingcap/tidb/pull/48251/components?src=pr&el=component&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=pingcap) | `52.9298% <ø> (ø)` | |

</details>
hawkingrei: /check-issue-triage-complete
ti-chi-bot[bot]: [APPROVALNOTIFIER] This PR is **NOT APPROVED**

This pull-request has been approved by: *<a href="https://github.com/pingcap/tidb/pull/48251#pullrequestreview-1710633080" title="Approved">hawkingrei</a>*
**Once this PR has been reviewed and has the lgtm label**, please ask for approval from [zimulala](https://github.com/zimulala). For more information see [the Kubernetes Code Review Process](https://book.prow.tidb.net/#/workflows/pr).

The full list of commands accepted by this bot can be found [here](https://prow.tidb.net/command-help?repo=pingcap%2Ftidb).

<details open>
Needs approval from an approver in each of these files:

- **[pkg/ddl/OWNERS](https://github.com/pingcap/tidb/blob/master/pkg/ddl/OWNERS)**

Approvers can indicate their approval by writing `/approve` in a comment
Approvers can cancel approval by writing `/approve cancel` in a comment
</details>
<!-- META={"approvers":["zimulala"]} -->
ti-chi-bot[bot]: [LGTM Timeline notifier]
---
Timeline:

- `2023-11-02 15:42:15.481126669 +0000 UTC m=+3140533.068236814`: :ballot_box_with_check: agreed by [hawkingrei](https://github.com/hawkingrei).

Patch:
diff --git a/pkg/disttask/framework/planner/plan.go b/pkg/disttask/framework/planner/plan.go
index fae9b2b9b9..e3d94a7c95 100644
--- a/pkg/disttask/framework/planner/plan.go
+++ b/pkg/disttask/framework/planner/plan.go
@@ -37,6 +37,7 @@ type PlanCtx struct {
 	PreviousSubtaskMetas map[proto.Step][][]byte
 	GlobalSort           bool
 	NextTaskStep         proto.Step
+	ExecuteNodesCnt      int
 }
 
 // LogicalPlan represents a logical plan in distribute framework.
diff --git a/pkg/disttask/importinto/dispatcher.go b/pkg/disttask/importinto/dispatcher.go
index b141f1e46c..ab8fd81bec 100644
--- a/pkg/disttask/importinto/dispatcher.go
+++ b/pkg/disttask/importinto/dispatcher.go
@@ -299,12 +299,18 @@ func (dsp *ImportDispatcherExt) OnNextSubtasksBatch(
 		return nil, errors.Errorf("unknown step %d", gTask.Step)
 	}
 
+	eligibleInstances, err := dsp.GetEligibleInstances(ctx, gTask)
+	if err != nil {
+		logger.Warn("failed to get eligible instances", zap.Error(err))
+	}
+
 	planCtx := planner.PlanCtx{
 		Ctx:                  ctx,
 		TaskID:               gTask.ID,
 		PreviousSubtaskMetas: previousSubtaskMetas,
 		GlobalSort:           dsp.GlobalSort,
 		NextTaskStep:         nextStep,
+		ExecuteNodesCnt:      len(eligibleInstances),
 	}
 	logicalPlan := &LogicalPlan{}
 	if err := logicalPlan.FromTaskMeta(gTask.Meta); err != nil {
diff --git a/pkg/disttask/importinto/planner.go b/pkg/disttask/importinto/planner.go
index 6ea23054b4..53ffe2e8c3 100644
--- a/pkg/disttask/importinto/planner.go
+++ b/pkg/disttask/importinto/planner.go
@@ -109,7 +109,7 @@ func (p *LogicalPlan) ToPhysicalPlan(planCtx planner.PlanCtx) (*planner.Physical
 	// we only generate needed plans for the next step.
 	switch planCtx.NextTaskStep {
 	case StepImport, StepEncodeAndSort:
-		specs, err := generateImportSpecs(planCtx.Ctx, p)
+		specs, err := generateImportSpecs(planCtx, p)
 		if err != nil {
 			return nil, err
 		}
@@ -249,7 +249,7 @@ func buildController(plan *importer.Plan, stmt string) (*importer.LoadDataContro
 	return controller, nil
 }
 
-func generateImportSpecs(ctx context.Context, p *LogicalPlan) ([]planner.PipelineSpec, error) {
+func generateImportSpecs(pCtx planner.PlanCtx, p *LogicalPlan) ([]planner.PipelineSpec, error) {
 	var chunkMap map[int32][]Chunk
 	if len(p.ChunkMap) > 0 {
 		chunkMap = p.ChunkMap
@@ -258,11 +258,12 @@ func generateImportSpecs(ctx context.Context, p *LogicalPlan) ([]planner.Pipelin
 		if err2 != nil {
 			return nil, err2
 		}
-		if err2 = controller.InitDataFiles(ctx); err2 != nil {
+		if err2 = controller.InitDataFiles(pCtx.Ctx); err2 != nil {
 			return nil, err2
 		}
 
-		engineCheckpoints, err2 := controller.PopulateChunks(ctx)
+		controller.SetExecuteNodeCnt(pCtx.ExecuteNodesCnt)
+		engineCheckpoints, err2 := controller.PopulateChunks(pCtx.Ctx)
 		if err2 != nil {
 			return nil, err2
 		}
diff --git a/pkg/executor/aggregate/agg_hash_base_worker.go b/pkg/executor/aggregate/agg_hash_base_worker.go
index 67fa32f356..b2ed73cd06 100644
--- a/pkg/executor/aggregate/agg_hash_base_worker.go
+++ b/pkg/executor/aggregate/agg_hash_base_worker.go
@@ -84,7 +84,7 @@ func (w *baseHashAggWorker) getPartialResult(_ *stmtctx.StatementContext, groupK
 
 func (w *baseHashAggWorker) getPartialResultSliceLenConsiderByteAlign() int {
 	length := len(w.aggFuncs)
-	if len(w.aggFuncs) == 1 {
+	if length == 1 {
 		return 1
 	}
 	return length + length&1
diff --git a/pkg/executor/aggregate/agg_hash_executor.go b/pkg/executor/aggregate/agg_hash_executor.go
index 9c4283dc69..a3bc512ece 100644
--- a/pkg/executor/aggregate/agg_hash_executor.go
+++ b/pkg/executor/aggregate/agg_hash_executor.go
@@ -104,7 +104,7 @@ type HashAggExec struct {
 
 	finishCh         chan struct{}
 	finalOutputCh    chan *AfFinalResult
-	partialOutputChs []chan *HashAggIntermData
+	partialOutputChs []chan *AggPartialResultMapper
 	inputCh          chan *HashAggInput
 	partialInputChs  []chan *chunk.Chunk
 	partialWorkers   []HashAggPartialWorker
@@ -264,9 +264,9 @@ func (e *HashAggExec) initForParallelExec(_ sessionctx.Context) {
 	for i := range e.partialInputChs {
 		e.partialInputChs[i] = make(chan *chunk.Chunk, 1)
 	}
-	e.partialOutputChs = make([]chan *HashAggIntermData, finalConcurrency)
+	e.partialOutputChs = make([]chan *AggPartialResultMapper, finalConcurrency)
 	for i := range e.partialOutputChs {
-		e.partialOutputChs[i] = make(chan *HashAggIntermData, partialConcurrency)
+		e.partialOutputChs[i] = make(chan *AggPartialResultMapper, partialConcurrency)
 	}
 
 	e.partialWorkers = make([]HashAggPartialWorker, partialConcurrency)
@@ -275,17 +275,30 @@ func (e *HashAggExec) initForParallelExec(_ sessionctx.Context) {
 
 	// Init partial workers.
 	for i := 0; i < partialConcurrency; i++ {
+		partialResultsMap := make([]AggPartialResultMapper, finalConcurrency)
+		for i := 0; i < finalConcurrency; i++ {
+			partialResultsMap[i] = make(AggPartialResultMapper)
+		}
+
 		w := HashAggPartialWorker{
-			baseHashAggWorker: newBaseHashAggWorker(e.Ctx(), e.finishCh, e.PartialAggFuncs, e.MaxChunkSize(), e.memTracker),
-			inputCh:           e.partialInputChs[i],
-			outputChs:         e.partialOutputChs,
-			giveBackCh:        e.inputCh,
-			globalOutputCh:    e.finalOutputCh,
-			partialResultsMap: make(AggPartialResultMapper),
-			groupByItems:      e.GroupByItems,
-			chk:               exec.TryNewCacheChunk(e.Children(0)),
-			groupKey:          make([][]byte, 0, 8),
+			baseHashAggWorker:    newBaseHashAggWorker(e.Ctx(), e.finishCh, e.PartialAggFuncs, e.MaxChunkSize(), e.memTracker),
+			inputCh:              e.partialInputChs[i],
+			outputChs:            e.partialOutputChs,
+			giveBackCh:           e.inputCh,
+			BInMaps:              make([]int, finalConcurrency),
+			partialResultsBuffer: make([][]aggfuncs.PartialResult, 0, 2048),
+			globalOutputCh:       e.finalOutputCh,
+			partialResultsMap:    partialResultsMap,
+			groupByItems:         e.GroupByItems,
+			chk:                  exec.TryNewCacheChunk(e.Children(0)),
+			groupKey:             make([][]byte, 0, 8),
 		}
+
+		w.partialResultNumInRow = w.getPartialResultSliceLenConsiderByteAlign()
+		for i := 0; i < finalConcurrency; i++ {
+			w.BInMaps[i] = 0
+		}
+
 		// There is a bucket in the empty partialResultsMap.
 		failpoint.Inject("ConsumeRandomPanic", nil)
 		e.memTracker.Consume(hack.DefBucketMemoryUsageForMapStrToSlice * (1 << w.BInMap))
@@ -309,6 +322,8 @@ func (e *HashAggExec) initForParallelExec(_ sessionctx.Context) {
 		w := HashAggFinalWorker{
 			baseHashAggWorker:   newBaseHashAggWorker(e.Ctx(), e.finishCh, e.FinalAggFuncs, e.MaxChunkSize(), e.memTracker),
 			partialResultMap:    make(AggPartialResultMapper),
+			BInMap:              0,
+			isFirstInput:        true,
 			groupSet:            groupSet,
 			inputCh:             e.partialOutputChs[i],
 			outputCh:            e.finalOutputCh,
diff --git a/pkg/executor/aggregate/agg_hash_final_worker.go b/pkg/executor/aggregate/agg_hash_final_worker.go
index a2e9f2edae..ab139f4915 100644
--- a/pkg/executor/aggregate/agg_hash_final_worker.go
+++ b/pkg/executor/aggregate/agg_hash_final_worker.go
@@ -19,10 +19,10 @@ import (
 	"time"
 
 	"github.com/pingcap/failpoint"
-	"github.com/pingcap/tidb/pkg/executor/aggfuncs"
 	"github.com/pingcap/tidb/pkg/sessionctx"
 	"github.com/pingcap/tidb/pkg/types"
 	"github.com/pingcap/tidb/pkg/util/chunk"
+	"github.com/pingcap/tidb/pkg/util/hack"
 	"github.com/pingcap/tidb/pkg/util/logutil"
 	"github.com/pingcap/tidb/pkg/util/set"
 	"go.uber.org/zap"
@@ -43,14 +43,16 @@ type HashAggFinalWorker struct {
 	rowBuffer           []types.Datum
 	mutableRow          chunk.MutRow
 	partialResultMap    AggPartialResultMapper
+	BInMap              int
+	isFirstInput        bool
 	groupSet            set.StringSetWithMemoryUsage
-	inputCh             chan *HashAggIntermData
+	inputCh             chan *AggPartialResultMapper
 	outputCh            chan *AfFinalResult
 	finalResultHolderCh chan *chunk.Chunk
 	groupKeys           [][]byte
 }
 
-func (w *HashAggFinalWorker) getPartialInput() (input *HashAggIntermData, ok bool) {
+func (w *HashAggFinalWorker) getPartialInput() (input *AggPartialResultMapper, ok bool) {
 	select {
 	case <-w.finishCh:
 		return nil, false
@@ -62,55 +64,60 @@ func (w *HashAggFinalWorker) getPartialInput() (input *HashAggIntermData, ok boo
 	return
 }
 
+func (w *HashAggFinalWorker) initBInMap() {
+	w.BInMap = 0
+	mapLen := len(w.partialResultMap)
+	for mapLen > (1<<w.BInMap)*hack.LoadFactorNum/hack.LoadFactorDen {
+		w.BInMap++
+	}
+}
+
 func (w *HashAggFinalWorker) consumeIntermData(sctx sessionctx.Context) (err error) {
-	var (
-		input            *HashAggIntermData
-		ok               bool
-		intermDataBuffer [][]aggfuncs.PartialResult
-		groupKeys        []string
-		sc               = sctx.GetSessionVars().StmtCtx
-	)
 	for {
 		waitStart := time.Now()
-		input, ok = w.getPartialInput()
+		input, ok := w.getPartialInput()
 		if w.stats != nil {
 			w.stats.WaitTime += int64(time.Since(waitStart))
 		}
 		if !ok {
 			return nil
 		}
-		execStart := time.Now()
-		if intermDataBuffer == nil {
-			intermDataBuffer = make([][]aggfuncs.PartialResult, 0, w.maxChunkSize)
+
+		// As the w.partialResultMap is empty when we get the first input.
+		// So it's better to directly assign the input to w.partialResultMap
+		if w.isFirstInput {
+			w.isFirstInput = false
+			w.partialResultMap = *input
+			w.initBInMap()
+			continue
 		}
-		// Consume input in batches, size of every batch is less than w.maxChunkSize.
-		for reachEnd := false; !reachEnd; {
-			intermDataBuffer, groupKeys, reachEnd = input.getPartialResultBatch(sc, intermDataBuffer[:0], w.aggFuncs, w.maxChunkSize)
-			groupKeysLen := len(groupKeys)
-			memSize := getGroupKeyMemUsage(w.groupKeys)
-			w.groupKeys = w.groupKeys[:0]
-			for i := 0; i < groupKeysLen; i++ {
-				w.groupKeys = append(w.groupKeys, []byte(groupKeys[i]))
-			}
-			failpoint.Inject("ConsumeRandomPanic", nil)
-			w.memTracker.Consume(getGroupKeyMemUsage(w.groupKeys) - memSize)
-			finalPartialResults := w.getPartialResult(sc, w.groupKeys, w.partialResultMap)
-			allMemDelta := int64(0)
-			for i, groupKey := range groupKeys {
-				if !w.groupSet.Exist(groupKey) {
-					allMemDelta += w.groupSet.Insert(groupKey)
+
+		failpoint.Inject("ConsumeRandomPanic", nil)
+
+		execStart := time.Now()
+		allMemDelta := int64(0)
+		for key, value := range *input {
+			dstVal, ok := w.partialResultMap[key]
+			if !ok {
+				// Map will expand when count > bucketNum * loadFactor. The memory usage will double.
+				if len(w.partialResultMap)+1 > (1<<w.BInMap)*hack.LoadFactorNum/hack.LoadFactorDen {
+					w.memTracker.Consume(hack.DefBucketMemoryUsageForMapStrToSlice * (1 << w.BInMap))
+					w.BInMap++
 				}
-				prs := intermDataBuffer[i]
-				for j, af := range w.aggFuncs {
-					memDelta, err := af.MergePartialResult(sctx, prs[j], finalPartialResults[i][j])
-					if err != nil {
-						return err
-					}
-					allMemDelta += memDelta
+				w.partialResultMap[key] = value
+				continue
+			}
+
+			for j, af := range w.aggFuncs {
+				memDelta, err := af.MergePartialResult(sctx, value[j], dstVal[j])
+				if err != nil {
+					return err
 				}
+				allMemDelta += memDelta
 			}
-			w.memTracker.Consume(allMemDelta)
 		}
+		w.memTracker.Consume(allMemDelta)
+
 		if w.stats != nil {
 			w.stats.ExecTime += int64(time.Since(execStart))
 			w.stats.TaskNum++
@@ -127,24 +134,21 @@ func (w *HashAggFinalWorker) loadFinalResult(sctx sessionctx.Context) {
 	if finished {
 		return
 	}
-	execStart := time.Now()
-	memSize := getGroupKeyMemUsage(w.groupKeys)
-	w.groupKeys = w.groupKeys[:0]
-	for groupKey := range w.groupSet.StringSet {
-		w.groupKeys = append(w.groupKeys, []byte(groupKey))
-	}
+
 	failpoint.Inject("ConsumeRandomPanic", nil)
-	w.memTracker.Consume(getGroupKeyMemUsage(w.groupKeys) - memSize)
-	partialResults := w.getPartialResult(sctx.GetSessionVars().StmtCtx, w.groupKeys, w.partialResultMap)
-	for i := 0; i < len(w.groupSet.StringSet); i++ {
+
+	execStart := time.Now()
+	for _, results := range w.partialResultMap {
 		for j, af := range w.aggFuncs {
-			if err := af.AppendFinalResult2Chunk(sctx, partialResults[i][j], result); err != nil {
+			if err := af.AppendFinalResult2Chunk(sctx, results[j], result); err != nil {
 				logutil.BgLogger().Error("HashAggFinalWorker failed to append final result to Chunk", zap.Error(err))
 			}
 		}
+
 		if len(w.aggFuncs) == 0 {
 			result.SetNumVirtualRows(result.NumRows() + 1)
 		}
+
 		if result.IsFull() {
 			w.outputCh <- &AfFinalResult{chk: result, giveBackCh: w.finalResultHolderCh}
 			result, finished = w.receiveFinalResultHolder()
@@ -153,6 +157,7 @@ func (w *HashAggFinalWorker) loadFinalResult(sctx sessionctx.Context) {
 			}
 		}
 	}
+
 	w.outputCh <- &AfFinalResult{chk: result, giveBackCh: w.finalResultHolderCh}
 	if w.stats != nil {
 		w.stats.ExecTime += int64(time.Since(execStart))
diff --git a/pkg/executor/aggregate/agg_hash_partial_worker.go b/pkg/executor/aggregate/agg_hash_partial_worker.go
index 1988d48eab..d41300bd2d 100644
--- a/pkg/executor/aggregate/agg_hash_partial_worker.go
+++ b/pkg/executor/aggregate/agg_hash_partial_worker.go
@@ -19,30 +19,30 @@ import (
 	"time"
 
 	"github.com/pingcap/failpoint"
+	"github.com/pingcap/tidb/pkg/executor/aggfuncs"
 	"github.com/pingcap/tidb/pkg/expression"
 	"github.com/pingcap/tidb/pkg/sessionctx"
-	"github.com/pingcap/tidb/pkg/sessionctx/stmtctx"
 	"github.com/pingcap/tidb/pkg/util/chunk"
+	"github.com/pingcap/tidb/pkg/util/hack"
 	"github.com/twmb/murmur3"
 )
 
-// HashAggIntermData indicates the intermediate data of aggregation execution.
-type HashAggIntermData struct {
-	groupKeys        []string
-	cursor           int
-	partialResultMap AggPartialResultMapper
-}
-
 // HashAggPartialWorker indicates the partial workers of parallel hash agg execution,
 // the number of the worker can be set by `tidb_hashagg_partial_concurrency`.
 type HashAggPartialWorker struct {
 	baseHashAggWorker
 
-	inputCh           chan *chunk.Chunk
-	outputChs         []chan *HashAggIntermData
-	globalOutputCh    chan *AfFinalResult
-	giveBackCh        chan<- *HashAggInput
-	partialResultsMap AggPartialResultMapper
+	inputCh        chan *chunk.Chunk
+	outputChs      []chan *AggPartialResultMapper
+	globalOutputCh chan *AfFinalResult
+	giveBackCh     chan<- *HashAggInput
+	BInMaps        []int
+
+	partialResultsBuffer  [][]aggfuncs.PartialResult
+	partialResultNumInRow int
+
+	// Length of this map is equal to the number of final workers
+	partialResultsMap []AggPartialResultMapper
 	groupByItems      []expression.Expression
 	groupKey          [][]byte
 	// chk stores the input data from child,
@@ -69,13 +69,13 @@ func (w *HashAggPartialWorker) getChildInput() bool {
 
 func (w *HashAggPartialWorker) run(ctx sessionctx.Context, waitGroup *sync.WaitGroup, finalConcurrency int) {
 	start := time.Now()
-	needShuffle, sc := false, ctx.GetSessionVars().StmtCtx
+	needShuffle := false
 	defer func() {
 		if r := recover(); r != nil {
 			recoveryHashAgg(w.globalOutputCh, r)
 		}
 		if needShuffle {
-			w.shuffleIntermData(sc, finalConcurrency)
+			w.shuffleIntermData(finalConcurrency)
 		}
 		w.memTracker.Consume(-w.chk.MemoryUsage())
 		if w.stats != nil {
@@ -93,7 +93,7 @@ func (w *HashAggPartialWorker) run(ctx sessionctx.Context, waitGroup *sync.WaitG
 			return
 		}
 		execStart := time.Now()
-		if err := w.updatePartialResult(ctx, sc, w.chk, len(w.partialResultsMap)); err != nil {
+		if err := w.updatePartialResult(ctx, w.chk, len(w.partialResultsMap)); err != nil {
 			w.globalOutputCh <- &AfFinalResult{err: err}
 			return
 		}
@@ -107,7 +107,48 @@ func (w *HashAggPartialWorker) run(ctx sessionctx.Context, waitGroup *sync.WaitG
 	}
 }
 
-func (w *HashAggPartialWorker) updatePartialResult(ctx sessionctx.Context, sc *stmtctx.StatementContext, chk *chunk.Chunk, _ int) (err error) {
+// If the group key has appeared before, reuse the partial result.
+// If the group key has not appeared before, create empty partial results.
+func (w *HashAggPartialWorker) getPartialResultsOfEachRow(groupKey [][]byte, finalConcurrency int) [][]aggfuncs.PartialResult {
+	mapper := w.partialResultsMap
+	numRows := len(groupKey)
+	allMemDelta := int64(0)
+	w.partialResultsBuffer = w.partialResultsBuffer[0:0]
+
+	for i := 0; i < numRows; i++ {
+		finalWorkerIdx := int(murmur3.Sum32(groupKey[i])) % finalConcurrency
+		tmp, ok := mapper[finalWorkerIdx][string(hack.String(groupKey[i]))]
+
+		// This group by key has appeared before, reuse the partial result.
+		if ok {
+			w.partialResultsBuffer = append(w.partialResultsBuffer, tmp)
+			continue
+		}
+
+		// It's the first time that this group by key appeared, create it
+		w.partialResultsBuffer = append(w.partialResultsBuffer, make([]aggfuncs.PartialResult, w.partialResultNumInRow))
+		lastIdx := len(w.partialResultsBuffer) - 1
+		for j, af := range w.aggFuncs {
+			partialResult, memDelta := af.AllocPartialResult()
+			w.partialResultsBuffer[lastIdx][j] = partialResult
+			allMemDelta += memDelta // the memory usage of PartialResult
+		}
+		allMemDelta += int64(w.partialResultNumInRow * 8)
+
+		// Map will expand when count > bucketNum * loadFactor. The memory usage will double.
+		if len(mapper[finalWorkerIdx])+1 > (1<<w.BInMaps[finalWorkerIdx])*hack.LoadFactorNum/hack.LoadFactorDen {
+			w.memTracker.Consume(hack.DefBucketMemoryUsageForMapStrToSlice * (1 << w.BInMaps[finalWorkerIdx]))
+			w.BInMaps[finalWorkerIdx]++
+		}
+
+		mapper[finalWorkerIdx][string(groupKey[i])] = w.partialResultsBuffer[lastIdx]
+		allMemDelta += int64(len(groupKey[i]))
+	}
+	w.memTracker.Consume(allMemDelta)
+	return w.partialResultsBuffer
+}
+
+func (w *HashAggPartialWorker) updatePartialResult(ctx sessionctx.Context, chk *chunk.Chunk, finalConcurrency int) (err error) {
 	memSize := getGroupKeyMemUsage(w.groupKey)
 	w.groupKey, err = GetGroupKey(w.ctx, chk, w.groupKey, w.groupByItems)
 	failpoint.Inject("ConsumeRandomPanic", nil)
@@ -116,14 +157,16 @@ func (w *HashAggPartialWorker) updatePartialResult(ctx sessionctx.Context, sc *s
 		return err
 	}
 
-	partialResults := w.getPartialResult(sc, w.groupKey, w.partialResultsMap)
+	partialResultOfEachRow := w.getPartialResultsOfEachRow(w.groupKey, finalConcurrency)
+
 	numRows := chk.NumRows()
 	rows := make([]chunk.Row, 1)
 	allMemDelta := int64(0)
 	for i := 0; i < numRows; i++ {
+		partialResult := partialResultOfEachRow[i]
+		rows[0] = chk.GetRow(i)
 		for j, af := range w.aggFuncs {
-			rows[0] = chk.GetRow(i)
-			memDelta, err := af.UpdatePartialResult(ctx, rows, partialResults[i][j])
+			memDelta, err := af.UpdatePartialResult(ctx, rows, partialResult[j])
 			if err != nil {
 				return err
 			}
@@ -134,25 +177,8 @@ func (w *HashAggPartialWorker) updatePartialResult(ctx sessionctx.Context, sc *s
 	return nil
 }
 
-// shuffleIntermData shuffles the intermediate data of partial workers to corresponded final workers.
-// We only support parallel execution for single-machine, so process of encode and decode can be skipped.
-func (w *HashAggPartialWorker) shuffleIntermData(_ *stmtctx.StatementContext, finalConcurrency int) {
-	groupKeysSlice := make([][]string, finalConcurrency)
-	for groupKey := range w.partialResultsMap {
-		finalWorkerIdx := int(murmur3.Sum32([]byte(groupKey))) % finalConcurrency
-		if groupKeysSlice[finalWorkerIdx] == nil {
-			groupKeysSlice[finalWorkerIdx] = make([]string, 0, len(w.partialResultsMap)/finalConcurrency)
-		}
-		groupKeysSlice[finalWorkerIdx] = append(groupKeysSlice[finalWorkerIdx], groupKey)
-	}
-
-	for i := range groupKeysSlice {
-		if groupKeysSlice[i] == nil {
-			continue
-		}
-		w.outputChs[i] <- &HashAggIntermData{
-			groupKeys:        groupKeysSlice[i],
-			partialResultMap: w.partialResultsMap,
-		}
+func (w *HashAggPartialWorker) shuffleIntermData(finalConcurrency int) {
+	for i := 0; i < finalConcurrency; i++ {
+		w.outputChs[i] <- &w.partialResultsMap[i]
 	}
 }
diff --git a/pkg/executor/aggregate/agg_util.go b/pkg/executor/aggregate/agg_util.go
index 9b7a75edb8..5508fa0b26 100644
--- a/pkg/executor/aggregate/agg_util.go
+++ b/pkg/executor/aggregate/agg_util.go
@@ -28,7 +28,6 @@ import (
 	"github.com/pingcap/tidb/pkg/parser/mysql"
 	"github.com/pingcap/tidb/pkg/parser/terror"
 	"github.com/pingcap/tidb/pkg/sessionctx"
-	"github.com/pingcap/tidb/pkg/sessionctx/stmtctx"
 	"github.com/pingcap/tidb/pkg/util"
 	"github.com/pingcap/tidb/pkg/util/chunk"
 	"github.com/pingcap/tidb/pkg/util/codec"
@@ -37,18 +36,6 @@ import (
 	"go.uber.org/zap"
 )
 
-// getPartialResultBatch fetches a batch of partial results from HashAggIntermData.
-func (d *HashAggIntermData) getPartialResultBatch(_ *stmtctx.StatementContext, prs [][]aggfuncs.PartialResult, _ []aggfuncs.AggFunc, maxChunkSize int) (_ [][]aggfuncs.PartialResult, groupKeys []string, reachEnd bool) {
-	keyStart := d.cursor
-	for ; d.cursor < len(d.groupKeys) && len(prs) < maxChunkSize; d.cursor++ {
-		prs = append(prs, d.partialResultMap[d.groupKeys[d.cursor]])
-	}
-	if d.cursor == len(d.groupKeys) {
-		reachEnd = true
-	}
-	return prs, d.groupKeys[keyStart:d.cursor], reachEnd
-}
-
 func closeBaseExecutor(b *exec.BaseExecutor) {
 	if r := recover(); r != nil {
 		// Release the resource, but throw the panic again and let the top level handle it.
@@ -122,7 +109,7 @@ func GetGroupKey(ctx sessionctx.Context, input *chunk.Chunk, groupKey [][]byte,
 		}
 		expression.PutColumn(buf)
 	}
-	return groupKey, nil
+	return groupKey[:numRows], nil
 }
 
 // HashAggRuntimeStats record the HashAggExec runtime stat
diff --git a/pkg/executor/importer/BUILD.bazel b/pkg/executor/importer/BUILD.bazel
index 07f6fb8622..5e1afc2d83 100644
--- a/pkg/executor/importer/BUILD.bazel
+++ b/pkg/executor/importer/BUILD.bazel
@@ -88,7 +88,7 @@ go_test(
     embed = [":importer"],
     flaky = True,
     race = "on",
-    shard_count = 19,
+    shard_count = 20,
     deps = [
         "//br/pkg/errors",
         "//br/pkg/lightning/backend/encode",
diff --git a/pkg/executor/importer/import.go b/pkg/executor/importer/import.go
index 4d4cc844f8..2dd931e828 100644
--- a/pkg/executor/importer/import.go
+++ b/pkg/executor/importer/import.go
@@ -264,6 +264,8 @@ type LoadDataController struct {
 	dataFiles []*mydump.SourceFileMeta
 	// GlobalSortStore is used to store sorted data when using global sort.
 	GlobalSortStore storage.ExternalStorage
+	// ExecuteNodesCnt is the count of execute nodes.
+	ExecuteNodesCnt int
 }
 
 func getImportantSysVars(sctx sessionctx.Context) map[string]string {
@@ -452,10 +454,11 @@ func NewLoadDataController(plan *Plan, tbl table.Table, astArgs *ASTArgs) (*Load
 	fullTableName := tbl.Meta().Name.String()
 	logger := log.L().With(zap.String("table", fullTableName))
 	c := &LoadDataController{
-		Plan:    plan,
-		ASTArgs: astArgs,
-		Table:   tbl,
-		logger:  logger,
+		Plan:            plan,
+		ASTArgs:         astArgs,
+		Table:           tbl,
+		logger:          logger,
+		ExecuteNodesCnt: 1,
 	}
 	if err := c.checkFieldParams(); err != nil {
 		return nil, err
diff --git a/pkg/executor/importer/table_import.go b/pkg/executor/importer/table_import.go
index ace1df765e..016c688aeb 100644
--- a/pkg/executor/importer/table_import.go
+++ b/pkg/executor/importer/table_import.go
@@ -299,12 +299,11 @@ func (ti *TableImporter) getKVEncoder(chunk *checkpoints.ChunkCheckpoint) (KVEnc
 	return NewTableKVEncoder(cfg, ti)
 }
 
-func (e *LoadDataController) getAdjustedMaxEngineSize() int64 {
+func (e *LoadDataController) calculateSubtaskCnt() int {
 	// we want to split data files into subtask of size close to MaxEngineSize to reduce range overlap,
 	// and evenly distribute them to subtasks.
-	// so we adjust MaxEngineSize to make sure each subtask has a similar amount of data to import.
-	// we calculate subtask count first by round(TotalFileSize / maxEngineSize), then adjust maxEngineSize
-	//
+	// we calculate subtask count first by round(TotalFileSize / maxEngineSize)
+
 	// AllocateEngineIDs is using ceil() to calculate subtask count, engine size might be too small in some case,
 	// such as 501G data, maxEngineSize will be about 250G, so we don't relay on it.
 	// see https://github.com/pingcap/tidb/blob/b4183e1dc9bb01fb81d3aa79ca4b5b74387c6c2a/br/pkg/lightning/mydump/region.go#L109
@@ -315,13 +314,33 @@ func (e *LoadDataController) getAdjustedMaxEngineSize() int64 {
 	// [750, 1250)            2    [375, 625)
 	// [1250, 1750)           3    [416, 583)
 	// [1750, 2250)           4    [437, 562)
-	maxEngineSize := int64(e.MaxEngineSize)
+	var (
+		subtaskCount  float64
+		maxEngineSize = int64(e.MaxEngineSize)
+	)
 	if e.TotalFileSize <= maxEngineSize {
-		return e.TotalFileSize
+		subtaskCount = 1
+	} else {
+		subtaskCount = math.Round(float64(e.TotalFileSize) / float64(e.MaxEngineSize))
+	}
+
+	// for global sort task, since there is no overlap,
+	// we make sure subtask count is a multiple of execute nodes count
+	if e.IsGlobalSort() && e.ExecuteNodesCnt > 0 {
+		subtaskCount = math.Ceil(subtaskCount/float64(e.ExecuteNodesCnt)) * float64(e.ExecuteNodesCnt)
 	}
-	subtaskCount := math.Round(float64(e.TotalFileSize) / float64(maxEngineSize))
-	adjusted := math.Ceil(float64(e.TotalFileSize) / subtaskCount)
-	return int64(adjusted)
+	return int(subtaskCount)
+}
+
+func (e *LoadDataController) getAdjustedMaxEngineSize() int64 {
+	subtaskCount := e.calculateSubtaskCnt()
+	// we adjust MaxEngineSize to make sure each subtask has a similar amount of data to import.
+	return int64(math.Ceil(float64(e.TotalFileSize) / float64(subtaskCount)))
+}
+
+// SetExecuteNodeCnt sets the execute node count.
+func (e *LoadDataController) SetExecuteNodeCnt(cnt int) {
+	e.ExecuteNodesCnt = cnt
 }
 
 // PopulateChunks populates chunks from table regions.
diff --git a/pkg/server/internal/BUILD.bazel b/pkg/server/internal/BUILD.bazel
index 9b7576c20c..05f0b1cf84 100644
--- a/pkg/server/internal/BUILD.bazel
+++ b/pkg/server/internal/BUILD.bazel
@@ -23,7 +23,7 @@ go_test(
     srcs = ["packetio_test.go"],
     embed = [":internal"],
     flaky = True,
-    shard_count = 4,
+    shard_count = 7,
     deps = [
         "//pkg/parser/mysql",
         "//pkg/server/internal/testutil",
diff --git a/pkg/server/internal/packetio.go b/pkg/server/internal/packetio.go
index 2fae04517b..cf1b10d559 100644
--- a/pkg/server/internal/packetio.go
+++ b/pkg/server/internal/packetio.go
@@ -46,7 +46,7 @@ import (
 	"github.com/pingcap/errors"
 	"github.com/pingcap/tidb/pkg/parser/mysql"
 	"github.com/pingcap/tidb/pkg/parser/terror"
-	"github.com/pingcap/tidb/pkg/server/err"
+	server_err "github.com/pingcap/tidb/pkg/server/err"
 	"github.com/pingcap/tidb/pkg/server/internal/util"
 	server_metrics "github.com/pingcap/tidb/pkg/server/metrics"
 	"github.com/pingcap/tidb/pkg/sessionctx/variable"
@@ -60,6 +60,7 @@ type PacketIO struct {
 	bufReadConn      *util.BufferedReadConn
 	bufWriter        *bufio.Writer
 	compressedWriter *compressedWriter
+	compressedReader *compressedReader
 	readTimeout      time.Duration
 	// maxAllowedPacket is the maximum size of one packet in ReadPacket.
 	maxAllowedPacket uint64
@@ -119,8 +120,10 @@ func (p *PacketIO) ResetBufWriter(w io.Writer) {
 // SetCompressionAlgorithm sets the compression algorithm of PacketIO.
 func (p *PacketIO) SetCompressionAlgorithm(ca int) {
 	p.compressionAlgorithm = ca
-	p.compressedWriter = newCompressedWriter(p.bufReadConn, ca)
+	p.compressedWriter = newCompressedWriter(p.bufReadConn, ca, &p.compressedSequence)
 	p.compressedWriter.zstdLevel = p.zstdLevel
+	p.compressedReader = newCompressedReader(p.bufReadConn, ca, &p.compressedSequence)
+	p.compressedReader.zstdLevel = p.zstdLevel
 	p.bufWriter.Flush()
 }
 
@@ -143,56 +146,30 @@ func (p *PacketIO) readOnePacket() ([]byte, error) {
 			return nil, err
 		}
 	}
-	if p.compressionAlgorithm != mysql.CompressionNone {
-		var compressedHeader [7]byte
-		if _, err := io.ReadFull(p.bufReadConn, compressedHeader[:]); err != nil {
+	if p.compressionAlgorithm == mysql.CompressionNone {
+		if _, err := io.ReadFull(r, header[:]); err != nil {
 			return nil, errors.Trace(err)
 		}
-		compressedSequence := compressedHeader[3]
-		if compressedSequence != p.compressedSequence {
-			return nil, err.ErrInvalidSequence.GenWithStack(
-				"invalid compressed sequence %d != %d", compressedSequence, p.compressedSequence)
-		}
-		p.compressedSequence++
-		p.compressedWriter.compressedSequence = p.compressedSequence
-		uncompressedLength := int(uint32(compressedHeader[4]) | uint32(compressedHeader[5])<<8 | uint32(compressedHeader[6])<<16)
-
-		if uncompressedLength > 0 {
-			switch p.compressionAlgorithm {
-			case mysql.CompressionZlib:
-				var err error
-				r, err = zlib.NewReader(p.bufReadConn)
-				if err != nil {
-					return nil, errors.Trace(err)
-				}
-			case mysql.CompressionZstd:
-				zstdReader, err := zstd.NewReader(p.bufReadConn, zstd.WithDecoderConcurrency(1))
-				if err != nil {
-					return nil, errors.Trace(err)
-				}
-				r = zstdReader.IOReadCloser()
-			default:
-				return nil, errors.New("Unknown compression algorithm")
-			}
+	} else {
+		if _, err := io.ReadFull(p.compressedReader, header[:]); err != nil {
+			return nil, errors.Trace(err)
 		}
 	}
-	if _, err := io.ReadFull(r, header[:]); err != nil {
-		return nil, errors.Trace(err)
-	}
 
+	length := int(uint32(header[0]) | uint32(header[1])<<8 | uint32(header[2])<<16)
 	sequence := header[3]
+
 	if sequence != p.sequence {
-		return nil, err.ErrInvalidSequence.GenWithStack("invalid sequence %d != %d", sequence, p.sequence)
+		return nil, server_err.ErrInvalidSequence.GenWithStack(
+			"invalid sequence, received %d while expecting %d", sequence, p.sequence)
 	}
 
 	p.sequence++
 
-	length := int(uint32(header[0]) | uint32(header[1])<<8 | uint32(header[2])<<16)
-
 	// Accumulated payload length exceeds the limit.
 	if p.accumulatedLength += uint64(length); p.accumulatedLength > p.maxAllowedPacket {
-		terror.Log(err.ErrNetPacketTooLarge)
-		return nil, err.ErrNetPacketTooLarge
+		terror.Log(server_err.ErrNetPacketTooLarge)
+		return nil, server_err.ErrNetPacketTooLarge
 	}
 
 	data := make([]byte, length)
@@ -201,8 +178,14 @@ func (p *PacketIO) readOnePacket() ([]byte, error) {
 			return nil, err
 		}
 	}
-	if _, err := io.ReadFull(r, data); err != nil {
-		return nil, errors.Trace(err)
+	if p.compressionAlgorithm == mysql.CompressionNone {
+		if _, err := io.ReadFull(r, data); err != nil {
+			return nil, errors.Trace(err)
+		}
+	} else {
+		if _, err := io.ReadFull(p.compressedReader, data); err != nil {
+			return nil, errors.Trace(err)
+		}
 	}
 	err := r.Close()
 	if err != nil {
@@ -256,11 +239,7 @@ func (p *PacketIO) ReadPacket() ([]byte, error) {
 func (p *PacketIO) WritePacket(data []byte) error {
 	length := len(data) - 4
 	server_metrics.WritePacketBytes.Add(float64(len(data)))
-
 	maxPayloadLen := mysql.MaxPayloadLen
-	if p.compressionAlgorithm != mysql.CompressionNone {
-		maxPayloadLen -= 4
-	}
 
 	for length >= maxPayloadLen {
 		data[3] = p.sequence
@@ -324,15 +303,19 @@ func (p *PacketIO) Flush() error {
 	if err != nil {
 		return errors.Trace(err)
 	}
+
+	if p.compressionAlgorithm != mysql.CompressionNone {
+		p.sequence = p.compressedSequence
+	}
 	return err
 }
 
-func newCompressedWriter(w io.Writer, ca int) *compressedWriter {
+func newCompressedWriter(w io.Writer, ca int, seq *uint8) *compressedWriter {
 	return &compressedWriter{
 		w,
 		new(bytes.Buffer),
+		seq,
 		ca,
-		0,
 		3,
 	}
 }
@@ -340,8 +323,8 @@ func newCompressedWriter(w io.Writer, ca int) *compressedWriter {
 type compressedWriter struct {
 	w                    io.Writer
 	buf                  *bytes.Buffer
+	compressedSequence   *uint8
 	compressionAlgorithm int
-	compressedSequence   uint8
 	zstdLevel            zstd.EncoderLevel
 }
 
@@ -421,7 +404,7 @@ func (cw *compressedWriter) Flush() error {
 	compressedHeader[0] = byte(compressedLength)
 	compressedHeader[1] = byte(compressedLength >> 8)
 	compressedHeader[2] = byte(compressedLength >> 16)
-	compressedHeader[3] = cw.compressedSequence
+	compressedHeader[3] = *cw.compressedSequence
 	compressedHeader[4] = byte(uncompressedLength)
 	compressedHeader[5] = byte(uncompressedLength >> 8)
 	compressedHeader[6] = byte(uncompressedLength >> 16)
@@ -429,7 +412,7 @@ func (cw *compressedWriter) Flush() error {
 	if err != nil {
 		return errors.Trace(err)
 	}
-	cw.compressedSequence++
+	*cw.compressedSequence++
 
 	if len(data) > minCompressLength {
 		_, err = compressedPacket.Write(payload.Bytes())
@@ -446,3 +429,89 @@ func (cw *compressedWriter) Flush() error {
 	}
 	return nil
 }
+
+func newCompressedReader(r io.Reader, ca int, seq *uint8) *compressedReader {
+	return &compressedReader{
+		r,
+		seq,
+		nil,
+		ca,
+		3,
+		0,
+	}
+}
+
+type compressedReader struct {
+	r                    io.Reader
+	compressedSequence   *uint8
+	data                 []byte
+	compressionAlgorithm int
+	zstdLevel            zstd.EncoderLevel
+	pos                  uint64
+}
+
+func (cr *compressedReader) Read(data []byte) (n int, err error) {
+	if cr.data == nil {
+		var compressedHeader [7]byte
+		if _, err = io.ReadFull(cr.r, compressedHeader[:]); err != nil {
+			return
+		}
+
+		compressedLength := int(uint32(compressedHeader[0]) | uint32(compressedHeader[1])<<8 | uint32(compressedHeader[2])<<16)
+		compressedSequence := compressedHeader[3]
+		uncompressedLength := int(uint32(compressedHeader[4]) | uint32(compressedHeader[5])<<8 | uint32(compressedHeader[6])<<16)
+
+		if compressedSequence != *cr.compressedSequence {
+			return n, server_err.ErrInvalidSequence.GenWithStack(
+				"invalid compressed sequence, received %d while expecting %d", compressedSequence, cr.compressedSequence)
+		}
+		*cr.compressedSequence++
+
+		r := io.NopCloser(cr.r)
+		if uncompressedLength > 0 {
+			switch cr.compressionAlgorithm {
+			case mysql.CompressionZlib:
+				var err error
+				lr := io.LimitReader(cr.r, int64(compressedLength))
+				r, err = zlib.NewReader(lr)
+				if err != nil {
+					return n, errors.Trace(err)
+				}
+			case mysql.CompressionZstd:
+				zstdReader, err := zstd.NewReader(cr.r, zstd.WithDecoderConcurrency(1))
+				if err != nil {
+					return n, errors.Trace(err)
+				}
+				r = zstdReader.IOReadCloser()
+			default:
+				return n, errors.New("Unknown compression algorithm")
+			}
+			cr.data = make([]byte, uncompressedLength)
+			if _, err := io.ReadFull(r, cr.data); err != nil {
+				return n, errors.Trace(err)
+			}
+			n = copy(data, cr.data)
+		} else {
+			cr.data = make([]byte, compressedLength)
+			if _, err := io.ReadFull(r, cr.data); err != nil {
+				return n, errors.Trace(err)
+			}
+			n = copy(data, cr.data)
+		}
+	} else {
+		if cr.pos > uint64(len(cr.data)) {
+			return n, io.EOF
+		}
+		n = copy(data, cr.data[cr.pos:])
+	}
+	cr.pos += uint64(n)
+	if cr.pos == uint64(len(cr.data)) {
+		cr.pos = 0
+		cr.data = nil
+	}
+	return
+}
+
+func (*compressedReader) Close() error {
+	return nil
+}


Test patch:
diff --git a/pkg/ddl/tests/fail/fail_db_test.go b/pkg/ddl/tests/fail/fail_db_test.go
index da4a9018c1..36ae8f8f31 100644
--- a/pkg/ddl/tests/fail/fail_db_test.go
+++ b/pkg/ddl/tests/fail/fail_db_test.go
@@ -126,6 +126,11 @@ func TestHalfwayCancelOperations(t *testing.T) {
 	tk.MustExec("use cancel_job_db")
 	tk.MustExec("select * from tx")
 	// test for exchanging partition
+	limit := variable.GetDDLErrorCountLimit()
+	variable.SetDDLErrorCountLimit(3)
+	defer func() {
+		variable.SetDDLErrorCountLimit(limit)
+	}()
 	require.NoError(t, failpoint.Enable("github.com/pingcap/tidb/pkg/ddl/exchangePartitionErr", `return(true)`))
 	defer func() {
 		require.NoError(t, failpoint.Disable("github.com/pingcap/tidb/pkg/ddl/exchangePartitionErr"))
diff --git a/pkg/executor/importer/import_test.go b/pkg/executor/importer/import_test.go
index 4cb202c735..35860ae3f3 100644
--- a/pkg/executor/importer/import_test.go
+++ b/pkg/executor/importer/import_test.go
@@ -19,6 +19,7 @@ import (
 	"fmt"
 	"net/url"
 	"os"
+	"os/user"
 	"path"
 	"path/filepath"
 	"runtime"
@@ -308,6 +309,11 @@ func TestGetBackendWorkerConcurrency(t *testing.T) {
 }
 
 func TestSupportedSuffixForServerDisk(t *testing.T) {
+	username, err := user.Current()
+	require.NoError(t, err)
+	if username.Name == "root" {
+		t.Skip("it cannot run as root")
+	}
 	tempDir := t.TempDir()
 	ctx := context.Background()
 
@@ -363,7 +369,7 @@ func TestSupportedSuffixForServerDisk(t *testing.T) {
 	require.ErrorContains(t, err2, "URI of data source is invalid")
 	// non-exist parent directory
 	c.Path = "/path/to/non/exists/file.csv"
-	err := c.InitDataFiles(ctx)
+	err = c.InitDataFiles(ctx)
 	require.ErrorIs(t, err, exeerrors.ErrLoadDataInvalidURI)
 	require.ErrorContains(t, err, "no such file or directory")
 	// without permission to parent dir
diff --git a/pkg/executor/importer/table_import_test.go b/pkg/executor/importer/table_import_test.go
index 3fc5688a7d..e72b6c009f 100644
--- a/pkg/executor/importer/table_import_test.go
+++ b/pkg/executor/importer/table_import_test.go
@@ -80,29 +80,92 @@ func TestPrepareSortDir(t *testing.T) {
 	require.Nil(t, info)
 }
 
+func TestCalculateSubtaskCnt(t *testing.T) {
+	tests := []struct {
+		totalSize       int64
+		maxEngineSize   config.ByteSize
+		executeNodeCnt  int
+		cloudStorageURL string
+		want            int
+	}{
+		{1, 500, 0, "", 1},
+		{499, 500, 1, "", 1},
+		{500, 500, 2, "", 1},
+		{749, 500, 3, "", 1},
+		{750, 500, 4, "", 2},
+		{1249, 500, 5, "", 2},
+		{1250, 500, 6, "", 3},
+		{100, 30, 7, "", 3},
+
+		{1, 500, 0, "url", 1},
+		{499, 500, 1, "url", 1},
+		{500, 500, 2, "url", 2},
+		{749, 500, 3, "url", 3},
+		{750, 500, 4, "url", 4},
+		{1249, 500, 5, "url", 5},
+		{1250, 500, 6, "url", 6},
+		{100, 30, 2, "url", 4},
+		{400, 99, 3, "url", 6},
+		{500, 100, 5, "url", 5},
+		{500, 200, 5, "url", 5},
+	}
+	for _, tt := range tests {
+		t.Run(fmt.Sprintf("%d/%d", tt.totalSize, tt.maxEngineSize), func(t *testing.T) {
+			e := &LoadDataController{
+				Plan: &Plan{
+					MaxEngineSize:   tt.maxEngineSize,
+					TotalFileSize:   tt.totalSize,
+					CloudStorageURI: tt.cloudStorageURL,
+				},
+				ExecuteNodesCnt: tt.executeNodeCnt,
+			}
+			if got := e.calculateSubtaskCnt(); got != tt.want {
+				t.Errorf("calculateSubtaskCnt() = %v, want %v", got, tt.want)
+			}
+		})
+	}
+}
+
 func TestLoadDataControllerGetAdjustedMaxEngineSize(t *testing.T) {
 	tests := []struct {
-		totalSize     int64
-		maxEngineSize config.ByteSize
-		want          int64
+		totalSize       int64
+		maxEngineSize   config.ByteSize
+		executeNodeCnt  int
+		cloudStorageURL string
+		want            int64
 	}{
-		{1, 500, 1},
-		{499, 500, 499},
-		{500, 500, 500},
-		{749, 500, 749},
-		{750, 500, 375},
-		{1249, 500, 625},
-		{1250, 500, 417},
+		{1, 500, 0, "", 1},
+		{499, 500, 1, "", 499},
+		{500, 500, 2, "", 500},
+		{749, 500, 3, "", 749},
+		{750, 500, 4, "", 375},
+		{1249, 500, 5, "", 625},
+		{1250, 500, 6, "", 417},
 		// ceil(100/3)
-		{100, 30, 34},
+		{100, 30, 7, "", 34},
+
+		{1, 500, 0, "url", 1},
+		{499, 500, 1, "url", 499},
+		{500, 500, 2, "url", 250},
+		{749, 500, 3, "url", 250},
+		{750, 500, 4, "url", 188},
+		{1249, 500, 5, "url", 250},
+		{1250, 500, 6, "url", 209},
+		{100, 30, 2, "url", 25},
+		{400, 99, 3, "url", 67},
+		{500, 100, 5, "url", 100},
+		{500, 200, 5, "url", 100},
+		{500, 100, 1, "url", 100},
 	}
 	for _, tt := range tests {
 		t.Run(fmt.Sprintf("%d/%d", tt.totalSize, tt.maxEngineSize), func(t *testing.T) {
 			e := &LoadDataController{
 				Plan: &Plan{
-					MaxEngineSize: tt.maxEngineSize,
-					TotalFileSize: tt.totalSize,
+					MaxEngineSize:   tt.maxEngineSize,
+					TotalFileSize:   tt.totalSize,
+					CloudStorageURI: tt.cloudStorageURL,
 				},
+				ExecuteNodesCnt: tt.executeNodeCnt,
 			}
 			if got := e.getAdjustedMaxEngineSize(); got != tt.want {
 				t.Errorf("getAdjustedMaxEngineSize() = %v, want %v", got, tt.want)
diff --git a/pkg/executor/test/aggregate/BUILD.bazel b/pkg/executor/test/aggregate/BUILD.bazel
index 575f1ccf53..efa11544af 100644
--- a/pkg/executor/test/aggregate/BUILD.bazel
+++ b/pkg/executor/test/aggregate/BUILD.bazel
@@ -9,7 +9,7 @@ go_test(
     ],
     data = glob(["testdata/**"]),
     flaky = True,
-    shard_count = 5,
+    shard_count = 6,
     deps = [
         "//pkg/config",
         "//pkg/executor/aggregate",
diff --git a/pkg/executor/test/aggregate/aggregate_test.go b/pkg/executor/test/aggregate/aggregate_test.go
index 83fcbdd73d..8f37c62089 100644
--- a/pkg/executor/test/aggregate/aggregate_test.go
+++ b/pkg/executor/test/aggregate/aggregate_test.go
@@ -15,6 +15,7 @@
 package aggregate
 
 import (
+	"bytes"
 	"context"
 	"fmt"
 	"math"
@@ -332,3 +333,125 @@ func TestRandomPanicConsume(t *testing.T) {
 		}
 	}
 }
+
+func checkResults(actualRes [][]interface{}, expectedRes map[string]string) bool {
+	if len(actualRes) != len(expectedRes) {
+		return false
+	}
+
+	var key string
+	var expectVal string
+	var actualVal string
+	var ok bool
+	for _, row := range actualRes {
+		if len(row) != 2 {
+			return false
+		}
+
+		key, ok = row[0].(string)
+		if !ok {
+			return false
+		}
+
+		expectVal, ok = expectedRes[key]
+		if !ok {
+			return false
+		}
+
+		actualVal, ok = row[1].(string)
+		if !ok {
+			return false
+		}
+
+		if expectVal != actualVal {
+			return false
+		}
+	}
+	return true
+}
+
+func genListPartition(begin, end int) string {
+	buf := &bytes.Buffer{}
+	buf.WriteString("(")
+	for i := begin; i < end-1; i++ {
+		buf.WriteString(fmt.Sprintf("%v, ", i))
+	}
+	buf.WriteString(fmt.Sprintf("%v)", end-1))
+	return buf.String()
+}
+
+func TestParallelHashAgg(t *testing.T) {
+	store, _ := testkit.CreateMockStoreAndDomain(t)
+	tk := testkit.NewTestKit(t, store)
+
+	tk.MustExec("use test")
+	tk.MustExec("drop table if exists test.parallel_hash_agg;")
+	tk.MustExec("create table test.parallel_hash_agg(k varchar(30), v int);")
+	for i := 0; i < 20; i++ {
+		tk.MustExec("insert into test.parallel_hash_agg (k, v) values ('aa', 1), ('AA', 1), ('aA', 1), ('Aa', 1), ('bb', 1), ('BB', 1), ('bB', 1), ('Bb', 1), ('cc', 1), ('CC', 1), ('cC', 1), ('Cc', 1), ('dd', 1), ('DD', 1), ('dD', 1), ('Dd', 1), ('ee', 1), ('EE', 1), ('eE', 1), ('Ee', 1);")
+	}
+
+	tk.MustExec("set @@tidb_max_chunk_size=32;")
+
+	expectedResult := make(map[string]string)
+	expectedResult["dd"] = "20"
+	expectedResult["AA"] = "20"
+	expectedResult["cc"] = "20"
+	expectedResult["eE"] = "20"
+	expectedResult["bb"] = "20"
+	expectedResult["Cc"] = "20"
+	expectedResult["EE"] = "20"
+	expectedResult["Aa"] = "20"
+	expectedResult["ee"] = "20"
+	expectedResult["Bb"] = "20"
+	expectedResult["dD"] = "20"
+	expectedResult["aa"] = "20"
+	expectedResult["cC"] = "20"
+	expectedResult["DD"] = "20"
+	expectedResult["BB"] = "20"
+	expectedResult["Dd"] = "20"
+	expectedResult["CC"] = "20"
+	expectedResult["bB"] = "20"
+	expectedResult["aA"] = "20"
+	expectedResult["Ee"] = "20"
+	res := tk.MustQuery("select k, sum(v) from parallel_hash_agg group by k;")
+	tk.RequireEqual(true, checkResults(res.Rows(), expectedResult))
+
+	tk.MustExec("create database list_partition_agg")
+	tk.MustExec("use list_partition_agg")
+	tk.MustExec("drop table if exists tlist")
+	tk.MustExec(`set tidb_enable_list_partition = 1`)
+	tk.MustExec(`create table tlist (a int, b int) partition by list(a) (` +
+		` partition p0 values in ` + genListPartition(0, 20) +
+		`, partition p1 values in ` + genListPartition(20, 40) +
+		`, partition p2 values in ` + genListPartition(40, 60) +
+		`, partition p3 values in ` + genListPartition(60, 80) +
+		`, partition p4 values in ` + genListPartition(80, 100) + `)`)
+	tk.MustExec(`create table tnormal (a int, b int)`)
+
+	vals := ""
+	for i := 0; i < 100; i++ {
+		if vals != "" {
+			vals += ", "
+		}
+		vals += fmt.Sprintf("(%v, %v)", rand.Intn(100), rand.Intn(100))
+	}
+	tk.MustExec(`insert into tnormal values ` + vals)
+	tk.MustExec(`insert into tlist values ` + vals)
+
+	for _, aggFunc := range []string{"min", "max", "sum", "count"} {
+		c1, c2 := "a", "b"
+		for i := 0; i < 2; i++ {
+			rs := tk.MustQuery(fmt.Sprintf(`select %v, %v(%v) from tnormal group by %v`, c1, aggFunc, c2, c1)).Sort()
+
+			tk.MustExec("set @@tidb_partition_prune_mode = 'dynamic'")
+			rsDynamic := tk.MustQuery(fmt.Sprintf(`select %v, %v(%v) from tlist group by %v`, c1, aggFunc, c2, c1)).Sort()
+
+			tk.MustExec("set @@tidb_partition_prune_mode = 'static'")
+			rsStatic := tk.MustQuery(fmt.Sprintf(`select %v, %v(%v) from tlist group by %v`, c1, aggFunc, c2, c1)).Sort()
+
+			rs.Check(rsDynamic.Rows())
+			rs.Check(rsStatic.Rows())
+		}
+	}
+}
diff --git a/pkg/server/internal/packetio_test.go b/pkg/server/internal/packetio_test.go
index 0953f573e0..7e4230dcdb 100644
--- a/pkg/server/internal/packetio_test.go
+++ b/pkg/server/internal/packetio_test.go
@@ -62,6 +62,32 @@ func TestPacketIOWrite(t *testing.T) {
 	require.Equal(t, byte(0), res[3])
 }
 
+func TestPacketIOWriteCompressed(t *testing.T) {
+	var testdata, outBuffer bytes.Buffer
+
+	seq := uint8(0)
+	pkt := &PacketIO{
+		bufWriter:            bufio.NewWriter(&outBuffer),
+		compressionAlgorithm: mysql.CompressionZlib,
+		compressedWriter:     newCompressedWriter(&testdata, mysql.CompressionZlib, &seq),
+	}
+
+	payload := bytes.Repeat([]byte{'A'}, 16*1024*1024)
+	err := pkt.WritePacket(payload)
+	require.NoError(t, err)
+
+	err = pkt.Flush()
+	require.NoError(t, err)
+
+	compressedLength := []byte{0x18, 0x4, 0x0} // 1048 bytes
+	packetNr := []byte{0x0}
+	uncompressedLength := []byte{0x0, 0x0, 0x10} // 1048576 bytes
+
+	require.Equal(t, compressedLength, testdata.Bytes()[:3])
+	require.Equal(t, packetNr, testdata.Bytes()[3:4])
+	require.Equal(t, uncompressedLength, testdata.Bytes()[4:7])
+}
+
 func TestPacketIORead(t *testing.T) {
 	t.Run("uncompressed", func(t *testing.T) {
 		var inBuffer bytes.Buffer
@@ -215,8 +241,9 @@ func TestPacketIORead(t *testing.T) {
 func TestCompressedWriterShort(t *testing.T) {
 	var testdata bytes.Buffer
 	payload := []byte("test_short")
+	seq := uint8(0)
 
-	cw := newCompressedWriter(&testdata, mysql.CompressionZlib)
+	cw := newCompressedWriter(&testdata, mysql.CompressionZlib, &seq)
 	cw.Write(payload)
 	cw.Flush()
 
@@ -236,8 +263,9 @@ func TestCompressedWriterLong(t *testing.T) {
 	t.Run("zlib", func(t *testing.T) {
 		var testdata, decoded bytes.Buffer
 		payload := []byte("test_zlib test_zlib test_zlib test_zlib test_zlib test_zlib test_zlib")
+		seq := uint8(0)
 
-		cw := newCompressedWriter(&testdata, mysql.CompressionZlib)
+		cw := newCompressedWriter(&testdata, mysql.CompressionZlib, &seq)
 		cw.Write(payload)
 		cw.Flush()
 
@@ -262,8 +290,9 @@ func TestCompressedWriterLong(t *testing.T) {
 	t.Run("zstd", func(t *testing.T) {
 		var testdata bytes.Buffer
 		payload := []byte("test_zstd test_zstd test_zstd test_zstd test_zstd test_zstd test_zstd")
+		seq := uint8(0)
 
-		cw := newCompressedWriter(&testdata, mysql.CompressionZstd)
+		cw := newCompressedWriter(&testdata, mysql.CompressionZstd, &seq)
 		cw.Write(payload)
 		cw.Flush()
 
@@ -286,3 +315,112 @@ func TestCompressedWriterLong(t *testing.T) {
 		require.Equal(t, payload, decoded)
 	})
 }
+
+// TestCompressedReaderShort test a compressed protocol packet that has an uncompressed
+// length of 0, which means the actual payload isn't compressed.
+func TestCompressedReaderShort(t *testing.T) {
+	// payload: 7 bytes compressed header, 37 bytes uncompressed payload
+	payload := []byte{0x25, 0x0, 0x0, 0x0, 0x0, 0x0, 0x0, 0x21, 0x0, 0x0, 0x0, 0x3,
+		0x73, 0x65, 0x6c, 0x65, 0x63, 0x74, 0x20, 0x40, 0x40, 0x76, 0x65, 0x72,
+		0x73, 0x69, 0x6f, 0x6e, 0x5f, 0x63, 0x6f, 0x6d, 0x6d, 0x65, 0x6e, 0x74,
+		0x20, 0x6c, 0x69, 0x6d, 0x69, 0x74, 0x20, 0x31}
+	r := bytes.NewReader(payload)
+	seq := uint8(0)
+	cr := newCompressedReader(r, mysql.CompressionZlib, &seq)
+
+	// Read 4 byte header from the payload. This is the regular packet header,
+	// not the compressed header.
+	header := make([]byte, 4)
+	_, err := cr.Read(header)
+	require.NoError(t, err)
+	require.Equal(t, []byte{0x21, 0x0, 0x0, 0x0}, header)
+
+	length := int(uint32(header[0]) | uint32(header[1])<<8 | uint32(header[2])<<16)
+	sequence := header[3]
+	require.Equal(t, 33, length)
+	require.Equal(t, uint8(0), sequence)
+
+	data := make([]byte, length)
+	_, err = cr.Read(data)
+	require.NoError(t, err)
+	expected := []byte{0x3, 0x73, 0x65, 0x6c, 0x65, 0x63, 0x74, 0x20, 0x40, 0x40,
+		0x76, 0x65, 0x72, 0x73, 0x69, 0x6f, 0x6e, 0x5f, 0x63, 0x6f, 0x6d, 0x6d,
+		0x65, 0x6e, 0x74, 0x20, 0x6c, 0x69, 0x6d, 0x69, 0x74, 0x20, 0x31}
+	require.Equal(t, expected, data)
+}
+
+func TestCompressedReaderLong(t *testing.T) {
+	t.Run("zlib", func(t *testing.T) {
+		payload := []byte{0x19, 0x0, 0x0, 0x0, 0x9c, 0x0, 0x0, 0x78, 0x5e, 0x9b,
+			0xc1, 0xc0, 0xc0, 0xc0, 0x1c, 0xec, 0xea, 0xe3, 0xea, 0x1c, 0xa2,
+			0xa0, 0xe4, 0x38, 0xa8, 0x80, 0x12, 0x0, 0xbe, 0xe6, 0x26, 0xce}
+		r := bytes.NewReader(payload)
+		seq := uint8(0)
+		cr := newCompressedReader(r, mysql.CompressionZlib, &seq)
+		header := make([]byte, 4)
+		_, err := cr.Read(header)
+		require.NoError(t, err)
+		require.Equal(t, []byte{0x98, 0x0, 0x0, 0x0}, header)
+
+		length := int(uint32(header[0]) | uint32(header[1])<<8 | uint32(header[2])<<16)
+		sequence := header[3]
+		require.Equal(t, 152, length)
+		require.Equal(t, uint8(0), sequence)
+
+		data := make([]byte, length)
+		_, err = cr.Read(data)
+		require.NoError(t, err)
+		expected := []byte{0x3, 0x53, 0x45, 0x4c, 0x45, 0x43, 0x54, 0x20, 0x22,
+			0x41, 0x41, 0x41, 0x41, 0x41, 0x41, 0x41, 0x41, 0x41, 0x41, 0x41,
+			0x41, 0x41, 0x41, 0x41, 0x41, 0x41, 0x41, 0x41, 0x41, 0x41, 0x41,
+			0x41, 0x41, 0x41, 0x41, 0x41, 0x41, 0x41, 0x41, 0x41, 0x41, 0x41,
+			0x41, 0x41, 0x41, 0x41, 0x41, 0x41, 0x41, 0x41, 0x41, 0x41, 0x41,
+			0x41, 0x41, 0x41, 0x41, 0x41, 0x41, 0x41, 0x41, 0x41, 0x41, 0x41,
+			0x41, 0x41, 0x41, 0x41, 0x41, 0x41, 0x41, 0x41, 0x41, 0x41, 0x41,
+			0x41, 0x41, 0x41, 0x41, 0x41, 0x41, 0x41, 0x41, 0x41, 0x41, 0x41,
+			0x41, 0x41, 0x41, 0x41, 0x41, 0x41, 0x41, 0x41, 0x41, 0x41, 0x41,
+			0x41, 0x41, 0x41, 0x41, 0x41, 0x41, 0x41, 0x41, 0x41, 0x41, 0x41,
+			0x41, 0x41, 0x41, 0x41, 0x41, 0x41, 0x41, 0x41, 0x41, 0x41, 0x41,
+			0x41, 0x41, 0x41, 0x41, 0x41, 0x41, 0x41, 0x41, 0x41, 0x41, 0x41,
+			0x41, 0x41, 0x41, 0x41, 0x41, 0x41, 0x41, 0x41, 0x41, 0x41, 0x41,
+			0x41, 0x41, 0x41, 0x41, 0x41, 0x41, 0x41, 0x41, 0x41, 0x41, 0x22}
+		require.Equal(t, expected, data)
+	})
+	t.Run("zstd", func(t *testing.T) {
+		payload := []byte{0x1f, 0x0, 0x0, 0x0, 0x9c, 0x0, 0x0, 0x28, 0xb5, 0x2f, 0xfd,
+			0x20, 0x9c, 0xb5, 0x0, 0x0, 0x78, 0x98, 0x0, 0x0, 0x0, 0x3, 0x53,
+			0x45, 0x4c, 0x45, 0x43, 0x54, 0x20, 0x22, 0x41, 0x22, 0x1, 0x0, 0xa,
+			0xa, 0x28, 0x1}
+		r := bytes.NewReader(payload)
+		seq := uint8(0)
+		cr := newCompressedReader(r, mysql.CompressionZstd, &seq)
+		header := make([]byte, 4)
+		_, err := cr.Read(header)
+		require.NoError(t, err)
+		require.Equal(t, []byte{0x98, 0x0, 0x0, 0x0}, header)
+
+		length := int(uint32(header[0]) | uint32(header[1])<<8 | uint32(header[2])<<16)
+		sequence := header[3]
+		require.Equal(t, 152, length)
+		require.Equal(t, uint8(0), sequence)
+
+		data := make([]byte, length)
+		_, err = cr.Read(data)
+		require.NoError(t, err)
+		expected := []byte{0x3, 0x53, 0x45, 0x4c, 0x45, 0x43, 0x54, 0x20, 0x22,
+			0x41, 0x41, 0x41, 0x41, 0x41, 0x41, 0x41, 0x41, 0x41, 0x41, 0x41,
+			0x41, 0x41, 0x41, 0x41, 0x41, 0x41, 0x41, 0x41, 0x41, 0x41, 0x41,
+			0x41, 0x41, 0x41, 0x41, 0x41, 0x41, 0x41, 0x41, 0x41, 0x41, 0x41,
+			0x41, 0x41, 0x41, 0x41, 0x41, 0x41, 0x41, 0x41, 0x41, 0x41, 0x41,
+			0x41, 0x41, 0x41, 0x41, 0x41, 0x41, 0x41, 0x41, 0x41, 0x41, 0x41,
+			0x41, 0x41, 0x41, 0x41, 0x41, 0x41, 0x41, 0x41, 0x41, 0x41, 0x41,
+			0x41, 0x41, 0x41, 0x41, 0x41, 0x41, 0x41, 0x41, 0x41, 0x41, 0x41,
+			0x41, 0x41, 0x41, 0x41, 0x41, 0x41, 0x41, 0x41, 0x41, 0x41, 0x41,
+			0x41, 0x41, 0x41, 0x41, 0x41, 0x41, 0x41, 0x41, 0x41, 0x41, 0x41,
+			0x41, 0x41, 0x41, 0x41, 0x41, 0x41, 0x41, 0x41, 0x41, 0x41, 0x41,
+			0x41, 0x41, 0x41, 0x41, 0x41, 0x41, 0x41, 0x41, 0x41, 0x41, 0x41,
+			0x41, 0x41, 0x41, 0x41, 0x41, 0x41, 0x41, 0x41, 0x41, 0x41, 0x41,
+			0x41, 0x41, 0x41, 0x41, 0x41, 0x41, 0x41, 0x41, 0x41, 0x41, 0x22}
+		require.Equal(t, expected, data)
+	})
+}