[SPARK-27839][SQL] Change UTF8String.replace() to operate on UTF8 bytes

common/unsafe/pom.xml

@@ -65,6 +65,10 @@
       <groupId>com.google.guava</groupId>
       <artifactId>guava</artifactId>
     </dependency>
+    <dependency>
+      <groupId>org.apache.commons</groupId>
+      <artifactId>commons-lang3</artifactId>
+    </dependency>
 
     <!-- Provided dependencies -->
     <dependency>
@@ -84,11 +88,6 @@
       <artifactId>scalacheck_${scala.binary.version}</artifactId>
       <scope>test</scope>
     </dependency>
-    <dependency>
-      <groupId>org.apache.commons</groupId>
-      <artifactId>commons-lang3</artifactId>
-      <scope>test</scope>
-    </dependency>
     <dependency>
       <groupId>org.apache.commons</groupId>
       <artifactId>commons-text</artifactId>

common/unsafe/src/main/java/org/apache/spark/unsafe/types/UTF8String.java

@@ -25,6 +25,7 @@
 import java.util.Arrays;
 import java.util.Map;
 
+import org.apache.commons.lang3.StringUtils;
 import com.esotericsoftware.kryo.Kryo;
 import com.esotericsoftware.kryo.KryoSerializable;
 import com.esotericsoftware.kryo.io.Input;
@@ -976,9 +977,28 @@ public UTF8String replace(UTF8String search, UTF8String replace) {
     if (EMPTY_UTF8.equals(search)) {
       return this;
     }
-    String replaced = toString().replace(
-      search.toString(), replace.toString());
-    return fromString(replaced);
+    return replace(search.toString(), replace.toString());
+  }
+
+  public UTF8String replace(String search, String replace) {
+    String before = toString();
+    String after;
+    if (search.length() == 1 && replace.length() == 1) {
+      // Use single-character-replacement fast path
+      after = before.replace(search.charAt(0), replace.charAt(0));
+    } else {
+      // In Java 8, String.replace() is implemented using a regex and is therefore
+      // somewhat inefficient (see https://bugs.openjdk.java.net/browse/JDK-8058779).
+      // This is fixed in Java 9, but in Java 8 we can use Commons StringUtils instead:
+      after = StringUtils.replace(before, search, replace);
+    }
+    // Use reference equality to cheaply detect whether the replacement had no effect,
+    // in which case we can simply return the original UTF8String and save some copying.
+    if (before == after) {
+      return this;
+    } else {
+      return fromString(after);
+    }
   }
 
   // TODO: Need to use `Code Point` here instead of Char in case the character longer than 2 bytes

sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/stringExpressions.scala

@@ -443,9 +443,40 @@ case class StringReplace(srcExpr: Expression, searchExpr: Expression, replaceExp
   }
 
   override def doGenCode(ctx: CodegenContext, ev: ExprCode): ExprCode = {
-    nullSafeCodeGen(ctx, ev, (src, search, replace) => {
-      s"""${ev.value} = $src.replace($search, $replace);"""
-    })
+    if (searchExpr.foldable && replaceExpr.foldable) {
+      // The search and replacement strings are constants, so we can use a more optimized path
+      // which avoids repeatedly decoding those UTF8Strings.
+      val search = searchExpr.eval()
+      val replace = replaceExpr.eval()
+      if (search == null || replace == null) {
+        // Either the search or replacement is null, so the entire expression evaluates to null
+        ev.copy(code = code"""
+          boolean ${ev.isNull} = true;
+          ${CodeGenerator.javaType(dataType)} ${ev.value} = ${CodeGenerator.defaultValue(dataType)};
+        """)
+      } else {
+        val searchStr: String = search.asInstanceOf[UTF8String].toString
+        val replaceStr: String = replace.asInstanceOf[UTF8String].toString
+        val searchStrRef = ctx.addReferenceObj("searchStr", searchStr, "String")
+        val replaceStrRef = ctx.addReferenceObj("searchStr", replaceStr, "String")
+        // We don't use nullSafeCodeGen here because we don't want to re-evaluate the search
+        // and replace expressions.
+        val eval = srcExpr.genCode(ctx)
+        ev.copy(code = code"""
+        ${eval.code}
+        boolean ${ev.isNull} = ${eval.isNull};
+        ${CodeGenerator.javaType(dataType)} ${ev.value} = ${CodeGenerator.defaultValue(dataType)};
+        if (!${ev.isNull}) {
+          ${ev.value} = ${eval.value}.replace($searchStrRef, $replaceStrRef);
+        }
+      """)
+      }
+    } else {
+      // Either the search or replace expression is non-foldable, so use a slower path:
+      nullSafeCodeGen(ctx, ev, (src, search, replace) => {
+        s"""${ev.value} = $src.replace($search, $replace);"""
+      })
+    }
   }
 
   override def dataType: DataType = StringType

sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/expressions/StringExpressionsSuite.scala

@@ -409,22 +409,23 @@ class StringExpressionsSuite extends SparkFunSuite with ExpressionEvalHelper {
   }
 
   test("replace") {
-    checkEvaluation(
-      StringReplace(Literal("replace"), Literal("pl"), Literal("123")), "re123ace")
+    val nullString = Literal.create(null, StringType)
+    checkEvaluation(StringReplace(Literal("replace"), Literal("pl"), Literal("123")), "re123ace")
     checkEvaluation(StringReplace(Literal("replace"), Literal("pl"), Literal("")), "reace")
     checkEvaluation(StringReplace(Literal("replace"), Literal(""), Literal("123")), "replace")
-    checkEvaluation(StringReplace(Literal.create(null, StringType),
-      Literal("pl"), Literal("123")), null)
-    checkEvaluation(StringReplace(Literal("replace"),
-      Literal.create(null, StringType), Literal("123")), null)
-    checkEvaluation(StringReplace(Literal("replace"),
-      Literal("pl"), Literal.create(null, StringType)), null)
+    checkEvaluation(StringReplace(nullString, Literal("pl"), Literal("123")), null)
+    checkEvaluation(StringReplace(Literal("replace"), nullString, Literal("123")), null)
+    checkEvaluation(StringReplace(Literal("replace"), Literal("pl"), nullString), null)
     // test for multiple replace
     checkEvaluation(StringReplace(Literal("abcabc"), Literal("b"), Literal("12")), "a12ca12c")
     checkEvaluation(StringReplace(Literal("abcdabcd"), Literal("bc"), Literal("")), "adad")
+    // tests for single character search and replacement strings
+    checkEvaluation(StringReplace(Literal("abcabc"), Literal("a"), Literal("A")), "AbcAbc")
+    checkEvaluation(StringReplace(Literal("abcabc"), Literal("Z"), Literal("A")), "abcabc")
     // scalastyle:off
     // non ascii characters are not allowed in the source code, so we disable the scalastyle.
     checkEvaluation(StringReplace(Literal("花花世界"), Literal("花世"), Literal("ab")), "花ab界")
+    checkEvaluation(StringReplace(Literal("火"), Literal("火"), Literal("水")), "水")
     // scalastyle:on
   }