[SPARK-58275][SQL][PYTHON] Add normalize Unicode normalization SQL function

Author: SreeramaYeshwanthGowd

common/utils/src/main/resources/error/error-conditions.json

@@ -4648,6 +4648,11 @@
           "expects a long literal, but got <invalidValue>."
         ]
       },
+      "NORMALIZE_FORM" : {
+        "message" : [
+          "expects one of the normalization forms 'NFC', 'NFD', 'NFKC', 'NFKD', but got <form>."
+        ]
+      },
       "NULL" : {
         "message" : [
           "expects a non-NULL value."

python/docs/source/reference/pyspark.sql/functions.rst

@@ -191,6 +191,7 @@ String Functions
     ltrim
     make_valid_utf8
     mask
+    normalize
     octet_length
     overlay
     position

python/pyspark/sql/connect/functions/builtin.py

@@ -2665,6 +2665,16 @@ def try_validate_utf8(str: "ColumnOrName") -> Column:
 try_validate_utf8.__doc__ = pysparkfuncs.try_validate_utf8.__doc__
 
 
+def normalize(str: "ColumnOrName", form: Optional["ColumnOrName"] = None) -> Column:
+    if form is None:
+        return _invoke_function_over_columns("normalize", str)
+    else:
+        return _invoke_function_over_columns("normalize", str, form)
+
+
+normalize.__doc__ = pysparkfuncs.normalize.__doc__
+
+
 def format_number(col: "ColumnOrName", d: int) -> Column:
     return _invoke_function("format_number", _to_col(col), lit(d))
 

python/pyspark/sql/functions/__init__.py

@@ -152,6 +152,7 @@
     "ltrim",
     "make_valid_utf8",
     "mask",
+    "normalize",
     "octet_length",
     "overlay",
     "position",

python/pyspark/sql/functions/builtin.py

@@ -15757,6 +15757,43 @@ def try_validate_utf8(str: "ColumnOrName") -> Column:
     return _invoke_function_over_columns("try_validate_utf8", str)
 
 
+@_try_remote_functions
+def normalize(str: "ColumnOrName", form: Optional["ColumnOrName"] = None) -> Column:
+    """
+    Returns the Unicode normalization of ``str`` using the given normalization ``form``.
+
+    .. versionadded:: 4.3.0
+
+    Parameters
+    ----------
+    str : :class:`~pyspark.sql.Column` or column name
+        the input string to normalize.
+    form : :class:`~pyspark.sql.Column` or column name, optional
+        the normalization form, one of 'NFC', 'NFD', 'NFKC', 'NFKD' (case-insensitive).
+        If omitted, 'NFC' is used.
+
+    Returns
+    -------
+    :class:`~pyspark.sql.Column`
+        the normalized string.
+
+    Examples
+    --------
+    >>> import pyspark.sql.functions as sf
+    >>> df = spark.createDataFrame([("\ufb01",)], ["s"])
+    >>> df.select(sf.normalize(df.s, sf.lit("NFKC"))).show()
+    +------------------+
+    |normalize(s, NFKC)|
+    +------------------+
+    |                fi|
+    +------------------+
+    """
+    if form is None:
+        return _invoke_function_over_columns("normalize", str)
+    else:
+        return _invoke_function_over_columns("normalize", str, form)
+
+
 @_try_remote_functions
 def format_number(col: "ColumnOrName", d: int) -> Column:
     """

python/pyspark/sql/tests/test_functions.py

@@ -1022,6 +1022,14 @@ def test_jaro_winkler_similarity_function(self):
         null_result = df.select(F.jaro_winkler_similarity(df.l, F.lit(None))).first()[0]
         self.assertIsNone(null_result)
 
+    def test_normalize_function(self):
+        df = self.spark.createDataFrame([("\ufb01",)], ["s"])
+        result = df.select(F.normalize(df.s, F.lit("NFKC"))).first()[0]
+        self.assertEqual(result, "fi")
+        # Null handling
+        null_result = df.select(F.normalize(F.lit(None), F.lit("NFC"))).first()[0]
+        self.assertIsNone(null_result)
+
     def test_between_function(self):
         df = self.spark.createDataFrame(
             [Row(a=1, b=2, c=3), Row(a=2, b=1, c=3), Row(a=4, b=1, c=4)]

sql/api/src/main/scala/org/apache/spark/sql/functions.scala

@@ -7407,6 +7407,32 @@ object functions {
   def try_validate_utf8(str: Column): Column =
     Column.fn("try_validate_utf8", str)
 
+  /**
+   * Returns the Unicode normalization of `str` using the given normalization `form`. Valid forms
+   * are 'NFC', 'NFD', 'NFKC', and 'NFKD'. The form name is case-insensitive.
+   *
+   * @param str
+   *   the input string to normalize.
+   * @param form
+   *   the normalization form: 'NFC', 'NFD', 'NFKC', or 'NFKD'.
+   * @group string_funcs
+   * @since 4.3.0
+   */
+  def normalize(str: Column, form: Column): Column =
+    Column.fn("normalize", str, form)
+
+  /**
+   * Returns the Unicode normalization of `str` using the default form 'NFC'. To use a different
+   * form, call the two-argument overload.
+   *
+   * @param str
+   *   the input string to normalize.
+   * @group string_funcs
+   * @since 4.3.0
+   */
+  def normalize(str: Column): Column =
+    Column.fn("normalize", str)
+
   /**
    * Formats numeric column x to a format like '#,###,###.##', rounded to d decimal places with
    * HALF_EVEN round mode, and returns the result as a string column.

sql/catalyst/src/main/java/org/apache/spark/sql/catalyst/expressions/ExpressionImplUtils.java

@@ -23,6 +23,7 @@
 import java.security.spec.AlgorithmParameterSpec;
 import java.text.BreakIterator;
 import java.text.DecimalFormat;
+import java.text.Normalizer;
 import java.util.ArrayList;
 import java.util.List;
 import java.util.Locale;
@@ -149,6 +150,24 @@ public static UTF8String tryValidateUTF8String(UTF8String utf8String) {
     else return null;
   }
 
+  /**
+   * Normalizes the given string using the given Unicode normalization form.
+   *
+   * @param input the input string to normalize.
+   * @param form the normalization form, one of NFC, NFD, NFKC, NFKD (case-insensitive).
+   * @return the normalized string.
+   */
+  public static UTF8String normalize(UTF8String input, UTF8String form) {
+    Normalizer.Form normalizerForm = switch (form.toString().toUpperCase(Locale.ROOT)) {
+      case "NFC" -> Normalizer.Form.NFC;
+      case "NFD" -> Normalizer.Form.NFD;
+      case "NFKC" -> Normalizer.Form.NFKC;
+      case "NFKD" -> Normalizer.Form.NFKD;
+      default -> throw QueryExecutionErrors.invalidNormalizeFormError(form.toString());
+    };
+    return UTF8String.fromString(Normalizer.normalize(input.toString(), normalizerForm));
+  }
+
   public static byte[] aesEncrypt(byte[] input,
                                   byte[] key,
                                   UTF8String mode,

sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/analysis/FunctionRegistry.scala

@@ -687,6 +687,7 @@ object FunctionRegistry {
     expression[ValidateUTF8]("validate_utf8"),
     expression[TryValidateUTF8]("try_validate_utf8"),
     expression[Quote]("quote"),
+    expression[Normalize]("normalize"),
 
     // url functions
     expression[UrlEncode]("url_encode"),

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

@@ -961,6 +961,62 @@ case class TryValidateUTF8(input: Expression) extends RuntimeReplaceable with Im
 
 }
 
+/**
+ * A function that returns the Unicode normalization of a string.
+ */
+// scalastyle:off
+@ExpressionDescription(
+  usage = """
+    _FUNC_(str[, form]) - Returns the Unicode normalization of `str` using the normalization `form`.
+      Valid forms are 'NFC' (default), 'NFD', 'NFKC', and 'NFKD'. The form name is case-insensitive.
+  """,
+  arguments = """
+    Arguments:
+      * str - a string expression to normalize.
+      * form - a string expression giving the normalization form: 'NFC', 'NFD', 'NFKC', or 'NFKD'.
+          If omitted, 'NFC' is used.
+  """,
+  examples = """
+    Examples:
+      > SELECT _FUNC_('fi', 'NFKC');
+       fi
+  """,
+  since = "4.3.0",
+  group = "string_funcs")
+// scalastyle:on
+case class Normalize(input: Expression, form: Expression)
+  extends RuntimeReplaceable with ImplicitCastInputTypes with BinaryLike[Expression] {
+  override def nullIntolerant: Boolean = true
+
+  override lazy val replacement: Expression =
+    StaticInvoke(
+      classOf[ExpressionImplUtils],
+      input.dataType,
+      "normalize",
+      Seq(input, form),
+      inputTypes)
+
+  def this(input: Expression) = this(input, Literal("NFC"))
+
+  override def inputTypes: Seq[AbstractDataType] =
+    Seq(StringTypeWithCollation(supportsTrimCollation = true),
+      StringTypeWithCollation(supportsTrimCollation = true))
+
+  override def nodeName: String = "normalize"
+
+  override def nullable: Boolean = true
+
+  override def left: Expression = input
+
+  override def right: Expression = form
+
+  override protected def withNewChildrenInternal(
+      newLeft: Expression, newRight: Expression): Normalize = {
+    copy(input = newLeft, form = newRight)
+  }
+
+}
+
 /**
  * Replace all occurrences with string.
  */