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:
Buqian Zheng
2026-03-31 14:15:32 +08:00
committed by GitHub
co-authored by Claude
parent d022c6fff0
commit e1aefb24b8
5 changed files with 208 additions and 10 deletions
@@ -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);
+66 -2
View File
@@ -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{
+1 -1
View File
@@ -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",