[YDF] Use RE2 for regexes during inference

PiperOrigin-RevId: 911192343

Author: rstz

CHANGELOG.md

@@ -5,6 +5,16 @@ Changelog under `yggdrasil_decision_forests/port/python/CHANGELOG.md`.
 
 ## HEAD
 
+### Breaking changes
+
+-   Replace std::regex by Google's RE2 during inference. Models using Regex
+    expressions for tokenization must now follow the RE2 syntax. See
+    https://github.com/google/re2/wiki/Syntax for the a detailed breakdown of
+    this syntax. Note that other regexes (notably for the data spec columns)
+    can still use the more permissive std::regex syntax.
+
+## HEAD
+
 ### Features
 
 -   "True" and "False" are now recognized as boolean in CSV files. This change

MODULE.bazel

@@ -21,6 +21,7 @@ bazel_dep(name = "rules_cc", version = "0.2.0")
 bazel_dep(name = "rules_proto", version = "7.1.0")
 bazel_dep(name = "rules_python", version = "1.6.0")
 bazel_dep(name = "zlib", version = "1.3.1.bcr.7")
+bazel_dep(name = "re2", version = "2025-11-05.bcr.1")
 
 bazel_dep(name = "emsdk", version = "4.0.13", dev_dependency = True)
 

yggdrasil_decision_forests/dataset/BUILD

@@ -59,6 +59,7 @@ cc_library_ydf(
     deps = [
         ":data_spec_cc_proto",
         ":example_cc_proto",
+        "//third_party/re2",
         "//yggdrasil_decision_forests/utils:hash",
         "//yggdrasil_decision_forests/utils:logging",
         "//yggdrasil_decision_forests/utils:regex",

yggdrasil_decision_forests/dataset/data_spec.cc

@@ -43,6 +43,7 @@
 #include "absl/strings/string_view.h"
 #include "absl/strings/substitute.h"
 #include "absl/types/span.h"
+#include "third_party/re2/re2.h"
 #include "yggdrasil_decision_forests/dataset/data_spec.pb.h"
 #include "yggdrasil_decision_forests/dataset/example.pb.h"
 #include "yggdrasil_decision_forests/utils/logging.h"
@@ -703,12 +704,23 @@ absl::Status Tokenize(const absl::string_view text,
           absl::StrSplit(cased_text, absl::ByAnyChar(tokenizer.separator()));
       break;
     case proto::Tokenizer::REGEX_MATCH: {
-      std::string remaining = cased_text;
-      std::regex re(tokenizer.regex());
-      std::smatch sm;
-      while (std::regex_search(remaining, sm, re)) {
-        unit_tokens.emplace_back(sm.str());
-        remaining = sm.suffix();
+      RE2 re(tokenizer.regex());
+      if (!re.ok()) {
+        return absl::InvalidArgumentError(
+            absl::StrCat("Invalid regular expression: ", tokenizer.regex()));
+      }
+      absl::string_view sp(cased_text);
+      size_t startpos = 0;
+      absl::string_view submatch;
+      while (re.Match(sp, startpos, sp.size(), RE2::UNANCHORED, &submatch, 1)) {
+        unit_tokens.emplace_back(submatch.data(), submatch.size());
+        startpos = (submatch.data() - sp.data()) + submatch.size();
+        if (submatch.empty()) {
+          startpos++;
+        }
+        if (startpos > sp.size()) {
+          break;
+        }
       }
     } break;
     case proto::Tokenizer::CHARACTER:

yggdrasil_decision_forests/port/python/ydf/learner/wrapper/BUILD

@@ -28,6 +28,7 @@ cc_library_ydf(
     srcs = ["wrapper_generator.cc"],
     hdrs = ["wrapper_generator.h"],
     deps = [
+        "//third_party/re2",
         "@com_google_absl//absl/container:flat_hash_map",
         "@com_google_absl//absl/log",
         "@com_google_absl//absl/status",

yggdrasil_decision_forests/port/python/ydf/learner/wrapper/wrapper_generator.cc

@@ -20,7 +20,6 @@
 #include <cstddef>
 #include <iterator>
 #include <memory>
-#include <regex>
 #include <string>
 #include <tuple>
 #include <utility>
@@ -39,6 +38,7 @@
 #include "absl/strings/str_split.h"
 #include "absl/strings/string_view.h"
 #include "absl/strings/substitute.h"
+#include "third_party/re2/re2.h"
 #include "yggdrasil_decision_forests/learner/abstract_learner.h"
 #include "yggdrasil_decision_forests/learner/abstract_learner.pb.h"
 #include "yggdrasil_decision_forests/learner/learner_library.h"
@@ -934,15 +934,15 @@ absl::Status FixDefaultFieldsDocumentation(std::string* fields_documentation,
       return absl::InvalidArgumentError(
           "Missing documentation for discretize_numerical_columns");
     }
-    std::regex end_pattern("\\n {4}[^ \\t\\n\\r\\f\\v]");
-    auto search_start_it =
-        fields_documentation->cbegin() + start_pos + start_marker.length();
-    std::smatch match;
-    size_t end_pos;
-
-    if (std::regex_search(search_start_it, fields_documentation->cend(), match,
-                          end_pattern)) {
-      end_pos = std::distance(fields_documentation->cbegin(), match[0].first);
+    static const LazyRE2 end_pattern(R"((\n {4}[^ \t\n\r\f\v]))");
+    STATUS_CHECK(end_pattern->ok());
+
+    absl::string_view subset = absl::string_view(*fields_documentation)
+                                   .substr(start_pos + start_marker.length());
+    absl::string_view match;
+
+    if (RE2::PartialMatch(subset, *end_pattern, &match)) {
+      size_t end_pos = match.data() - fields_documentation->data();
       size_t length = end_pos - start_pos;
       fields_documentation->replace(start_pos, length, replacement_string);
     } else {

yggdrasil_decision_forests/utils/BUILD

@@ -107,6 +107,7 @@ cc_library_ydf(
         ":logging_default",
         ":protobuf",
         ":status_macros",
+        "//third_party/re2",
         "@com_google_absl//absl/base:core_headers",
         "@com_google_absl//absl/log:check",
         "@com_google_absl//absl/status",
@@ -127,6 +128,7 @@ cc_library_ydf(
         ":filesystem_interface",
         ":logging",
         ":status_macros",
+        "//third_party/re2",
         "@com_google_absl//absl/status",
         "@com_google_absl//absl/strings",
         "@google_cloud_cpp//:storage",
@@ -209,6 +211,7 @@ cc_library_ydf(
         ":logging",
         ":protobuf",
         ":status_macros",
+        "//third_party/re2",
         "@com_google_absl//absl/functional:any_invocable",
         "@com_google_absl//absl/log",
         "@com_google_absl//absl/log:check",

yggdrasil_decision_forests/utils/filesystem_default.cc

@@ -23,7 +23,6 @@
 #include <initializer_list>
 #include <ios>
 #include <memory>
-#include <regex>  // NOLINT
 #include <string>
 #include <utility>
 #include <vector>
@@ -36,6 +35,7 @@
 #include "absl/strings/str_replace.h"
 #include "absl/strings/string_view.h"
 #include "absl/types/optional.h"
+#include "third_party/re2/re2.h"
 #include "yggdrasil_decision_forests/utils/filesystem_interface.h"
 #include "yggdrasil_decision_forests/utils/logging.h"
 #include "yggdrasil_decision_forests/utils/status_macros.h"
@@ -164,18 +164,12 @@ std::string JoinPathList(std::initializer_list<absl::string_view> paths) {
 
 bool GenerateShardedFilenames(absl::string_view spec,
                               std::vector<std::string>* names) {
-  std::regex num_shard_pattern(R"((.*)\@(\*|[0-9]+)(?:(\..+))?)");
-  std::smatch match;
+  static const LazyRE2 num_shard_pattern = {R"((.*)\@(\*|[0-9]+)(?:(\..+))?)"};
+  std::string prefix, count, suffix;
   std::string str_spec(spec);
-  if (!std::regex_match(str_spec, match, num_shard_pattern)) {
+  if (!RE2::FullMatch(str_spec, *num_shard_pattern, &prefix, &count, &suffix)) {
     return false;
   }
-  if (match.size() != 4) {
-    return false;
-  }
-  const auto prefix = match[1].str();
-  const auto count = match[2].str();
-  const auto suffix = match[3].str();
 
   int int_count;
   if (count == "*") {
@@ -212,7 +206,8 @@ absl::Status Match(absl::string_view pattern, std::vector<std::string>* results,
     const auto filename = fs::path(SV_ABSL_TO_STD(pattern)).filename().string();
     std::string regexp_filename =
         absl::StrReplaceAll(filename, {{".", "\\."}, {"*", ".*"}, {"?", "."}});
-    std::regex regexp_pattern(regexp_filename);
+    RE2 regexp_pattern(regexp_filename);
+    STATUS_CHECK(regexp_pattern.ok());
     std::error_code error;
 
     const fs::directory_iterator path_end;
@@ -221,7 +216,7 @@ absl::Status Match(absl::string_view pattern, std::vector<std::string>* results,
       if (!fs::is_regular_file(path->path())) {
         continue;
       }
-      if (std::regex_match(path->path().filename().string(), regexp_pattern)) {
+      if (RE2::FullMatch(path->path().filename().string(), regexp_pattern)) {
         results->push_back(path->path().string());
       }
     }

yggdrasil_decision_forests/utils/filesystem_default_gcs.cc

@@ -18,7 +18,6 @@
 #include <cstring>
 #include <initializer_list>
 #include <memory>
-#include <regex>  // NOLINT
 #include <string>
 #include <utility>
 #include <vector>
@@ -28,6 +27,7 @@
 #include "absl/strings/str_replace.h"
 #include "absl/strings/string_view.h"
 #include "google/cloud/storage/client.h"
+#include "third_party/re2/re2.h"
 #include "yggdrasil_decision_forests/utils/bytestream.h"
 #include "yggdrasil_decision_forests/utils/filesystem_default.h"
 #include "yggdrasil_decision_forests/utils/filesystem_interface.h"
@@ -198,11 +198,15 @@ class FileSystemImplementation : public FileSystemInterface {
 
     std::string regexp_filename = absl::StrReplaceAll(
         cloud_path.object, {{".", "\\."}, {"*", ".*"}, {"?", "."}});
-    std::regex regexp_pattern(regexp_filename);
+    RE2 regexp_pattern(regexp_filename);
+    if (!regexp_pattern.ok()) {
+      return absl::InvalidArgumentError(
+          "Invalid regular expression generated from pattern.");
+    }
 
     for (const auto& candidate : candidates) {
       GCS_RETURN_IF_ERROR(candidate.status());
-      if (!std::regex_match(candidate->name(), regexp_pattern)) {
+      if (!RE2::FullMatch(candidate->name(), regexp_pattern)) {
         continue;
       }
       results->push_back(

yggdrasil_decision_forests/utils/filesystem_tensorflow_impl.cc

@@ -18,7 +18,6 @@
 #include <cstring>
 #include <initializer_list>
 #include <memory>
-#include <regex>  // NOLINT
 #include <string>
 #include <vector>
 
@@ -31,6 +30,7 @@
 #include "src/google/protobuf/message.h"
 #include "src/google/protobuf/message_lite.h"
 #include "src/google/protobuf/text_format.h"
+#include "third_party/re2/re2.h"
 #include "tensorflow/core/platform/env.h"
 #include "tensorflow/core/platform/file_system.h"
 #include "tensorflow/core/platform/path.h"
@@ -225,18 +225,14 @@ class FileSystemImplementation : public FileSystemInterface {
 
   bool GenerateShardedFilenames(absl::string_view spec,
                                 std::vector<std::string>* names) override {
-    std::regex num_shard_pattern(R"((.*)\@(\*|[0-9]+)(?:(\..+))?)");
-    std::smatch match;
+    static const LazyRE2 num_shard_pattern = {
+        R"((.*)\@(\*|[0-9]+)(?:(\..+))?)"};
+    std::string prefix, count, suffix;
     std::string str_spec(spec);
-    if (!std::regex_match(str_spec, match, num_shard_pattern)) {
+    if (!RE2::FullMatch(str_spec, *num_shard_pattern, &prefix, &count,
+                        &suffix)) {
       return false;
     }
-    if (match.size() != 4) {
-      return false;
-    }
-    const auto prefix = match[1].str();
-    const auto count = match[2].str();
-    const auto suffix = match[3].str();
 
     int int_count;
     if (count == "*") {