[SPARK-12567] [SQL] Add aes_{encrypt,decrypt} UDFs

Author: Kai Jiang <jiangkai@gmail.com>

Closes #10527 from vectorijk/spark-12567.

2 files changed, 60 insertions, 0 deletions

diff --git a/sql/core/src/main/scala/org/apache/spark/sql/functions.scala b/sql/core/src/main/scala/org/apache/spark/sql/functions.scala
index 97c6992e18..8da50bedfc 100644
--- a/sql/core/src/main/scala/org/apache/spark/sql/functions.scala
+++ b/sql/core/src/main/scala/org/apache/spark/sql/functions.scala
@@ -1982,6 +1982,42 @@ object functions extends LegacyFunctions {
     new Murmur3Hash(cols.map(_.expr))
   }
 
+  /**
+   * Encrypts input using AES and Returns the result as a binary column.
+   * Key lengths of 128, 192 or 256 bits can be used. 192 and 256 bits keys can be used if Java
+   * Cryptography Extension (JCE) Unlimited Strength Jurisdiction Policy Files are installed. If
+   * either argument is NULL, the result will also be null. If input is invalid, key length is not
+   * one of the permitted values or using 192/256 bits key before installing JCE, an exception will
+   * be thrown.
+   *
+   * @param input binary column to encrypt input
+   * @param key binary column of 128, 192 or 256 bits key
+   *
+   * @group misc_funcs
+   * @since 2.0.0
+   */
+  def aes_encrypt(input: Column, key: Column): Column = withExpr {
+    AesEncrypt(input.expr, key.expr)
+  }
+
+  /**
+   * Decrypts input using AES and Returns the result as a string column.
+   * Key lengths of 128, 192 or 256 bits can be used. 192 and 256 bits keys can be used if Java
+   * Cryptography Extension (JCE) Unlimited Strength Jurisdiction Policy Files are installed. If
+   * either argument is NULL, the result will also be null. If input is invalid, key length is not
+   * one of the permitted values or using 192/256 bits key before installing JCE, an exception will
+   * be thrown.
+   *
+   * @param input binary column to decrypt input
+   * @param key binary column of 128, 192 or 256 bits key
+   *
+   * @group misc_funcs
+   * @since 2.0.0
+   */
+  def aes_decrypt(input: Column, key: Column): Column = withExpr {
+    AesDecrypt(input.expr, key.expr)
+  }
+
   //////////////////////////////////////////////////////////////////////////////////////////////
   // String functions
   //////////////////////////////////////////////////////////////////////////////////////////////
diff --git a/sql/core/src/test/scala/org/apache/spark/sql/DataFrameFunctionsSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/DataFrameFunctionsSuite.scala
index aff9efe4b2..0381d57280 100644
--- a/sql/core/src/test/scala/org/apache/spark/sql/DataFrameFunctionsSuite.scala
+++ b/sql/core/src/test/scala/org/apache/spark/sql/DataFrameFunctionsSuite.scala
@@ -206,6 +206,30 @@ class DataFrameFunctionsSuite extends QueryTest with SharedSQLContext {
       Row(2743272264L, 2180413220L))
   }
 
+  test("misc aes encrypt function") {
+    val df = Seq(("ABC", "1234567890123456")).toDF("input", "key")
+    checkAnswer(
+      df.select(base64(aes_encrypt($"input", $"key"))),
+      Row("y6Ss+zCYObpCbgfWfyNWTw==")
+    )
+    checkAnswer(
+      sql("SELECT base64(aes_encrypt('', '1234567890123456'))"),
+      Row("BQGHoM3lqYcsurCRq3PlUw==")
+    )
+  }
+
+  test("misc aes decrypt function") {
+    val df = Seq(("y6Ss+zCYObpCbgfWfyNWTw==", "1234567890123456")).toDF("input", "key")
+    checkAnswer(
+      df.select((aes_decrypt(unbase64($"input"), $"key"))),
+      Row("ABC")
+    )
+    checkAnswer(
+      sql("SELECT aes_decrypt(unbase64('BQGHoM3lqYcsurCRq3PlUw=='), '1234567890123456')"),
+      Row("")
+    )
+  }
+
   test("string function find_in_set") {
     val df = Seq(("abc,b,ab,c,def", "abc,b,ab,c,def")).toDF("a", "b")