add getPkExpr unit test (matrixorigin#14232)

add getPkExpr unit test Approved by: @triump2020, @ouyuanning, @aunjgr
w-zr · Jan 16, 2024 · 5049505 · 5049505
1 parent 40a9a47
commit 5049505
Show file tree

Hide file tree

Showing 2 changed files with 145 additions and 0 deletions.
diff --git 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
@@ -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())
+}