wehub-resource-sync
498b235461
Build and test / Build and test AMD64 Ubuntu 22.04 (push) Failing after 0s
Publish Builder / amazonlinux2023 (push) Failing after 1s
Build and test / UT for Go (push) Has been skipped
Publish KRTE Images / KRTE (push) Failing after 1s
Build and test / Integration Test (push) Has been skipped
Build and test / Upload Code Coverage (push) Has been skipped
Publish Builder / rockylinux9 (push) Failing after 1s
Publish Builder / ubuntu22.04 (push) Failing after 0s
Publish Builder / ubuntu24.04 (push) Failing after 0s
Publish Gpu Builder / publish-gpu-builder (push) Failing after 1s
Publish Test Images / PyTest (push) Failing after 0s
Build and test / UT for Cpp (push) Has been cancelled
433 lines
16 KiB
Go
433 lines
16 KiB
Go
package exprutil
|
|
|
|
import (
|
|
"testing"
|
|
|
|
"github.com/stretchr/testify/assert"
|
|
"github.com/stretchr/testify/require"
|
|
|
|
"github.com/milvus-io/milvus-proto/go-api/v3/commonpb"
|
|
"github.com/milvus-io/milvus-proto/go-api/v3/schemapb"
|
|
"github.com/milvus-io/milvus/internal/parser/planparserv2"
|
|
"github.com/milvus-io/milvus/internal/util/testutil"
|
|
"github.com/milvus-io/milvus/pkg/v3/common"
|
|
"github.com/milvus-io/milvus/pkg/v3/proto/planpb"
|
|
"github.com/milvus-io/milvus/pkg/v3/util/funcutil"
|
|
"github.com/milvus-io/milvus/pkg/v3/util/typeutil"
|
|
)
|
|
|
|
func mkColumnInfo(isPK bool) *planpb.ColumnInfo {
|
|
return &planpb.ColumnInfo{IsPrimaryKey: isPK, DataType: schemapb.DataType_Int64}
|
|
}
|
|
|
|
func mkTermExpr(isPK bool) *planpb.Expr {
|
|
return &planpb.Expr{Expr: &planpb.Expr_TermExpr{TermExpr: &planpb.TermExpr{
|
|
ColumnInfo: mkColumnInfo(isPK),
|
|
Values: []*planpb.GenericValue{{Val: &planpb.GenericValue_Int64Val{Int64Val: 1}}},
|
|
}}}
|
|
}
|
|
|
|
func mkUnaryRangeExpr(isPK bool) *planpb.Expr {
|
|
return &planpb.Expr{Expr: &planpb.Expr_UnaryRangeExpr{UnaryRangeExpr: &planpb.UnaryRangeExpr{
|
|
ColumnInfo: mkColumnInfo(isPK),
|
|
Op: planpb.OpType_Equal,
|
|
Value: &planpb.GenericValue{Val: &planpb.GenericValue_Int64Val{Int64Val: 5}},
|
|
}}}
|
|
}
|
|
|
|
func mkBinaryRangeExpr(isPK bool) *planpb.Expr {
|
|
return &planpb.Expr{Expr: &planpb.Expr_BinaryRangeExpr{BinaryRangeExpr: &planpb.BinaryRangeExpr{
|
|
ColumnInfo: mkColumnInfo(isPK),
|
|
LowerValue: &planpb.GenericValue{Val: &planpb.GenericValue_Int64Val{Int64Val: 1}},
|
|
UpperValue: &planpb.GenericValue{Val: &planpb.GenericValue_Int64Val{Int64Val: 10}},
|
|
LowerInclusive: true,
|
|
UpperInclusive: false,
|
|
}}}
|
|
}
|
|
|
|
func mkBinaryExpr(op planpb.BinaryExpr_BinaryOp, left, right *planpb.Expr) *planpb.Expr {
|
|
return &planpb.Expr{Expr: &planpb.Expr_BinaryExpr{BinaryExpr: &planpb.BinaryExpr{
|
|
Op: op, Left: left, Right: right,
|
|
}}}
|
|
}
|
|
|
|
func mkUnaryExpr(child *planpb.Expr) *planpb.Expr {
|
|
return &planpb.Expr{Expr: &planpb.Expr_UnaryExpr{UnaryExpr: &planpb.UnaryExpr{
|
|
Op: planpb.UnaryExpr_Not,
|
|
Child: child,
|
|
}}}
|
|
}
|
|
|
|
func TestHasOptimizablePkPredicate(t *testing.T) {
|
|
tests := []struct {
|
|
name string
|
|
expr *planpb.Expr
|
|
expect bool
|
|
}{
|
|
{"nil", nil, false},
|
|
{"term on pk", mkTermExpr(true), true},
|
|
{"term on non-pk", mkTermExpr(false), false},
|
|
{"unary range on pk", mkUnaryRangeExpr(true), true},
|
|
{"unary range on non-pk", mkUnaryRangeExpr(false), false},
|
|
{"binary range on pk", mkBinaryRangeExpr(true), true},
|
|
{"binary range on non-pk", mkBinaryRangeExpr(false), false},
|
|
{"AND: pk left, non-pk right", mkBinaryExpr(planpb.BinaryExpr_LogicalAnd, mkTermExpr(true), mkTermExpr(false)), true},
|
|
{"AND: non-pk left, pk right", mkBinaryExpr(planpb.BinaryExpr_LogicalAnd, mkTermExpr(false), mkTermExpr(true)), true},
|
|
{"AND: neither pk", mkBinaryExpr(planpb.BinaryExpr_LogicalAnd, mkTermExpr(false), mkTermExpr(false)), false},
|
|
{"OR: both pk", mkBinaryExpr(planpb.BinaryExpr_LogicalOr, mkTermExpr(true), mkTermExpr(true)), true},
|
|
{"OR: one non-pk", mkBinaryExpr(planpb.BinaryExpr_LogicalOr, mkTermExpr(true), mkTermExpr(false)), false},
|
|
{"NOT wrapping pk", mkUnaryExpr(mkTermExpr(true)), false},
|
|
}
|
|
for _, tc := range tests {
|
|
t.Run(tc.name, func(t *testing.T) {
|
|
assert.Equal(t, tc.expect, HasOptimizablePkPredicate(tc.expr))
|
|
})
|
|
}
|
|
}
|
|
|
|
func TestParsePartitionKeys(t *testing.T) {
|
|
prefix := "TestParsePartitionKeys"
|
|
collectionName := prefix + funcutil.GenRandomStr()
|
|
|
|
fieldName2Type := make(map[string]schemapb.DataType)
|
|
fieldName2Type["int64_field"] = schemapb.DataType_Int64
|
|
fieldName2Type["varChar_field"] = schemapb.DataType_VarChar
|
|
fieldName2Type["fvec_field"] = schemapb.DataType_FloatVector
|
|
schema := testutil.ConstructCollectionSchemaByDataType(collectionName, fieldName2Type,
|
|
"int64_field", false, 8)
|
|
partitionKeyField := &schemapb.FieldSchema{
|
|
Name: "partition_key_field",
|
|
DataType: schemapb.DataType_Int64,
|
|
IsPartitionKey: true,
|
|
}
|
|
schema.Fields = append(schema.Fields, partitionKeyField)
|
|
|
|
fieldID := common.StartOfUserFieldID
|
|
for _, field := range schema.Fields {
|
|
field.FieldID = int64(fieldID)
|
|
fieldID++
|
|
}
|
|
schemaHelper, err := typeutil.CreateSchemaHelper(schema)
|
|
require.NoError(t, err)
|
|
|
|
queryInfo := &planpb.QueryInfo{
|
|
Topk: 10,
|
|
MetricType: "L2",
|
|
SearchParams: "",
|
|
RoundDecimal: -1,
|
|
}
|
|
|
|
type testCase struct {
|
|
name string
|
|
expr string
|
|
expected int
|
|
validPartitionKeys []int64
|
|
invalidPartitionKeys []int64
|
|
}
|
|
cases := []testCase{
|
|
{
|
|
name: "binary_expr_and with term",
|
|
expr: "partition_key_field in [7, 8] && int64_field >= 10",
|
|
expected: 2,
|
|
validPartitionKeys: []int64{7, 8},
|
|
invalidPartitionKeys: []int64{},
|
|
},
|
|
{
|
|
name: "binary_expr_and with equal",
|
|
expr: "partition_key_field == 7 && int64_field >= 10",
|
|
expected: 1,
|
|
validPartitionKeys: []int64{7},
|
|
invalidPartitionKeys: []int64{},
|
|
},
|
|
{
|
|
name: "binary_expr_and with term2",
|
|
expr: "partition_key_field in [7, 8] && int64_field == 10",
|
|
expected: 2,
|
|
validPartitionKeys: []int64{7, 8},
|
|
invalidPartitionKeys: []int64{10},
|
|
},
|
|
{
|
|
name: "binary_expr_and with partition key in range",
|
|
expr: "partition_key_field in [1, 2, 3, 4, 5, 6, 7, 8, 9, 10] && partition_key_field > 11",
|
|
expected: 0,
|
|
validPartitionKeys: []int64{},
|
|
invalidPartitionKeys: []int64{},
|
|
},
|
|
{
|
|
name: "binary_expr_and with partition key in range2",
|
|
expr: "int64_field == 10 && partition_key_field > 9",
|
|
expected: 0,
|
|
validPartitionKeys: []int64{},
|
|
invalidPartitionKeys: []int64{},
|
|
},
|
|
{
|
|
name: "binary_expr_and with term and not",
|
|
expr: "partition_key_field in [7, 8] && partition_key_field not in [10, 20]",
|
|
expected: 2,
|
|
validPartitionKeys: []int64{7, 8},
|
|
invalidPartitionKeys: []int64{10, 20},
|
|
},
|
|
{
|
|
name: "binary_expr_or with term and not",
|
|
expr: "partition_key_field in [7, 8] or partition_key_field not in [10, 20]",
|
|
expected: 0,
|
|
validPartitionKeys: []int64{},
|
|
invalidPartitionKeys: []int64{},
|
|
},
|
|
{
|
|
name: "binary_expr_or with term and not 2",
|
|
expr: "partition_key_field in [7, 8] or int64_field not in [10, 20]",
|
|
expected: 0,
|
|
validPartitionKeys: []int64{},
|
|
invalidPartitionKeys: []int64{},
|
|
},
|
|
}
|
|
|
|
for _, tc := range cases {
|
|
idx := 0
|
|
t.Run(tc.name, func(t *testing.T) {
|
|
idx++
|
|
t.Log(idx, tc.name, tc.expr)
|
|
// test search plan
|
|
searchPlan, err := planparserv2.CreateSearchPlan(schemaHelper, tc.expr, "fvec_field", queryInfo, nil, nil)
|
|
assert.NoError(t, err)
|
|
expr, err := ParseExprFromPlan(searchPlan)
|
|
assert.NoError(t, err)
|
|
partitionKeys := ParseKeys(expr, PartitionKey)
|
|
assert.Equal(t, tc.expected, len(partitionKeys))
|
|
for _, key := range partitionKeys {
|
|
int64Val := key.Val.(*planpb.GenericValue_Int64Val).Int64Val
|
|
assert.Contains(t, tc.validPartitionKeys, int64Val)
|
|
assert.NotContains(t, tc.invalidPartitionKeys, int64Val)
|
|
}
|
|
|
|
// test query plan
|
|
queryPlan, err := planparserv2.CreateRetrievePlan(schemaHelper, tc.expr, nil)
|
|
assert.NoError(t, err)
|
|
expr, err = ParseExprFromPlan(queryPlan)
|
|
assert.NoError(t, err)
|
|
partitionKeys = ParseKeys(expr, PartitionKey)
|
|
assert.Equal(t, tc.expected, len(partitionKeys))
|
|
for _, key := range partitionKeys {
|
|
int64Val := key.Val.(*planpb.GenericValue_Int64Val).Int64Val
|
|
assert.Contains(t, tc.validPartitionKeys, int64Val)
|
|
assert.NotContains(t, tc.invalidPartitionKeys, int64Val)
|
|
}
|
|
})
|
|
}
|
|
}
|
|
|
|
func TestValidatePartitionKeyIsolation(t *testing.T) {
|
|
prefix := "TestValidatePartitionKeyIsolation"
|
|
collectionName := prefix + funcutil.GenRandomStr()
|
|
|
|
fieldName2Type := make(map[string]schemapb.DataType)
|
|
fieldName2Type["int64_field"] = schemapb.DataType_Int64
|
|
fieldName2Type["varChar_field"] = schemapb.DataType_VarChar
|
|
fieldName2Type["fvec_field"] = schemapb.DataType_FloatVector
|
|
schema := testutil.ConstructCollectionSchemaByDataType(collectionName, fieldName2Type,
|
|
"int64_field", false, 8)
|
|
schema.Properties = append(schema.Properties, &commonpb.KeyValuePair{
|
|
Key: common.PartitionKeyIsolationKey,
|
|
Value: "true",
|
|
})
|
|
partitionKeyField := &schemapb.FieldSchema{
|
|
Name: "key_field",
|
|
DataType: schemapb.DataType_Int64,
|
|
IsPartitionKey: true,
|
|
}
|
|
schema.Fields = append(schema.Fields, partitionKeyField)
|
|
fieldID := common.StartOfUserFieldID
|
|
for _, field := range schema.Fields {
|
|
field.FieldID = int64(fieldID)
|
|
fieldID++
|
|
}
|
|
schemaHelper, err := typeutil.CreateSchemaHelper(schema)
|
|
require.NoError(t, err)
|
|
|
|
type testCase struct {
|
|
name string
|
|
expr string
|
|
expectedErrorString string
|
|
}
|
|
cases := []testCase{
|
|
{
|
|
name: "partition key isolation equal",
|
|
expr: "key_field == 10",
|
|
expectedErrorString: "",
|
|
},
|
|
{
|
|
name: "partition key isolation equal AND with same field equal",
|
|
expr: "key_field == 10 && key_field == 10",
|
|
expectedErrorString: "",
|
|
},
|
|
{
|
|
name: "partition key isolation equal AND with same field equal diff",
|
|
expr: "key_field == 10 && key_field == 20",
|
|
expectedErrorString: "",
|
|
},
|
|
{
|
|
name: "partition key isolation equal AND with same field equal 3",
|
|
expr: "key_field == 10 && key_field == 11 && key_field == 12",
|
|
expectedErrorString: "",
|
|
},
|
|
{
|
|
name: "partition key isolation equal AND with varchar field equal",
|
|
expr: "key_field == 10 && varChar_field == 'a'",
|
|
expectedErrorString: "",
|
|
},
|
|
{
|
|
name: "partition key isolation equal AND with varchar field not equal",
|
|
expr: "key_field == 10 && varChar_field != 'a'",
|
|
expectedErrorString: "",
|
|
},
|
|
{
|
|
name: "partition key isolation equal AND with varchar field in",
|
|
expr: "key_field == 10 && varChar_field in ['a', 'b']",
|
|
expectedErrorString: "",
|
|
},
|
|
{
|
|
name: "partition key isolation equal AND with varchar field in Reversed",
|
|
expr: "varChar_field in ['a', 'b'] && key_field == 10",
|
|
expectedErrorString: "",
|
|
},
|
|
{
|
|
name: "partition key isolation equal AND with varchar field OR",
|
|
expr: "key_field == 10 && (varChar_field == 'a' || varChar_field == 'b')",
|
|
expectedErrorString: "",
|
|
},
|
|
{
|
|
name: "partition key isolation equal AND with varchar field OR Reversed",
|
|
expr: "(varChar_field == 'a' || varChar_field == 'b') && key_field == 10",
|
|
expectedErrorString: "",
|
|
},
|
|
{
|
|
name: "partition key isolation equal to arithmic operations",
|
|
expr: "key_field == (1+1)",
|
|
expectedErrorString: "",
|
|
},
|
|
{
|
|
name: "partition key isolation empty",
|
|
expr: "",
|
|
expectedErrorString: "partition key not found in expr or the expr is invalid when validating partition key isolation",
|
|
},
|
|
{
|
|
name: "partition key isolation not equal",
|
|
expr: "key_field != 10",
|
|
expectedErrorString: "partition key isolation does not support NotEqual",
|
|
},
|
|
{
|
|
name: "partition key isolation term",
|
|
expr: "key_field in [1, 2, 3, 4, 5, 6, 7, 8, 9, 10]",
|
|
expectedErrorString: "partition key isolation does not support IN",
|
|
},
|
|
{
|
|
name: "partition key isolation term multiple",
|
|
expr: "key_field in [1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16, 17, 18, 19, 20]",
|
|
expectedErrorString: "partition key isolation does not support IN",
|
|
},
|
|
{
|
|
name: "partition key isolation NOT term",
|
|
expr: "key_field not in [1, 2, 3, 4, 5, 6, 7, 8, 9, 10]",
|
|
expectedErrorString: "partition key isolation does not support IN",
|
|
},
|
|
{
|
|
name: "partition key isolation less",
|
|
expr: "key_field < 10",
|
|
expectedErrorString: "partition key isolation does not support LessThan",
|
|
},
|
|
{
|
|
name: "partition key isolation less or equal",
|
|
expr: "key_field <= 10",
|
|
expectedErrorString: "partition key isolation does not support LessEq",
|
|
},
|
|
{
|
|
name: "partition key isolation greater",
|
|
expr: "key_field > 10",
|
|
expectedErrorString: "partition key isolation does not support GreaterThan",
|
|
},
|
|
{
|
|
name: "partition key isolation equal greator or equal",
|
|
expr: "key_field >= 10",
|
|
expectedErrorString: "partition key isolation does not support GreaterEqual",
|
|
},
|
|
{
|
|
name: "partition key isolation binary range",
|
|
expr: "1 < key_field < 10",
|
|
expectedErrorString: "partition key isolation does not support BinaryRange",
|
|
},
|
|
{
|
|
name: "partition key isolation NOT equal",
|
|
expr: "not(key_field == 10)",
|
|
expectedErrorString: "partition key isolation does not support NotEqual",
|
|
},
|
|
{
|
|
name: "partition key isolation equal AND with same field term",
|
|
expr: "key_field == 10 && key_field in [10]",
|
|
expectedErrorString: "",
|
|
},
|
|
{
|
|
name: "partition key isolation equal OR with same field equal",
|
|
expr: "key_field == 10 || key_field == 11",
|
|
expectedErrorString: "partition key isolation does not support",
|
|
},
|
|
{
|
|
name: "partition key isolation equal OR with same field equal Reversed",
|
|
expr: "key_field == 11 || key_field == 10",
|
|
expectedErrorString: "partition key isolation does not support",
|
|
},
|
|
{
|
|
name: "partition key isolation equal OR with other field equal",
|
|
expr: "key_field == 10 || varChar_field == 'a'",
|
|
expectedErrorString: "partition key isolation does not support",
|
|
},
|
|
{
|
|
name: "partition key isolation equal OR with other field equal Reversed",
|
|
expr: "varChar_field == 'a' || key_field == 10",
|
|
expectedErrorString: "partition key isolation does not support",
|
|
},
|
|
{
|
|
name: "partition key isolation equal OR with other field equal",
|
|
expr: "key_field == 10 || varChar_field == 'a'",
|
|
expectedErrorString: "partition key isolation does not support",
|
|
},
|
|
{
|
|
// Rewriter merges OR into IN, then AND+IN+Equal simplifies:
|
|
// key_field == 10 && key_field in [10, 11] → key_field == 10
|
|
// Final plan is just ==, so isolation validation passes.
|
|
name: "partition key isolation equal AND",
|
|
expr: "key_field == 10 && (key_field == 10 || key_field == 11)",
|
|
expectedErrorString: "",
|
|
},
|
|
{
|
|
name: "partition key isolation other field equal",
|
|
expr: "varChar_field == 'a'",
|
|
expectedErrorString: "partition key not found in expr or the expr is invalid when validating partition key isolation",
|
|
},
|
|
{
|
|
name: "partition key isolation other field equal AND",
|
|
expr: "varChar_field == 'a' && int64_field == 1",
|
|
expectedErrorString: "partition key not found in expr or the expr is invalid when validating partition key isolation",
|
|
},
|
|
{
|
|
name: "partition key isolation complex OR",
|
|
expr: "(key_field == 10 and int64_field == 11) or (key_field == 10 and varChar_field == 'a')",
|
|
expectedErrorString: "partition key isolation does not support",
|
|
},
|
|
}
|
|
|
|
for _, tc := range cases {
|
|
t.Run(tc.name, func(t *testing.T) {
|
|
queryPlan, err := planparserv2.CreateRetrievePlan(schemaHelper, tc.expr, nil)
|
|
assert.NoError(t, err)
|
|
planExpr, err := ParseExprFromPlan(queryPlan)
|
|
assert.NoError(t, err)
|
|
if tc.expectedErrorString != "" {
|
|
assert.ErrorContains(t, ValidatePartitionKeyIsolation(planExpr), tc.expectedErrorString)
|
|
} else {
|
|
assert.NoError(t, ValidatePartitionKeyIsolation(planExpr))
|
|
}
|
|
})
|
|
}
|
|
}
|