checkSum leads to NPE

Author: scwf

core/src/main/scala/org/apache/spark/broadcast/BroadcastFactory.scala

@@ -48,10 +48,11 @@ private[spark] trait BroadcastFactory {
   def newExecutorBroadcast[T: ClassTag](
       value: T,
       id: Long,
-      nBlocks: Int): Broadcast[T]
+      nBlocks: Int,
+      cSums: Array[Int]): Broadcast[T]
 
   // Called from executor to put broadcast data to blockmanager.
-  def uploadBroadcast[T: ClassTag](value_ : T, id: Long): Int
+  def uploadBroadcast[T: ClassTag](value_ : T, id: Long): Seq[Int]
 
   def unbroadcast(id: Long, removeFromDriver: Boolean, blocking: Boolean): Unit
 

core/src/main/scala/org/apache/spark/broadcast/BroadcastManager.scala

@@ -93,16 +93,17 @@ private[spark] class BroadcastManager(
   def uploadBroadcast[T: ClassTag](
       value_ : T,
       id: Long
-     ): Int = {
+     ): Seq[Int] = {
     broadcastFactory.uploadBroadcast[T](value_, id)
   }
 
   // Called from driver to create broadcast with specified id
   def newExecutorBroadcast[T: ClassTag](
       value_ : T,
       id: Long,
-      nBlocks: Int): Broadcast[T] = {
-    broadcastFactory.newExecutorBroadcast[T](value_, id, nBlocks)
+      nBlocks: Int,
+      cSums: Array[Int]): Broadcast[T] = {
+    broadcastFactory.newExecutorBroadcast[T](value_, id, nBlocks, cSums)
   }
 
   def unbroadcast(id: Long, removeFromDriver: Boolean, blocking: Boolean) {

core/src/main/scala/org/apache/spark/broadcast/TorrentBroadcast.scala

@@ -74,7 +74,8 @@ private[spark] class TorrentBroadcast[T: ClassTag](
     obj: T,
     id: Long,
     isExecutorSide: Boolean = false,
-    nBlocks: Option[Int] = None)
+    nBlocks: Option[Int] = None,
+    cSums: Option[Array[Int]] = None)
   extends Broadcast[T](id) with Logging with Serializable {
 
   /**
@@ -90,6 +91,11 @@ private[spark] class TorrentBroadcast[T: ClassTag](
   /** Size of each block. Default value is 4MB.  This value is only read by the broadcaster. */
   @transient private var blockSize: Int = _
 
+  /** Whether to generate checksum for blocks or not. */
+  private var checksumEnabled: Boolean = false
+  /** The checksum for all the blocks. */
+  private var checksums: Array[Int] = cSums.getOrElse(null)
+
   private def setConf(conf: SparkConf) {
     compressionCodec = if (conf.getBoolean("spark.broadcast.compress", true)) {
       Some(CompressionCodec.createCodec(conf))
@@ -104,15 +110,14 @@ private[spark] class TorrentBroadcast[T: ClassTag](
 
   private val broadcastId = BroadcastBlockId(id)
 
-  def getNumBlocks: Int = numBlocks
+  def getNumBlocksAndChecksums: Seq[Int] = if (checksumEnabled) {
+    Seq(numBlocks) ++ checksums
+  } else {
+    Seq(numBlocks)
+  }
 
   /** Total number of blocks this broadcast variable contains. */
-  private val numBlocks: Int = nBlocks.getOrElse(writeBlocks(obj))
-
-  /** Whether to generate checksum for blocks or not. */
-  private var checksumEnabled: Boolean = false
-  /** The checksum for all the blocks. */
-  private var checksums: Array[Int] = _
+  private val numBlocks: Int = nBlocks.getOrElse(writeBlocks(obj)) // this must be after checkSums
 
   override protected def getValue() = {
     _value

core/src/main/scala/org/apache/spark/broadcast/TorrentBroadcastFactory.scala

@@ -37,12 +37,14 @@ private[spark] class TorrentBroadcastFactory extends BroadcastFactory {
   override def newExecutorBroadcast[T: ClassTag](
       value: T,
       id: Long,
-      nBlocks: Int): Broadcast[T] = {
-    new TorrentBroadcast[T](value, id, true, Option(nBlocks))
+      nBlocks: Int,
+      cSums: Array[Int]): Broadcast[T] = {
+    new TorrentBroadcast[T](value, id, true, Option(nBlocks), Option(cSums))
   }
 
-  override def uploadBroadcast[T: ClassTag](value_ : T, id: Long): Int = {
-    new TorrentBroadcast[T](value_, id, true).getNumBlocks
+  override def uploadBroadcast[T: ClassTag](value_ : T, id: Long): Seq[Int] = {
+    val executorBroadcast = new TorrentBroadcast[T](value_, id, true)
+    executorBroadcast.getNumBlocksAndChecksums
   }
 
   override def stop() { }

core/src/main/scala/org/apache/spark/rdd/RDD.scala

@@ -949,15 +949,17 @@ abstract class RDD[T: ClassTag](
       val id = sc.env.broadcastManager.newBroadcastId
 
       // first: write blocks to block manager from executor.
-      val nBlocks = coalesce(1).mapPartitions { iter =>
-        val numBlocks =
-          SparkEnv.get.broadcastManager.uploadBroadcast(transFunc.transform(iter.toArray), id)
-        Seq(numBlocks).iterator
-      }.collect().head
+      val numBlocksAndChecksums = coalesce(1).mapPartitions { iter =>
+          SparkEnv.get.broadcastManager
+            .uploadBroadcast(transFunc.transform(iter.toArray), id).iterator
+      }.collect()
 
       // then: create broadcast from driver, this will not write blocks
       val res = SparkEnv.get.broadcastManager.newExecutorBroadcast(
-        transFunc.transform(Array.empty[T]), id, nBlocks)
+        transFunc.transform(Array.empty[T]),
+        id,
+        numBlocksAndChecksums.head,
+        numBlocksAndChecksums.tail)
 
       val callSite = sc.getCallSite
       logInfo("Created executor side broadcast " + res.id + " from " + callSite.shortForm)