feat(index): accelerate regex and infix LIKE with the ngram index

Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>

Author: Vova Kolmakov

Cargo.lock

@@ -180,6 +180,7 @@ rand_distr = { version = "0.5.1" }
 rand_xoshiro = "0.7.0"
 rangemap = { version = "1.0" }
 rayon = "1.10"
+regex-syntax = "0.8.10"
 roaring = "0.11.4"
 rstest = "0.26.1"
 serde = { version = "^1" }

Cargo.toml

@@ -29,4 +29,10 @@ The N-gram index provides inexact results for the following query types:
 
 | Query Type     | Description              | Operation                                             | Result Type |
 |----------------|--------------------------|-------------------------------------------------------|-------------|
-| **contains**   | Substring search in text | Finds all trigrams in query, intersects posting lists | AtMost      |
+| **contains**   | Substring search in text | Finds all trigrams in query, intersects posting lists | AtMost      |
+| **regexp_like** / **regexp_match** | Regular-expression match | Derives a necessary trigram condition from the pattern (AND of intersections, OR of unions), then rechecks the true regex | AtMost |
+| **LIKE** (infix) | Wildcard match such as `%foo%bar%` | Uses the literal segments of the pattern as a trigram condition, then rechecks the LIKE | AtMost |
+
+Patterns from which no trigram can be derived - for example `a.b`, `.*`,
+case-insensitive matches, or literal runs shorter than three characters - fall
+back to rechecking every row. This is always correct, just not accelerated.

docs/src/format/index/scalar/ngram.md

@@ -56,6 +56,7 @@ object_store.workspace = true
 prost.workspace = true
 prost-types.workspace = true
 rand.workspace = true
+regex-syntax.workspace = true
 roaring.workspace = true
 rayon.workspace = true
 serde_json.workspace = true

python/Cargo.lock

@@ -8,6 +8,7 @@ use arrow_array::{BooleanArray, ListArray, RecordBatch, UInt64Array};
 use arrow_schema::{Field, Schema};
 use async_trait::async_trait;
 use bytes::Bytes;
+use datafusion::functions::regex::regexplike::RegexpLikeFunc;
 use datafusion::functions::string::contains::ContainsFunc;
 use datafusion::functions_nested::array_has;
 use datafusion::physical_plan::SendableRecordBatchStream;
@@ -628,9 +629,15 @@ impl AnyQuery for LabelListQuery {
 pub enum TextQuery {
     /// Retrieve all row ids where the text contains the given string
     StringContains(String),
-    // TODO: In the future we should be able to do string-insensitive contains
-    // as well as partial matches (e.g. LIKE 'foo%') and potentially even
-    // some regular expressions
+    /// Retrieve all row ids whose text matches the given regular expression.
+    ///
+    /// The pattern is a full regular expression (as accepted by `regexp_like`).
+    /// The index returns a candidate superset that the scan rechecks, so any
+    /// pattern is sound; patterns with no usable trigram structure simply fall
+    /// back to rechecking every row.
+    Regex(String),
+    // TODO: In the future we should be able to do case-insensitive contains
+    // as well as partial matches (e.g. LIKE 'foo%').
 }
 
 impl AnyQuery for TextQuery {
@@ -651,6 +658,17 @@ impl AnyQuery for TextQuery {
                     Expr::Literal(ScalarValue::Utf8(Some(substr.clone())), None),
                 ],
             }),
+            // `regexp_like` returns Boolean directly, so the reconstructed
+            // expression can be used as-is for the recheck filter (no IsNotNull
+            // wrapper, unlike `regexp_match`). It is the semantic equivalent of
+            // the original predicate for the "does it match" question.
+            Self::Regex(pattern) => Expr::ScalarFunction(ScalarFunction {
+                func: Arc::new(RegexpLikeFunc::new().into()),
+                args: vec![
+                    Expr::Column(Column::new_unqualified(col)),
+                    Expr::Literal(ScalarValue::Utf8(Some(pattern.clone())), None),
+                ],
+            }),
         }
     }
 

rust/lance-index/Cargo.toml

@@ -781,20 +781,28 @@ impl ScalarQueryParser for LabelListQueryParser {
     }
 }
 
-/// A parser for indices that handle string contains queries
+/// A parser for indices that handle string `contains` queries, and -- when
+/// `supports_regex` is set -- `regexp_like` / `regexp_match` queries.
 #[derive(Debug, Clone)]
 pub struct TextQueryParser {
     index_name: String,
     index_type: String,
     needs_recheck: bool,
+    supports_regex: bool,
 }
 
 impl TextQueryParser {
-    pub fn new(index_name: String, index_type: String, needs_recheck: bool) -> Self {
+    pub fn new(
+        index_name: String,
+        index_type: String,
+        needs_recheck: bool,
+        supports_regex: bool,
+    ) -> Self {
         Self {
             index_name,
             index_type,
             needs_recheck,
+            supports_regex,
         }
     }
 }
@@ -837,31 +845,156 @@ impl ScalarQueryParser for TextQueryParser {
         func: &ScalarUDF,
         args: &[Expr],
     ) -> Option<IndexedExpression> {
-        if args.len() != 2 {
+        // The first argument is the indexed column; the second is the substring
+        // / pattern. `contains` takes exactly two arguments; the regex functions
+        // optionally take a third flags argument.
+        if args.len() < 2 {
             return None;
         }
-        let scalar = maybe_scalar(&args[1], data_type)?;
-        match scalar {
-            ScalarValue::Utf8(Some(scalar_str)) | ScalarValue::LargeUtf8(Some(scalar_str)) => {
-                if func.name() == "contains" {
-                    let query = TextQuery::StringContains(scalar_str);
-                    Some(IndexedExpression::index_query_with_recheck(
-                        column.to_string(),
-                        self.index_name.clone(),
-                        self.index_type.clone(),
-                        Arc::new(query),
-                        self.needs_recheck,
-                    ))
-                } else {
+        // A non-string pattern cannot be handled.
+        let (ScalarValue::Utf8(Some(pattern)) | ScalarValue::LargeUtf8(Some(pattern))) =
+            maybe_scalar(&args[1], data_type)?
+        else {
+            return None;
+        };
+
+        let query = match func.name() {
+            "contains" if args.len() == 2 => TextQuery::StringContains(pattern),
+            "regexp_like" | "regexp_match" if self.supports_regex => {
+                let pattern = match args.get(2) {
+                    Some(flags_expr) => apply_regex_flags(&pattern, flags_expr)?,
+                    None => pattern,
+                };
+                // If the pattern yields no usable trigram (e.g. `a.b`), leave it
+                // to a full scan instead of routing it to the index, which could
+                // only answer with an unsupported "recheck everything" result.
+                if !crate::scalar::ngram::regex_can_use_index(&pattern) {
+                    return None;
+                }
+                TextQuery::Regex(pattern)
+            }
+            _ => return None,
+        };
+
+        Some(IndexedExpression::index_query_with_recheck(
+            column.to_string(),
+            self.index_name.clone(),
+            self.index_type.clone(),
+            Arc::new(query),
+            self.needs_recheck,
+        ))
+    }
+
+    fn visit_like(
+        &self,
+        column: &str,
+        like: &Like,
+        pattern: &ScalarValue,
+    ) -> Option<IndexedExpression> {
+        // Infix LIKE is accelerated only by the ngram index (via its regex
+        // machinery). A plain-literal `regexp_like(col, 'foo')` is rewritten to
+        // `col LIKE '%foo%'` before it reaches the index, so this is the path
+        // that accelerates those. ILIKE is skipped because its case folding does
+        // not match the index's normalization.
+        if !self.supports_regex || like.case_insensitive {
+            return None;
+        }
+        let pattern_str = match pattern {
+            ScalarValue::Utf8(Some(s)) | ScalarValue::LargeUtf8(Some(s)) => s.as_str(),
+            _ => return None,
+        };
+        // Translate the LIKE pattern into a loose regex used only for candidate
+        // generation; the original LIKE stays as the recheck filter, so the
+        // regex only needs to be a sound superset.
+        let regex = like_to_regex(pattern_str, like.escape_char)?;
+        if !crate::scalar::ngram::regex_can_use_index(&regex) {
+            return None;
+        }
+        Some(IndexedExpression {
+            scalar_query: Some(ScalarIndexExpr::Query(ScalarIndexSearch {
+                column: column.to_string(),
+                index_name: self.index_name.clone(),
+                index_type: self.index_type.clone(),
+                query: Arc::new(TextQuery::Regex(regex)),
+                needs_recheck: self.needs_recheck,
+                fragment_bitmap: None,
+            })),
+            refine_expr: Some(Expr::Like(like.clone())),
+        })
+    }
+}
+
+/// Translate a LIKE pattern into a regular expression used purely for ngram
+/// candidate generation: `%` becomes `.*`, `_` becomes `.`, and literal
+/// characters are regex-escaped. Returns `None` when no literal run is long
+/// enough to yield a trigram (the index could not help, so a full scan is left
+/// to handle it).
+fn like_to_regex(pattern: &str, escape: Option<char>) -> Option<String> {
+    let mut regex = String::new();
+    let mut run = 0usize;
+    let mut longest_run = 0usize;
+    let mut chars = pattern.chars();
+    while let Some(c) = chars.next() {
+        let literal = if Some(c) == escape {
+            // The next character is escaped, i.e. a literal.
+            chars.next()
+        } else {
+            match c {
+                '%' => {
+                    regex.push_str(".*");
+                    run = 0;
                     None
                 }
+                '_' => {
+                    regex.push('.');
+                    run = 0;
+                    None
+                }
+                other => Some(other),
             }
-            _ => {
-                // If the scalar is not a string, we cannot handle it
-                None
+        };
+        if let Some(lit) = literal {
+            if regex_syntax::is_meta_character(lit) {
+                regex.push('\\');
+            }
+            regex.push(lit);
+            // Only runs of alphanumeric characters can produce a trigram.
+            if lit.is_alphanumeric() {
+                run += 1;
+                longest_run = longest_run.max(run);
+            } else {
+                run = 0;
             }
         }
     }
+    (longest_run >= 3).then_some(regex)
+}
+
+/// Fold the supported `regexp_like` / `regexp_match` flags into an inline prefix
+/// on the pattern (e.g. flags `"i"` -> `"(?i)pattern"`). Returns `None` for a
+/// non-literal flags argument or an unrecognized flag, so the caller leaves the
+/// predicate to a full recheck rather than risk changing its semantics.
+fn apply_regex_flags(pattern: &str, flags_expr: &Expr) -> Option<String> {
+    let (Expr::Literal(ScalarValue::Utf8(Some(flags)), _)
+    | Expr::Literal(ScalarValue::LargeUtf8(Some(flags)), _)) = flags_expr
+    else {
+        return None;
+    };
+    let mut inline = String::new();
+    for flag in flags.chars() {
+        // Only flags expressible as an inline `(?...)` group in the regex crate
+        // (which the recheck uses) are safe to fold.
+        if ['i', 's', 'm', 'x'].contains(&flag) {
+            inline.push(flag);
+        } else {
+            return None;
+        }
+    }
+    if inline.is_empty() {
+        Some(pattern.to_string())
+    } else {
+        Some(format!("(?{inline}){pattern}"))
+    }
 }
 
 /// A parser for indices that handle queries with the contains_tokens function
@@ -1813,7 +1946,18 @@ fn visit_node(
         Expr::IsFalse(expr) => Ok(visit_is_bool(expr.as_ref(), index_info, false)),
         Expr::IsTrue(expr) => Ok(visit_is_bool(expr.as_ref(), index_info, true)),
         Expr::IsNull(expr) => Ok(visit_is_null(expr.as_ref(), index_info, false)),
-        Expr::IsNotNull(expr) => Ok(visit_is_null(expr.as_ref(), index_info, true)),
+        Expr::IsNotNull(expr) => {
+            // `regexp_match(col, pat)` returns a list and is coerced to
+            // `IsNotNull(regexp_match(...))` before it reaches here. Unwrap that
+            // so the regex acceleration applies; everything else is a genuine
+            // IS NOT NULL check.
+            if let Expr::ScalarFunction(scalar_fn) = expr.as_ref()
+                && scalar_fn.func.name() == "regexp_match"
+            {
+                return Ok(visit_scalar_fn(scalar_fn, index_info));
+            }
+            Ok(visit_is_null(expr.as_ref(), index_info, true))
+        }
         Expr::Not(expr) => visit_not(expr.as_ref(), index_info, depth),
         Expr::BinaryExpr(binary_expr) => visit_binary_expr(binary_expr, index_info, depth),
         Expr::ScalarFunction(scalar_fn) => Ok(visit_scalar_fn(scalar_fn, index_info)),
@@ -2690,6 +2834,59 @@ mod tests {
         assert!(matches!(negated.upper, NullableRowAddrMask::BlockList(_)));
     }
 
+    #[test]
+    fn test_like_to_regex() {
+        // `%` -> `.*`, `_` -> `.`, with a literal run of at least three chars.
+        assert_eq!(like_to_regex("%foo%", None).as_deref(), Some(".*foo.*"));
+        assert_eq!(like_to_regex("foo%bar", None).as_deref(), Some("foo.*bar"));
+        assert_eq!(like_to_regex("foo_bar", None).as_deref(), Some("foo.bar"));
+        assert_eq!(like_to_regex("foobar", None).as_deref(), Some("foobar"));
+
+        // Regex metacharacters in the literal portion are escaped.
+        assert_eq!(
+            like_to_regex("%a.bcd%", None).as_deref(),
+            Some(".*a\\.bcd.*")
+        );
+
+        // No literal run of three alphanumeric characters -> no index help.
+        assert_eq!(like_to_regex("%ab%", None), None);
+        assert_eq!(like_to_regex("%a%b%c%", None), None);
+        assert_eq!(like_to_regex("%", None), None);
+
+        // The escape character makes the following character a literal.
+        assert_eq!(
+            like_to_regex(r"%foo\%bar%", Some('\\')).as_deref(),
+            Some(".*foo%bar.*")
+        );
+    }
+
+    #[test]
+    fn test_apply_regex_flags() {
+        fn flags(s: &str) -> Expr {
+            Expr::Literal(ScalarValue::Utf8(Some(s.to_string())), None)
+        }
+
+        // Empty flags leave the pattern untouched (no inline group emitted).
+        assert_eq!(apply_regex_flags("foo", &flags("")).as_deref(), Some("foo"));
+        // Supported flags are folded into an inline `(?...)` prefix.
+        assert_eq!(
+            apply_regex_flags("foo", &flags("i")).as_deref(),
+            Some("(?i)foo")
+        );
+        assert_eq!(
+            apply_regex_flags("foo", &flags("is")).as_deref(),
+            Some("(?is)foo")
+        );
+        // An unrecognized flag bails out so the caller leaves the predicate to a
+        // full recheck rather than risk changing its semantics.
+        assert_eq!(apply_regex_flags("foo", &flags("g")), None);
+        // A non-string (hence non-literal-flags) argument cannot be folded.
+        assert_eq!(
+            apply_regex_flags("foo", &Expr::Literal(ScalarValue::Int32(Some(1)), None)),
+            None
+        );
+    }
+
     #[test]
     fn test_extract_like_leading_prefix() {
         // Simple prefix patterns (no recheck needed)

rust/lance-index/src/scalar.rs

@@ -1311,6 +1311,12 @@ impl ScalarIndex for FMIndexScalarIndex {
                     Default::default(),
                 )))
             }
+            // Regex queries are routed only to the ngram index (the FM-index's
+            // query parser advertises `supports_regex = false`), so this is
+            // unreachable in practice; reject it explicitly rather than silently.
+            TextQuery::Regex(_) => Err(Error::invalid_input(
+                "FMIndex does not support regular expression queries",
+            )),
         }
     }
     fn can_remap(&self) -> bool {
@@ -1599,6 +1605,9 @@ impl ScalarIndexPlugin for FMIndexPlugin {
         Some(Box::new(TextQueryParser::new(
             index_name,
             self.name().to_string(),
+            // needs_recheck: the FM-index returns exact substring matches.
+            false,
+            // supports_regex: regex acceleration is only implemented for ngram.
             false,
         )))
     }