Github user cloud-fan commented on a diff in the pull request:
https://github.com/apache/spark/pull/16989#discussion_r114954964
--- Diff:
core/src/test/scala/org/apache/spark/storage/ShuffleBlockFetcherIteratorSuite.scala
---
@@ -401,4 +424,74 @@ class ShuffleBlockFetcherIteratorSuite extends
SparkFunSuite with PrivateMethodT
assert(id3 === ShuffleBlockId(0, 2, 0))
}
+ test("Blocks should be shuffled to disk when size is above the
threshold, otherwise to memory.") {
+ val blockManager = mock(classOf[BlockManager])
+ val localBmId = BlockManagerId("test-client", "test-client", 1)
+ doReturn(localBmId).when(blockManager).blockManagerId
+
+ val diskBlockManager = mock(classOf[DiskBlockManager])
+ doReturn(new
File("shuffle-read-file")).when(diskBlockManager).getFile(any(classOf[String]))
+ doReturn(diskBlockManager).when(blockManager).diskBlockManager
+
+ val remoteBmId = BlockManagerId("test-client-1", "test-client-1", 2)
+ val remoteBlocks = Map[BlockId, ManagedBuffer](
+ ShuffleBlockId(0, 0, 0) -> createMockManagedBuffer(),
+ ShuffleBlockId(0, 1, 0) -> createMockManagedBuffer(),
+ ShuffleBlockId(0, 2, 0) -> createMockManagedBuffer())
+ val transfer = mock(classOf[BlockTransferService])
+ var shuffleFilesOpt: Option[Array[File]] = None
+ when(transfer.fetchBlocks(any(), any(), any(), any(), any(), any()))
+ .thenAnswer(new Answer[Unit] {
+ override def answer(invocation: InvocationOnMock): Unit = {
+ val listener =
invocation.getArguments()(4).asInstanceOf[BlockFetchingListener]
+ shuffleFilesOpt =
invocation.getArguments()(5).asInstanceOf[Option[Array[File]]]
+ Future {
+ listener.onBlockFetchSuccess(
+ ShuffleBlockId(0, 0, 0).toString,
remoteBlocks(ShuffleBlockId(0, 0, 0)))
+ listener.onBlockFetchSuccess(
+ ShuffleBlockId(0, 1, 0).toString,
remoteBlocks(ShuffleBlockId(0, 1, 0)))
+ listener.onBlockFetchSuccess(
+ ShuffleBlockId(0, 2, 0).toString,
remoteBlocks(ShuffleBlockId(0, 2, 0)))
+ }
+ }
+ })
+ val taskMemoryManager = mock(classOf[TaskMemoryManager])
+ when(taskMemoryManager.acquireExecutionMemory(any(), any()))
+ .thenAnswer(new Answer[Long] {
+ override def answer(invocationOnMock: InvocationOnMock): Long =
500L
+ })
+
+ val tc = new TaskContextImpl(0, 0, 0, 0, taskMemoryManager, new
Properties, null)
+
+ val blocksByAddress1 = Seq[(BlockManagerId, Seq[(BlockId, Long)])](
+ (remoteBmId, remoteBlocks.keys.map(blockId => (blockId,
100.asInstanceOf[Long])).toSeq))
+
+ val iterator1 = new ShuffleBlockFetcherIterator(
+ tc,
+ transfer,
+ blockManager,
+ blocksByAddress1,
+ (_, in) => in,
+ 48 * 1024 * 1024,
+ Int.MaxValue,
+ true,
+ taskMemoryManager)
+
+ assert(shuffleFilesOpt.isEmpty)
+ val blocksByAddress2 = Seq[(BlockManagerId, Seq[(BlockId, Long)])](
+ (remoteBmId, remoteBlocks.keys.map(blockId => (blockId,
200.asInstanceOf[Long])).toSeq)
--- End diff --
same here.
---
If your project is set up for it, you can reply to this email and have your
reply appear on GitHub as well. If your project does not have this feature
enabled and wishes so, or if the feature is enabled but not working, please
contact infrastructure at [email protected] or file a JIRA ticket
with INFRA.
---
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
