executor: write multi-valued index

executor/admin.go

@@ -163,7 +163,7 @@ func (e *CheckIndexRangeExec) constructIndexScanPB() *tipb.Executor {
 	idxExec := &tipb.IndexScan{
 		TableId: e.table.ID,
 		IndexId: e.index.ID,
-		Columns: util.ColumnsToProto(e.cols, e.table.PKIsHandle),
+		Columns: util.ColumnsToProto(e.cols, e.table.PKIsHandle, true),
 	}
 	return &tipb.Executor{Tp: tipb.ExecType_TypeIndexScan, IdxScan: idxExec}
 }
@@ -814,7 +814,7 @@ func (e *CleanupIndexExec) constructIndexScanPB() *tipb.Executor {
 	idxExec := &tipb.IndexScan{
 		TableId:          e.physicalID,
 		IndexId:          e.index.Meta().ID,
-		Columns:          util.ColumnsToProto(e.columns, e.table.Meta().PKIsHandle),
+		Columns:          util.ColumnsToProto(e.columns, e.table.Meta().PKIsHandle, true),
 		PrimaryColumnIds: tables.TryGetCommonPkColumnIds(e.table.Meta()),
 	}
 	return &tipb.Executor{Tp: tipb.ExecType_TypeIndexScan, IdxScan: idxExec}

executor/builder.go

@@ -2580,7 +2580,7 @@ func (b *executorBuilder) buildAnalyzeSamplingPushdown(task plannercore.AnalyzeC
 		SampleSize:   int64(opts[ast.AnalyzeOptNumSamples]),
 		SampleRate:   sampleRate,
 		SketchSize:   maxSketchSize,
-		ColumnsInfo:  util.ColumnsToProto(task.ColsInfo, task.TblInfo.PKIsHandle),
+		ColumnsInfo:  util.ColumnsToProto(task.ColsInfo, task.TblInfo.PKIsHandle, false),
 		ColumnGroups: colGroups,
 	}
 	if task.TblInfo != nil {
@@ -2741,7 +2741,7 @@ func (b *executorBuilder) buildAnalyzeColumnsPushdown(task plannercore.AnalyzeCo
 		BucketSize:    int64(opts[ast.AnalyzeOptNumBuckets]),
 		SampleSize:    MaxRegionSampleSize,
 		SketchSize:    maxSketchSize,
-		ColumnsInfo:   util.ColumnsToProto(cols, task.HandleCols != nil && task.HandleCols.IsInt()),
+		ColumnsInfo:   util.ColumnsToProto(cols, task.HandleCols != nil && task.HandleCols.IsInt(), false),
 		CmsketchDepth: &depth,
 		CmsketchWidth: &width,
 	}

expression/multi_valued_index_test.go

@@ -15,11 +15,20 @@
 package expression_test
 
 import (
+	"context"
 	"fmt"
 	"testing"
 
 	"github.com/pingcap/tidb/errno"
+	"github.com/pingcap/tidb/kv"
+	"github.com/pingcap/tidb/parser/model"
+	"github.com/pingcap/tidb/sessiontxn"
+	"github.com/pingcap/tidb/table"
+	"github.com/pingcap/tidb/tablecodec"
 	"github.com/pingcap/tidb/testkit"
+	"github.com/pingcap/tidb/types"
+	"github.com/pingcap/tidb/util/codec"
+	"github.com/stretchr/testify/require"
 )
 
 func TestMultiValuedIndexDDL(t *testing.T) {
@@ -225,3 +234,251 @@ func TestMultiValuedIndexDML(t *testing.T) {
 		tk.MustGetErrCode(`insert into t values (json_array(cast('{"a":1}' as json)));`, errno.ErrInvalidJSONValueForFuncIndex)
 	}
 }
+
+func TestWriteMultiValuedIndex(t *testing.T) {
+	store, dom := testkit.CreateMockStoreAndDomain(t)
+	tk := testkit.NewTestKit(t, store)
+	tk.MustExec("use test")
+	tk.MustExec("create table t1(pk int primary key, a json, index idx((cast(a as signed array))))")
+	tk.MustExec("insert into t1 values (1, '[1,2,2,3]')")
+	tk.MustExec("insert into t1 values (2, '[1,2,3]')")
+	tk.MustExec("insert into t1 values (3, '[]')")
+	tk.MustExec("insert into t1 values (4, '[2,3,4]')")
+	tk.MustExec("insert into t1 values (5, null)")
+
+	t1, err := dom.InfoSchema().TableByName(model.NewCIStr("test"), model.NewCIStr("t1"))
+	require.NoError(t, err)
+	for _, index := range t1.Indices() {
+		if index.Meta().MVIndex {
+			checkCount(t, t1.IndexPrefix(), index, store, 10)
+			checkKey(t, t1.IndexPrefix(), index, store, [][]types.Datum{
+				{types.NewDatum(nil), types.NewIntDatum(5)},
+				{types.NewIntDatum(1), types.NewIntDatum(1)},
+				{types.NewIntDatum(1), types.NewIntDatum(2)},
+				{types.NewIntDatum(2), types.NewIntDatum(1)},
+				{types.NewIntDatum(2), types.NewIntDatum(2)},
+				{types.NewIntDatum(2), types.NewIntDatum(4)},
+				{types.NewIntDatum(3), types.NewIntDatum(1)},
+				{types.NewIntDatum(3), types.NewIntDatum(2)},
+				{types.NewIntDatum(3), types.NewIntDatum(4)},
+				{types.NewIntDatum(4), types.NewIntDatum(4)},
+			})
+		}
+	}
+	tk.MustExec("delete from t1")
+	for _, index := range t1.Indices() {
+		if index.Meta().MVIndex {
+			checkCount(t, t1.IndexPrefix(), index, store, 0)
+		}
+	}
+
+	tk.MustExec("drop table t1")
+	tk.MustExec("create table t1(pk int primary key, a json, index idx((cast(a as char(5) array))))")
+	tk.MustExec("insert into t1 values (1, '[\"abc\", \"abc \"]')")
+	tk.MustExec("insert into t1 values (2, '[\"b\"]')")
+	tk.MustExec("insert into t1 values (3, '[\"b   \"]')")
+
+	t1, err = dom.InfoSchema().TableByName(model.NewCIStr("test"), model.NewCIStr("t1"))
+	require.NoError(t, err)
+	for _, index := range t1.Indices() {
+		if index.Meta().MVIndex {
+			checkCount(t, t1.IndexPrefix(), index, store, 4)
+			checkKey(t, t1.IndexPrefix(), index, store, [][]types.Datum{
+				{types.NewBytesDatum([]byte("abc")), types.NewIntDatum(1)},
+				{types.NewBytesDatum([]byte("abc ")), types.NewIntDatum(1)},
+				{types.NewBytesDatum([]byte("b")), types.NewIntDatum(2)},
+				{types.NewBytesDatum([]byte("b   ")), types.NewIntDatum(3)},
+			})
+		}
+	}
+
+	tk.MustExec("update t1 set a = json_array_append(a, '$', 'bcd') where pk = 1")
+	tk.MustExec("update t1 set a = '[]' where pk = 2")
+	tk.MustExec("update t1 set a = '[\"abc\"]' where pk = 3")
+
+	for _, index := range t1.Indices() {
+		if index.Meta().MVIndex {
+			checkCount(t, t1.IndexPrefix(), index, store, 4)
+			checkKey(t, t1.IndexPrefix(), index, store, [][]types.Datum{
+				{types.NewBytesDatum([]byte("abc")), types.NewIntDatum(1)},
+				{types.NewBytesDatum([]byte("abc")), types.NewIntDatum(3)},
+				{types.NewBytesDatum([]byte("abc ")), types.NewIntDatum(1)},
+				{types.NewBytesDatum([]byte("bcd")), types.NewIntDatum(1)},
+			})
+		}
+	}
+
+	tk.MustExec("delete from t1")
+	for _, index := range t1.Indices() {
+		if index.Meta().MVIndex {
+			checkCount(t, t1.IndexPrefix(), index, store, 0)
+		}
+	}
+}
+
+func TestWriteMultiValuedIndexPartitionTable(t *testing.T) {
+	store, dom := testkit.CreateMockStoreAndDomain(t)
+	tk := testkit.NewTestKit(t, store)
+	tk.MustExec("use test")
+	tk.MustExec(`create table t1
+(
+    pk int primary key,
+    a  json,
+    index idx ((cast(a as signed array)))
+) partition by range columns (pk) (partition p0 values less than (10), partition p1 values less than (20));`)
+	tk.MustExec("insert into t1 values (1, '[1,2,2,3]')")
+	tk.MustExec("insert into t1 values (11, '[1,2,3]')")
+	tk.MustExec("insert into t1 values (2, '[]')")
+	tk.MustExec("insert into t1 values (12, '[2,3,4]')")
+	tk.MustExec("insert into t1 values (3, null)")
+	tk.MustExec("insert into t1 values (13, null)")
+
+	t1, err := dom.InfoSchema().TableByName(model.NewCIStr("test"), model.NewCIStr("t1"))
+	require.NoError(t, err)
+
+	expect := map[string]struct {
+		count int
+		vals  [][]types.Datum
+	}{
+		"p0": {4, [][]types.Datum{
+			{types.NewDatum(nil), types.NewIntDatum(3)},
+			{types.NewIntDatum(1), types.NewIntDatum(1)},
+			{types.NewIntDatum(2), types.NewIntDatum(1)},
+			{types.NewIntDatum(3), types.NewIntDatum(1)},
+		}},
+		"p1": {7, [][]types.Datum{
+			{types.NewDatum(nil), types.NewIntDatum(13)},
+			{types.NewIntDatum(1), types.NewIntDatum(11)},
+			{types.NewIntDatum(2), types.NewIntDatum(11)},
+			{types.NewIntDatum(2), types.NewIntDatum(12)},
+			{types.NewIntDatum(3), types.NewIntDatum(11)},
+			{types.NewIntDatum(3), types.NewIntDatum(12)},
+			{types.NewIntDatum(4), types.NewIntDatum(12)},
+		}},
+	}
+
+	for _, def := range t1.Meta().GetPartitionInfo().Definitions {
+		partition := t1.(table.PartitionedTable).GetPartition(def.ID)
+		for _, index := range partition.Indices() {
+			if index.Meta().MVIndex {
+				checkCount(t, partition.IndexPrefix(), index, store, expect[def.Name.L].count)
+				checkKey(t, partition.IndexPrefix(), index, store, expect[def.Name.L].vals)
+			}
+		}
+	}
+
+	tk.MustExec("delete from t1")
+	for _, def := range t1.Meta().GetPartitionInfo().Definitions {
+		partition := t1.(table.PartitionedTable).GetPartition(def.ID)
+		for _, index := range partition.Indices() {
+			if index.Meta().MVIndex {
+				checkCount(t, partition.IndexPrefix(), index, store, 0)
+			}
+		}
+	}
+}
+
+func TestWriteMultiValuedIndexUnique(t *testing.T) {
+	store, dom := testkit.CreateMockStoreAndDomain(t)
+	tk := testkit.NewTestKit(t, store)
+	tk.MustExec("use test")
+	tk.MustExec("create table t1(pk int primary key, a json, unique index idx((cast(a as signed array))))")
+	tk.MustExec("insert into t1 values (1, '[1,2,2]')")
+	tk.MustGetErrCode("insert into t1 values (2, '[1]')", errno.ErrDupEntry)
+	tk.MustExec("insert into t1 values (3, '[3,3,4]')")
+
+	t1, err := dom.InfoSchema().TableByName(model.NewCIStr("test"), model.NewCIStr("t1"))
+	require.NoError(t, err)
+	for _, index := range t1.Indices() {
+		if index.Meta().MVIndex {
+			checkCount(t, t1.IndexPrefix(), index, store, 4)
+			checkKey(t, t1.IndexPrefix(), index, store, [][]types.Datum{
+				{types.NewIntDatum(1)},
+				{types.NewIntDatum(2)},
+				{types.NewIntDatum(3)},
+				{types.NewIntDatum(4)},
+			})
+		}
+	}
+}
+
+func TestWriteMultiValuedIndexComposite(t *testing.T) {
+	store, dom := testkit.CreateMockStoreAndDomain(t)
+	tk := testkit.NewTestKit(t, store)
+	tk.MustExec("use test")
+	tk.MustExec("create table t1(pk int primary key, a json, c int, d int, index idx(c, (cast(a as signed array)), d))")
+	tk.MustExec("insert into t1 values (1, '[1,2,2]', 1, 1)")
+	tk.MustExec("insert into t1 values (2, '[2,2,2]', 2, 2)")
+	tk.MustExec("insert into t1 values (3, '[3,3,4]', 3, 3)")
+	tk.MustExec("insert into t1 values (4, null, 4, 4)")
+	tk.MustExec("insert into t1 values (5, '[]', 5, 5)")
+
+	t1, err := dom.InfoSchema().TableByName(model.NewCIStr("test"), model.NewCIStr("t1"))
+	require.NoError(t, err)
+	for _, index := range t1.Indices() {
+		if index.Meta().MVIndex {
+			checkCount(t, t1.IndexPrefix(), index, store, 6)
+			checkKey(t, t1.IndexPrefix(), index, store, [][]types.Datum{
+				{types.NewIntDatum(1), types.NewIntDatum(1), types.NewIntDatum(1), types.NewIntDatum(1)},
+				{types.NewIntDatum(1), types.NewIntDatum(2), types.NewIntDatum(1), types.NewIntDatum(1)},
+				{types.NewIntDatum(2), types.NewIntDatum(2), types.NewIntDatum(2), types.NewIntDatum(2)},
+				{types.NewIntDatum(3), types.NewIntDatum(3), types.NewIntDatum(3), types.NewIntDatum(3)},
+				{types.NewIntDatum(3), types.NewIntDatum(4), types.NewIntDatum(3), types.NewIntDatum(3)},
+				{types.NewIntDatum(4), types.NewDatum(nil), types.NewIntDatum(4), types.NewIntDatum(4)},
+			})
+		}
+	}
+}
+
+func checkCount(t *testing.T, prefix kv.Key, index table.Index, store kv.Storage, except int) {
+	c := 0
+	checkIndex(t, prefix, index, store, func(it kv.Iterator) {
+		c++
+	})
+	require.Equal(t, except, c)
+}
+
+func checkKey(t *testing.T, prefix kv.Key, index table.Index, store kv.Storage, except [][]types.Datum) {
+	idx := 0
+	checkIndex(t, prefix, index, store, func(it kv.Iterator) {
+		indexKey := decodeIndexKey(t, it.Key())
+		require.Equal(t, except[idx], indexKey)
+		idx++
+	})
+}
+
+func checkIndex(t *testing.T, prefix kv.Key, index table.Index, store kv.Storage, fn func(kv.Iterator)) {
+	startKey := codec.EncodeInt(prefix, index.Meta().ID)
+	prefix.Next()
+	se := testkit.NewTestKit(t, store).Session()
+	err := sessiontxn.NewTxn(context.Background(), se)
+	require.NoError(t, err)
+	txn, err := se.Txn(true)
+	require.NoError(t, err)
+	it, err := txn.Iter(startKey, prefix.PrefixNext())
+	require.NoError(t, err)
+	for it.Valid() && it.Key().HasPrefix(prefix) {
+		fn(it)
+		err = it.Next()
+		require.NoError(t, err)
+	}
+	it.Close()
+	se.Close()
+}
+
+func decodeIndexKey(t *testing.T, key kv.Key) []types.Datum {
+	var idLen = 8
+	var prefixLen = 1 + idLen /*tableID*/ + 2
+	_, _, isRecord, err := tablecodec.DecodeKeyHead(key)
+	require.NoError(t, err)
+	require.False(t, isRecord)
+	indexKey := key[prefixLen+idLen:]
+	var datumValues []types.Datum
+	for len(indexKey) > 0 {
+		remain, d, err := codec.DecodeOne(indexKey)
+		require.NoError(t, err)
+		datumValues = append(datumValues, d)
+		indexKey = remain
+	}
+	return datumValues
+}

planner/core/logical_plan_builder.go

@@ -5739,7 +5739,10 @@ func (b *PlanBuilder) buildUpdateLists(ctx context.Context, tableList []*ast.Tab
 				}
 			}
 
+			o := b.allowBuildCastArray
+			b.allowBuildCastArray = true
 			newExpr, np, err = b.rewriteWithPreprocess(ctx, assign.Expr, p, nil, nil, false, rewritePreprocess(assign))
+			b.allowBuildCastArray = o
 			if err != nil {
 				return nil, nil, false, err
 			}

planner/core/plan_to_pb.go

@@ -370,7 +370,7 @@ func (p *PhysicalIndexScan) ToPB(_ sessionctx.Context, _ kv.StoreType) (*tipb.Ex
 	idxExec := &tipb.IndexScan{
 		TableId:          p.Table.ID,
 		IndexId:          p.Index.ID,
-		Columns:          util.ColumnsToProto(columns, p.Table.PKIsHandle),
+		Columns:          util.ColumnsToProto(columns, p.Table.PKIsHandle, true),
 		Desc:             p.Desc,
 		PrimaryColumnIds: pkColIds,
 	}

planner/core/plan_to_pb_test.go

@@ -35,16 +35,16 @@ func TestColumnToProto(t *testing.T) {
 	col := &model.ColumnInfo{
 		FieldType: *tp,
 	}
-	pc := util.ColumnToProto(col)
+	pc := util.ColumnToProto(col, false)
 	expect := &tipb.ColumnInfo{ColumnId: 0, Tp: 3, Collation: 83, ColumnLen: 11, Decimal: 0, Flag: 10, Elems: []string(nil), DefaultVal: []uint8(nil), PkHandle: false, XXX_unrecognized: []uint8(nil)}
 	require.Equal(t, expect, pc)
 
 	cols := []*model.ColumnInfo{col, col}
-	pcs := util.ColumnsToProto(cols, false)
+	pcs := util.ColumnsToProto(cols, false, false)
 	for _, v := range pcs {
 		require.Equal(t, int32(10), v.GetFlag())
 	}
-	pcs = util.ColumnsToProto(cols, true)
+	pcs = util.ColumnsToProto(cols, true, false)
 	for _, v := range pcs {
 		require.Equal(t, int32(10), v.GetFlag())
 	}
@@ -56,19 +56,19 @@ func TestColumnToProto(t *testing.T) {
 	col1 := &model.ColumnInfo{
 		FieldType: *tp,
 	}
-	pc = util.ColumnToProto(col1)
+	pc = util.ColumnToProto(col1, false)
 	require.Equal(t, int32(8), pc.Collation)
 
 	collate.SetNewCollationEnabledForTest(true)
 
-	pc = util.ColumnToProto(col)
+	pc = util.ColumnToProto(col, false)
 	expect = &tipb.ColumnInfo{ColumnId: 0, Tp: 3, Collation: -83, ColumnLen: 11, Decimal: 0, Flag: 10, Elems: []string(nil), DefaultVal: []uint8(nil), PkHandle: false, XXX_unrecognized: []uint8(nil)}
 	require.Equal(t, expect, pc)
-	pcs = util.ColumnsToProto(cols, true)
+	pcs = util.ColumnsToProto(cols, true, false)
 	for _, v := range pcs {
 		require.Equal(t, int32(-83), v.Collation)
 	}
-	pc = util.ColumnToProto(col1)
+	pc = util.ColumnToProto(col1, false)
 	require.Equal(t, int32(-8), pc.Collation)
 
 	tp = types.NewFieldType(mysql.TypeEnum)
@@ -77,6 +77,6 @@ func TestColumnToProto(t *testing.T) {
 	col2 := &model.ColumnInfo{
 		FieldType: *tp,
 	}
-	pc = util.ColumnToProto(col2)
+	pc = util.ColumnToProto(col2, false)
 	require.Len(t, pc.Elems, 2)
 }