address comment

distsql/distsql_test.go

@@ -70,7 +70,7 @@ func (s *testSuite) TestSelectNormal(c *C) {
 	response.Fetch(context.TODO())
 
 	// Test Next.
-	chk := chunk.New(colTypes, 32, 32*3)
+	chk := chunk.New(colTypes, 32, 32)
 	numAllRows := 0
 	for {
 		err = response.Next(context.TODO(), chk)

executor/executor.go

@@ -135,14 +135,12 @@ func newBaseExecutor(ctx sessionctx.Context, schema *expression.Schema, id strin
 	return e
 }
 
-func (e baseExecutor) withInitCap(initCap int) (ret baseExecutor) {
+func (e *baseExecutor) setInitCap(initCap int) {
 	e.initCap = initCap
-	return e
 }
 
-func (e baseExecutor) withMaxChunkSize(maxChunkSize int) (ret baseExecutor) {
+func (e *baseExecutor) setMaxChunkSize(maxChunkSize int) {
 	e.maxChunkSize = maxChunkSize
-	return e
 }
 
 // Executor is the physical implementation of a algebra operator.

executor/prepared.go

@@ -89,8 +89,10 @@ type PrepareExec struct {
 
 // NewPrepareExec creates a new PrepareExec.
 func NewPrepareExec(ctx sessionctx.Context, is infoschema.InfoSchema, sqlTxt string) *PrepareExec {
+	base := newBaseExecutor(ctx, nil, "PrepareStmt")
+	base.setInitCap(chunk.NoDataChunkCap)
 	return &PrepareExec{
-		baseExecutor: newBaseExecutor(ctx, nil, "PrepareStmt").withInitCap(chunk.NoDataChunkCap),
+		baseExecutor: base,
 		is:           is,
 		sqlText:      sqlTxt,
 	}

session/session.go

@@ -555,7 +555,7 @@ func (s *session) ExecRestrictedSQL(sctx sessionctx.Context, sql string) ([]chun
 	)
 	// Execute all recordset, take out the first one as result.
 	for i, rs := range recordSets {
-		tmp, err := drainRecordSet(ctx, sctx, rs)
+		tmp, err := drainRecordSet(ctx, se, rs)
 		if err != nil {
 			return nil, nil, errors.Trace(err)
 		}
@@ -604,7 +604,7 @@ func createSessionWithDomainFunc(store kv.Storage) func(*domain.Domain) (pools.R
 	}
 }
 
-func drainRecordSet(ctx context.Context, sctx sessionctx.Context, rs ast.RecordSet) ([]chunk.Row, error) {
+func drainRecordSet(ctx context.Context, se *session, rs ast.RecordSet) ([]chunk.Row, error) {
 	var rows []chunk.Row
 	chk := rs.NewChunk()
 	for {
@@ -616,12 +616,7 @@ func drainRecordSet(ctx context.Context, sctx sessionctx.Context, rs ast.RecordS
 		for r := iter.Begin(); r != iter.End(); r = iter.Next() {
 			rows = append(rows, r)
 		}
-		if sctx == nil {
-			// sadness, statistic will pass nil sctx
-			chk = chunk.Renew(chk, variable.DefMaxChunkSize)
-			continue
-		}
-		chk = chunk.Renew(chk, sctx.GetSessionVars().MaxChunkSize)
+		chk = chunk.Renew(chk, se.sessionVars.MaxChunkSize)
 	}
 }