*: add a session variable to control max chunk size (pingcap#5233)

executor/distsql.go

@@ -551,7 +551,7 @@ func (e *IndexLookUpExecutor) startIndexWorker(goCtx goctx.Context, kvRanges []k
 		finished:     e.finished,
 		resultCh:     e.resultCh,
 		keepOrder:    e.keepOrder,
-		batchSize:    chunk.MaxCapacity,
+		batchSize:    e.ctx.GetSessionVars().MaxChunkSize,
 		maxBatchSize: e.ctx.GetSessionVars().IndexLookupSize,
 	}
 	if worker.batchSize > worker.maxBatchSize {
@@ -795,7 +795,7 @@ func (e *IndexLookUpExecutor) NextChunk(chk *chunk.Chunk) error {
 		for resultTask.cursor < len(resultTask.rows) {
 			chk.AppendRow(0, resultTask.rows[resultTask.cursor])
 			resultTask.cursor++
-			if chk.NumRows() >= chunk.MaxCapacity {
+			if chk.NumRows() >= e.ctx.GetSessionVars().MaxChunkSize {
 				return nil
 			}
 		}

executor/sort.go

@@ -174,7 +174,7 @@ func (e *SortExec) NextChunk(chk *chunk.Chunk) error {
 		}
 		e.fetched = true
 	}
-	for chk.NumRows() < chunk.MaxCapacity {
+	for chk.NumRows() < e.ctx.GetSessionVars().MaxChunkSize {
 		if e.Idx >= len(e.rowPointers) {
 			return nil
 		}

sessionctx/variable/session.go

@@ -226,6 +226,9 @@ type SessionVars struct {
 
 	// MaxRowCountForINLJ defines max row count that the outer table of index nested loop join could be without force hint.
 	MaxRowCountForINLJ int
+
+	// MaxChunkSize defines max row count of a Chunk during query execution.
+	MaxChunkSize int
 }
 
 // NewSessionVars creates a session vars object.
@@ -249,6 +252,7 @@ func NewSessionVars() *SessionVars {
 		IndexSerialScanConcurrency: DefIndexSerialScanConcurrency,
 		DistSQLScanConcurrency:     DefDistSQLScanConcurrency,
 		MaxRowCountForINLJ:         DefMaxRowCountForINLJ,
+		MaxChunkSize:               DefMaxChunkSize,
 	}
 }
 

sessionctx/variable/sysvar.go

@@ -617,6 +617,7 @@ var defaultSysVars = []*SysVar{
 	{ScopeSession, TiDBBatchInsert, boolToIntStr(DefBatchInsert)},
 	{ScopeSession, TiDBBatchDelete, boolToIntStr(DefBatchDelete)},
 	{ScopeSession, TiDBCurrentTS, strconv.Itoa(DefCurretTS)},
+	{ScopeSession, TiDBMaxChunkSize, strconv.Itoa(DefMaxChunkSize)},
 }
 
 // SetNamesVariables is the system variable names related to set names statements.

sessionctx/variable/tidb_vars.go

@@ -100,6 +100,9 @@ const (
 	// It controls the max row count of outer table when do index nested loop join without hint.
 	// After the row count of the inner table is accurate, this variable will be removed.
 	TiDBMaxRowCountForINLJ = "tidb_max_row_count_for_inlj"
+
+	// tidb_max_chunk_capacity is used to control the max chunk size during query execution.
+	TiDBMaxChunkSize = "tidb_max_chunk_size"
 )
 
 // Default TiDB system variable values.
@@ -117,4 +120,5 @@ const (
 	DefBatchInsert                = false
 	DefBatchDelete                = false
 	DefCurretTS                   = 0
+	DefMaxChunkSize               = 1024
 )

sessionctx/varsutil/varsutil.go

@@ -148,6 +148,8 @@ func SetSessionSystemVar(vars *variable.SessionVars, name string, value types.Da
 		vars.MaxRowCountForINLJ = tidbOptPositiveInt(sVal, variable.DefMaxRowCountForINLJ)
 	case variable.TiDBCurrentTS:
 		return variable.ErrReadOnly
+	case variable.TiDBMaxChunkSize:
+		vars.MaxChunkSize = tidbOptPositiveInt(sVal, variable.DefMaxChunkSize)
 	}
 	vars.Systems[name] = sVal
 	return nil

sessionctx/varsutil/varsutil_test.go

@@ -161,6 +161,10 @@ func (s *testVarsutilSuite) TestVarsutil(c *C) {
 	c.Assert(v.MaxRowCountForINLJ, Equals, 128)
 	SetSessionSystemVar(v, variable.TiDBMaxRowCountForINLJ, types.NewStringDatum("127"))
 	c.Assert(v.MaxRowCountForINLJ, Equals, 127)
+
+	c.Assert(v.MaxChunkSize, Equals, 1024)
+	SetSessionSystemVar(v, variable.TiDBMaxChunkSize, types.NewStringDatum("2"))
+	c.Assert(v.MaxChunkSize, Equals, 2)
 }
 
 type mockGlobalAccessor struct {

util/chunk/chunk.go

@@ -35,7 +35,6 @@ type Chunk struct {
 // Capacity constants.
 const (
 	InitialCapacity = 32
-	MaxCapacity     = 1024
 )
 
 // NewChunk creates a new chunk with field types.