[SPARK-59403][CORE] Fix remote fetch of empty deserialized RDD blocks

core/src/main/scala/org/apache/spark/storage/BlockManager.scala

@@ -1408,14 +1408,25 @@ private[spark] class BlockManager(
     var totalFailureCount = 0
     val locations = sortLocations(locationsAndStatus.locations)
     val maxFetchFailures = locations.size
+    // For deserialized blocks, memSize is the estimated size of the JVM objects rather than
+    // the size of their serialized representation. An empty iterator can occupy memory while
+    // serializing to an empty buffer, so this size cannot be used to reject an empty response.
+    val isDeserializedMemoryBlock =
+      locationsAndStatus.status.storageLevel.useMemory &&
+        locationsAndStatus.status.storageLevel.deserialized
     var locationIterator = locations.iterator
     while (locationIterator.hasNext) {
       val loc = locationIterator.next()
       logDebug(s"Getting remote block $blockId from $loc")
       val data = try {
         val buf = blockTransferService.fetchBlockSync(loc.host, loc.port, loc.executorId,
           blockId.toString, tempFileManager)
-        if (blockSize > 0 && buf.size() == 0) {
+        // A different location can have a different storage level, including the external shuffle
+        // service which only serves disk files. Accept an empty response only from the location
+        // whose status reports a deserialized memory block.
+        val canServeEmptyDeserializedMemoryBlock =
+          isDeserializedMemoryBlock && loc == locationsAndStatus.statusLocation
+        if (!canServeEmptyDeserializedMemoryBlock && blockSize > 0 && buf.size() == 0) {
           throw SparkException.internalError("Empty buffer received for non empty block " +
             s"when fetching remote block $blockId from $loc", category = "STORAGE")
         }

core/src/main/scala/org/apache/spark/storage/BlockManagerMasterEndpoint.scala

@@ -1056,9 +1056,10 @@ class BlockManagerMasterEndpoint(
           BlockLocationsAndStatus(
             allLocations,
             blockStatus,
+            bmId,
             Option(executorIdToLocalDirs.getIfPresent(bmId.executorId)))
         } else {
-          BlockLocationsAndStatus(allLocations, blockStatus, None)
+          BlockLocationsAndStatus(allLocations, blockStatus, bmId, None)
         }
       }
       .orElse(None)

core/src/main/scala/org/apache/spark/storage/BlockManagerMessages.scala

@@ -133,13 +133,15 @@ private[spark] object BlockManagerMessages {
   /**
    * The response message of `GetLocationsAndStatus` request.
    *
+   * @param statusLocation the block manager location which owns `status`.
    * @param localDirs if it is persisted-to-disk on the same host as the requester executor is
    *                  running on then localDirs will be Some and the cached data will be in a file
    *                  in one of those dirs, otherwise it is None.
    */
   case class BlockLocationsAndStatus(
       locations: Seq[BlockManagerId],
       status: BlockStatus,
+      statusLocation: BlockManagerId,
       localDirs: Option[Array[String]]) {
     assert(locations.nonEmpty)
   }

core/src/test/scala/org/apache/spark/storage/BlockManagerSuite.scala

@@ -1832,7 +1832,8 @@ class BlockManagerSuite extends SparkFunSuite with Matchers with PrivateMethodTe
     val blockManagerIds = (0 to maxFailuresBeforeLocationRefresh)
       .map { i => BlockManagerId(s"id-$i", s"host-$i", i + 1) }
     when(mockBlockManagerMaster.getLocationsAndStatus(mc.any[BlockId], mc.any[String])).thenReturn(
-      Option(BlockLocationsAndStatus(blockManagerIds, BlockStatus.empty, None)))
+      Option(BlockLocationsAndStatus(
+        blockManagerIds, BlockStatus.empty, blockManagerIds.head, None)))
     when(mockBlockManagerMaster.getLocations(mc.any[BlockId])).thenReturn(
       blockManagerIds)
 
@@ -1905,6 +1906,104 @@ class BlockManagerSuite extends SparkFunSuite with Matchers with PrivateMethodTe
     assert(store2.getRemoteBytes("item").isEmpty)
   }
 
+  test("allow empty executor response for deserialized memory block") {
+    val mockBlockManagerMaster = mock(classOf[BlockManagerMaster])
+    val executor = BlockManagerId("executor1", "localhost", 8000)
+    val deserializedMemoryStatus =
+      BlockStatus(StorageLevel.MEMORY_ONLY, memSize = 1L, diskSize = 0L)
+    when(mockBlockManagerMaster.getLocationsAndStatus(
+      mc.any[BlockId], mc.any[String])).thenReturn(
+      Option(BlockLocationsAndStatus(
+        Seq(executor), deserializedMemoryStatus, executor, None)))
+    val blockFetcher = new MockBlockTransferService(0) {
+      override def fetchBlockSync(
+          host: String,
+          port: Int,
+          execId: String,
+          blockId: String,
+          tempFileManager: DownloadFileManager): ManagedBuffer = {
+        new NioManagedBuffer(ByteBuffer.allocate(0))
+      }
+    }
+    val store = makeBlockManager(
+      8000,
+      "executor2",
+      mockBlockManagerMaster,
+      transferService = Some(blockFetcher))
+
+    val remoteBytes = store.getRemoteBytes(rdd(0, 0))
+    assert(remoteBytes.isDefined)
+    assert(remoteBytes.get.size === 0)
+  }
+
+  test("reject empty shuffle-service response for deserialized memory block") {
+    val mockBlockManagerMaster = mock(classOf[BlockManagerMaster])
+    val externalShuffleServicePort = StorageUtils.externalShuffleServicePort(conf)
+    val executor = BlockManagerId("executor1", "localhost", 8000)
+    val shuffleService = BlockManagerId("executor2", "localhost", externalShuffleServicePort)
+    val deserializedMemoryStatus =
+      BlockStatus(StorageLevel.MEMORY_ONLY, memSize = 1L, diskSize = 0L)
+    when(mockBlockManagerMaster.getLocationsAndStatus(
+      mc.any[BlockId], mc.any[String])).thenReturn(
+      Option(BlockLocationsAndStatus(
+        Seq(executor, shuffleService), deserializedMemoryStatus, executor, None)))
+    val blockFetcher = new MockBlockTransferService(0) {
+      override def fetchBlockSync(
+          host: String,
+          port: Int,
+          execId: String,
+          blockId: String,
+          tempFileManager: DownloadFileManager): ManagedBuffer = {
+        if (port == externalShuffleServicePort) {
+          val transConf = SparkTransportConf.fromSparkConf(conf, "shuffle", numUsableCores = 1)
+          new FileSegmentManagedBuffer(transConf, new File("missing.file"), 0, 0)
+        } else {
+          throw new RuntimeException("simulated executor fetch failure")
+        }
+      }
+    }
+    val store = makeBlockManager(
+      8000,
+      "executor3",
+      mockBlockManagerMaster,
+      transferService = Some(blockFetcher))
+
+    assert(store.getRemoteBytes("item").isEmpty)
+  }
+
+  test("reject empty executor response when status belongs to another location") {
+    val mockBlockManagerMaster = mock(classOf[BlockManagerMaster])
+    val staleExecutor = BlockManagerId("executor1", "MockBlockTransferServiceHost", 8000)
+    val statusExecutor = BlockManagerId("executor2", "other-host", 8000)
+    val deserializedMemoryStatus =
+      BlockStatus(StorageLevel.MEMORY_ONLY, memSize = 1L, diskSize = 0L)
+    when(mockBlockManagerMaster.getLocationsAndStatus(
+      mc.any[BlockId], mc.any[String])).thenReturn(
+      Option(BlockLocationsAndStatus(
+        Seq(staleExecutor, statusExecutor), deserializedMemoryStatus, statusExecutor, None)))
+    val blockFetcher = new MockBlockTransferService(0) {
+      override def fetchBlockSync(
+          host: String,
+          port: Int,
+          execId: String,
+          blockId: String,
+          tempFileManager: DownloadFileManager): ManagedBuffer = {
+        if (execId == staleExecutor.executorId) {
+          new NioManagedBuffer(ByteBuffer.allocate(0))
+        } else {
+          throw new RuntimeException("simulated status executor fetch failure")
+        }
+      }
+    }
+    val store = makeBlockManager(
+      8000,
+      "executor3",
+      mockBlockManagerMaster,
+      transferService = Some(blockFetcher))
+
+    assert(store.getRemoteBytes("item").isEmpty)
+  }
+
   test("test sorting of block locations") {
     val localHost = "localhost"
     val otherHost = "otherHost"
@@ -2022,7 +2121,7 @@ class BlockManagerSuite extends SparkFunSuite with Matchers with PrivateMethodTe
     val blockStatus = BlockStatus(StorageLevel.DISK_ONLY, 0L, 2000L)
 
     when(mockBlockManagerMaster.getLocationsAndStatus(mc.any[BlockId], mc.any[String])).thenReturn(
-      Option(BlockLocationsAndStatus(blockLocations, blockStatus, None)))
+      Option(BlockLocationsAndStatus(blockLocations, blockStatus, blockLocations.head, None)))
     when(mockBlockManagerMaster.getLocations(mc.any[BlockId])).thenReturn(blockLocations)
 
     val store = makeBlockManager(8000, "executor1", mockBlockManagerMaster,