mirror of
https://github.com/milvus-io/milvus.git
synced 2026-07-21 10:15:43 +00:00
enhance: Optimize bool IN/NOT IN expressions with nullable field handling (#48459)
issue: https://github.com/milvus-io/milvus/issues/45525 ## Summary This PR adds optimizations for boolean IN and NOT IN expressions, with special handling for nullable boolean fields. The changes improve query performance by rewriting certain bool IN expressions to simpler, faster operations. ## Key Changes - **Bool IN expression optimizations:** - `BoolField IN [true, false]` → `AlwaysTrueExpr` (non-nullable) or `IS NOT NULL` (nullable) - `BoolField IN [true]` → `BoolField == true` (uses fast SIMD path) - `BoolField IN [false]` → `BoolField == false` - Single-value deduplication: `BoolField IN [true, true]` → `BoolField == true` - **Bool NOT IN expression optimizations:** - `BoolField NOT IN [true, false]` → `AlwaysFalseExpr` (non-nullable) or `IS NULL` (nullable) - `BoolField NOT IN [true]` → `BoolField != true` - `BoolField NOT IN [false]` → `BoolField != false` - **Unary expression optimizations:** - `NOT AlwaysTrueExpr` → `AlwaysFalseExpr` - `NOT (IS NOT NULL)` → `IS NULL` (and vice versa) - `NOT (col == val)` → `col != val` - **Helper functions and test infrastructure:** - Added `newNullExpr()` utility to create NULL check expressions - Added `allBoolVals()` helper to validate boolean value lists - Added `buildSchemaHelperForRewriteNullableT()` for testing nullable fields - Comprehensive test coverage for all bool IN/NOT IN optimization paths ## Implementation Details The optimization logic is implemented in the `visitTermExpr()` method, which: 1. Detects when a TermExpr operates on a boolean column 2. Analyzes the values to determine if both true and false are present 3. Rewrites to appropriate expressions based on nullability and value set 4. Handles single-value cases by converting to equality comparisons The unary expression visitor was enhanced to handle NOT operations on the newly created expressions, enabling proper composition of optimizations (e.g., NOT IN becomes NOT of an optimized expression). https://claude.ai/code/session_01GkYztNzb75zgjuevzqVcwb --------- Signed-off-by: Claude <noreply@anthropic.com> Signed-off-by: Buqian Zheng <zhengbuqian@gmail.com> Co-authored-by: Claude <noreply@anthropic.com>
This commit is contained in:
@@ -429,10 +429,11 @@ TEST_F(ExprMaterializedViewTest, TestUnaryRangeCompareExpr) {
|
||||
TEST_F(ExprMaterializedViewTest, TestInMultipleExpr) {
|
||||
constexpr int kNumValues = 15;
|
||||
for (const auto& data_type : GetNumericAndVarcharScalarDataTypes()) {
|
||||
// Skip BOOL: only 2 unique values (true/false), always below
|
||||
// the IN threshold so it gets split into OR equals.
|
||||
if (data_type == DataType::BOOL)
|
||||
// Skip Bool: "BoolField in [false, true]" covers the full domain
|
||||
// and is optimized to AlwaysTrueExpr by the expression rewriter.
|
||||
if (data_type == DataType::BOOL) {
|
||||
continue;
|
||||
}
|
||||
std::string field_name = GetFieldName(data_type);
|
||||
std::string vals;
|
||||
for (int i = 0; i < kNumValues; i++) {
|
||||
@@ -456,6 +457,11 @@ TEST_F(ExprMaterializedViewTest, TestInMultipleExpr) {
|
||||
// Test numeric and varchar expr: F0 not in [A]
|
||||
TEST_F(ExprMaterializedViewTest, TestUnaryLogicalNotInMultipleExpr) {
|
||||
for (const auto& data_type : GetNumericAndVarcharScalarDataTypes()) {
|
||||
// Skip Bool: "BoolField not in [false, true]" covers the full domain
|
||||
// and is optimized to AlwaysFalseExpr by the expression rewriter.
|
||||
if (data_type == DataType::BOOL) {
|
||||
continue;
|
||||
}
|
||||
std::string field_name = GetFieldName(data_type);
|
||||
std::string val0 = GetTestValue(data_type, 0);
|
||||
std::string val1 = GetTestValue(data_type, 1);
|
||||
|
||||
@@ -1,6 +1,7 @@
|
||||
package rewriter
|
||||
|
||||
import (
|
||||
"github.com/milvus-io/milvus-proto/go-api/v2/schemapb"
|
||||
"github.com/milvus-io/milvus/pkg/v2/proto/planpb"
|
||||
"github.com/milvus-io/milvus/pkg/v2/util/paramtable"
|
||||
)
|
||||
@@ -85,10 +86,13 @@ func (v *visitor) visitBinaryExpr(expr *planpb.BinaryExpr) interface{} {
|
||||
|
||||
func (v *visitor) visitUnaryExpr(expr *planpb.UnaryExpr) interface{} {
|
||||
// Handle NOT(TermExpr) before visiting child, so we can split not in → != and
|
||||
// Skip bool types here — they are handled via visitTermExpr bool optimization + NOT simplification.
|
||||
if expr.GetOp() == planpb.UnaryExpr_Not {
|
||||
if te := expr.GetChild().GetTermExpr(); te != nil {
|
||||
sortTermValues(te)
|
||||
if !shouldUseInExprWithPK(te.GetColumnInfo().GetDataType(), len(te.GetValues()), te.GetColumnInfo().GetIsPrimaryKey()) {
|
||||
if v.optimizeEnabled && effectiveDataType(te.GetColumnInfo()) == schemapb.DataType_Bool {
|
||||
// Let bool NOT IN flow through to visitTermExpr for bool-specific optimization
|
||||
} else if !shouldUseInExprWithPK(te.GetColumnInfo().GetDataType(), len(te.GetValues()), te.GetColumnInfo().GetIsPrimaryKey()) {
|
||||
return splitTermToAndNotEquals(te)
|
||||
}
|
||||
}
|
||||
@@ -96,11 +100,31 @@ func (v *visitor) visitUnaryExpr(expr *planpb.UnaryExpr) interface{} {
|
||||
|
||||
child := v.visitExpr(expr.GetChild()).(*planpb.Expr)
|
||||
|
||||
// Optimize double negation: NOT (NOT AlwaysTrue) → AlwaysTrue
|
||||
if expr.GetOp() == planpb.UnaryExpr_Not {
|
||||
// NOT (NOT AlwaysTrue) → AlwaysTrue
|
||||
if IsAlwaysFalseExpr(child) {
|
||||
return newAlwaysTrueExpr()
|
||||
}
|
||||
// NOT AlwaysTrueExpr → AlwaysFalseExpr
|
||||
// Handles: non-nullable bool NOT IN [true, false] → AlwaysFalse
|
||||
if IsAlwaysTrueExpr(child) {
|
||||
return newAlwaysFalseExpr()
|
||||
}
|
||||
// NOT (IS NOT NULL) → IS NULL
|
||||
// Handles: nullable bool NOT IN [true, false] → IS NULL
|
||||
if ne := child.GetNullExpr(); ne != nil {
|
||||
if ne.GetOp() == planpb.NullExpr_IsNotNull {
|
||||
return newNullExpr(ne.GetColumnInfo(), planpb.NullExpr_IsNull)
|
||||
}
|
||||
if ne.GetOp() == planpb.NullExpr_IsNull {
|
||||
return newNullExpr(ne.GetColumnInfo(), planpb.NullExpr_IsNotNull)
|
||||
}
|
||||
}
|
||||
// NOT (col == val) → col != val
|
||||
// Handles: bool NOT IN [true] → != true, bool NOT IN [false] → != false
|
||||
if ure := child.GetUnaryRangeExpr(); ure != nil && ure.GetOp() == planpb.OpType_Equal {
|
||||
return newUnaryRangeExpr(ure.GetColumnInfo(), planpb.OpType_NotEqual, ure.GetValue())
|
||||
}
|
||||
}
|
||||
|
||||
return &planpb.Expr{
|
||||
@@ -116,6 +140,33 @@ func (v *visitor) visitUnaryExpr(expr *planpb.UnaryExpr) interface{} {
|
||||
func (v *visitor) visitTermExpr(expr *planpb.TermExpr) interface{} {
|
||||
sortTermValues(expr)
|
||||
|
||||
// Optimize bool IN expressions:
|
||||
// - in [true, false] → AlwaysTrueExpr (non-nullable) or IS NOT NULL (nullable)
|
||||
// - in [true] → == true (uses fast SIMD path instead of slow scalar loop)
|
||||
// - in [false] → == false
|
||||
if v.optimizeEnabled && effectiveDataType(expr.GetColumnInfo()) == schemapb.DataType_Bool {
|
||||
values := expr.GetValues()
|
||||
if allBoolVals(values) {
|
||||
hasFalse, hasTrue := false, false
|
||||
for _, val := range values {
|
||||
if val.GetBoolVal() {
|
||||
hasTrue = true
|
||||
} else {
|
||||
hasFalse = true
|
||||
}
|
||||
}
|
||||
if hasTrue && hasFalse {
|
||||
if expr.GetColumnInfo().GetNullable() {
|
||||
return newNullExpr(expr.GetColumnInfo(), planpb.NullExpr_IsNotNull)
|
||||
}
|
||||
return newAlwaysTrueExpr()
|
||||
}
|
||||
if len(values) == 1 {
|
||||
return newUnaryRangeExpr(expr.GetColumnInfo(), planpb.OpType_Equal, values[0])
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// Split in → == or when shouldUseInExpr returns false
|
||||
if !shouldUseInExprWithPK(expr.GetColumnInfo().GetDataType(), len(expr.GetValues()), expr.GetColumnInfo().GetIsPrimaryKey()) {
|
||||
return splitTermToOrEquals(expr)
|
||||
@@ -124,6 +175,19 @@ func (v *visitor) visitTermExpr(expr *planpb.TermExpr) interface{} {
|
||||
return &planpb.Expr{Expr: &planpb.Expr_TermExpr{TermExpr: expr}}
|
||||
}
|
||||
|
||||
// allBoolVals returns true if all values in the slice are BoolVal type.
|
||||
func allBoolVals(values []*planpb.GenericValue) bool {
|
||||
if len(values) == 0 {
|
||||
return false
|
||||
}
|
||||
for _, v := range values {
|
||||
if _, ok := v.GetVal().(*planpb.GenericValue_BoolVal); !ok {
|
||||
return false
|
||||
}
|
||||
}
|
||||
return true
|
||||
}
|
||||
|
||||
// visitValueExpr converts constant boolean ValueExpr to AlwaysTrueExpr/AlwaysFalseExpr.
|
||||
// This handles cases like "1==1" which the parser constant-folds into ValueExpr(bool=true),
|
||||
// normalizing them to the canonical AlwaysTrueExpr/AlwaysFalseExpr representation.
|
||||
|
||||
@@ -40,6 +40,21 @@ func buildSchemaHelperForRewriteT(t *testing.T) *typeutil.SchemaHelper {
|
||||
return helper
|
||||
}
|
||||
|
||||
func buildSchemaHelperForRewriteNullableT(t *testing.T) *typeutil.SchemaHelper {
|
||||
fields := []*schemapb.FieldSchema{
|
||||
{FieldID: 101, Name: "Int64Field", DataType: schemapb.DataType_Int64},
|
||||
{FieldID: 106, Name: "NullableBoolField", DataType: schemapb.DataType_Bool, Nullable: true},
|
||||
}
|
||||
schema := &schemapb.CollectionSchema{
|
||||
Name: "rewrite_nullable_test",
|
||||
AutoID: false,
|
||||
Fields: fields,
|
||||
}
|
||||
helper, err := typeutil.CreateSchemaHelper(schema)
|
||||
require.NoError(t, err)
|
||||
return helper
|
||||
}
|
||||
|
||||
// --- shouldUseInExpr threshold tests ---
|
||||
|
||||
func TestRewrite_OREquals_ToIN_VarChar_AboveThreshold(t *testing.T) {
|
||||
@@ -196,13 +211,115 @@ func TestRewrite_Term_SortAndDedup_Int(t *testing.T) {
|
||||
|
||||
func TestRewrite_In_SortAndDedup_Bool(t *testing.T) {
|
||||
helper := buildSchemaHelperForRewriteT(t)
|
||||
// BoolField in [true, false] covers all possible bool values → AlwaysTrueExpr
|
||||
expr, err := parser.ParseExpr(helper, `BoolField in [true,false,false,true]`, nil)
|
||||
require.NoError(t, err)
|
||||
require.NotNil(t, expr)
|
||||
// bool: 2 unique values < threshold 3 → split to == or
|
||||
be := expr.GetBinaryExpr()
|
||||
require.NotNil(t, be, "bool in [true,false] should split to == or")
|
||||
require.Equal(t, planpb.BinaryExpr_LogicalOr, be.GetOp())
|
||||
require.True(t, rewriter.IsAlwaysTrueExpr(expr),
|
||||
"bool IN [true, false] should be rewritten to AlwaysTrueExpr")
|
||||
}
|
||||
|
||||
func TestRewrite_Bool_In_SingleTrue_ToEqual(t *testing.T) {
|
||||
helper := buildSchemaHelperForRewriteT(t)
|
||||
expr, err := parser.ParseExpr(helper, `BoolField in [true]`, nil)
|
||||
require.NoError(t, err)
|
||||
require.NotNil(t, expr)
|
||||
ure := expr.GetUnaryRangeExpr()
|
||||
require.NotNil(t, ure, "bool IN [true] should be rewritten to == true")
|
||||
require.Equal(t, planpb.OpType_Equal, ure.GetOp())
|
||||
require.Equal(t, true, ure.GetValue().GetBoolVal())
|
||||
}
|
||||
|
||||
func TestRewrite_Bool_In_SingleFalse_ToEqual(t *testing.T) {
|
||||
helper := buildSchemaHelperForRewriteT(t)
|
||||
expr, err := parser.ParseExpr(helper, `BoolField in [false]`, nil)
|
||||
require.NoError(t, err)
|
||||
require.NotNil(t, expr)
|
||||
ure := expr.GetUnaryRangeExpr()
|
||||
require.NotNil(t, ure, "bool IN [false] should be rewritten to == false")
|
||||
require.Equal(t, planpb.OpType_Equal, ure.GetOp())
|
||||
require.Equal(t, false, ure.GetValue().GetBoolVal())
|
||||
}
|
||||
|
||||
func TestRewrite_Bool_In_DedupedSingleTrue_ToEqual(t *testing.T) {
|
||||
helper := buildSchemaHelperForRewriteT(t)
|
||||
// After dedup, [true, true] becomes [true] → == true
|
||||
expr, err := parser.ParseExpr(helper, `BoolField in [true, true]`, nil)
|
||||
require.NoError(t, err)
|
||||
require.NotNil(t, expr)
|
||||
ure := expr.GetUnaryRangeExpr()
|
||||
require.NotNil(t, ure, "bool IN [true, true] should dedup then rewrite to == true")
|
||||
require.Equal(t, planpb.OpType_Equal, ure.GetOp())
|
||||
require.Equal(t, true, ure.GetValue().GetBoolVal())
|
||||
}
|
||||
|
||||
func TestRewrite_Bool_In_TrueFalse_Nullable_ToIsNotNull(t *testing.T) {
|
||||
// For nullable bool, in [true, false] should become IS NOT NULL (not AlwaysTrueExpr)
|
||||
// because null values should not match.
|
||||
helper := buildSchemaHelperForRewriteNullableT(t)
|
||||
expr, err := parser.ParseExpr(helper, `NullableBoolField in [true,false]`, nil)
|
||||
require.NoError(t, err)
|
||||
require.NotNil(t, expr)
|
||||
nullExpr := expr.GetNullExpr()
|
||||
require.NotNil(t, nullExpr, "nullable bool IN [true, false] should be rewritten to IS NOT NULL")
|
||||
require.Equal(t, planpb.NullExpr_IsNotNull, nullExpr.GetOp())
|
||||
}
|
||||
|
||||
func TestRewrite_Bool_In_SingleTrue_Nullable_ToEqual(t *testing.T) {
|
||||
// For nullable bool, in [true] should still become == true
|
||||
helper := buildSchemaHelperForRewriteNullableT(t)
|
||||
expr, err := parser.ParseExpr(helper, `NullableBoolField in [true]`, nil)
|
||||
require.NoError(t, err)
|
||||
require.NotNil(t, expr)
|
||||
ure := expr.GetUnaryRangeExpr()
|
||||
require.NotNil(t, ure, "nullable bool IN [true] should be rewritten to == true")
|
||||
require.Equal(t, planpb.OpType_Equal, ure.GetOp())
|
||||
require.Equal(t, true, ure.GetValue().GetBoolVal())
|
||||
}
|
||||
|
||||
func TestRewrite_Bool_NotIn_TrueFalse_ToAlwaysFalse(t *testing.T) {
|
||||
helper := buildSchemaHelperForRewriteT(t)
|
||||
// not in [true, false] on non-nullable → nothing can match → AlwaysFalse
|
||||
expr, err := parser.ParseExpr(helper, `BoolField not in [true,false]`, nil)
|
||||
require.NoError(t, err)
|
||||
require.NotNil(t, expr)
|
||||
require.True(t, rewriter.IsAlwaysFalseExpr(expr),
|
||||
"bool NOT IN [true, false] should be rewritten to AlwaysFalseExpr")
|
||||
}
|
||||
|
||||
func TestRewrite_Bool_NotIn_TrueFalse_Nullable_ToIsNull(t *testing.T) {
|
||||
helper := buildSchemaHelperForRewriteNullableT(t)
|
||||
// not in [true, false] on nullable → only NULL matches → IS NULL
|
||||
expr, err := parser.ParseExpr(helper, `NullableBoolField not in [true,false]`, nil)
|
||||
require.NoError(t, err)
|
||||
require.NotNil(t, expr)
|
||||
nullExpr := expr.GetNullExpr()
|
||||
require.NotNil(t, nullExpr, "nullable bool NOT IN [true, false] should be rewritten to IS NULL")
|
||||
require.Equal(t, planpb.NullExpr_IsNull, nullExpr.GetOp())
|
||||
}
|
||||
|
||||
func TestRewrite_Bool_NotIn_SingleTrue_ToNotEqual(t *testing.T) {
|
||||
helper := buildSchemaHelperForRewriteT(t)
|
||||
// not in [true] → != true
|
||||
expr, err := parser.ParseExpr(helper, `BoolField not in [true]`, nil)
|
||||
require.NoError(t, err)
|
||||
require.NotNil(t, expr)
|
||||
ure := expr.GetUnaryRangeExpr()
|
||||
require.NotNil(t, ure, "bool NOT IN [true] should be rewritten to != true")
|
||||
require.Equal(t, planpb.OpType_NotEqual, ure.GetOp())
|
||||
require.Equal(t, true, ure.GetValue().GetBoolVal())
|
||||
}
|
||||
|
||||
func TestRewrite_Bool_NotIn_SingleFalse_ToNotEqual(t *testing.T) {
|
||||
helper := buildSchemaHelperForRewriteT(t)
|
||||
// not in [false] → != false
|
||||
expr, err := parser.ParseExpr(helper, `BoolField not in [false]`, nil)
|
||||
require.NoError(t, err)
|
||||
require.NotNil(t, expr)
|
||||
ure := expr.GetUnaryRangeExpr()
|
||||
require.NotNil(t, ure, "bool NOT IN [false] should be rewritten to != false")
|
||||
require.Equal(t, planpb.OpType_NotEqual, ure.GetOp())
|
||||
require.Equal(t, false, ure.GetValue().GetBoolVal())
|
||||
}
|
||||
|
||||
func TestRewrite_Flatten_Then_OR_ToIN(t *testing.T) {
|
||||
|
||||
@@ -200,6 +200,17 @@ func newBoolConstExpr(v bool) *planpb.Expr {
|
||||
}
|
||||
}
|
||||
|
||||
func newNullExpr(col *planpb.ColumnInfo, op planpb.NullExpr_NullOp) *planpb.Expr {
|
||||
return &planpb.Expr{
|
||||
Expr: &planpb.Expr_NullExpr{
|
||||
NullExpr: &planpb.NullExpr{
|
||||
ColumnInfo: col,
|
||||
Op: op,
|
||||
},
|
||||
},
|
||||
}
|
||||
}
|
||||
|
||||
func newAlwaysTrueExpr() *planpb.Expr {
|
||||
return &planpb.Expr{
|
||||
Expr: &planpb.Expr_AlwaysTrueExpr{
|
||||
|
||||
@@ -290,7 +290,7 @@ func TestValidatePartitionKeyIsolation(t *testing.T) {
|
||||
{
|
||||
name: "partition key isolation NOT equal",
|
||||
expr: "not(key_field == 10)",
|
||||
expectedErrorString: "partition key isolation does not support NOT",
|
||||
expectedErrorString: "partition key isolation does not support NotEqual",
|
||||
},
|
||||
{
|
||||
name: "partition key isolation equal AND with same field term",
|
||||
|
||||
Reference in New Issue
Block a user