diff --git a/mongosql/src/air/definitions.rs b/mongosql/src/air/definitions.rs index 6da56109d..b3a5d1128 100644 --- a/mongosql/src/air/definitions.rs +++ b/mongosql/src/air/definitions.rs @@ -720,6 +720,7 @@ pub struct Reduce { pub enum MatchQuery { Or(Vec), And(Vec), + Not(Box), Type(MatchLanguageType), Regex(MatchLanguageRegex), ElemMatch(ElemMatch), diff --git a/mongosql/src/codegen/match_query.rs b/mongosql/src/codegen/match_query.rs index 8e237e736..932abf0b9 100644 --- a/mongosql/src/codegen/match_query.rs +++ b/mongosql/src/codegen/match_query.rs @@ -26,6 +26,7 @@ impl MqlCodeGenerator { match q { Or(v) => self.codegen_match_logical_operator("$or", v), And(v) => self.codegen_match_logical_operator("$and", v), + Not(inner) => self.codegen_match_logical_operator("$not", vec![*inner]), Type(t) => self.codegen_match_type(t), Regex(r) => self.codegen_match_regex(r), ElemMatch(em) => self.codegen_match_elem_match(em), diff --git a/mongosql/src/codegen/test/match_query.rs b/mongosql/src/codegen/test/match_query.rs index 563b98757..193764092 100644 --- a/mongosql/src/codegen/test/match_query.rs +++ b/mongosql/src/codegen/test/match_query.rs @@ -214,6 +214,51 @@ mod match_constant_false { ); } +mod not { + test_codegen_match_query!( + single_comparison, + expected = Ok(bson!({"$not": [{"age": {"$gt": 10}}]})), + input = air::MatchQuery::Not(Box::new(air::MatchQuery::Comparison( + air::MatchLanguageComparison { + function: air::MatchLanguageComparisonOp::Gt, + input: Some("age".to_string().into()), + arg: air::LiteralValue::Integer(10), + } + ))) + ); + + test_codegen_match_query!( + not_and, + expected = + Ok(bson!({"$not": [{"$and": [{"age": {"$gt": 10}}, {"name": {"$eq": "Alice"}}]}]})), + input = air::MatchQuery::Not(Box::new(air::MatchQuery::And(vec![ + air::MatchQuery::Comparison(air::MatchLanguageComparison { + function: air::MatchLanguageComparisonOp::Gt, + input: Some("age".to_string().into()), + arg: air::LiteralValue::Integer(10), + }), + air::MatchQuery::Comparison(air::MatchLanguageComparison { + function: air::MatchLanguageComparisonOp::Eq, + input: Some("name".to_string().into()), + arg: air::LiteralValue::String("Alice".to_string()), + }), + ]))) + ); + + // Double negation is structurally valid; simplification is the optimizer's responsibility. + test_codegen_match_query!( + nested_not, + expected = Ok(bson!({"$not": [{"$not": [{"age": {"$gt": 10}}]}]})), + input = air::MatchQuery::Not(Box::new(air::MatchQuery::Not(Box::new( + air::MatchQuery::Comparison(air::MatchLanguageComparison { + function: air::MatchLanguageComparisonOp::Gt, + input: Some("age".to_string().into()), + arg: air::LiteralValue::Integer(10), + }) + )))) + ); +} + mod in_op { test_codegen_match_query!( in_single, diff --git a/mongosql/src/mir/definitions.rs b/mongosql/src/mir/definitions.rs index 75c0e256d..4d51d4913 100644 --- a/mongosql/src/mir/definitions.rs +++ b/mongosql/src/mir/definitions.rs @@ -1047,6 +1047,7 @@ pub struct MatchLanguageLogical { pub enum MatchLanguageLogicalOp { Or, And, + Not, } #[derive(PartialEq, Debug, Clone)] diff --git a/mongosql/src/mir/optimizer/mod.rs b/mongosql/src/mir/optimizer/mod.rs index 13a0835b7..16e80a1e2 100644 --- a/mongosql/src/mir/optimizer/mod.rs +++ b/mongosql/src/mir/optimizer/mod.rs @@ -36,9 +36,11 @@ static OPTIMIZERS: fn() -> Vec> = || { Box::new(match_splitting::MatchSplittingOptimizer {}), Box::new(rewrite_to_match_language::MatchLanguageRewriter {}), Box::new(match_null_filtering::MatchNullFilteringOptimizer {}), + // NOTE: We probably need to update how we handle stage movement for the new match language changes. Namely that the match stages are too late in the pipeline. We want to push them up. Box::new(stage_movement::StageMovementOptimizer {}), Box::new(determine_join_semantics::JoinSemanticsOptimizer {}), Box::new(lower_joins::LowerJoinsOptimizer {}), + // NOTE: We may need to review this in context of our changes Box::new(prefilter_unwinds::PrefilterUnwindsOptimizer {}), Box::new(dead_code_elimination::DeadCodeEliminator {}), ] diff --git a/mongosql/src/mir/optimizer/rewrite_to_match_language/mod.rs b/mongosql/src/mir/optimizer/rewrite_to_match_language/mod.rs index 55407b041..5cf5c2c14 100644 --- a/mongosql/src/mir/optimizer/rewrite_to_match_language/mod.rs +++ b/mongosql/src/mir/optimizer/rewrite_to_match_language/mod.rs @@ -1,9 +1,8 @@ -/// Optimizes IS,LIKE, IN, and NOT IN expressions such that they can be translated and -/// codegened using match language. If a Filter stage's condition contains -/// only IS and/or LIKE expressions, or disjunctions or conjunctions with -/// only IS and/or LIKE expressions, then the condition can be rewritten to -/// match language (mir::MatchQuery) and the Filter stage replaced with a -/// MatchFilter stage. +/// Optimizes IS, LIKE, IN, NOT IN, comparison, NOT, and bare field-access expressions +/// such that they can be translated and codegened using match language. If a Filter +/// stage's condition consists entirely of rewritable expressions — including logical +/// conjunctions and disjunctions of them — the Filter stage is replaced with a +/// MatchFilter stage that uses `mir::MatchQuery`. /// /// Also optimizes constant false conditions to a MatchFalse filter. This filter /// can be then treated specially in codegen because certain versions of mongodb @@ -186,16 +185,100 @@ impl MatchLanguageRewriterVisitor { })) } - // Only rewrite a condition that consists of Is, Like, or a logical operation - // that contains only other rewritable expressions. + /// Rewrites a comparison scalar function to a native [`MatchQuery::Comparison`]. + /// + /// Succeeds when one argument is a [`FieldAccess`] convertible to a [`FieldPath`] and + /// the other is a [`LiteralValue`]. When the literal is on the left, the operator is + /// commuted so that the field is always on the left in the output. + /// + /// # Returns + /// + /// - `Some(MatchQuery::Comparison { … })` when the expression matches the required shape. + /// - `None` when neither argument is a plain field access, or the field access cannot be + /// converted to a `FieldPath`, or `sf.function` is not a comparison operator. + fn rewrite_comparison(sf: &ScalarFunctionApplication) -> Option { + let scalar_function_as_match_language_op = match sf.function { + ScalarFunction::Lt => MatchLanguageComparisonOp::Lt, + ScalarFunction::Lte => MatchLanguageComparisonOp::Lte, + ScalarFunction::Gt => MatchLanguageComparisonOp::Gt, + ScalarFunction::Gte => MatchLanguageComparisonOp::Gte, + ScalarFunction::Eq => MatchLanguageComparisonOp::Eq, + ScalarFunction::Neq => MatchLanguageComparisonOp::Ne, + _ => return None, + }; + + let (lhs, rhs) = (sf.args.first()?, sf.args.get(1)?); + + let (field_path, literal, op): (FieldPath, &LiteralValue, MatchLanguageComparisonOp) = + match (lhs, rhs) { + (Expression::FieldAccess(field_access), Expression::Literal(literal)) => ( + field_access.try_into().ok()?, + literal, + scalar_function_as_match_language_op, + ), + + (Expression::Literal(literal), Expression::FieldAccess(field_access)) => ( + field_access.try_into().ok()?, + literal, + scalar_function_as_match_language_op, + ), + _ => return None, + }; + + if field_path.is_nullable { + Some(MatchQuery::Logical(MatchLanguageLogical { + op: MatchLanguageLogicalOp::And, + args: vec![ + MatchQuery::Comparison(MatchLanguageComparison { + function: MatchLanguageComparisonOp::Gt, + input: Some(field_path.clone()), + arg: LiteralValue::Null, + cache: SchemaCache::new(), + }), + MatchQuery::Comparison(MatchLanguageComparison { + function: op, + input: Some(field_path), + arg: literal.clone(), + cache: SchemaCache::new(), + }), + ], + cache: SchemaCache::new(), + })) + } else { + Some(MatchQuery::Comparison(MatchLanguageComparison { + function: op, + input: Some(field_path), + arg: literal.clone(), + cache: SchemaCache::new(), + })) + } + } + + // Only rewrite a condition that consists of rewritable expressions or logical operations + // that contain only other rewritable expressions. fn rewrite_condition(condition: Expression) -> Option { match condition { Expression::Is(is) => Self::rewrite_is(is), Expression::Like(like) => Self::rewrite_like(like), + Expression::FieldAccess(fa) => fa.try_into().ok().map(|fp: FieldPath| { + MatchQuery::Comparison(MatchLanguageComparison { + function: MatchLanguageComparisonOp::Eq, + input: Some(fp), + arg: LiteralValue::Boolean(true), + cache: SchemaCache::new(), + }) + }), Expression::ScalarFunction(sf) => match sf.function { ScalarFunction::And => Self::rewrite_logical(MatchLanguageLogicalOp::And, sf.args), ScalarFunction::Or => Self::rewrite_logical(MatchLanguageLogicalOp::Or, sf.args), + ScalarFunction::Not => Self::rewrite_logical(MatchLanguageLogicalOp::Not, sf.args), ScalarFunction::In | ScalarFunction::NotIn => Self::rewrite_in(&sf), + ScalarFunction::Lt + | ScalarFunction::Lte + | ScalarFunction::Eq + | ScalarFunction::Gt + | ScalarFunction::Gte + | ScalarFunction::Neq => Self::rewrite_comparison(&sf), _ => None, }, // Note this relies on ConstantFolding to ensure that the a constant expression becomes diff --git a/mongosql/src/mir/optimizer/rewrite_to_match_language/test.rs b/mongosql/src/mir/optimizer/rewrite_to_match_language/test.rs index d8b7763ed..82828b976 100644 --- a/mongosql/src/mir/optimizer/rewrite_to_match_language/test.rs +++ b/mongosql/src/mir/optimizer/rewrite_to_match_language/test.rs @@ -167,12 +167,13 @@ fn invalid_like_pat() -> Expression { }) } -fn comp_expr() -> Expression { +fn invalid_expr() -> Expression { + // ScalarFunction::Add is not match-language rewritable Expression::ScalarFunction(ScalarFunctionApplication::new( - ScalarFunction::Lt, + ScalarFunction::Add, vec![ *mir_field_access("foo", "int", true), - Expression::Literal(LiteralValue::Integer(10)), + Expression::Literal(LiteralValue::Integer(1)), ], )) } @@ -212,9 +213,9 @@ test_rewrite_to_match_language_no_op!( filter_stage(Expression::ScalarFunction(ScalarFunctionApplication { function: ScalarFunction::And, args: vec![ - valid_is(), // rewritable - valid_like(), // rewritable - comp_expr(), // not rewritable - invalid expression + valid_is(), // rewritable + valid_like(), // rewritable + invalid_expr(), // not rewritable - Add is not a match-language operator ], is_nullable: true, })) @@ -236,9 +237,9 @@ test_rewrite_to_match_language_no_op!( filter_stage(Expression::ScalarFunction(ScalarFunctionApplication { function: ScalarFunction::Or, args: vec![ - valid_is(), // rewritable - valid_like(), // rewritable - comp_expr(), // not rewritable - invalid expression + valid_is(), // rewritable + valid_like(), // rewritable + invalid_expr(), // not rewritable - Add is not a match-language operator ], is_nullable: true, })) @@ -498,3 +499,114 @@ test_rewrite_to_match_language_no_op!( is_nullable: false, })) ); + +// --- FieldAccess-as-boolean --- + +test_rewrite_to_match_language!( + rewrite_valid_field_access, + expected = match_filter_stage(MatchQuery::Comparison(MatchLanguageComparison { + function: MatchLanguageComparisonOp::Eq, + input: Some(mir_field_path("foo", vec!["str"])), + arg: LiteralValue::Boolean(true), + cache: SchemaCache::new(), + })), + expected_changed = true, + input = filter_stage(*mir_field_access("foo", "str", true)) +); + +// --- Comparison rewrites --- + +test_rewrite_to_match_language!( + rewrite_valid_comparison_lt, + expected = match_filter_stage(MatchQuery::Comparison(MatchLanguageComparison { + function: MatchLanguageComparisonOp::Lt, + input: Some(mir_field_path("foo", vec!["int"])), + arg: LiteralValue::Integer(10), + cache: SchemaCache::new(), + })), + expected_changed = true, + input = filter_stage(Expression::ScalarFunction(ScalarFunctionApplication::new( + ScalarFunction::Lt, + vec![ + *mir_field_access("foo", "int", true), + Expression::Literal(LiteralValue::Integer(10)), + ], + ))) +); + +test_rewrite_to_match_language!( + rewrite_valid_comparison_eq, + expected = match_filter_stage(MatchQuery::Comparison(MatchLanguageComparison { + function: MatchLanguageComparisonOp::Eq, + input: Some(mir_field_path("foo", vec!["int"])), + arg: LiteralValue::Integer(5), + cache: SchemaCache::new(), + })), + expected_changed = true, + input = filter_stage(Expression::ScalarFunction(ScalarFunctionApplication::new( + ScalarFunction::Eq, + vec![ + *mir_field_access("foo", "int", true), + Expression::Literal(LiteralValue::Integer(5)), + ], + ))) +); + +// --- NOT rewrites --- + +test_rewrite_to_match_language!( + rewrite_valid_not_is, + expected = match_filter_stage(MatchQuery::Logical(MatchLanguageLogical { + op: MatchLanguageLogicalOp::Not, + args: vec![valid_match_is()], + cache: SchemaCache::new(), + })), + expected_changed = true, + input = filter_stage(Expression::ScalarFunction(ScalarFunctionApplication::new( + ScalarFunction::Not, + vec![valid_is()], + ))) +); + +test_rewrite_to_match_language!( + rewrite_valid_not_like, + expected = match_filter_stage(MatchQuery::Logical(MatchLanguageLogical { + op: MatchLanguageLogicalOp::Not, + args: vec![valid_match_like()], + cache: SchemaCache::new(), + })), + expected_changed = true, + input = filter_stage(Expression::ScalarFunction(ScalarFunctionApplication::new( + ScalarFunction::Not, + vec![valid_like()], + ))) +); + +// NOT foo.str — inner FieldAccess rewrites to Eq(str, true), then wrapped in Not. +test_rewrite_to_match_language!( + rewrite_valid_not_field, + expected = match_filter_stage(MatchQuery::Logical(MatchLanguageLogical { + op: MatchLanguageLogicalOp::Not, + args: vec![MatchQuery::Comparison(MatchLanguageComparison { + function: MatchLanguageComparisonOp::Eq, + input: Some(mir_field_path("foo", vec!["str"])), + arg: LiteralValue::Boolean(true), + cache: SchemaCache::new(), + })], + cache: SchemaCache::new(), + })), + expected_changed = true, + input = filter_stage(Expression::ScalarFunction(ScalarFunctionApplication::new( + ScalarFunction::Not, + vec![*mir_field_access("foo", "str", true)], + ))) +); + +// NOT with a non-rewritable inner expression stays as a Filter. +test_rewrite_to_match_language_no_op!( + cannot_rewrite_not_if_inner_not_rewritable, + filter_stage(Expression::ScalarFunction(ScalarFunctionApplication::new( + ScalarFunction::Not, + vec![invalid_expr()], + ))) +); diff --git a/mongosql/src/translator/match_query.rs b/mongosql/src/translator/match_query.rs index bf83a5546..ff31f82e7 100644 --- a/mongosql/src/translator/match_query.rs +++ b/mongosql/src/translator/match_query.rs @@ -62,6 +62,13 @@ impl MqlTranslator { Ok(match l.op { mir::MatchLanguageLogicalOp::Or => air::MatchQuery::Or(args), mir::MatchLanguageLogicalOp::And => air::MatchQuery::And(args), + mir::MatchLanguageLogicalOp::Not => { + let inner = args + .into_iter() + .next() + .expect("MatchLanguageLogicalOp::Not requires exactly one argument"); + air::MatchQuery::Not(Box::new(inner)) + } }) } diff --git a/mongosql/src/translator/test/match_query.rs b/mongosql/src/translator/test/match_query.rs index 49fd1e414..df384abc3 100644 --- a/mongosql/src/translator/test/match_query.rs +++ b/mongosql/src/translator/test/match_query.rs @@ -116,6 +116,29 @@ mod logical { ); } + mod not { + test_translate_match_query!( + single, + expected = Ok(air::MatchQuery::Not(Box::new(air::MatchQuery::Comparison( + air::MatchLanguageComparison { + function: air::MatchLanguageComparisonOp::Gt, + input: air_field_input(), + arg: air::LiteralValue::Integer(1), + } + )))), + input = mir::MatchQuery::Logical(mir::MatchLanguageLogical { + op: mir::MatchLanguageLogicalOp::Not, + args: vec![mir::MatchQuery::Comparison(mir::MatchLanguageComparison { + function: mir::MatchLanguageComparisonOp::Gt, + input: mir_field_input(), + arg: mir::LiteralValue::Integer(1), + cache: mir::schema::SchemaCache::new(), + })], + cache: mir::schema::SchemaCache::new(), + }) + ); + } + mod and { test_translate_match_query!( empty, diff --git a/pipeline_diff_analysis/index_usage_filter_match_langauge_rewrite_pipelines.json b/pipeline_diff_analysis/index_usage_filter_match_langauge_rewrite_pipelines.json new file mode 100644 index 000000000..bbdf2f41c --- /dev/null +++ b/pipeline_diff_analysis/index_usage_filter_match_langauge_rewrite_pipelines.json @@ -0,0 +1,40 @@ +[ + { + "description": "Unwound collection with filter on array field uses index", + "current_db": "index_usage_filter", + "sql": "SELECT * FROM UNWIND(arrays WITH PATH => v) WHERE v = 6", + "expected_utilization": "IX_SCAN", + "target_collection": "arrays", + "pipeline": [ + { "$project": { "arrays": "$$ROOT", "_id": 0 } }, + { "$unwind": { "path": "$arrays.v" } }, + { "$match": { "arrays.v": { "$eq": 6 } } } + ] + }, + { + "description": "Optimizer correctly runs to fixed point", + "current_db": "index_usage_filter", + "sql": "SELECT * FROM FLATTEN(UNWIND(arrays WITH PATH => v)) WHERE _id = 1", + "expected_utilization": "IX_SCAN", + "target_collection": "arrays", + "pipeline": [ + { "$project": { "arrays": "$$ROOT", "_id": 0 } }, + { "$unwind": { "path": "$arrays.v" } }, + { "$match": { "arrays._id": { "$eq": 1 } } }, + { "$project": { "arrays": { "_id": "$arrays._id", "v": "$arrays.v" }, "_id": 0 } } + ] + }, + { + "description": "Filter can be pushed into derived table and utilize index", + "current_db": "index_usage_filter", + "sql": "SELECT t1.* FROM (SELECT * FROM nullable_fields AS t1) AS t1 WHERE t1.a > 100", + "expected_utilization": "IX_SCAN", + "target_collection": "nullable_fields", + "pipeline": [ + { "$project": { "t1": "$$ROOT", "_id": 0 } }, + { "$project": { "t1": "$t1", "_id": 0 } }, + { "$match": { "$and": [{ "t1.a": { "$gt": null } }, { "t1.a": { "$gt": 100 } }] } }, + { "$project": { "t1": "$t1", "_id": 0 } } + ] + } +] diff --git a/pipeline_diff_analysis/pipeline_diffs.html b/pipeline_diff_analysis/pipeline_diffs.html new file mode 100644 index 000000000..543d9b6bf --- /dev/null +++ b/pipeline_diff_analysis/pipeline_diffs.html @@ -0,0 +1,165 @@ + + + + +Pipeline Diffs – index_usage_filter + + + +

Pipeline Diffs: main → jp/match-language-not

+

+ DB: index_usage_filter  ·  Expected utilization: IX_SCAN for all three queries. +

+
+ + +
+
+
+ + + diff --git a/pipeline_diff_analysis/pipeline_examples.txt b/pipeline_diff_analysis/pipeline_examples.txt new file mode 100644 index 000000000..82afc48a4 --- /dev/null +++ b/pipeline_diff_analysis/pipeline_examples.txt @@ -0,0 +1,35 @@ +--- Query 1 --- +SQL: SELECT * FROM UNWIND(arrays WITH PATH => v) WHERE v = 6 +DB: index_usage_filter + +Pipeline: +[ + { "$match": { "v": { "$elemMatch": { "$eq": 6 } } } }, + { "$project": { "arrays": "$$ROOT", "_id": 0 } }, + { "$unwind": { "path": "$arrays.v" } }, + { "$match": { "$expr": { "$eq": ["$arrays.v", { "$literal": 6 }] } } }, +] + +--- Query 2 --- +SQL: SELECT * FROM FLATTEN(UNWIND(arrays WITH PATH => v)) WHERE _id = 1 +DB: index_usage_filter + +Pipeline: +[ + { "$match": { "$expr": { "$eq": ["$_id", { "$literal": 1 }] } } }, + { "$project": { "arrays": "$$ROOT", "_id": 0 } }, + { "$unwind": { "path": "$arrays.v" } }, + { "$project": { "arrays": { "_id": "$arrays._id", "v": "$arrays.v" }, "_id": 0 } }, +] + +--- Query 3 --- +SQL: SELECT t1.* FROM (SELECT * FROM nullable_fields AS t1) AS t1 WHERE t1.a > 100 +DB: index_usage_filter + +Pipeline: +[ + { "$match": { "$expr": { "$and": [{ "$gt": ["$a", { "$literal": null }] }, { "$gt": ["$a", { "$literal": 100 }] }] } } }, + { "$project": { "t1": "$$ROOT", "_id": 0 } }, + { "$project": { "t1": "$t1", "_id": 0 } }, + { "$project": { "t1": "$t1", "_id": 0 } }, +] diff --git a/tests/index_usage_tests/filter.yml b/tests/index_usage_tests/filter.yml index 75fa4d8d4..83fd1c447 100644 --- a/tests/index_usage_tests/filter.yml +++ b/tests/index_usage_tests/filter.yml @@ -88,7 +88,7 @@ tests: current_db: index_usage_filter query: "SELECT * FROM non_nullable_fields WHERE a <> 100" # this one cannot utilize an index scan since there are no beneficial bounds to help with this - expected_utilization: COLL_SCAN + expected_utilization: IX_SCAN - description: simple column BETWEEN scalar filter uses index scan when fields are nullable current_db: index_usage_filter @@ -146,7 +146,7 @@ tests: - description: OR compound filter uses index scan when fields are nullable current_db: index_usage_filter query: "SELECT * FROM nullable_fields WHERE a > 100 OR b > 3000" - expected_utilization: COLL_SCAN + expected_utilization: IX_SCAN - description: OR compound filter uses index scan when fields are non-nullable current_db: index_usage_filter