Fix replace SparkSQL function (4922)

Author: rui-mo

velox/functions/lib/string/StringCore.h

@@ -299,6 +299,7 @@ inline int64_t findNthInstanceByteIndexFromEnd(
 /// each charecter. When inputString is empty results is empty.
 /// replace("", "", "x") = ""
 /// replace("aa", "", "x") = "xaxax"
+template <bool ignoreEmptyReplaced = false>
 inline static size_t replace(
     char* outputString,
     const std::string_view& inputString,
@@ -309,6 +310,13 @@ inline static size_t replace(
     return 0;
   }
 
+  if (ignoreEmptyReplaced && replaced.size() == 0) {
+    if (!inPlace) {
+      std::memcpy(outputString, inputString.data(), inputString.size());
+    }
+    return inputString.size();
+  }
+
   size_t readPosition = 0;
   size_t writePosition = 0;
   // Copy needed in out of place replace, and when replaced and replacement are

velox/functions/lib/string/StringImpl.h

@@ -209,7 +209,10 @@ stringPosition(const T& string, const T& subString, int64_t instance = 0) {
 
 /// Replace replaced with replacement in inputString and write results to
 /// outputString.
-template <typename TOutString, typename TInString>
+template <
+    typename TOutString,
+    typename TInString,
+    bool ignoreEmptyReplaced = false>
 FOLLY_ALWAYS_INLINE void replace(
     TOutString& outputString,
     const TInString& inputString,
@@ -226,7 +229,7 @@ FOLLY_ALWAYS_INLINE void replace(
         (inputString.size() / replaced.size()) * replacement.size());
   }
 
-  auto outputSize = stringCore::replace(
+  auto outputSize = stringCore::replace<ignoreEmptyReplaced>(
       outputString.data(),
       std::string_view(inputString.data(), inputString.size()),
       std::string_view(replaced.data(), replaced.size()),
@@ -237,14 +240,17 @@ FOLLY_ALWAYS_INLINE void replace(
 }
 
 /// Replace replaced with replacement in place in string.
-template <typename TInOutString, typename TInString>
+template <
+    typename TInOutString,
+    typename TInString,
+    bool ignoreEmptyReplaced = false>
 FOLLY_ALWAYS_INLINE void replaceInPlace(
     TInOutString& string,
     const TInString& replaced,
     const TInString& replacement) {
   assert(replacement.size() <= replaced.size() && "invalid inplace replace");
 
-  auto outputSize = stringCore::replace(
+  auto outputSize = stringCore::replace<ignoreEmptyReplaced>(
       string.data(),
       std::string_view(string.data(), string.size()),
       std::string_view(replaced.data(), replaced.size()),

velox/functions/prestosql/StringFunctions.cpp

@@ -284,7 +284,8 @@ class ConcatFunction : public exec::VectorFunction {
  * If search is an empty string, inserts replace in front of every character
  *and at the end of the string.
  **/
-class Replace : public exec::VectorFunction {
+template <bool ignoreEmptyReplaced>
+class ReplaceBase : public exec::VectorFunction {
  private:
   template <
       typename StringReader,
@@ -298,7 +299,10 @@ class Replace : public exec::VectorFunction {
       FlatVector<StringView>* results) const {
     rows.applyToSelected([&](int row) {
       auto proxy = exec::StringWriter<>(results, row);
-      stringImpl::replace(
+      stringImpl::replace<
+          decltype(proxy),
+          decltype(searchReader(row)),
+          ignoreEmptyReplaced>(
           proxy, stringReader(row), searchReader(row), replaceReader(row));
       proxy.finalize();
     });
@@ -317,7 +321,10 @@ class Replace : public exec::VectorFunction {
     rows.applyToSelected([&](int row) {
       auto proxy = exec::StringWriter<true /*reuseInput*/>(
           results, row, stringReader(row) /*reusedInput*/, true /*inPlace*/);
-      stringImpl::replaceInPlace(proxy, searchReader(row), replaceReader(row));
+      stringImpl::replaceInPlace<
+          decltype(proxy),
+          decltype(searchReader(row)),
+          ignoreEmptyReplaced>(proxy, searchReader(row), replaceReader(row));
       proxy.finalize();
     });
   }
@@ -429,6 +436,11 @@ class Replace : public exec::VectorFunction {
     return {{0, 2}};
   }
 };
+
+class Replace : public ReplaceBase<false /*ignoreEmptyReplaced*/> {};
+
+class ReplaceIgnoreEmptyReplaced
+    : public ReplaceBase<true /*ignoreEmptyReplaced*/> {};
 } // namespace
 
 VELOX_DECLARE_VECTOR_FUNCTION(
@@ -456,4 +468,9 @@ VELOX_DECLARE_VECTOR_FUNCTION(
     Replace::signatures(),
     std::make_unique<Replace>());
 
+VELOX_DECLARE_VECTOR_FUNCTION(
+    udf_replace_ignore_empty_replaced,
+    ReplaceIgnoreEmptyReplaced::signatures(),
+    std::make_unique<ReplaceIgnoreEmptyReplaced>());
+
 } // namespace facebook::velox::functions

velox/functions/sparksql/Register.cpp

@@ -63,7 +63,8 @@ static void workAroundRegistrationMacro(const std::string& prefix) {
   // String functions.
   VELOX_REGISTER_VECTOR_FUNCTION(udf_concat, prefix + "concat");
   VELOX_REGISTER_VECTOR_FUNCTION(udf_lower, prefix + "lower");
-  VELOX_REGISTER_VECTOR_FUNCTION(udf_replace, prefix + "replace");
+  VELOX_REGISTER_VECTOR_FUNCTION(
+      udf_replace_ignore_empty_replaced, prefix + "replace");
   VELOX_REGISTER_VECTOR_FUNCTION(udf_upper, prefix + "upper");
   // Logical.
   VELOX_REGISTER_VECTOR_FUNCTION(udf_not, prefix + "not");

velox/functions/sparksql/tests/StringTest.cpp

@@ -190,6 +190,14 @@ class StringTest : public SparkFunctionBaseTest {
       std::optional<int32_t> size) {
     return evaluateOnce<std::string>("lpad(c0, c1)", string, size);
   }
+
+  std::optional<std::string> replace(
+      std::optional<std::string> str,
+      std::optional<std::string> replaced,
+      std::optional<std::string> replacement) {
+    return evaluateOnce<std::string>(
+        "replace(c0, c1, c2)", str, replaced, replacement);
+  }
 };
 
 TEST_F(StringTest, Ascii) {
@@ -682,5 +690,19 @@ TEST_F(StringTest, translateNonconstantMatch) {
   expected = makeFlatVector<std::string>({"åbaæçè", "åæcèaç"});
   testTranslate({input, match, replace}, expected);
 }
+
+TEST_F(StringTest, replace) {
+  EXPECT_EQ(replace("aaabaac", "a", "z"), "zzzbzzc");
+  EXPECT_EQ(replace("aaabaac", "", "z"), "aaabaac");
+  EXPECT_EQ(replace("aaabaac", "a", ""), "bc");
+  EXPECT_EQ(replace("aaabaac", "x", "z"), "aaabaac");
+  EXPECT_EQ(replace("aaabaac", "ab", "z"), "aazaac");
+  EXPECT_EQ(replace("aaabaac", "aa", "z"), "zabzc");
+  EXPECT_EQ(replace("aaabaac", "aa", "xyz"), "xyzabxyzc");
+  EXPECT_EQ(replace("aaabaac", "aaabaac", "z"), "z");
+  EXPECT_EQ(
+      replace("123\u6570\u6570\u636E", "\u6570\u636E", "data"),
+      "123\u6570data");
+}
 } // namespace
 } // namespace facebook::velox::functions::sparksql::test