diff --git a/pkg/ddl/backfilling_operators.go b/pkg/ddl/backfilling_operators.go index 80d8ac9f285e8..e6ea2774ef2bb 100644 --- a/pkg/ddl/backfilling_operators.go +++ b/pkg/ddl/backfilling_operators.go @@ -915,7 +915,7 @@ func (w *indexIngestWorker) WriteChunk(rs *IndexRecordChunk) (count int, bytes i indexConditionCheckers = nil } cnt, kvBytes, err := writeChunk(w.ctx, w.writers, w.indexes, indexConditionCheckers, w.copCtx, - sc.TimeZone(), sc.ErrCtx(), vars.GetWriteStmtBufs(), rs.Chunk, w.tbl.Meta(), w.tbl.UseNewCollate()) + sc.TimeZone(), sc.ErrCtx(), vars.GetWriteStmtBufs(), rs.Chunk, w.tbl.GetPhysicalID(), w.tbl.Meta(), w.tbl.UseNewCollate()) if err != nil || cnt == 0 { return 0, 0, err } diff --git a/pkg/ddl/index.go b/pkg/ddl/index.go index f675c705921dd..2e8d4014039d2 100644 --- a/pkg/ddl/index.go +++ b/pkg/ddl/index.go @@ -2697,6 +2697,58 @@ func getLocalWriterConfig(indexCnt, writerCnt int) *backend.LocalWriterConfig { return writerCfg } +func generateIndexKVsForRow( + indexes []table.Index, + loc *time.Location, + errCtx errctx.Context, + handle kv.Handle, + physicalID int64, + idxDataBuf []types.Datum, + checkPartialCondition func(int, table.Index) (bool, error), + fetchIndexValues func(int, table.Index, []types.Datum) ([]types.Datum, error), + restoreData func(int, table.Index) []types.Datum, + sink func(int, table.Index, []types.Datum, table.IndexKVGenerator) (int64, error), +) (int64, error) { + var totalBytes int64 + for i, index := range indexes { + if index.Meta().HasCondition() && checkPartialCondition != nil { + ok, err := checkPartialCondition(i, index) + if err != nil { + return totalBytes, errors.Trace(err) + } + if !ok { + continue + } + } + indexedValues, err := fetchIndexValues(i, index, idxDataBuf) + if err != nil { + return totalBytes, errors.Trace(err) + } + var rsData []types.Datum + if restoreData != nil { + rsData = restoreData(i, index) + } + indexHandle := indexKVHandleForPhysicalTable(index, physicalID, handle) + iter := index.GenIndexKVIter(errCtx, loc, indexedValues, indexHandle, rsData) + kvBytes, err := sink(i, index, indexedValues, iter) + if err != nil { + return totalBytes, errors.Trace(err) + } + totalBytes += kvBytes + } + return totalBytes, nil +} + +func indexKVHandleForPhysicalTable(index table.Index, physicalID int64, handle kv.Handle) kv.Handle { + if index.Meta().Global && index.Meta().GlobalIndexVersion >= model.GlobalIndexVersionV1 { + if ph, ok := handle.(kv.PartitionHandle); ok { + handle = ph.Handle + } + return kv.NewPartitionHandle(physicalID, handle) + } + return handle +} + func writeChunk( ctx context.Context, writers []ingest.Writer, @@ -2707,6 +2759,7 @@ func writeChunk( errCtx errctx.Context, writeStmtBufs *variable.WriteStmtBufs, copChunk *chunk.Chunk, + physicalID int64, tblInfo *model.TableInfo, useNewCollate bool, ) (rowCnt int, bytes int, err error) { @@ -2757,34 +2810,42 @@ func writeChunk( if err != nil { return 0, totalBytes, errors.Trace(err) } - for i, index := range indexes { - // If the `IndexRecordChunk.conditionPushed` is true and we have only 1 index, the `indexConditionCheckers` - // will not be initialized. - if index.Meta().HasCondition() && indexConditionCheckers != nil { + var checkPartialCondition func(int, table.Index) (bool, error) + if indexConditionCheckers != nil { + checkPartialCondition = func(i int, _ table.Index) (bool, error) { ok, err := indexConditionCheckers[i](row) if err != nil { - return 0, 0, errors.Trace(err) - } - if !ok { - continue + return false, errors.Trace(err) } + return ok, nil } - - idxID := index.Meta().ID - idxDataBuf = ExtractDatumByOffsets(ectx, - row, copCtx.IndexColumnOutputOffsets(idxID), c.ExprColumnInfos, idxDataBuf) - idxData := idxDataBuf[:len(index.Meta().Columns)] - var rsData []types.Datum - if needRestoreForIndexes[i] { - rsData = getRestoreData(useNewCollate, c.TableInfo, copCtx.IndexInfo(idxID), c.PrimaryKeyInfo, restoreDataBuf) - } - kvBytes, err := writeOneKV(ctx, writers[i], index, loc, errCtx, writeStmtBufs, idxData, rsData, h) - if err != nil { - err = ingest.TryConvertToKeyExistsErr(err, index.Meta(), tblInfo) - return 0, totalBytes, errors.Trace(err) - } - totalBytes += int(kvBytes) } + kvBytes, err := generateIndexKVsForRow( + indexes, loc, errCtx, h, physicalID, idxDataBuf, checkPartialCondition, + func(_ int, index table.Index, buf []types.Datum) ([]types.Datum, error) { + idxID := index.Meta().ID + idxData := ExtractDatumByOffsets(ectx, row, copCtx.IndexColumnOutputOffsets(idxID), c.ExprColumnInfos, buf) + return idxData[:len(index.Meta().Columns)], nil + }, + func(i int, index table.Index) []types.Datum { + if !needRestoreForIndexes[i] { + return nil + } + return getRestoreData(useNewCollate, c.TableInfo, copCtx.IndexInfo(index.Meta().ID), c.PrimaryKeyInfo, restoreDataBuf) + }, + func(i int, _ table.Index, _ []types.Datum, iter table.IndexKVGenerator) (int64, error) { + kvBytes, err := writeOneKV(ctx, writers[i], writeStmtBufs, iter, h) + if err != nil { + err = ingest.TryConvertToKeyExistsErr(err, indexes[i].Meta(), tblInfo) + return 0, errors.Trace(err) + } + return kvBytes, nil + }, + ) + if err != nil { + return 0, totalBytes + int(kvBytes), errors.Trace(err) + } + totalBytes += int(kvBytes) count++ } return count, totalBytes, nil @@ -2804,15 +2865,11 @@ func maxIndexColumnCount(indexes []table.Index) int { func writeOneKV( ctx context.Context, writer ingest.Writer, - index table.Index, - loc *time.Location, - errCtx errctx.Context, writeBufs *variable.WriteStmtBufs, - idxDt, rsData []types.Datum, + iter table.IndexKVGenerator, handle kv.Handle, ) (int64, error) { var kvBytes int64 - iter := index.GenIndexKVIter(errCtx, loc, idxDt, handle, rsData) for iter.Valid() { key, idxVal, _, err := iter.Next(writeBufs.IndexKeyBuf, writeBufs.RowValBuf) if err != nil { diff --git a/pkg/ddl/reorg_util_test.go b/pkg/ddl/reorg_util_test.go index d637ed6d6aa9e..a2a3b6d9a5566 100644 --- a/pkg/ddl/reorg_util_test.go +++ b/pkg/ddl/reorg_util_test.go @@ -19,6 +19,7 @@ import ( "testing" "github.com/docker/go-units" + "github.com/pingcap/tidb/pkg/kv" "github.com/pingcap/tidb/pkg/meta/model" "github.com/pingcap/tidb/pkg/store/helper" "github.com/pingcap/tidb/pkg/table/tables" @@ -84,6 +85,23 @@ func expectedRegionRange(tableID int64) ([]byte, []byte) { return mockCodec{}.EncodeRegionRange(tableStart, tableEnd) } +func TestIndexKVHandleForPhysicalTableRewrapsGlobalIndexHandle(t *testing.T) { + idx, err := tables.NewIndex(10, &model.TableInfo{ID: 1}, &model.IndexInfo{ + ID: 2, + Global: true, + GlobalIndexVersion: model.GlobalIndexVersionV1, + }) + require.NoError(t, err) + + handle := indexKVHandleForPhysicalTable(idx, 20, kv.NewPartitionHandle(10, kv.IntHandle(42))) + partitionHandle, ok := handle.(kv.PartitionHandle) + require.True(t, ok) + require.Equal(t, int64(20), partitionHandle.PartitionID) + require.True(t, partitionHandle.Handle.Equal(kv.IntHandle(42))) + _, nested := partitionHandle.Handle.(kv.PartitionHandle) + require.False(t, nested) +} + func TestEstimateTableSizeByIDUsesMaxApproximateSizes(t *testing.T) { pdCli := &mockPDHTTPClient{ regionInfos: []*pdhttp.RegionsInfo{ diff --git a/pkg/table/tables/BUILD.bazel b/pkg/table/tables/BUILD.bazel index 9775792b76383..e9b4d421d24cf 100644 --- a/pkg/table/tables/BUILD.bazel +++ b/pkg/table/tables/BUILD.bazel @@ -80,7 +80,7 @@ go_test( ], embed = [":tables"], flaky = True, - shard_count = 44, + shard_count = 45, deps = [ "//pkg/ddl", "//pkg/domain", diff --git a/pkg/table/tables/index.go b/pkg/table/tables/index.go index 5d433e5d656c8..6887d16308b55 100644 --- a/pkg/table/tables/index.go +++ b/pkg/table/tables/index.go @@ -206,13 +206,14 @@ func (c *index) GenIndexValue(ec errctx.Context, loc *time.Location, distinct, u return idx, err } -// getIndexedValue will produce the result like: +// BuildMultiValueIndexValueGroups expands indexed values for a multi-valued index. +// It will produce the result like: // 1. If not multi-valued index, return directly. // 2. (i1, [m1,m2], i2, ...) ==> [(i1, m1, i2, ...), (i1, m2, i2, ...)] // 3. (i1, null, i2, ...) ==> [(i1, null, i2, ...)] // 4. (i1, [], i2, ...) ==> nothing. -func (c *index) getIndexedValue(indexedValues []types.Datum) [][]types.Datum { - if !c.idxInfo.MVIndex { +func BuildMultiValueIndexValueGroups(tblInfo *model.TableInfo, idxInfo *model.IndexInfo, indexedValues []types.Datum) [][]types.Datum { + if idxInfo == nil || !idxInfo.MVIndex { return [][]types.Datum{indexedValues} } @@ -224,7 +225,7 @@ func (c *index) getIndexedValue(indexedValues []types.Datum) [][]types.Datum { for !jsonIsNull { val := make([]types.Datum, 0, len(indexedValues)) for i, v := range indexedValues { - if !c.tblInfo.Columns[c.idxInfo.Columns[i].Offset].FieldType.IsArray() { + if !tblInfo.Columns[idxInfo.Columns[i].Offset].FieldType.IsArray() { val = append(val, v) } else { // if the datum type is not JSON, it must come from cleanup index. @@ -258,6 +259,35 @@ out: return vals } +// EncodeRawIndexKeyValues encodes one raw index-key suffix per generated index entry. +// For multi-valued indexes, it expands the indexed values in the same order as index +// entry generation, then truncates and encodes each group. +func EncodeRawIndexKeyValues( + useNewCollate bool, + loc *time.Location, + tblInfo *model.TableInfo, + idxInfo *model.IndexInfo, + indexedValues []types.Datum, +) ([][]byte, error) { + rawValueGroups := BuildMultiValueIndexValueGroups(tblInfo, idxInfo, indexedValues) + rawKeys := make([][]byte, 0, len(rawValueGroups)) + encoder := codec.NewEncoder(useNewCollate) + for _, rawValues := range rawValueGroups { + clonedValues := append([]types.Datum(nil), rawValues...) + tablecodec.TruncateIndexValues(tblInfo, idxInfo, clonedValues) + rawKey, err := encoder.EncodeKey(loc, nil, clonedValues...) + if err != nil { + return nil, err + } + rawKeys = append(rawKeys, rawKey) + } + return rawKeys, nil +} + +func (c *index) getIndexedValue(indexedValues []types.Datum) [][]types.Datum { + return BuildMultiValueIndexValueGroups(c.tblInfo, c.idxInfo, indexedValues) +} + // MeetPartialCondition checks whether the row meets the partial index condition of the index. func (c *index) MeetPartialCondition(row []types.Datum) (meet bool, err error) { if c.conditionExpr == nil { diff --git a/pkg/table/tables/index_test.go b/pkg/table/tables/index_test.go index 3d5e7d8a0e7ed..a29650252982c 100644 --- a/pkg/table/tables/index_test.go +++ b/pkg/table/tables/index_test.go @@ -183,6 +183,56 @@ func buildTableInfo(t *testing.T, sql string) *model.TableInfo { return tblInfo } +func TestEncodeRawIndexKeyValues(t *testing.T) { + loc := time.UTC + + t.Run("non multi-valued index truncates prefix columns", func(t *testing.T) { + tblInfo := buildTableInfo(t, "create table t (a varchar(10), key idx(a(2)))") + idxInfo := tblInfo.Indices[0] + + rawKeys, err := tables.EncodeRawIndexKeyValues(true, loc, tblInfo, idxInfo, types.MakeDatums("abcd")) + require.NoError(t, err) + require.Len(t, rawKeys, 1) + + remain, datum, err := codec.DecodeOne(rawKeys[0]) + require.NoError(t, err) + require.Empty(t, remain) + require.Equal(t, "ab", datum.GetString()) + }) + + t.Run("multi-valued index expands unique array elements in order", func(t *testing.T) { + jsonType := types.NewFieldType(mysql.TypeJSON) + jsonType.SetArray(true) + tblInfo := &model.TableInfo{ + Columns: []*model.ColumnInfo{ + {ID: 1, Offset: 0, FieldType: *jsonType}, + }, + } + idxInfo := &model.IndexInfo{ + MVIndex: true, + Columns: []*model.IndexColumn{ + {Offset: 0}, + }, + } + + arrayDatum, err := types.ParseBinaryJSONFromString(`[2, 1, 2]`) + require.NoError(t, err) + rawKeys, err := tables.EncodeRawIndexKeyValues(true, loc, tblInfo, idxInfo, types.MakeDatums(arrayDatum)) + require.NoError(t, err) + require.Len(t, rawKeys, 2) + + remain, datum, err := codec.DecodeOne(rawKeys[0]) + require.NoError(t, err) + require.Empty(t, remain) + require.EqualValues(t, 2, datum.GetInt64()) + + remain, datum, err = codec.DecodeOne(rawKeys[1]) + require.NoError(t, err) + require.Empty(t, remain) + require.EqualValues(t, 1, datum.GetInt64()) + }) +} + func TestGenIndexValueFromIndex(t *testing.T) { tblInfo := buildTableInfo(t, "create table a (a int primary key, b int not null, c text, unique key key_b(b));") tblInfo.State = model.StatePublic