lightning: refactor to reuse in load data, part 3

br/pkg/lightning/backend/kv/BUILD.bazel

@@ -4,6 +4,7 @@ go_library(
     name = "kv",
     srcs = [
         "allocator.go",
+        "base.go",
         "kv2sql.go",
         "session.go",
         "sql2kv.go",

br/pkg/lightning/backend/kv/base.go

@@ -0,0 +1,330 @@
+// Copyright 2023 PingCAP, Inc.
+//
+// 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 kv
+
+import (
+	"context"
+	"math/rand"
+
+	"github.com/pingcap/errors"
+	"github.com/pingcap/tidb/br/pkg/lightning/backend/encode"
+	"github.com/pingcap/tidb/br/pkg/lightning/common"
+	"github.com/pingcap/tidb/br/pkg/lightning/log"
+	"github.com/pingcap/tidb/br/pkg/logutil"
+	"github.com/pingcap/tidb/br/pkg/redact"
+	"github.com/pingcap/tidb/expression"
+	"github.com/pingcap/tidb/meta/autoid"
+	"github.com/pingcap/tidb/parser/model"
+	"github.com/pingcap/tidb/parser/mysql"
+	"github.com/pingcap/tidb/table"
+	"github.com/pingcap/tidb/table/tables"
+	"github.com/pingcap/tidb/types"
+	"github.com/pingcap/tidb/util/chunk"
+	"go.uber.org/zap"
+	"go.uber.org/zap/zapcore"
+)
+
+// ExtraHandleColumnInfo is the column info of extra handle column.
+var ExtraHandleColumnInfo = model.NewExtraHandleColInfo()
+
+// GeneratedCol generated column info.
+type GeneratedCol struct {
+	// index of the column in the table
+	Index int
+	Expr  expression.Expression
+}
+
+// AutoIDConverterFn is a function to convert auto id.
+type AutoIDConverterFn func(int64) int64
+
+// RowArrayMarshaller wraps a slice of types.Datum for logging the content into zap.
+type RowArrayMarshaller []types.Datum
+
+var kindStr = [...]string{
+	types.KindNull:          "null",
+	types.KindInt64:         "int64",
+	types.KindUint64:        "uint64",
+	types.KindFloat32:       "float32",
+	types.KindFloat64:       "float64",
+	types.KindString:        "string",
+	types.KindBytes:         "bytes",
+	types.KindBinaryLiteral: "binary",
+	types.KindMysqlDecimal:  "decimal",
+	types.KindMysqlDuration: "duration",
+	types.KindMysqlEnum:     "enum",
+	types.KindMysqlBit:      "bit",
+	types.KindMysqlSet:      "set",
+	types.KindMysqlTime:     "time",
+	types.KindInterface:     "interface",
+	types.KindMinNotNull:    "min",
+	types.KindMaxValue:      "max",
+	types.KindRaw:           "raw",
+	types.KindMysqlJSON:     "json",
+}
+
+// MarshalLogArray implements the zapcore.ArrayMarshaler interface
+func (row RowArrayMarshaller) MarshalLogArray(encoder zapcore.ArrayEncoder) error {
+	for _, datum := range row {
+		kind := datum.Kind()
+		var str string
+		var err error
+		switch kind {
+		case types.KindNull:
+			str = "NULL"
+		case types.KindMinNotNull:
+			str = "-inf"
+		case types.KindMaxValue:
+			str = "+inf"
+		default:
+			str, err = datum.ToString()
+			if err != nil {
+				return err
+			}
+		}
+		if err := encoder.AppendObject(zapcore.ObjectMarshalerFunc(func(enc zapcore.ObjectEncoder) error {
+			enc.AddString("kind", kindStr[kind])
+			enc.AddString("val", redact.String(str))
+			return nil
+		})); err != nil {
+			return err
+		}
+	}
+	return nil
+}
+
+// BaseKVEncoder encodes a row into a KV pair.
+type BaseKVEncoder struct {
+	GenCols         []GeneratedCol
+	SessionCtx      *Session
+	Table           table.Table
+	Columns         []*table.Column
+	AutoRandomColID int64
+	// convert auto id for shard rowid or auto random id base on row id generated by lightning
+	AutoIDFn AutoIDConverterFn
+
+	logger      *zap.Logger
+	recordCache []types.Datum
+}
+
+// NewBaseKVEncoder creates a new BaseKVEncoder.
+func NewBaseKVEncoder(config *encode.EncodingConfig) (*BaseKVEncoder, error) {
+	meta := config.Table.Meta()
+	cols := config.Table.Cols()
+	se := NewSession(&config.SessionOptions, config.Logger)
+	// Set CommonAddRecordCtx to session to reuse the slices and BufStore in AddRecord
+	recordCtx := tables.NewCommonAddRecordCtx(len(cols))
+	tables.SetAddRecordCtx(se, recordCtx)
+
+	var autoRandomColID int64
+	autoIDFn := func(id int64) int64 { return id }
+	if meta.ContainsAutoRandomBits() {
+		col := common.GetAutoRandomColumn(meta)
+		autoRandomColID = col.ID
+
+		shardFmt := autoid.NewShardIDFormat(&col.FieldType, meta.AutoRandomBits, meta.AutoRandomRangeBits)
+		shard := rand.New(rand.NewSource(config.AutoRandomSeed)).Int63()
+		autoIDFn = func(id int64) int64 {
+			return shardFmt.Compose(shard, id)
+		}
+	} else if meta.ShardRowIDBits > 0 {
+		rd := rand.New(rand.NewSource(config.AutoRandomSeed)) // nolint:gosec
+		mask := int64(1)<<meta.ShardRowIDBits - 1
+		shift := autoid.RowIDBitLength - meta.ShardRowIDBits - 1
+		autoIDFn = func(id int64) int64 {
+			rd.Seed(id)
+			shardBits := (int64(rd.Uint32()) & mask) << shift
+			return shardBits | id
+		}
+	}
+
+	// collect expressions for evaluating stored generated columns
+	genCols, err := CollectGeneratedColumns(se, meta, cols)
+	if err != nil {
+		return nil, errors.Annotate(err, "failed to parse generated column expressions")
+	}
+	return &BaseKVEncoder{
+		GenCols:         genCols,
+		SessionCtx:      se,
+		Table:           config.Table,
+		Columns:         cols,
+		AutoRandomColID: autoRandomColID,
+		AutoIDFn:        autoIDFn,
+		logger:          config.Logger.Logger,
+	}, nil
+}
+
+// GetOrCreateRecord returns a record slice from the cache if possible, otherwise creates a new one.
+func (e *BaseKVEncoder) GetOrCreateRecord() []types.Datum {
+	if e.recordCache != nil {
+		return e.recordCache
+	}
+	return make([]types.Datum, 0, len(e.Columns)+1)
+}
+
+// Record2KV converts a row into a KV pair.
+func (e *BaseKVEncoder) Record2KV(record, originalRow []types.Datum, rowID int64) (*KvPairs, error) {
+	_, err := e.Table.AddRecord(e.SessionCtx, record)
+	if err != nil {
+		e.logger.Error("kv encode failed",
+			zap.Array("originalRow", RowArrayMarshaller(originalRow)),
+			zap.Array("convertedRow", RowArrayMarshaller(record)),
+			log.ShortError(err),
+		)
+		return nil, errors.Trace(err)
+	}
+	kvPairs := e.SessionCtx.TakeKvPairs()
+	for i := 0; i < len(kvPairs.Pairs); i++ {
+		var encoded [9]byte // The max length of encoded int64 is 9.
+		kvPairs.Pairs[i].RowID = common.EncodeIntRowIDToBuf(encoded[:0], rowID)
+	}
+	e.recordCache = record[:0]
+	return kvPairs, nil
+}
+
+// ProcessColDatum processes the datum of a column.
+func (e *BaseKVEncoder) ProcessColDatum(col *table.Column, rowID int64, inputDatum *types.Datum) (types.Datum, error) {
+	value, err := e.getActualDatum(col, rowID, inputDatum)
+	if err != nil {
+		return value, err
+	}
+
+	if e.IsAutoRandomCol(col.ToInfo()) {
+		meta := e.Table.Meta()
+		shardFmt := autoid.NewShardIDFormat(&col.FieldType, meta.AutoRandomBits, meta.AutoRandomRangeBits)
+		// this allocator is the same as the allocator in table importer, i.e. PanickingAllocators. below too.
+		alloc := e.Table.Allocators(e.SessionCtx).Get(autoid.AutoRandomType)
+		if err := alloc.Rebase(context.Background(), value.GetInt64()&shardFmt.IncrementalMask(), false); err != nil {
+			return value, errors.Trace(err)
+		}
+	}
+	if IsAutoIncCol(col.ToInfo()) {
+		alloc := e.Table.Allocators(e.SessionCtx).Get(autoid.AutoIncrementType)
+		if err := alloc.Rebase(context.Background(), GetAutoRecordID(value, &col.FieldType), false); err != nil {
+			return value, errors.Trace(err)
+		}
+	}
+	return value, nil
+}
+
+func (e *BaseKVEncoder) getActualDatum(col *table.Column, rowID int64, inputDatum *types.Datum) (types.Datum, error) {
+	var (
+		value types.Datum
+		err   error
+	)
+
+	isBadNullValue := false
+	if inputDatum != nil {
+		value, err = table.CastValue(e.SessionCtx, *inputDatum, col.ToInfo(), false, false)
+		if err != nil {
+			return value, err
+		}
+		if err := col.CheckNotNull(&value, 0); err == nil {
+			return value, nil // the most normal case
+		}
+		isBadNullValue = true
+	}
+	// handle special values
+	switch {
+	case IsAutoIncCol(col.ToInfo()):
+		// we still need a conversion, e.g. to catch overflow with a TINYINT column.
+		value, err = table.CastValue(e.SessionCtx, types.NewIntDatum(rowID), col.ToInfo(), false, false)
+	case e.IsAutoRandomCol(col.ToInfo()):
+		var val types.Datum
+		realRowID := e.AutoIDFn(rowID)
+		if mysql.HasUnsignedFlag(col.GetFlag()) {
+			val = types.NewUintDatum(uint64(realRowID))
+		} else {
+			val = types.NewIntDatum(realRowID)
+		}
+		value, err = table.CastValue(e.SessionCtx, val, col.ToInfo(), false, false)
+	case col.IsGenerated():
+		// inject some dummy value for gen col so that MutRowFromDatums below sees a real value instead of nil.
+		// if MutRowFromDatums sees a nil it won't initialize the underlying storage and cause SetDatum to panic.
+		value = types.GetMinValue(&col.FieldType)
+	case isBadNullValue:
+		err = col.HandleBadNull(&value, e.SessionCtx.Vars.StmtCtx, 0)
+	default:
+		value, err = table.GetColDefaultValue(e.SessionCtx, col.ToInfo())
+	}
+	return value, err
+}
+
+// IsAutoRandomCol checks if the column is auto random column.
+func (e *BaseKVEncoder) IsAutoRandomCol(col *model.ColumnInfo) bool {
+	return e.Table.Meta().ContainsAutoRandomBits() && col.ID == e.AutoRandomColID
+}
+
+// EvalGeneratedColumns evaluates the generated columns.
+func (e *BaseKVEncoder) EvalGeneratedColumns(record []types.Datum, cols []*table.Column) (errCol *model.ColumnInfo, err error) {
+	return evalGeneratedColumns(e.SessionCtx, record, cols, e.GenCols)
+}
+
+// LogKVConvertFailed logs the error when converting a row to KV pair failed.
+func (e *BaseKVEncoder) LogKVConvertFailed(row []types.Datum, j int, colInfo *model.ColumnInfo, err error) error {
+	var original types.Datum
+	if 0 <= j && j < len(row) {
+		original = row[j]
+		row = row[j : j+1]
+	}
+
+	e.logger.Error("kv convert failed",
+		zap.Array("original", RowArrayMarshaller(row)),
+		zap.Int("originalCol", j),
+		zap.String("colName", colInfo.Name.O),
+		zap.Stringer("colType", &colInfo.FieldType),
+		log.ShortError(err),
+	)
+
+	e.logger.Error("failed to convert kv value", logutil.RedactAny("origVal", original.GetValue()),
+		zap.Stringer("fieldType", &colInfo.FieldType), zap.String("column", colInfo.Name.O),
+		zap.Int("columnID", j+1))
+	return errors.Annotatef(
+		err,
+		"failed to cast value as %s for column `%s` (#%d)", &colInfo.FieldType, colInfo.Name.O, j+1,
+	)
+}
+
+// LogEvalGenExprFailed logs the error when evaluating the generated column expression failed.
+func (e *BaseKVEncoder) LogEvalGenExprFailed(row []types.Datum, colInfo *model.ColumnInfo, err error) error {
+	e.logger.Error("kv convert failed: cannot evaluate generated column expression",
+		zap.Array("original", RowArrayMarshaller(row)),
+		zap.String("colName", colInfo.Name.O),
+		log.ShortError(err),
+	)
+
+	return errors.Annotatef(
+		err,
+		"failed to evaluate generated column expression for column `%s`",
+		colInfo.Name.O,
+	)
+}
+
+func evalGeneratedColumns(se *Session, record []types.Datum, cols []*table.Column, genCols []GeneratedCol) (errCol *model.ColumnInfo, err error) {
+	mutRow := chunk.MutRowFromDatums(record)
+	for _, gc := range genCols {
+		col := cols[gc.Index].ToInfo()
+		evaluated, err := gc.Expr.Eval(mutRow.ToRow())
+		if err != nil {
+			return col, err
+		}
+		value, err := table.CastValue(se, evaluated, col, false, false)
+		if err != nil {
+			return col, err
+		}
+		mutRow.SetDatum(gc.Index, value)
+		record[gc.Index] = value
+	}
+	return nil, nil
+}

br/pkg/lightning/backend/kv/kv2sql.go

@@ -32,7 +32,7 @@ type TableKVDecoder struct {
 	se  *Session
 	// tableName is the unique table name in the form "`db`.`tbl`".
 	tableName string
-	genCols   []genCol
+	genCols   []GeneratedCol
 }
 
 func (t *TableKVDecoder) Name() string {
@@ -77,7 +77,7 @@ func (t *TableKVDecoder) IterRawIndexKeys(h kv.Handle, rawRow []byte, fn func([]
 				row[i] = types.GetMinValue(&col.FieldType)
 			}
 		}
-		if _, err := evaluateGeneratedColumns(t.se, row, t.tbl.Cols(), t.genCols); err != nil {
+		if _, err := evalGeneratedColumns(t.se, row, t.tbl.Cols(), t.genCols); err != nil {
 			return err
 		}
 	}
@@ -121,7 +121,7 @@ func NewTableKVDecoder(
 	recordCtx := tables.NewCommonAddRecordCtx(len(cols))
 	tables.SetAddRecordCtx(se, recordCtx)
 
-	genCols, err := collectGeneratedColumns(se, tbl.Meta(), cols)
+	genCols, err := CollectGeneratedColumns(se, tbl.Meta(), cols)
 	if err != nil {
 		return nil, err
 	}