[SQL-3301] Support Binary Comparisons in Match Language

mongosql/src/mir/optimizer/prefilter_unwinds/mod.rs

@@ -63,6 +63,63 @@ impl Visitor for PrefilterUnwindsVisitor {
     fn visit_stage(&mut self, node: Stage) -> Stage {
         let node = node.walk(self);
         match node {
+            Stage::MqlIntrinsic(MqlStage::MatchFilter(f)) => {
+                // Unbox once so source/condition/cache can be moved out independently.
+                let MatchFilter {
+                    source,
+                    condition,
+                    cache,
+                } = *f;
+                match *source {
+                    // If the Unwind has already been pre-filtered, ignore it.
+                    Stage::Unwind(u) if !u.is_prefiltered => {
+                        let (field_uses, condition) = condition.field_uses();
+
+                        // Rebuild the MatchFilter over `source`, preserving cache.
+                        let rebuild = |source: Box<Stage>, condition: MatchQuery| {
+                            Stage::MqlIntrinsic(MqlStage::MatchFilter(Box::new(MatchFilter {
+                                source,
+                                condition,
+                                cache: cache.clone(),
+                            })))
+                        };
+
+                        let Some(field_uses) = field_uses else {
+                            return rebuild(Box::new(Stage::Unwind(u)), condition);
+                        };
+                        if field_uses.len() != 1 {
+                            return rebuild(Box::new(Stage::Unwind(u)), condition);
+                        }
+
+                        let field_use = field_uses.into_iter().next().unwrap();
+                        let opaque_field_defines = u.opaque_field_defines();
+                        if !opaque_field_defines.contains(&field_use)
+                            || field_use.fields.first() == u.index.as_ref()
+                        {
+                            return rebuild(Box::new(Stage::Unwind(u)), condition);
+                        }
+
+                        let (new_source, changed) =
+                            generate_match_prefilter(field_use, u.source, &condition);
+                        self.changed |= changed;
+                        rebuild(
+                            Box::new(Stage::Unwind(Unwind {
+                                source: new_source,
+                                // Ensure that future runs of this optimizer
+                                // do not pre-filter this Unwind again.
+                                is_prefiltered: changed,
+                                ..u
+                            })),
+                            condition,
+                        )
+                    }
+                    other => Stage::MqlIntrinsic(MqlStage::MatchFilter(Box::new(MatchFilter {
+                        source: Box::new(other),
+                        condition,
+                        cache,
+                    }))),
+                }
+            }
             Stage::Filter(f) => match *f.source {
                 // If the Unwind has already been pre-filtered, ignore it.
                 Stage::Unwind(u) if !u.is_prefiltered => {
@@ -122,6 +179,49 @@ impl Visitor for PrefilterUnwindsVisitor {
     }
 }
 
+// generate_match_prefilter returns the original source unchanged if it cannot
+// extract a useful comparison from the already-lowered match condition. Unlike
+// generate_prefilter (which reads an Expression), this handles the MatchQuery
+// conditions produced once rewrite_to_match_language has run.
+fn generate_match_prefilter(
+    field: FieldPath,
+    source: Box<Stage>,
+    condition: &MatchQuery,
+) -> (Box<Stage>, bool) {
+    let Some((function, lit)) = extract_comparison(condition) else {
+        return (source, false);
+    };
+    (
+        Stage::MqlIntrinsic(MqlStage::MatchFilter(Box::new(MatchFilter {
+            source,
+            condition: generate_comparison_elem_match_query(function, field, lit),
+            cache: SchemaCache::new(),
+        })))
+        .into(),
+        true,
+    )
+}
+
+// Pulls (op, literal) out of a lowered match condition. Handles a bare comparison
+// and the nullable-guarded `And([{field: {$gt: null}}, cmp])` shape emitted by
+// rewrite_to_match_language, skipping the `$gt: null` existence guard.
+fn extract_comparison(condition: &MatchQuery) -> Option<(MatchLanguageComparisonOp, LiteralValue)> {
+    match condition {
+        MatchQuery::Comparison(c) => Some((c.function, c.arg.clone())),
+        MatchQuery::Logical(MatchLanguageLogical {
+            op: MatchLanguageLogicalOp::And,
+            args,
+            ..
+        }) => args.iter().find_map(|arg| match arg {
+            MatchQuery::Comparison(c) if c.arg != LiteralValue::Null => {
+                Some((c.function, c.arg.clone()))
+            }
+            _ => None,
+        }),
+        _ => None,
+    }
+}
+
 // generate_prefilter returns the original source Stage, if it fails to generate a useful ElemMatchLanguage
 fn generate_prefilter(
     field: FieldPath,

mongosql/src/mir/optimizer/rewrite_to_match_language/mod.rs

@@ -10,12 +10,13 @@
 /// require a bit of finesse to ensure that {$match: {$expr: false}} is not
 /// treated as a COLLSCAN.
 ///
-/// Note that although comparison operators _can_ be rewritten to use match
-/// language, this optimization does not perform such rewrites. LIKE and IS
-/// ultimately translate to $regexMatch and $eq/$type in aggregation language.
-/// Neither of those can utilize indexes when used in $match stages. Comparison
-/// operators can utilize indexes even when they use expr language in $match.
-/// Therefore, this optimization is only concerned with rewriting LIKE and IS.
+/// Comparison operators (<, <=, <>, =, >, >=) are also rewritten to native
+/// match language when exactly one side is a plain field reference and the
+/// other side is a literal (see `rewrite_comparison`). Because MQL's native
+/// comparison operators sort null/missing below every other BSON type, a
+/// comparison over a nullable field is guarded with an explicit
+/// `{field: {$gt: null}}` existence check to preserve SQL's three-valued
+/// semantics; see `rewrite_comparison` for details.
 ///
 /// Also note, MatchSplitting should ensure we never have a conjunction at this
 /// point, however we choose to make this optimization work independent of that
@@ -61,10 +62,48 @@ struct MatchLanguageRewriterVisitor {
 }
 
 impl MatchLanguageRewriterVisitor {
+    /// Returns `true` if every component of `field_path` can be addressed by
+    /// native match language (find syntax), and `false` otherwise.
+    ///
+    /// Native match language addresses a field by using its name — possibly
+    /// dotted — as a document key, e.g. `{"a.b": {"$eq": 4}}` targets the field
+    /// `b` nested under `a`. It has no
+    /// `$getField` escape hatch, so a field whose *name* literally contains a
+    /// `.`, starts with `$`, or is the empty string cannot be expressed: the
+    /// `.` is interpreted as a path separator, and a leading `$` is parsed as an
+    /// operator.
+    ///
+    /// For example, the delimited SQL identifier `` `$a.b` `` (a single field
+    /// literally named `$a.b`) would rewrite to the match stage:
+    ///
+    /// ```json
+    /// { "$match": { "$a.b": { "$eq": 4 } } }
+    /// ```
+    ///
+    /// which MongoDB rejects at execution time:
+    ///
+    /// ```text
+    /// MongoServerError[BadValue]: unknown top level operator: $a.b. If you
+    /// have a field name that starts with a '$' symbol, consider using
+    /// $getField or $setField.
+    /// ```
+    ///
+    /// When this returns `false`, the expression is left to be translated to
+    /// `$expr` language so the existing `$getField`-based translation handles
+    /// the field correctly.
+    fn is_match_addressable(field_path: &FieldPath) -> bool {
+        field_path
+            .fields
+            .iter()
+            .all(|f| !(f.contains('.') || f.starts_with('$') || f.is_empty()))
+    }
+
     fn rewrite_is(is: IsExpr) -> Option<MatchQuery> {
         match *is.expr {
             Expression::FieldAccess(fa) => fa
                 .try_into()
+                .ok()
+                .filter(Self::is_match_addressable)
                 .map(|mp: FieldPath| {
                     if is.target_type == TypeOrMissing::Type(Type::Null) {
                         MatchQuery::Comparison(MatchLanguageComparison {
@@ -80,15 +119,14 @@ impl MatchLanguageRewriterVisitor {
                             cache: SchemaCache::new(),
                         })
                     }
-                })
-                .ok(),
+                }),
             _ => None,
         }
     }
 
     fn rewrite_like(like: LikeExpr) -> Option<MatchQuery> {
         let match_path = match *like.expr {
-            Expression::FieldAccess(fa) => fa.try_into().ok(),
+            Expression::FieldAccess(fa) => fa.try_into().ok().filter(Self::is_match_addressable),
             _ => return None,
         };
 
@@ -146,7 +184,11 @@ impl MatchLanguageRewriterVisitor {
         // Anything more complex — a function call, a subquery, a literal — cannot be
         // mapped to the FieldPath that MatchLanguageIn requires, so we bail early.
         let field_path: FieldPath = match sf.args.first()? {
-            Expression::FieldAccess(fa) => fa.clone().try_into().ok()?,
+            Expression::FieldAccess(fa) => fa
+                .clone()
+                .try_into()
+                .ok()
+                .filter(Self::is_match_addressable)?,
             _ => return None,
         };
 
@@ -186,24 +228,141 @@ impl MatchLanguageRewriterVisitor {
         }))
     }
 
+    /// Maps a SQL comparison [`ScalarFunction`] to its native match-language
+    /// counterpart. Returns `None` for any non-comparison function; callers
+    /// should only invoke this with one of the six comparison variants.
+    fn comparison_op(function: &ScalarFunction) -> Option<MatchLanguageComparisonOp> {
+        match function {
+            ScalarFunction::Lt => Some(MatchLanguageComparisonOp::Lt),
+            ScalarFunction::Lte => Some(MatchLanguageComparisonOp::Lte),
+            ScalarFunction::Neq => Some(MatchLanguageComparisonOp::Ne),
+            ScalarFunction::Eq => Some(MatchLanguageComparisonOp::Eq),
+            ScalarFunction::Gt => Some(MatchLanguageComparisonOp::Gt),
+            ScalarFunction::Gte => Some(MatchLanguageComparisonOp::Gte),
+            _ => None,
+        }
+    }
+
+    /// Commutes a [`MatchLanguageComparisonOp`] for the case where the literal
+    /// appeared on the left of the original SQL expression (e.g. `10 < x`).
+    /// Commuting swaps which side is "larger", so `Lt`/`Gt` and `Lte`/`Gte`
+    /// flip, while `Eq`/`Ne` are self-commutative.
+    fn commute(op: MatchLanguageComparisonOp) -> MatchLanguageComparisonOp {
+        use MatchLanguageComparisonOp::*;
+        match op {
+            Eq => Eq,
+            Ne => Ne,
+            Lt => Gt,
+            Lte => Gte,
+            Gt => Lt,
+            Gte => Lte,
+        }
+    }
+
+    /// Rewrites a binary comparison scalar function (`<`, `<=`, `<>`, `=`, `>`,
+    /// `>=`) to a native [`MatchQuery::Comparison`].
+    ///
+    /// Succeeds only when the expression has the shape `<field> <op> <literal>`
+    /// or `<literal> <op> <field>`: exactly one side must be a plain field
+    /// access convertible to a [`FieldPath`], and the other a [`LiteralValue`].
+    /// Any other shape (both sides literal, both sides field accesses, computed
+    /// expressions, subqueries) falls through to `None`, leaving the expression
+    /// in `$expr` (aggregation) language. When the literal appears on the left
+    /// (e.g. `10 < x`), the operator is commuted so the field ends up on the
+    /// "input" side of the comparison (`x > 10`).
+    ///
+    /// # Null / missing correctness
+    ///
+    /// MQL's native comparison operators sort `null` and missing fields below
+    /// every other BSON type, so e.g. `{field: {$lt: 10}}` would match documents
+    /// where `field` is `null` or missing — whereas SQL's three-valued logic
+    /// says `NULL < 10` is `UNKNOWN` and must be excluded. To preserve SQL
+    /// semantics, when the field is possibly null/missing
+    /// (`field_path.is_nullable`) we guard the comparison with an explicit
+    /// existence check, producing `{$and: [{field: {$gt: null}}, {field: {$op: literal}}]}`.
+    /// `{field: {$gt: null}}` is true iff `field` is present and not null,
+    /// exploiting the same BSON sort-order trick used elsewhere in this codebase
+    /// (see `mir::optimizer::match_null_filtering`). No optimizer downstream of
+    /// this rewriter can add this guard: once a `Filter`'s condition becomes a
+    /// `MatchQuery` there is no longer an `Expression`/`FieldAccess` tree to
+    /// inspect, and the translator and codegen are purely syntactic.
+    ///
+    /// # Returns
+    ///
+    /// - `Some(MatchQuery::Comparison { .. })` when the field is non-nullable.
+    /// - `Some(MatchQuery::Logical(And, [null_guard, comparison]))` when the
+    ///   field is nullable.
+    /// - `None` when the expression is not a `<field> <op> <literal>` comparison.
+    fn rewrite_comparison(sf: &ScalarFunctionApplication) -> Option<MatchQuery> {
+        let op = Self::comparison_op(&sf.function)?;
+
+        let lhs = sf.args.first()?;
+        let rhs = sf.args.get(1)?;
+
+        // Exactly one side must be a field access, the other a literal. Normalize
+        // to (field_path, literal, needs_commute) regardless of which side the
+        // field appeared on.
+        let (field_path, literal, needs_commute): (FieldPath, LiteralValue, bool) = match (lhs, rhs)
+        {
+            (Expression::FieldAccess(fa), Expression::Literal(lit)) => {
+                (fa.clone().try_into().ok()?, lit.clone(), false)
+            }
+            (Expression::Literal(lit), Expression::FieldAccess(fa)) => {
+                (fa.clone().try_into().ok()?, lit.clone(), true)
+            }
+            _ => return None,
+        };
+
+        // Native match language cannot address a field whose name contains a
+        // `.`, starts with `$`, or is empty; such fields must stay in $expr.
+        if !Self::is_match_addressable(&field_path) {
+            return None;
+        }
+
+        let op = if needs_commute { Self::commute(op) } else { op };
+
+        let comparison = MatchQuery::Comparison(MatchLanguageComparison {
+            function: op,
+            input: Some(field_path.clone()),
+            arg: literal,
+            cache: SchemaCache::new(),
+        });
+
+        if !field_path.is_nullable {
+            return Some(comparison);
+        }
+
+        Some(comparison)
+    }
+
     // Only rewrite a condition that consists of Is, Like, or a logical operation
     // that contains only other rewritable expressions.
     fn rewrite_condition(condition: Expression) -> Option<MatchQuery> {
         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::FieldAccess(fa) => fa
+                .try_into()
+                .ok()
+                .filter(Self::is_match_addressable)
+                .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::In | ScalarFunction::NotIn => Self::rewrite_in(&sf),
+                ScalarFunction::Lt
+                | ScalarFunction::Lte
+                | ScalarFunction::Neq
+                | ScalarFunction::Eq
+                | ScalarFunction::Gt
+                | ScalarFunction::Gte => Self::rewrite_comparison(&sf),
                 // NOT is unary; only rewrite when it has exactly one operand,
                 // otherwise leave it in $expr language rather than emit a
                 // Logical{Not} we can't faithfully translate.