[SPARK-53154][SQL][PYTHON] Add hmac SQL function

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

@@ -4577,6 +4577,16 @@
           "Invalid extension: <invalidValue>. Extension is limited to exactly 3 letters (e.g. csv, tsv, etc...)"
         ]
       },
+      "HMAC_ALGORITHM" : {
+        "message" : [
+          "expects one of the algorithms 'SHA-224', 'SHA-256', 'SHA-384', 'SHA-512', 'SHA-1', 'MD5', but got <algorithm>."
+        ]
+      },
+      "HMAC_CRYPTO_ERROR" : {
+        "message" : [
+          "detail message: <detailMessage>"
+        ]
+      },
       "INTEGER" : {
         "message" : [
           "expects an integer literal, but got <invalidValue>."

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

@@ -659,6 +659,7 @@ Misc Functions
     current_path
     current_schema
     current_user
+    hmac
     input_file_block_length
     input_file_block_start
     input_file_name

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

@@ -5373,6 +5373,20 @@ def zeroifnull(col: "ColumnOrName") -> Column:
 zeroifnull.__doc__ = pysparkfuncs.zeroifnull.__doc__
 
 
+def hmac(
+    key: "ColumnOrName",
+    message: "ColumnOrName",
+    algorithm: Optional["ColumnOrName"] = None,
+) -> Column:
+    if algorithm is None:
+        return _invoke_function_over_columns("hmac", key, message)
+    else:
+        return _invoke_function_over_columns("hmac", key, message, algorithm)
+
+
+hmac.__doc__ = pysparkfuncs.hmac.__doc__
+
+
 def aes_encrypt(
     input: "ColumnOrName",
     key: "ColumnOrName",

python/pyspark/sql/functions/__init__.py

@@ -519,6 +519,7 @@
     "current_path",
     "current_schema",
     "current_user",
+    "hmac",
     "input_file_block_length",
     "input_file_block_start",
     "input_file_name",

python/pyspark/sql/functions/builtin.py

@@ -30189,6 +30189,65 @@ def zeroifnull(col: "ColumnOrName") -> Column:
     return _invoke_function_over_columns("zeroifnull", col)
 
 
+@_try_remote_functions
+def hmac(
+    key: "ColumnOrName",
+    message: "ColumnOrName",
+    algorithm: Optional["ColumnOrName"] = None,
+) -> Column:
+    """
+    Returns the keyed-hash message authentication code (HMAC) of `message` using `key` and the
+    given hash `algorithm`. The result is returned as raw MAC bytes; wrap it with :func:`hex` or
+    :func:`base64` for a textual value. The default algorithm is 'SHA-256'.
+
+    .. versionadded:: 4.3.0
+
+    Parameters
+    ----------
+    key : :class:`~pyspark.sql.Column` or column name
+        The secret key, as a binary value.
+    message : :class:`~pyspark.sql.Column` or column name
+        The message to authenticate, as a binary value.
+    algorithm : :class:`~pyspark.sql.Column` or column name, optional
+        The hash algorithm. Valid values: SHA-224, SHA-256, SHA-384, SHA-512, SHA-1, MD5.
+        The default is SHA-256.
+
+    Returns
+    -------
+    :class:`~pyspark.sql.Column`
+        A new column that contains the raw HMAC bytes.
+
+    Examples
+    --------
+
+    Example 1: Compute the HMAC with the default SHA-256 algorithm.
+
+    >>> import pyspark.sql.functions as sf
+    >>> df = spark.createDataFrame([("key", "message")], ["key", "message"])
+    >>> df.select(sf.hex(sf.hmac(df.key, df.message))).show(truncate=False)
+    +----------------------------------------------------------------+
+    |hex(hmac(key, message, SHA-256))                                |
+    +----------------------------------------------------------------+
+    |6E9EF29B75FFFC5B7ABAE527D58FDADB2FE42E7219011976917343065F58ED4A|
+    +----------------------------------------------------------------+
+
+    Example 2: Compute the HMAC with an explicit algorithm.
+
+    >>> import pyspark.sql.functions as sf
+    >>> df = spark.createDataFrame([("key", "message")], ["key", "message"])
+    >>> df.select(sf.hex(sf.hmac(df.key, df.message, sf.lit("SHA-1")))).show(truncate=False)
+    +----------------------------------------+
+    |hex(hmac(key, message, SHA-1))          |
+    +----------------------------------------+
+    |2088DF74D5F2146B48146CAF4965377E9D0BE3A4|
+    +----------------------------------------+
+    """
+    if algorithm is None:
+        return _invoke_function_over_columns("hmac", key, message)
+    else:
+        return _invoke_function_over_columns("hmac", key, message, algorithm)
+
+
 @_try_remote_functions
 def aes_encrypt(
     input: "ColumnOrName",

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

@@ -4608,6 +4608,40 @@ object functions {
    */
   def uuid(seed: Column): Column = Column.fn("uuid", seed)
 
+  /**
+   * Returns the keyed-hash message authentication code (HMAC) of `message` using `key` and the
+   * given hash `algorithm`. The result is returned as raw MAC bytes; wrap it with `hex` or
+   * `base64` for a textual value.
+   *
+   * @param key
+   *   The secret key, as a binary value.
+   * @param message
+   *   The message to authenticate, as a binary value.
+   * @param algorithm
+   *   The hash algorithm. Valid values: SHA-224, SHA-256, SHA-384, SHA-512, SHA-1, MD5.
+   *
+   * @group misc_funcs
+   * @since 4.3.0
+   */
+  def hmac(key: Column, message: Column, algorithm: Column): Column =
+    Column.fn("hmac", key, message, algorithm)
+
+  /**
+   * Returns the keyed-hash message authentication code (HMAC) of `message` using `key` and
+   * SHA-256. The result is returned as raw MAC bytes; wrap it with `hex` or `base64` for a
+   * textual value. To use a different algorithm, call the three-argument overload.
+   *
+   * @param key
+   *   The secret key, as a binary value.
+   * @param message
+   *   The message to authenticate, as a binary value.
+   *
+   * @group misc_funcs
+   * @since 4.3.0
+   */
+  def hmac(key: Column, message: Column): Column =
+    Column.fn("hmac", key, message)
+
   /**
    * Returns an encrypted value of `input` using AES in given `mode` with the specified `padding`.
    * Key lengths of 16, 24 and 32 bits are supported. Supported combinations of (`mode`,

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

@@ -30,6 +30,7 @@
 import java.util.regex.PatternSyntaxException;
 import java.util.zip.CRC32;
 import javax.crypto.Cipher;
+import javax.crypto.Mac;
 import javax.crypto.spec.GCMParameterSpec;
 import javax.crypto.spec.IvParameterSpec;
 import javax.crypto.spec.SecretKeySpec;
@@ -181,6 +182,42 @@ public static byte[] aesDecrypt(byte[] input,
     );
   }
 
+  /**
+   * Computes a keyed-hash message authentication code (HMAC) of the given message using the
+   * given key and hash algorithm.
+   * @param key The secret key.
+   * @param message The message to authenticate.
+   * @param algorithm The hash algorithm. Supported values (case-insensitive):
+   *                  SHA-224, SHA-256, SHA-384, SHA-512, SHA-1, MD5.
+   * @return The raw HMAC bytes.
+   */
+  public static byte[] hmac(byte[] key, byte[] message, UTF8String algorithm) {
+    String macName = hmacName(algorithm.toString());
+    try {
+      Mac mac = Mac.getInstance(macName);
+      mac.init(new SecretKeySpec(key, macName));
+      return mac.doFinal(message);
+    } catch (GeneralSecurityException | IllegalArgumentException e) {
+      throw QueryExecutionErrors.hmacCryptoError(e.getMessage());
+    }
+  }
+
+  /**
+   * Maps a user-facing hash algorithm name to its JCA {@link Mac} algorithm name. Supported
+   * algorithms are SHA-224, SHA-256, SHA-384, SHA-512, SHA-1 and MD5 (case-insensitive).
+   */
+  private static String hmacName(String algorithm) {
+    return switch (algorithm.toUpperCase(Locale.ROOT)) {
+      case "SHA-224", "SHA224" -> "HmacSHA224";
+      case "SHA-256", "SHA256" -> "HmacSHA256";
+      case "SHA-384", "SHA384" -> "HmacSHA384";
+      case "SHA-512", "SHA512" -> "HmacSHA512";
+      case "SHA-1", "SHA1" -> "HmacSHA1";
+      case "MD5" -> "HmacMD5";
+      default -> throw QueryExecutionErrors.hmacUnsupportedAlgorithmError(algorithm);
+    };
+  }
+
   /**
    * Function to return the Spark version.
    * @return

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

@@ -851,6 +851,7 @@ object FunctionRegistry {
     expression[Sha2]("sha2"),
     expression[AesEncrypt]("aes_encrypt"),
     expression[AesDecrypt]("aes_decrypt"),
+    expression[Hmac]("hmac"),
     expression[SparkPartitionID]("spark_partition_id"),
     expression[InputFileName]("input_file_name"),
     expression[InputFileBlockStart]("input_file_block_start"),

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

@@ -600,3 +600,56 @@ case class TryAesDecrypt(
     this.copy(replacement = newChild)
 }
 // scalastyle:on line.size.limit
+
+/**
+ * A function that computes a keyed-hash message authentication code (HMAC) of the `message`
+ * using the given `key` and hash `algorithm`. The result is returned as raw MAC bytes, which
+ * can be fed back in as the `key` of a subsequent HMAC to chain derivations, or wrapped with
+ * `hex`/`base64` for a textual representation.
+ */
+// scalastyle:off line.size.limit
+@ExpressionDescription(
+  usage = """
+    _FUNC_(key, message[, algorithm]) - Returns the keyed-hash message authentication code (HMAC) of `message` using `key` and the hash `algorithm`. The result is returned as raw MAC bytes; wrap it with `hex` or `base64` for a textual value. The default algorithm is 'SHA-256'.
+  """,
+  arguments = """
+    Arguments:
+      * key - The secret key, as a binary value.
+      * message - The message to authenticate, as a binary value.
+      * algorithm - Optional hash algorithm. Valid values: SHA-224, SHA-256, SHA-384, SHA-512, SHA-1, MD5. The default is SHA-256.
+  """,
+  examples = """
+    Examples:
+      > SELECT hex(_FUNC_('key', 'message'));
+       6E9EF29B75FFFC5B7ABAE527D58FDADB2FE42E7219011976917343065F58ED4A
+      > SELECT hex(_FUNC_('key', 'message', 'SHA-512'));
+       E477384D7CA229DD1426E64B63EBF2D36EBD6D7E669A6735424E72EA6C01D3F8B56EB39C36D8232F5427999B8D1A3F9CD1128FC69F4D75B434216810FA367E98
+  """,
+  since = "4.3.0",
+  group = "misc_funcs")
+// scalastyle:on line.size.limit
+case class Hmac(key: Expression, message: Expression, algorithm: Expression)
+  extends RuntimeReplaceable with ImplicitCastInputTypes {
+
+  override lazy val replacement: Expression = StaticInvoke(
+    classOf[ExpressionImplUtils],
+    BinaryType,
+    "hmac",
+    Seq(key, message, algorithm),
+    inputTypes)
+
+  def this(key: Expression, message: Expression) =
+    this(key, message, Literal("SHA-256"))
+
+  override def prettyName: String = "hmac"
+
+  override def inputTypes: Seq[AbstractDataType] =
+    Seq(BinaryType, BinaryType, StringTypeWithCollation(supportsTrimCollation = true))
+
+  override def children: Seq[Expression] = Seq(key, message, algorithm)
+
+  override protected def withNewChildrenInternal(
+      newChildren: IndexedSeq[Expression]): Expression = {
+    copy(newChildren(0), newChildren(1), newChildren(2))
+  }
+}

sql/catalyst/src/main/scala/org/apache/spark/sql/errors/QueryExecutionErrors.scala

@@ -2526,6 +2526,24 @@ private[sql] object QueryExecutionErrors extends QueryErrorsBase with ExecutionE
         "functionName" -> toSQLId("aes_encrypt")))
   }
 
+  def hmacUnsupportedAlgorithmError(algorithm: String): RuntimeException = {
+    new SparkRuntimeException(
+      errorClass = "INVALID_PARAMETER_VALUE.HMAC_ALGORITHM",
+      messageParameters = Map(
+        "parameter" -> toSQLId("algorithm"),
+        "functionName" -> toSQLId("hmac"),
+        "algorithm" -> toSQLValue(algorithm, StringType)))
+  }
+
+  def hmacCryptoError(detailMessage: String): RuntimeException = {
+    new SparkRuntimeException(
+      errorClass = "INVALID_PARAMETER_VALUE.HMAC_CRYPTO_ERROR",
+      messageParameters = Map(
+        "parameter" -> toSQLId("key"),
+        "functionName" -> toSQLId("hmac"),
+        "detailMessage" -> detailMessage))
+  }
+
   def hiveTableWithAnsiIntervalsError(
       table: TableIdentifier): SparkUnsupportedOperationException = {
     new SparkUnsupportedOperationException(

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

@@ -403,6 +403,65 @@ class ExpressionImplUtilsSuite extends SparkFunSuite {
     // scalastyle:on nonascii
   }
 
+  // RFC 4231 test case 1 (HMAC-SHA-224/256/384/512) and RFC 2202 test case 1
+  // (HMAC-SHA-1 / HMAC-MD5): key is 0x0b repeated, data is "Hi There".
+  private val hmacKey20 = Array.fill[Byte](20)(0x0b)
+  private val hmacKey16 = Array.fill[Byte](16)(0x0b)
+  private val hmacData = "Hi There".getBytes("UTF-8")
+
+  private def hmacHex(key: Array[Byte], message: Array[Byte], algorithm: String): String = {
+    ExpressionImplUtils.hmac(key, message, UTF8String.fromString(algorithm))
+      .map(b => f"$b%02x").mkString
+  }
+
+  test("Hmac with supported algorithms") {
+    assert(hmacHex(hmacKey20, hmacData, "SHA-224") ==
+      "896fb1128abbdf196832107cd49df33f47b4b1169912ba4f53684b22")
+    assert(hmacHex(hmacKey20, hmacData, "SHA-256") ==
+      "b0344c61d8db38535ca8afceaf0bf12b881dc200c9833da726e9376c2e32cff7")
+    assert(hmacHex(hmacKey20, hmacData, "SHA-384") ==
+      "afd03944d84895626b0825f4ab46907f15f9dadbe4101ec682aa034c7cebc59c" +
+      "faea9ea9076ede7f4af152e8b2fa9cb6")
+    assert(hmacHex(hmacKey20, hmacData, "SHA-512") ==
+      "87aa7cdea5ef619d4ff0b4241a1d6cb02379f4e2ce4ec2787ad0b30545e17cde" +
+      "daa833b7d6b8a702038b274eaea3f4e4be9d914eeb61f1702e696c203a126854")
+    assert(hmacHex(hmacKey20, hmacData, "SHA-1") ==
+      "b617318655057264e28bc0b6fb378c8ef146be00")
+    assert(hmacHex(hmacKey16, hmacData, "MD5") ==
+      "9294727a3638bb1c13f48ef8158bfc9d")
+  }
+
+  test("Hmac algorithm name is case-insensitive and accepts the compact form") {
+    val expected = "b0344c61d8db38535ca8afceaf0bf12b881dc200c9833da726e9376c2e32cff7"
+    assert(hmacHex(hmacKey20, hmacData, "sha-256") == expected)
+    assert(hmacHex(hmacKey20, hmacData, "SHA256") == expected)
+  }
+
+  test("Hmac unsupported algorithm error") {
+    checkError(
+      exception = intercept[SparkRuntimeException] {
+        ExpressionImplUtils.hmac(hmacKey20, hmacData, UTF8String.fromString("SHA-3"))
+      },
+      condition = "INVALID_PARAMETER_VALUE.HMAC_ALGORITHM",
+      parameters = Map(
+        "parameter" -> "`algorithm`",
+        "functionName" -> "`hmac`",
+        "algorithm" -> "'SHA-3'"))
+  }
+
+  test("Hmac empty key error") {
+    checkError(
+      exception = intercept[SparkRuntimeException] {
+        ExpressionImplUtils.hmac(Array.empty[Byte], hmacData, UTF8String.fromString("SHA-256"))
+      },
+      condition = "INVALID_PARAMETER_VALUE.HMAC_CRYPTO_ERROR",
+      parameters = Map(
+        "parameter" -> "`key`",
+        "functionName" -> "`hmac`",
+        "detailMessage" -> ".*"),
+      matchPVals = true)
+  }
+
   test("TryValidate UTF8 string") {
     def tryValidateUTF8(str: UTF8String, expected: UTF8String): Unit = {
       assert(ExpressionImplUtils.tryValidateUTF8String(str) == expected)