From 425f21653af0df05fa7609d408cce42f13aee753 Mon Sep 17 00:00:00 2001
From: Lynn <zimu_xia@126.com>
Date: Thu, 3 Mar 2022 11:45:46 +0800
Subject: [PATCH] cherry pick #32621 to release-5.2

Signed-off-by: ti-srebot <ti-srebot@pingcap.com>
---
 table/tables/partition.go | 53 +++++++++++++++++++++++++++++++--------
 1 file changed, 42 insertions(+), 11 deletions(-)

diff --git a/table/tables/partition.go b/table/tables/partition.go
index 74c473268dec4..cad2971184f73 100644
--- a/table/tables/partition.go
+++ b/table/tables/partition.go
@@ -252,6 +252,14 @@ type ForListColumnPruning struct {
 	valueTp  *types.FieldType
 	valueMap map[string]ListPartitionLocation
 	sorted   *btree.BTree
+
+	// To deal with the location partition failure caused by inconsistent NewCollationEnabled values(see issue #32416).
+	// The following fields are used to delay building valueMap.
+	ctx     sessionctx.Context
+	tblInfo *model.TableInfo
+	schema  *expression.Schema
+	names   types.NameSlice
+	colIdx  int
 }
 
 // ListPartitionGroup indicate the group index of the column value in a partition.
@@ -648,7 +656,6 @@ func (lp *ForListPruning) buildListColumnsPruner(ctx sessionctx.Context, tblInfo
 	columns []*expression.Column, names types.NameSlice) error {
 	pi := tblInfo.GetPartitionInfo()
 	schema := expression.NewSchema(columns...)
-	p := parser.New()
 	colPrunes := make([]*ForListColumnPruning, 0, len(pi.Columns))
 	for colIdx := range pi.Columns {
 		colInfo := model.FindColumnInfo(tblInfo.Columns, pi.Columns[colIdx].L)
@@ -660,15 +667,17 @@ func (lp *ForListPruning) buildListColumnsPruner(ctx sessionctx.Context, tblInfo
 			return table.ErrUnknownColumn.GenWithStackByArgs(pi.Columns[colIdx].L)
 		}
 		colPrune := &ForListColumnPruning{
+			ctx:      ctx,
+			tblInfo:  tblInfo,
+			schema:   schema,
+			names:    names,
+			colIdx:   colIdx,
 			ExprCol:  columns[idx],
 			valueTp:  &colInfo.FieldType,
 			valueMap: make(map[string]ListPartitionLocation),
 			sorted:   btree.New(btreeDegree),
 		}
-		err := colPrune.buildPartitionValueMapAndSorted(ctx, tblInfo, colIdx, schema, names, p)
-		if err != nil {
-			return err
-		}
+
 		colPrunes = append(colPrunes, colPrune)
 	}
 	lp.ColPrunes = colPrunes
@@ -752,13 +761,13 @@ func (lp *ForListPruning) locateListColumnsPartitionByRow(ctx sessionctx.Context
 // buildListPartitionValueMapAndSorted builds list columns partition value map for the specified column.
 // it also builds list columns partition value btree for the specified column.
 // colIdx is the specified column index in the list columns.
-func (lp *ForListColumnPruning) buildPartitionValueMapAndSorted(ctx sessionctx.Context, tblInfo *model.TableInfo, colIdx int,
-	schema *expression.Schema, names types.NameSlice, p *parser.Parser) error {
-	pi := tblInfo.GetPartitionInfo()
-	sc := ctx.GetSessionVars().StmtCtx
+func (lp *ForListColumnPruning) buildPartitionValueMapAndSorted() error {
+	p := parser.New()
+	pi := lp.tblInfo.GetPartitionInfo()
+	sc := lp.ctx.GetSessionVars().StmtCtx
 	for partitionIdx, def := range pi.Definitions {
 		for groupIdx, vs := range def.InValues {
-			keyBytes, err := lp.genConstExprKey(ctx, sc, vs[colIdx], schema, names, p)
+			keyBytes, err := lp.genConstExprKey(lp.ctx, sc, vs[lp.colIdx], lp.schema, lp.names, p)
 			if err != nil {
 				return errors.Trace(err)
 			}
@@ -804,11 +813,20 @@ func (lp *ForListColumnPruning) genKey(sc *stmtctx.StatementContext, v types.Dat
 	if err != nil {
 		return nil, errors.Trace(err)
 	}
-	return codec.EncodeKey(sc, nil, v)
+	valByte, err := codec.EncodeKey(sc, nil, v)
+	return valByte, err
 }
 
 // LocatePartition locates partition by the column value
 func (lp *ForListColumnPruning) LocatePartition(sc *stmtctx.StatementContext, v types.Datum) (ListPartitionLocation, error) {
+	// To deal with the location partition failure caused by inconsistent NewCollationEnabled values(see issue #32416).
+	if len(lp.valueMap) == 0 {
+		err := lp.buildPartitionValueMapAndSorted()
+		if err != nil {
+			return nil, err
+		}
+	}
+
 	key, err := lp.genKey(sc, v)
 	if err != nil {
 		return nil, errors.Trace(err)
@@ -822,6 +840,19 @@ func (lp *ForListColumnPruning) LocatePartition(sc *stmtctx.StatementContext, v
 
 // LocateRanges locates partition ranges by the column range
 func (lp *ForListColumnPruning) LocateRanges(sc *stmtctx.StatementContext, r *ranger.Range) ([]ListPartitionLocation, error) {
+<<<<<<< HEAD
+=======
+	// To deal with the location partition failure caused by inconsistent NewCollationEnabled values(see issue #32416).
+	if len(lp.valueMap) == 0 {
+		err := lp.buildPartitionValueMapAndSorted()
+		if err != nil {
+			return nil, err
+		}
+	}
+
+	var err error
+	var lowKey, highKey []byte
+>>>>>>> 0e4084c02... table/tables: fix buildListColumnsPruner issue in the list partition (#32621)
 	lowVal := r.LowVal[0]
 	if r.LowVal[0].Kind() == types.KindMinNotNull {
 		lowVal = types.GetMinValue(lp.ExprCol.GetType())