Skip to content

Commit

Permalink
executor: write multi-valued index (#40172)
Browse files Browse the repository at this point in the history 
close #40207
  • Loading branch information
@xiongjiwei
xiongjiwei authored Dec 30, 2022
1 parent 0134d0b commit d3b952a
Show file tree
Hide file tree
Showing 14 changed files with 570 additions and 192 deletions.
4 changes: 2 additions & 2 deletions executor/admin.go
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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}
}
Expand Down Expand Up @@ -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}
Expand Down
4 changes: 2 additions & 2 deletions executor/builder.go
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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 {
Expand Down Expand Up @@ -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,
}
Expand Down
257 changes: 257 additions & 0 deletions expression/multi_valued_index_test.go
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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) {
Expand Down Expand Up @@ -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
}
3 changes: 3 additions & 0 deletions planner/core/logical_plan_builder.go
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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
}
Expand Down
2 changes: 1 addition & 1 deletion planner/core/plan_to_pb.go
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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,
}
Expand Down
16 changes: 8 additions & 8 deletions planner/core/plan_to_pb_test.go
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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())
}
Expand All @@ -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)
Expand All @@ -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)
}
Loading

0 comments on commit d3b952a

Please sign in to comment.