調用 udf 以觸發數據幀時任務不可序列化錯誤

Question

我有一個Scala的功能進行加密，然后創建一個udf了出來，並把它傳遞給其中一列，我als_embeddings數據幀被添加到我的數據幀新列。

import java.util.Base64
import javax.crypto.Cipher
import javax.crypto.spec.{IvParameterSpec, SecretKeySpec}

val Algorithm = "AES/CBC/PKCS5Padding"
val Key = new SecretKeySpec(Base64.getDecoder.decode("BiwHeIqzQa8X6MXtdg/hhQ=="), "AES")
val IvSpec = new IvParameterSpec(new Array[Byte](16))

def encrypt(text: String): String = {
  val cipher = Cipher.getInstance(Algorithm)
  cipher.init(Cipher.ENCRYPT_MODE, Key, IvSpec)

  new String(Base64.getEncoder.encode(cipher.doFinal(text.getBytes("utf-8"))), "utf-8")
}


val encryptUDF = udf((uid : String) => encrypt(uid))

將encryptUDF傳遞給我的 spark 數據encryptUDF以創建一個帶有加密uid的新列

val als_encrypt_embeddings = als_embeddings.withColumn("encrypt_uid",encryptUDF(col("uid")))
als_encrypt_embeddings.show()

但是當我這樣做時，它給了我以下錯誤：

線程“main”org.apache.spark.SparkException 中的異常：任務不可序列化

我在這里錯過了什么。

Answer 1

錯誤消息Task not serializable是正確的，但不是很清楚。 在堆棧跟蹤的進一步下方，有更詳細的說明出了什么問題：

Exception in thread "main" org.apache.spark.SparkException: Task not serializable
    at org.apache.spark.util.ClosureCleaner$.ensureSerializable(ClosureCleaner.scala:403)
    at org.apache.spark.util.ClosureCleaner$.org$apache$spark$util$ClosureCleaner$$clean(ClosureCleaner.scala:393)
    at org.apache.spark.util.ClosureCleaner$.clean(ClosureCleaner.scala:162)
[...]
Caused by: java.io.NotSerializableException: javax.crypto.spec.IvParameterSpec
Serialization stack:
    - object not serializable (class: javax.crypto.spec.IvParameterSpec, value: javax.crypto.spec.IvParameterSpec@7d4d65f5)
    - field (class: Starter$$anonfun$1, name: IvSpec$1, type: class javax.crypto.spec.IvParameterSpec)
    - object (class Starter$$anonfun$1, <function1>)
    - element of array (index: 2)
    - array (class [Ljava.lang.Object;, size 3)
    - field (class: org.apache.spark.sql.execution.WholeStageCodegenExec$$anonfun$13, name: references$1, type: class [Ljava.lang.Object;)
    - object (class org.apache.spark.sql.execution.WholeStageCodegenExec$$anonfun$13, <function2>)
    at org.apache.spark.serializer.SerializationDebugger$.improveException(SerializationDebugger.scala:40)
    at org.apache.spark.serializer.JavaSerializationStream.writeObject(JavaSerializer.scala:46)
    at org.apache.spark.serializer.JavaSerializerInstance.serialize(JavaSerializer.scala:100)
    at org.apache.spark.util.ClosureCleaner$.ensureSerializable(ClosureCleaner.scala:400)
    ... 48 more

在Caused by部分堆棧跟蹤Caused by Spark 報告中，它無法序列化javax.crypto.spec.IvParameterSpec的實例。

ParameterSpec 已在驅動程序 JVM 中創建，同時 udf 在其中一個執行程序中執行。 因此，必須對對象進行序列化才能將其移動到執行程序的 VM。 由於對象不可序列化，移動它的嘗試失敗。

解決問題的最簡單方法是通過將代碼塊移動到 udf 的閉包中，直接在執行器的 VM 中創建加密所需的對象：

val encryptUDF = udf((uid : String) => {
  val Algorithm = "AES/CBC/PKCS5Padding"
  val Key = new SecretKeySpec(Base64.getDecoder.decode("BiwHeIqzQa8X6MXtdg/hhQ=="), "AES")
  val IvSpec = new IvParameterSpec(new Array[Byte](16))

  def encrypt(text: String): String = {
    val cipher = Cipher.getInstance(Algorithm)
    cipher.init(Cipher.ENCRYPT_MODE, Key, IvSpec)

    new String(Base64.getEncoder.encode(cipher.doFinal(text.getBytes("utf-8"))), "utf-8")
  }
  encrypt(uid)
})

這樣，所有對象都將直接在 executors VM 中創建。

這種方法的缺點是每次調用 udf 都會創建一組加密對象。 如果這些對象的實例化開銷很大，這可能會導致性能問題。 一種選擇是使用mapPartitions而不是udf 。 在這個答案中， mapPartitions 用於避免在迭代數據幀時創建太多昂貴的數據庫連接。 這種方法也可以在這里使用。

Answer 2

我們可以將該函數定義為沒有對不可序列化值的引用的獨立對象的一部分。

object EncryptUtils extends Serializable {
  ...
  ...
  ...
}