GH-50988: [C++][Python] Substrait: add mappings for starts_with, ends_with and match_substring

cpp/src/arrow/engine/substrait/extension_set.cc

@@ -26,13 +26,15 @@
 
 #include "arrow/compute/api_scalar.h"
 #include "arrow/engine/substrait/options.h"
+#include "arrow/scalar.h"
 #include "arrow/type.h"
 #include "arrow/type_fwd.h"
 #include "arrow/util/checked_cast.h"
 #include "arrow/util/hash_util.h"
 #include "arrow/util/hashing.h"
 #include "arrow/util/logging_internal.h"
 #include "arrow/util/string.h"
+#include "arrow/util/utf8.h"
 
 namespace arrow {
 
@@ -757,6 +759,19 @@ static std::vector<std::string> kRoundModes = {
     "TIE_UP", "TIE_TOWARDS_ZERO", "TIE_AWAY_FROM_ZERO", "TIE_TO_EVEN",    "TIE_TO_ODD"};
 static EnumParser<compute::RoundMode> kRoundModeParser(kRoundModes);
 
+// The case_sensitivity option used by the Substrait string functions
+// (starts_with / ends_with / contains).  Arrow's MatchSubstringOptions only
+// distinguishes case sensitive from case insensitive, so CASE_INSENSITIVE_ASCII
+// is parsed but not implemented.
+enum class CaseSensitivity {
+  kCaseSensitive = 0,
+  kCaseInsensitive,
+  kCaseInsensitiveAscii
+};
+static std::vector<std::string> kCaseSensitivityOptions = {
+    "CASE_SENSITIVE", "CASE_INSENSITIVE", "CASE_INSENSITIVE_ASCII"};
+static EnumParser<CaseSensitivity> kCaseSensitivityParser(kCaseSensitivityOptions);
+
 template <typename Enum>
 Result<Enum> ParseOptionOrElse(const SubstraitCall& call, std::string_view option_name,
                                const EnumParser<Enum>& parser,
@@ -961,6 +976,74 @@ ExtensionIdRegistry::SubstraitCallToArrow DecodeConcatMapping() {
   };
 }
 
+// Substrait's starts_with / ends_with / contains take the pattern as a second
+// value argument.  The matching Arrow kernels (starts_with / ends_with /
+// match_substring) are unary and carry the pattern in MatchSubstringOptions, so
+// the second argument must be a string literal.
+ExtensionIdRegistry::SubstraitCallToArrow DecodeMatchSubstringMapping(
+    const std::string& function_name) {
+  return [function_name](const SubstraitCall& call) -> Result<compute::Expression> {
+    if (call.size() != 2) {
+      return Status::NotImplemented("Acero does not have a kernel for ", function_name,
+                                    " that receives ", call.size(), " arguments");
+    }
+    ARROW_ASSIGN_OR_RAISE(
+        CaseSensitivity case_sensitivity,
+        ParseOptionOrElse(
+            call, "case_sensitivity", kCaseSensitivityParser,
+            {CaseSensitivity::kCaseSensitive, CaseSensitivity::kCaseInsensitive},
+            CaseSensitivity::kCaseSensitive));
+    ARROW_ASSIGN_OR_RAISE(compute::Expression input, call.GetValueArg(0));
+    ARROW_ASSIGN_OR_RAISE(compute::Expression substring, call.GetValueArg(1));
+    const Datum* pattern = substring.literal();
+    if (pattern == nullptr || !pattern->is_scalar() ||
+        pattern->type()->id() != Type::STRING || !pattern->scalar()->is_valid) {
+      return Status::NotImplemented(
+          "The Arrow ", function_name,
+          " kernel requires the substring argument to be a non-null string literal");
+    }
+    auto options = std::make_shared<compute::MatchSubstringOptions>(
+        std::string(pattern->scalar_as<StringScalar>().view()),
+        /*ignore_case=*/case_sensitivity == CaseSensitivity::kCaseInsensitive);
+    return compute::call(function_name, {std::move(input)}, std::move(options));
+  };
+}
+
+ExtensionIdRegistry::ArrowToSubstraitCall EncodeMatchSubstring(Id substrait_fn_id) {
+  return
+      [substrait_fn_id](const compute::Expression::Call& call) -> Result<SubstraitCall> {
+        if (call.options == nullptr) {
+          return Status::Invalid("The call to ", call.function_name,
+                                 " is missing its MatchSubstringOptions");
+        }
+        // Substrait's starts_with / ends_with / contains only accept strings while
+        // the Arrow kernels also accept binary-like input and non-UTF-8 patterns
+        const DataType* input_type = call.arguments[0].type();
+        if (input_type == nullptr || input_type->id() != Type::STRING) {
+          return Status::NotImplemented(
+              "Substrait only supports ", substrait_fn_id.name,
+              " on string input but the input to ", call.function_name, " is ",
+              input_type == nullptr ? "unbound" : input_type->ToString());
+        }
+        auto match_options =
+            checked_pointer_cast<compute::MatchSubstringOptions>(call.options);
+        util::InitializeUTF8();
+        if (!util::ValidateUTF8(match_options->pattern)) {
+          return Status::NotImplemented("Substrait only supports ", substrait_fn_id.name,
+                                        " with a valid UTF-8 pattern");
+        }
+        // nullable=true errs on the side of caution
+        SubstraitCall substrait_call(substrait_fn_id, call.type.GetSharedPtr(),
+                                     /*nullable=*/true);
+        substrait_call.SetValueArg(0, call.arguments[0]);
+        substrait_call.SetValueArg(1, compute::literal(match_options->pattern));
+        substrait_call.SetOption(
+            "case_sensitivity",
+            {match_options->ignore_case ? "CASE_INSENSITIVE" : "CASE_SENSITIVE"});
+        return substrait_call;
+      };
+}
+
 ExtensionIdRegistry::SubstraitAggregateToArrow DecodeBasicAggregate(
     const std::string& arrow_function_name) {
   return [arrow_function_name](const SubstraitCall& call) -> Result<compute::Aggregate> {
@@ -1147,6 +1230,13 @@ struct DefaultExtensionIdRegistry : ExtensionIdRegistryImpl {
                                       DecodeTemporalExtractionMapping()));
     DCHECK_OK(AddSubstraitCallToArrow({kSubstraitStringFunctionsUri, "concat"},
                                       DecodeConcatMapping()));
+    for (const auto& fn_pair : std::vector<std::pair<std::string_view, std::string>>{
+             {"starts_with", "starts_with"},
+             {"ends_with", "ends_with"},
+             {"contains", "match_substring"}}) {
+      DCHECK_OK(AddSubstraitCallToArrow({kSubstraitStringFunctionsUri, fn_pair.first},
+                                        DecodeMatchSubstringMapping(fn_pair.second)));
+    }
     DCHECK_OK(
         AddSubstraitCallToArrow({kSubstraitComparisonFunctionsUri, "is_null"},
                                 DecodeOptionlessBasicMapping("is_null", /*max_args=*/1)));
@@ -1239,6 +1329,15 @@ struct DefaultExtensionIdRegistry : ExtensionIdRegistryImpl {
 
     DCHECK_OK(AddArrowToSubstraitCall(
         "is_null", EncodeIsNull({kSubstraitComparisonFunctionsUri, "is_null"})));
+
+    for (const auto& fn_pair : std::vector<std::pair<std::string, std::string_view>>{
+             {"starts_with", "starts_with"},
+             {"ends_with", "ends_with"},
+             {"match_substring", "contains"}}) {
+      DCHECK_OK(AddArrowToSubstraitCall(
+          fn_pair.first,
+          EncodeMatchSubstring({kSubstraitStringFunctionsUri, fn_pair.second})));
+    }
   }
 };
 

cpp/src/arrow/engine/substrait/serde_test.cc

@@ -6377,5 +6377,77 @@ TEST(Substrait, ExtendedExpressionInvalidPlans) {
               Raises(StatusCode::Invalid, testing::HasSubstr("Ambiguous plan")));
 }
 
+TEST(Substrait, StringMatchExpressionSerialization) {
+  std::shared_ptr<Schema> test_schema = schema({field("cat", utf8())});
+  for (const auto& fn_name : {"starts_with", "ends_with", "match_substring"}) {
+    for (bool ignore_case : {false, true}) {
+      CheckExpressionRoundTrip(
+          *test_schema, compute::call(fn_name, {compute::field_ref(0)},
+                                      compute::MatchSubstringOptions("al", ignore_case)));
+    }
+  }
+
+  // Substrait only accepts string input and UTF-8 patterns
+  auto serialize = [](const std::shared_ptr<DataType>& type,
+                      std::string pattern) -> Result<std::shared_ptr<Buffer>> {
+    Schema schema({field("s", type)});
+    ExtensionSet ext_set;
+    ARROW_ASSIGN_OR_RAISE(
+        compute::Expression bound,
+        compute::call("starts_with", {compute::field_ref(0)},
+                      compute::MatchSubstringOptions(std::move(pattern)))
+            .Bind(schema));
+    return SerializeExpression(bound, &ext_set);
+  };
+  ASSERT_THAT(serialize(binary(), "al"),
+              Raises(StatusCode::NotImplemented, HasSubstr("on string input")));
+  ASSERT_THAT(serialize(large_utf8(), "al"),
+              Raises(StatusCode::NotImplemented, HasSubstr("on string input")));
+  ASSERT_THAT(serialize(utf8(), "\xff"),
+              Raises(StatusCode::NotImplemented, HasSubstr("valid UTF-8 pattern")));
+}
+
+TEST(Substrait, StringMatchExpressionDeserialization) {
+  ExtensionSet ext_set;
+  ASSERT_OK_AND_ASSIGN(
+      uint32_t anchor,
+      ext_set.EncodeFunction({kSubstraitStringFunctionsUri, "starts_with"}));
+  // A starts_with call on field 0 with the given pattern argument and options
+  auto deserialize = [&](std::string_view pattern_arg,
+                         std::string_view options) -> Result<compute::Expression> {
+    std::string json = R"({"scalarFunction":{"functionReference":)" +
+                       std::to_string(anchor) +
+                       R"(,"outputType":{"bool":{}},"arguments":[
+        {"value":{"selection":{"directReference":{"structField":{"field":0}},
+                               "rootReference":{}}}},
+        {"value":)" + std::string(pattern_arg) +
+                       R"(}],"options":)" + std::string(options) + "}}";
+    ARROW_ASSIGN_OR_RAISE(std::shared_ptr<Buffer> buf,
+                          internal::SubstraitFromJSON("Expression", json));
+    return DeserializeExpression(*buf, ext_set);
+  };
+  constexpr std::string_view kLiteral = R"({"literal":{"string":"al"}})";
+
+  // No case_sensitivity option means case sensitive
+  ASSERT_OK_AND_ASSIGN(compute::Expression no_option, deserialize(kLiteral, "[]"));
+  ASSERT_EQ(compute::call("starts_with", {compute::field_ref(0)},
+                          compute::MatchSubstringOptions("al")),
+            no_option);
+
+  // CASE_INSENSITIVE_ASCII has no Arrow equivalent
+  ASSERT_THAT(
+      deserialize(
+          kLiteral,
+          R"([{"name":"case_sensitivity","preference":["CASE_INSENSITIVE_ASCII"]}])"),
+      Raises(StatusCode::NotImplemented, HasSubstr("the only supported options are")));
+
+  // The pattern has to be a literal because the Arrow kernel takes it as an option
+  ASSERT_THAT(
+      deserialize(R"({"selection":{"directReference":{"structField":{"field":0}},
+                                   "rootReference":{}}})",
+                  "[]"),
+      Raises(StatusCode::NotImplemented, HasSubstr("substring argument to be a")));
+}
+
 }  // namespace engine
 }  // namespace arrow

docs/source/cpp/acero/substrait.rst

@@ -227,6 +227,12 @@ Functions
   * Acero does not support the SATURATE option for overflow
   * Acero does not support kernels that take more than two arguments
     for the functions ``and``, ``or``, ``xor``
+  * The functions ``starts_with``, ``ends_with``, and ``contains`` map onto Acero
+    kernels that take the substring as a function option.  The second argument
+    must therefore be a non-null string literal.  The ``case_sensitivity`` option
+    supports ``CASE_SENSITIVE`` and ``CASE_INSENSITIVE`` but not
+    ``CASE_INSENSITIVE_ASCII``.  Only calls on ``string`` input with a valid
+    UTF-8 pattern can be serialized to Substrait
 
 * Substrait has not yet clearly identified the form that URIs should take for
   standard functions.  Acero will look for the URIs to the ``main`` GitHub branch.

python/pyarrow/tests/test_substrait.py

@@ -946,6 +946,19 @@ def test_serializing_expressions(expr):
     assert "test_expr" in returned.expressions
 
 
+@pytest.mark.parametrize("func", [pc.starts_with, pc.ends_with, pc.match_substring])
+def test_serializing_string_match_expressions(func):
+    # GH-50988: these map onto Substrait's starts_with / ends_with / contains
+    schema = pa.schema([pa.field("s", pa.string())])
+    expr = func(pc.field("s"), "al", ignore_case=True)
+
+    buf = pa.substrait.serialize_expressions([expr], ["test_expr"], schema)
+    returned = pa.substrait.deserialize_expressions(buf)
+    # Substrait refers to fields by index, so compare against the normalized form
+    assert str(returned.expressions["test_expr"]) == str(
+        func(pc.field(0), "al", ignore_case=True))
+
+
 def test_arrow_specific_types():
     fields = {
         "time_seconds": (pa.time32("s"), 0),