Add ends_with Presto function

facebookincubator · Oct 5, 2023 · 2293c5f · 2293c5f
1 parent b49b453
commit 2293c5f
Show file tree

Hide file tree

Showing 4 changed files with 42 additions and 0 deletions.
diff --git a/velox/docs/functions/presto/string.rst b/velox/docs/functions/presto/string.rst
@@ -33,6 +33,10 @@ String Functions
     This function provides the same functionality as the
     SQL-standard concatenation operator (``||``).
 
+.. function:: ends_with(string, substring) -> boolean
+
+    Returns whether ``string`` ends_with with ``substring``.
+
 .. function:: from_utf8(binary) -> varchar
 
     Decodes a UTF-8 encoded string from ``binary``. Invalid UTF-8 sequences

diff --git a/velox/functions/prestosql/StringFunctions.h b/velox/functions/prestosql/StringFunctions.h
@@ -239,6 +239,24 @@ struct StartsWithFunction {
   }
 };
 
+template <typename T>
+struct EndsWithFunction {
+  VELOX_DEFINE_FUNCTION_TYPES(T);
+
+  FOLLY_ALWAYS_INLINE void call(
+      out_type<bool>& result,
+      const arg_type<Varchar>& x,
+      const arg_type<Varchar>& y) {
+    if (x.size() < y.size()) {
+      result = false;
+      return;
+    }
+
+    result =
+        (memcmp(x.data() + (x.size() - y.size()), y.data(), y.size()) == 0);
+  }
+};
+
 /// Pad functions
 /// lpad(string, size, padString) → varchar
 ///     Left pads string to size characters with padString.  If size is

diff --git a/velox/functions/prestosql/registration/StringFunctionsRegistration.cpp b/velox/functions/prestosql/registration/StringFunctionsRegistration.cpp
@@ -46,6 +46,8 @@ void registerSimpleFunctions(const std::string& prefix) {
 
   registerFunction<StartsWithFunction, bool, Varchar, Varchar>(
       {prefix + "starts_with"});
+  registerFunction<EndsWithFunction, bool, Varchar, Varchar>(
+      {prefix + "ends_with"});
 
   registerFunction<SubstrFunction, Varchar, Varchar, int64_t>(
       {prefix + "substr"});

diff --git a/velox/functions/prestosql/tests/StringFunctionsTest.cpp b/velox/functions/prestosql/tests/StringFunctionsTest.cpp
@@ -870,6 +870,24 @@ TEST_F(StringFunctionsTest, startsWith) {
   ASSERT_FALSE(startsWith("", " "));
 }
 
+TEST_F(StringFunctionsTest, endsWith) {
+  auto endsWith = [&](const std::string& x, const std::string& y) {
+    return evaluateOnce<bool>(
+               "ends_with(c0, c1)", std::optional(x), std::optional(y))
+        .value();
+  };
+
+  ASSERT_TRUE(endsWith("", ""));
+  ASSERT_TRUE(endsWith("Hello world!", ""));
+  ASSERT_TRUE(endsWith("Hello world!", "world!"));
+  ASSERT_TRUE(endsWith("Hello world!", "lo world!"));
+  ASSERT_TRUE(endsWith("Hello world!", "Hello world!"));
+
+  ASSERT_FALSE(endsWith("Hello world!", " Hello world!"));
+  ASSERT_FALSE(endsWith("Hello world!", "hello"));
+  ASSERT_FALSE(endsWith("", " "));
+}
+
 // Test strpos function
 template <typename TInstance>
 void StringFunctionsTest::testStringPositionAllFlatVector(