diff --git a/cpp/src/arrow/engine/substrait/extension_set.cc b/cpp/src/arrow/engine/substrait/extension_set.cc index 4fb9b19645b1..78188dd8d72e 100644 --- a/cpp/src/arrow/engine/substrait/extension_set.cc +++ b/cpp/src/arrow/engine/substrait/extension_set.cc @@ -26,6 +26,7 @@ #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" @@ -33,6 +34,7 @@ #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 kRoundModes = { "TIE_UP", "TIE_TOWARDS_ZERO", "TIE_AWAY_FROM_ZERO", "TIE_TO_EVEN", "TIE_TO_ODD"}; static EnumParser 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 kCaseSensitivityOptions = { + "CASE_SENSITIVE", "CASE_INSENSITIVE", "CASE_INSENSITIVE_ASCII"}; +static EnumParser kCaseSensitivityParser(kCaseSensitivityOptions); + template Result ParseOptionOrElse(const SubstraitCall& call, std::string_view option_name, const EnumParser& 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 { + 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( + std::string(pattern->scalar_as().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 { + 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(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 { @@ -1147,6 +1230,13 @@ struct DefaultExtensionIdRegistry : ExtensionIdRegistryImpl { DecodeTemporalExtractionMapping())); DCHECK_OK(AddSubstraitCallToArrow({kSubstraitStringFunctionsUri, "concat"}, DecodeConcatMapping())); + for (const auto& fn_pair : std::vector>{ + {"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>{ + {"starts_with", "starts_with"}, + {"ends_with", "ends_with"}, + {"match_substring", "contains"}}) { + DCHECK_OK(AddArrowToSubstraitCall( + fn_pair.first, + EncodeMatchSubstring({kSubstraitStringFunctionsUri, fn_pair.second}))); + } } }; diff --git a/cpp/src/arrow/engine/substrait/serde_test.cc b/cpp/src/arrow/engine/substrait/serde_test.cc index 4de5480ef3cd..311876379830 100644 --- a/cpp/src/arrow/engine/substrait/serde_test.cc +++ b/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 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& type, + std::string pattern) -> Result> { + 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 { + 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 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 diff --git a/docs/source/cpp/acero/substrait.rst b/docs/source/cpp/acero/substrait.rst index a5532733627c..168cd9b5af37 100644 --- a/docs/source/cpp/acero/substrait.rst +++ b/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. diff --git a/python/pyarrow/tests/test_substrait.py b/python/pyarrow/tests/test_substrait.py index fcd1c8d48c5f..430bbeb65248 100644 --- a/python/pyarrow/tests/test_substrait.py +++ b/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),