diff --git a/pkg/sql/plan/make.go b/pkg/sql/plan/make.go index a52b813bd63de..6f904b8487feb 100644 --- a/pkg/sql/plan/make.go +++ b/pkg/sql/plan/make.go @@ -170,6 +170,23 @@ func makePlan2Int64ConstExprWithType(v int64) *plan.Expr { } } +var MakePlan2Int64VecExprWithType = makePlan2Int64VecExprWithType + +func makePlan2Int64VecExprWithType(vals ...int64) *plan.Expr { + data := types.EncodeSlice[int64](vals) + return &plan.Expr{ + Typ: &plan.Type{ + Id: int32(types.T_tuple), + }, + Expr: &plan.Expr_Vec{ + Vec: &plan.LiteralVec{ + Len: int32(len(vals)), + Data: data, + }, + }, + } +} + func makePlan2Uint64ConstExpr(v uint64) *plan.Expr_Lit { return &plan.Expr_Lit{Lit: &plan.Literal{ Isnull: false, diff --git a/pkg/vm/engine/disttae/util_test.go b/pkg/vm/engine/disttae/util_test.go index 922ff6f8b8850..610733ab28fc7 100644 --- a/pkg/vm/engine/disttae/util_test.go +++ b/pkg/vm/engine/disttae/util_test.go @@ -78,6 +78,37 @@ func makeFunctionExprForTest(name string, args []*plan.Expr) *plan.Expr { } } +func makeInExprForTest[T any](arg0 *plan.Expr, vals []T) *plan.Expr { + return &plan.Expr{ + Typ: &plan.Type{ + Id: int32(types.T_bool), + NotNullable: true, + }, + Expr: &plan.Expr_F{ + F: &plan.Function{ + Func: &plan.ObjectRef{ + Obj: function.InFunctionEncodedID, + ObjName: function.InFunctionName, + }, + Args: []*plan.Expr{ + arg0, + { + Typ: &plan.Type{ + Id: int32(types.T_tuple), + }, + Expr: &plan.Expr_Vec{ + Vec: &plan.LiteralVec{ + Len: int32(len(vals)), + Data: types.EncodeSlice[T](vals), + }, + }, + }, + }, + }, + }, + } +} + func TestBlockMetaMarshal(t *testing.T) { location := []byte("test") var info objectio.BlockInfo @@ -867,3 +898,100 @@ func TestForeachBlkInObjStatsList(t *testing.T) { require.Equal(t, count, 0) } + +func TestGetPKExpr(t *testing.T) { + m := mpool.MustNewNoFixed(t.Name()) + proc := testutil.NewProcessWithMPool(m) + type myCase struct { + desc []string + exprs []*plan.Expr + valExprs []*plan.Expr + } + tc := myCase{ + desc: []string{ + "a=10", + "a=20 and a=10", + "30=a and 20=a", + "a in (1,2)", + "b=40 and a=50", + "a=60 or b=70", + "b=80 and c=90", + }, + exprs: []*plan.Expr{ + makeFunctionExprForTest("=", []*plan.Expr{ + makeColExprForTest(0, types.T_int64), + plan2.MakePlan2Int64ConstExprWithType(10), + }), + makeFunctionExprForTest("and", []*plan.Expr{ + makeFunctionExprForTest("=", []*plan.Expr{ + makeColExprForTest(0, types.T_int64), + plan2.MakePlan2Int64ConstExprWithType(20), + }), + makeFunctionExprForTest("=", []*plan.Expr{ + makeColExprForTest(0, types.T_int64), + plan2.MakePlan2Int64ConstExprWithType(10), + }), + }), + makeFunctionExprForTest("and", []*plan.Expr{ + makeFunctionExprForTest("=", []*plan.Expr{ + plan2.MakePlan2Int64ConstExprWithType(30), + makeColExprForTest(0, types.T_int64), + }), + makeFunctionExprForTest("=", []*plan.Expr{ + plan2.MakePlan2Int64ConstExprWithType(20), + makeColExprForTest(0, types.T_int64), + }), + }), + makeInExprForTest[int64]( + makeColExprForTest(0, types.T_int64), + []int64{1, 2}, + ), + makeFunctionExprForTest("and", []*plan.Expr{ + makeFunctionExprForTest("=", []*plan.Expr{ + makeColExprForTest(1, types.T_int64), + plan2.MakePlan2Int64ConstExprWithType(40), + }), + makeFunctionExprForTest("=", []*plan.Expr{ + makeColExprForTest(0, types.T_int64), + plan2.MakePlan2Int64ConstExprWithType(50), + }), + }), + makeFunctionExprForTest("or", []*plan.Expr{ + makeFunctionExprForTest("=", []*plan.Expr{ + makeColExprForTest(0, types.T_int64), + plan2.MakePlan2Int64ConstExprWithType(60), + }), + makeFunctionExprForTest("=", []*plan.Expr{ + makeColExprForTest(1, types.T_int64), + plan2.MakePlan2Int64ConstExprWithType(70), + }), + }), + makeFunctionExprForTest("and", []*plan.Expr{ + makeFunctionExprForTest("=", []*plan.Expr{ + makeColExprForTest(1, types.T_int64), + plan2.MakePlan2Int64ConstExprWithType(80), + }), + makeFunctionExprForTest("=", []*plan.Expr{ + makeColExprForTest(2, types.T_int64), + plan2.MakePlan2Int64ConstExprWithType(90), + }), + }), + }, + valExprs: []*plan.Expr{ + plan2.MakePlan2Int64ConstExprWithType(10), + plan2.MakePlan2Int64ConstExprWithType(20), + plan2.MakePlan2Int64ConstExprWithType(30), + plan2.MakePlan2Int64VecExprWithType(int64(1), int64(2)), + plan2.MakePlan2Int64ConstExprWithType(50), + nil, + nil, + }, + } + pkName := "a" + for i, expr := range tc.exprs { + rExpr := getPkExpr(expr, pkName, proc) + // t.Log(plan2.FormatExpr(rExpr)) + require.Equal(t, tc.valExprs[i], rExpr) + } + require.Zero(t, m.CurrNB()) +}