beliefer commented on code in PR #41689:
URL: https://github.com/apache/spark/pull/41689#discussion_r1239309130
##########
python/pyspark/sql/functions.py:
##########
@@ -14394,6 +14394,403 @@ def nvl2(col1: "ColumnOrName", col2: "ColumnOrName",
col3: "ColumnOrName") -> Co
return _invoke_function_over_columns("nvl2", col1, col2, col3)
+@try_remote_functions
+def uuid() -> Column:
+ """
+ Returns an universally unique identifier (UUID) string. The value is
returned as a canonical
+ UUID 36-character string.
+
+ .. versionadded:: 3.5.0
+
+ Examples
+ --------
+ >>> df = spark.range(1)
+ >>> df.select(uuid()).show(truncate=False) # doctest: +SKIP
+ +------------------------------------+
+ |uuid() |
+ +------------------------------------+
+ |3dcc5174-9da9-41ca-815f-34c05c6d3926|
+ +------------------------------------+
+ """
+ return _invoke_function_over_columns("uuid")
+
+
+@try_remote_functions
+def aes_encrypt(
+ input: "ColumnOrName",
+ key: "ColumnOrName",
+ mode: Optional["ColumnOrName"] = None,
+ padding: Optional["ColumnOrName"] = None,
+ iv: Optional["ColumnOrName"] = None,
+ aad: Optional["ColumnOrName"] = None,
+) -> Column:
+ """
+ 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`,
+ `padding`) are ('ECB', 'PKCS'), ('GCM', 'NONE') and ('CBC', 'PKCS').
Optional initialization
+ vectors (IVs) are only supported for CBC and GCM modes. These must be 16
bytes for CBC and 12
+ bytes for GCM. If not provided, a random vector will be generated and
prepended to the
+ output. Optional additional authenticated data (AAD) is only supported for
GCM. If provided
+ for encryption, the identical AAD value must be provided for decryption.
The default mode is
+ GCM.
+
+ .. versionadded:: 3.5.0
+
+ Parameters
+ ----------
+ input : :class:`~pyspark.sql.Column` or str
+ The binary value to encrypt.
+ key : :class:`~pyspark.sql.Column` or str
+ The passphrase to use to encrypt the data.
+ mode : :class:`~pyspark.sql.Column` or str, optional
+ Specifies which block cipher mode should be used to encrypt messages.
Valid modes: ECB,
+ GCM, CBC.
+ padding : :class:`~pyspark.sql.Column` or str, optional
+ Specifies how to pad messages whose length is not a multiple of the
block size. Valid
+ values: PKCS, NONE, DEFAULT. The DEFAULT padding means PKCS for ECB,
NONE for GCM and PKCS
+ for CBC.
+ iv : :class:`~pyspark.sql.Column` or str, optional
+ Optional initialization vector. Only supported for CBC and GCM modes.
Valid values: None or
+ "". 16-byte array for CBC mode. 12-byte array for GCM mode.
+ aad : :class:`~pyspark.sql.Column` or str, optional
+ Optional additional authenticated data. Only supported for GCM mode.
This can be any
+ free-form input and must be provided for both encryption and
decryption.
+
+ Examples
+ --------
+ >>> df = spark.createDataFrame([(
+ ... "Spark", "abcdefghijklmnop12345678ABCDEFGH", "GCM", "DEFAULT",
+ ... "000000000000000000000000", "This is an AAD mixed into the
input",)],
+ ... ["input", "key", "mode", "padding", "iv", "aad"]
+ ... )
+ >>> df.select(base64(aes_encrypt(
+ ... df.input, df.key, df.mode, df.padding, to_binary(df.iv,
lit("hex")), df.aad)
+ ... ).alias('r')).collect()
+ [Row(r='AAAAAAAAAAAAAAAAQiYi+sTLm7KD9UcZ2nlRdYDe/PX4')]
+
+ >>> df.select(base64(aes_encrypt(
+ ... df.input, df.key, df.mode, df.padding, to_binary(df.iv,
lit("hex")))
+ ... ).alias('r')).collect()
+ [Row(r='AAAAAAAAAAAAAAAAQiYi+sRNYDAOTjdSEcYBFsAWPL1f')]
+
+ >>> df = spark.createDataFrame([(
+ ... "Spark SQL", "1234567890abcdef", "ECB", "PKCS",)],
+ ... ["input", "key", "mode", "padding"]
+ ... )
+ >>> df.select(aes_decrypt(aes_encrypt(df.input, df.key, df.mode,
df.padding),
Review Comment:
Personally, I think we only need display the usage for `aes_encrypt` here.
cc @zhengruifeng
##########
sql/core/src/test/scala/org/apache/spark/sql/MiscFunctionsSuite.scala:
##########
@@ -84,6 +88,152 @@ class MiscFunctionsSuite extends QueryTest with
SharedSparkSession {
}
}
}
+
+ test("uuid") {
+ val df = Seq((1, 2)).toDF("a", "b")
+ assert(df.selectExpr("uuid()").collect() != null)
+ assert(df.select(uuid()).collect() != null)
+ }
+
+ test("aes_encrypt") {
+ val iv = Hex.unhex("000000000000000000000000".getBytes())
+ val df = Seq(("Spark", "abcdefghijklmnop12345678ABCDEFGH",
+ "GCM", "DEFAULT", iv, "This is an AAD mixed into the input")).
+ toDF("input", "key", "mode", "padding", "iv", "aad")
+
+ checkAnswer(
+ df.selectExpr("aes_encrypt(input, key, mode, padding, iv, aad)"),
+ df.select(aes_encrypt(col("input"), col("key"), col("mode"),
+ col("padding"), col("iv"), col("aad"))))
+
+ checkAnswer(
+ df.selectExpr("aes_encrypt(input, key, mode, padding, iv)"),
+ df.select(aes_encrypt(col("input"), col("key"), col("mode"),
+ col("padding"), col("iv"))))
+
+ val df1 = Seq(("Spark SQL", "1234567890abcdef", "ECB", "PKCS")).
+ toDF("input", "key", "mode", "padding")
+
+ checkAnswer(
+ df1.selectExpr("base64(aes_encrypt(input, key, mode, padding))"),
+ df1.select(base64(aes_encrypt(col("input"), col("key"), col("mode"),
col("padding")))))
+
+ val df2 = Seq(("Spark SQL", "0000111122223333", "ECB")).toDF("input",
"key", "mode")
+
+ checkAnswer(
+ df2.selectExpr("hex(aes_encrypt(input, key, mode))"),
+ df2.select(hex(aes_encrypt(col("input"), col("key"), col("mode")))))
+
+ val df3 = Seq(("Spark", "abcdefghijklmnop")).toDF("input", "key")
+ checkAnswer(
+ df3.selectExpr("cast(aes_decrypt(unbase64(base64(" +
+ "aes_encrypt(input, key))), key) AS STRING)"),
+ Seq(Row("Spark")))
+ checkAnswer(
+ df3.select(aes_decrypt(unbase64(base64(
+ aes_encrypt(col("input"), col("key")))), col("key")).cast(StringType)),
+ Seq(Row("Spark")))
+ }
+
+ test("aes_decrypt") {
Review Comment:
Please merge the test cases for `aes_encrypt` and `aes_decrypt`.
The new test code could compare the origin input with
`aes_decrypt(aes_encrypt(input))` and simplify the code.
##########
python/pyspark/sql/functions.py:
##########
@@ -14394,6 +14394,403 @@ def nvl2(col1: "ColumnOrName", col2: "ColumnOrName",
col3: "ColumnOrName") -> Co
return _invoke_function_over_columns("nvl2", col1, col2, col3)
+@try_remote_functions
+def uuid() -> Column:
+ """
+ Returns an universally unique identifier (UUID) string. The value is
returned as a canonical
+ UUID 36-character string.
+
+ .. versionadded:: 3.5.0
+
+ Examples
+ --------
+ >>> df = spark.range(1)
+ >>> df.select(uuid()).show(truncate=False) # doctest: +SKIP
+ +------------------------------------+
+ |uuid() |
+ +------------------------------------+
+ |3dcc5174-9da9-41ca-815f-34c05c6d3926|
+ +------------------------------------+
+ """
+ return _invoke_function_over_columns("uuid")
+
+
+@try_remote_functions
+def aes_encrypt(
+ input: "ColumnOrName",
+ key: "ColumnOrName",
+ mode: Optional["ColumnOrName"] = None,
+ padding: Optional["ColumnOrName"] = None,
+ iv: Optional["ColumnOrName"] = None,
+ aad: Optional["ColumnOrName"] = None,
+) -> Column:
+ """
+ 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`,
+ `padding`) are ('ECB', 'PKCS'), ('GCM', 'NONE') and ('CBC', 'PKCS').
Optional initialization
+ vectors (IVs) are only supported for CBC and GCM modes. These must be 16
bytes for CBC and 12
+ bytes for GCM. If not provided, a random vector will be generated and
prepended to the
+ output. Optional additional authenticated data (AAD) is only supported for
GCM. If provided
+ for encryption, the identical AAD value must be provided for decryption.
The default mode is
+ GCM.
+
+ .. versionadded:: 3.5.0
+
+ Parameters
+ ----------
+ input : :class:`~pyspark.sql.Column` or str
+ The binary value to encrypt.
+ key : :class:`~pyspark.sql.Column` or str
+ The passphrase to use to encrypt the data.
+ mode : :class:`~pyspark.sql.Column` or str, optional
+ Specifies which block cipher mode should be used to encrypt messages.
Valid modes: ECB,
+ GCM, CBC.
+ padding : :class:`~pyspark.sql.Column` or str, optional
+ Specifies how to pad messages whose length is not a multiple of the
block size. Valid
+ values: PKCS, NONE, DEFAULT. The DEFAULT padding means PKCS for ECB,
NONE for GCM and PKCS
+ for CBC.
+ iv : :class:`~pyspark.sql.Column` or str, optional
+ Optional initialization vector. Only supported for CBC and GCM modes.
Valid values: None or
+ "". 16-byte array for CBC mode. 12-byte array for GCM mode.
+ aad : :class:`~pyspark.sql.Column` or str, optional
+ Optional additional authenticated data. Only supported for GCM mode.
This can be any
+ free-form input and must be provided for both encryption and
decryption.
+
+ Examples
+ --------
+ >>> df = spark.createDataFrame([(
+ ... "Spark", "abcdefghijklmnop12345678ABCDEFGH", "GCM", "DEFAULT",
+ ... "000000000000000000000000", "This is an AAD mixed into the
input",)],
+ ... ["input", "key", "mode", "padding", "iv", "aad"]
+ ... )
+ >>> df.select(base64(aes_encrypt(
+ ... df.input, df.key, df.mode, df.padding, to_binary(df.iv,
lit("hex")), df.aad)
+ ... ).alias('r')).collect()
+ [Row(r='AAAAAAAAAAAAAAAAQiYi+sTLm7KD9UcZ2nlRdYDe/PX4')]
+
+ >>> df.select(base64(aes_encrypt(
+ ... df.input, df.key, df.mode, df.padding, to_binary(df.iv,
lit("hex")))
+ ... ).alias('r')).collect()
+ [Row(r='AAAAAAAAAAAAAAAAQiYi+sRNYDAOTjdSEcYBFsAWPL1f')]
+
+ >>> df = spark.createDataFrame([(
+ ... "Spark SQL", "1234567890abcdef", "ECB", "PKCS",)],
+ ... ["input", "key", "mode", "padding"]
+ ... )
+ >>> df.select(aes_decrypt(aes_encrypt(df.input, df.key, df.mode,
df.padding),
+ ... df.key, df.mode, df.padding).alias('r')
+ ... ).collect()
+ [Row(r=bytearray(b'Spark SQL'))]
+
+ >>> df = spark.createDataFrame([(
+ ... "Spark SQL", "0000111122223333", "ECB",)],
+ ... ["input", "key", "mode"]
+ ... )
+ >>> df.select(aes_decrypt(aes_encrypt(df.input, df.key, df.mode),
+ ... df.key, df.mode).alias('r')
+ ... ).collect()
+ [Row(r=bytearray(b'Spark SQL'))]
+
+ >>> df = spark.createDataFrame([(
+ ... "Spark SQL", "abcdefghijklmnop",)],
+ ... ["input", "key"]
+ ... )
+ >>> df.select(aes_decrypt(
+ ... unbase64(base64(aes_encrypt(df.input, df.key))), df.key
+ ... ).cast("STRING").alias('r')).collect()
+ [Row(r='Spark SQL')]
+ """
+ _mode = lit("GCM") if mode is None else mode
+ _padding = lit("DEFAULT") if padding is None else padding
+ _iv = lit("") if iv is None else iv
+ _aad = lit("") if aad is None else aad
+ return _invoke_function_over_columns("aes_encrypt", input, key, _mode,
_padding, _iv, _aad)
Review Comment:
If we passed `padding` and without `mode`, what happens here ?
##########
python/pyspark/sql/functions.py:
##########
@@ -14394,6 +14394,403 @@ def nvl2(col1: "ColumnOrName", col2: "ColumnOrName",
col3: "ColumnOrName") -> Co
return _invoke_function_over_columns("nvl2", col1, col2, col3)
+@try_remote_functions
+def uuid() -> Column:
+ """
+ Returns an universally unique identifier (UUID) string. The value is
returned as a canonical
+ UUID 36-character string.
+
+ .. versionadded:: 3.5.0
+
+ Examples
+ --------
+ >>> df = spark.range(1)
+ >>> df.select(uuid()).show(truncate=False) # doctest: +SKIP
+ +------------------------------------+
+ |uuid() |
+ +------------------------------------+
+ |3dcc5174-9da9-41ca-815f-34c05c6d3926|
+ +------------------------------------+
+ """
+ return _invoke_function_over_columns("uuid")
+
+
+@try_remote_functions
+def aes_encrypt(
+ input: "ColumnOrName",
+ key: "ColumnOrName",
+ mode: Optional["ColumnOrName"] = None,
+ padding: Optional["ColumnOrName"] = None,
+ iv: Optional["ColumnOrName"] = None,
+ aad: Optional["ColumnOrName"] = None,
+) -> Column:
+ """
+ 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`,
+ `padding`) are ('ECB', 'PKCS'), ('GCM', 'NONE') and ('CBC', 'PKCS').
Optional initialization
+ vectors (IVs) are only supported for CBC and GCM modes. These must be 16
bytes for CBC and 12
+ bytes for GCM. If not provided, a random vector will be generated and
prepended to the
+ output. Optional additional authenticated data (AAD) is only supported for
GCM. If provided
+ for encryption, the identical AAD value must be provided for decryption.
The default mode is
+ GCM.
+
+ .. versionadded:: 3.5.0
+
+ Parameters
+ ----------
+ input : :class:`~pyspark.sql.Column` or str
+ The binary value to encrypt.
+ key : :class:`~pyspark.sql.Column` or str
+ The passphrase to use to encrypt the data.
+ mode : :class:`~pyspark.sql.Column` or str, optional
+ Specifies which block cipher mode should be used to encrypt messages.
Valid modes: ECB,
+ GCM, CBC.
+ padding : :class:`~pyspark.sql.Column` or str, optional
+ Specifies how to pad messages whose length is not a multiple of the
block size. Valid
+ values: PKCS, NONE, DEFAULT. The DEFAULT padding means PKCS for ECB,
NONE for GCM and PKCS
+ for CBC.
+ iv : :class:`~pyspark.sql.Column` or str, optional
+ Optional initialization vector. Only supported for CBC and GCM modes.
Valid values: None or
+ "". 16-byte array for CBC mode. 12-byte array for GCM mode.
+ aad : :class:`~pyspark.sql.Column` or str, optional
+ Optional additional authenticated data. Only supported for GCM mode.
This can be any
+ free-form input and must be provided for both encryption and
decryption.
+
+ Examples
+ --------
+ >>> df = spark.createDataFrame([(
+ ... "Spark", "abcdefghijklmnop12345678ABCDEFGH", "GCM", "DEFAULT",
+ ... "000000000000000000000000", "This is an AAD mixed into the
input",)],
+ ... ["input", "key", "mode", "padding", "iv", "aad"]
+ ... )
+ >>> df.select(base64(aes_encrypt(
+ ... df.input, df.key, df.mode, df.padding, to_binary(df.iv,
lit("hex")), df.aad)
+ ... ).alias('r')).collect()
+ [Row(r='AAAAAAAAAAAAAAAAQiYi+sTLm7KD9UcZ2nlRdYDe/PX4')]
+
+ >>> df.select(base64(aes_encrypt(
+ ... df.input, df.key, df.mode, df.padding, to_binary(df.iv,
lit("hex")))
+ ... ).alias('r')).collect()
+ [Row(r='AAAAAAAAAAAAAAAAQiYi+sRNYDAOTjdSEcYBFsAWPL1f')]
+
+ >>> df = spark.createDataFrame([(
+ ... "Spark SQL", "1234567890abcdef", "ECB", "PKCS",)],
+ ... ["input", "key", "mode", "padding"]
+ ... )
+ >>> df.select(aes_decrypt(aes_encrypt(df.input, df.key, df.mode,
df.padding),
+ ... df.key, df.mode, df.padding).alias('r')
+ ... ).collect()
+ [Row(r=bytearray(b'Spark SQL'))]
+
+ >>> df = spark.createDataFrame([(
+ ... "Spark SQL", "0000111122223333", "ECB",)],
+ ... ["input", "key", "mode"]
+ ... )
+ >>> df.select(aes_decrypt(aes_encrypt(df.input, df.key, df.mode),
+ ... df.key, df.mode).alias('r')
+ ... ).collect()
+ [Row(r=bytearray(b'Spark SQL'))]
+
+ >>> df = spark.createDataFrame([(
+ ... "Spark SQL", "abcdefghijklmnop",)],
+ ... ["input", "key"]
+ ... )
+ >>> df.select(aes_decrypt(
+ ... unbase64(base64(aes_encrypt(df.input, df.key))), df.key
+ ... ).cast("STRING").alias('r')).collect()
+ [Row(r='Spark SQL')]
+ """
+ _mode = lit("GCM") if mode is None else mode
+ _padding = lit("DEFAULT") if padding is None else padding
+ _iv = lit("") if iv is None else iv
+ _aad = lit("") if aad is None else aad
+ return _invoke_function_over_columns("aes_encrypt", input, key, _mode,
_padding, _iv, _aad)
+
+
+@try_remote_functions
+def aes_decrypt(
+ input: "ColumnOrName",
+ key: "ColumnOrName",
+ mode: Optional["ColumnOrName"] = None,
+ padding: Optional["ColumnOrName"] = None,
+ aad: Optional["ColumnOrName"] = None,
+) -> Column:
+ """
+ Returns a decrypted value of `input` using AES in `mode` with `padding`.
Key lengths of 16,
+ 24 and 32 bits are supported. Supported combinations of (`mode`,
`padding`) are ('ECB',
+ 'PKCS'), ('GCM', 'NONE') and ('CBC', 'PKCS'). Optional additional
authenticated data (AAD) is
+ only supported for GCM. If provided for encryption, the identical AAD
value must be provided
+ for decryption. The default mode is GCM.
+
+ .. versionadded:: 3.5.0
+
+ Parameters
+ ----------
+ input : :class:`~pyspark.sql.Column` or str
+ The binary value to decrypt.
+ key : :class:`~pyspark.sql.Column` or str
+ The passphrase to use to decrypt the data.
+ mode : :class:`~pyspark.sql.Column` or str, optional
+ Specifies which block cipher mode should be used to decrypt messages.
Valid modes: ECB,
+ GCM, CBC.
+ padding : :class:`~pyspark.sql.Column` or str, optional
+ Specifies how to pad messages whose length is not a multiple of the
block size. Valid
+ values: PKCS, NONE, DEFAULT. The DEFAULT padding means PKCS for ECB,
NONE for GCM and PKCS
+ for CBC.
+ aad : :class:`~pyspark.sql.Column` or str, optional
+ Optional additional authenticated data. Only supported for GCM mode.
This can be any
+ free-form input and must be provided for both encryption and
decryption.
+
+ Examples
+ --------
+ >>> df = spark.createDataFrame([(
+ ... "AAAAAAAAAAAAAAAAQiYi+sTLm7KD9UcZ2nlRdYDe/PX4",
+ ... "abcdefghijklmnop12345678ABCDEFGH", "GCM", "DEFAULT",
+ ... "This is an AAD mixed into the input",)],
+ ... ["input", "key", "mode", "padding", "aad"]
+ ... )
+ >>> df.select(aes_decrypt(
+ ... unbase64(df.input), df.key, df.mode, df.padding, df.aad).alias('r')
+ ... ).collect()
+ [Row(r=bytearray(b'Spark'))]
+
+ >>> df = spark.createDataFrame([(
+ ... "AAAAAAAAAAAAAAAAAAAAAPSd4mWyMZ5mhvjiAPQJnfg=",
+ ... "abcdefghijklmnop12345678ABCDEFGH", "CBC", "DEFAULT",)],
+ ... ["input", "key", "mode", "padding"]
+ ... )
+ >>> df.select(aes_decrypt(
+ ... unbase64(df.input), df.key, df.mode, df.padding).alias('r')
+ ... ).collect()
+ [Row(r=bytearray(b'Spark'))]
+
+ >>> df.select(aes_decrypt(unbase64(df.input), df.key,
df.mode).alias('r')).collect()
+ [Row(r=bytearray(b'Spark'))]
+
+ >>> df = spark.createDataFrame([(
+ ...
"83F16B2AA704794132802D248E6BFD4E380078182D1544813898AC97E709B28A94",
+ ... "0000111122223333",)],
+ ... ["input", "key"]
+ ... )
+ >>> df.select(aes_decrypt(unhex(df.input), df.key).alias('r')).collect()
+ [Row(r=bytearray(b'Spark'))]
+ """
+ _mode = lit("GCM") if mode is None else mode
+ _padding = lit("DEFAULT") if padding is None else padding
+ _aad = lit("") if aad is None else aad
+ return _invoke_function_over_columns("aes_decrypt", input, key, _mode,
_padding, _aad)
Review Comment:
ditto
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
