Github user JoshRosen commented on a diff in the pull request:
https://github.com/apache/spark/pull/6707#discussion_r52853940
--- Diff:
streaming/src/test/scala/org/apache/spark/streaming/ReceivedBlockHandlerSuite.scala
---
@@ -174,6 +176,130 @@ class ReceivedBlockHandlerSuite
}
}
+ test("Test Block - count messages") {
+ // Test count with BlockManagedBasedBlockHandler
+ testCountWithBlockManagerBasedBlockHandler(true)
+ // Test count with WriteAheadLogBasedBlockHandler
+ testCountWithBlockManagerBasedBlockHandler(false)
+ }
+
+ test("Test Block - isFullyConsumed") {
+ val sparkConf = new SparkConf()
+ sparkConf.set("spark.storage.unrollMemoryThreshold", "512")
+ // spark.storage.unrollFraction set to 0.4 for BlockManager
+ sparkConf.set("spark.storage.unrollFraction", "0.4")
+ // Block Manager with 12000 * 0.4 = 4800 bytes of free space for unroll
+ blockManager = createBlockManager(12000, sparkConf)
+
+ // there is not enough space to store this block in MEMORY,
+ // But BlockManager will be able to sereliaze this block to WAL
+ // and hence count returns correct value.
+ testRecordcount(false, StorageLevel.MEMORY_ONLY,
+ IteratorBlock((List.fill(70)(new Array[Byte](100))).iterator),
blockManager, Some(70))
+
+ // there is not enough space to store this block in MEMORY,
+ // But BlockManager will be able to sereliaze this block to DISK
+ // and hence count returns correct value.
+ testRecordcount(true, StorageLevel.MEMORY_AND_DISK,
+ IteratorBlock((List.fill(70)(new Array[Byte](100))).iterator),
blockManager, Some(70))
+
+ // there is not enough space to store this block With MEMORY_ONLY
StorageLevel.
+ // BlockManager will not be able to unroll this block
+ // and hence it will not tryToPut this block, resulting the
SparkException
+ storageLevel = StorageLevel.MEMORY_ONLY
+ withBlockManagerBasedBlockHandler { handler =>
+ val thrown = intercept[SparkException] {
+ storeSingleBlock(handler, IteratorBlock((List.fill(70)(new
Array[Byte](100))).iterator))
+ }
+ }
+ }
+
+ private def
testCountWithBlockManagerBasedBlockHandler(isBlockManagerBasedBlockHandler:
Boolean) {
+ // ByteBufferBlock-MEMORY_ONLY
+ testRecordcount(isBlockManagerBasedBlockHandler,
StorageLevel.MEMORY_ONLY,
+ ByteBufferBlock(ByteBuffer.wrap(Array.tabulate(100)(i =>
i.toByte))), blockManager, None)
+ // ByteBufferBlock-MEMORY_ONLY_SER
+ testRecordcount(isBlockManagerBasedBlockHandler,
StorageLevel.MEMORY_ONLY_SER,
+ ByteBufferBlock(ByteBuffer.wrap(Array.tabulate(100)(i =>
i.toByte))), blockManager, None)
+ // ArrayBufferBlock-MEMORY_ONLY
+ testRecordcount(isBlockManagerBasedBlockHandler,
StorageLevel.MEMORY_ONLY,
+ ArrayBufferBlock(ArrayBuffer.fill(25)(0)), blockManager, Some(25))
+ // ArrayBufferBlock-MEMORY_ONLY_SER
+ testRecordcount(isBlockManagerBasedBlockHandler,
StorageLevel.MEMORY_ONLY_SER,
+ ArrayBufferBlock(ArrayBuffer.fill(25)(0)), blockManager, Some(25))
+ // ArrayBufferBlock-DISK_ONLY
+ testRecordcount(isBlockManagerBasedBlockHandler,
StorageLevel.DISK_ONLY,
+ ArrayBufferBlock(ArrayBuffer.fill(50)(0)), blockManager, Some(50))
+ // ArrayBufferBlock-MEMORY_AND_DISK
+ testRecordcount(isBlockManagerBasedBlockHandler,
StorageLevel.MEMORY_AND_DISK,
+ ArrayBufferBlock(ArrayBuffer.fill(75)(0)), blockManager, Some(75))
+ // IteratorBlock-MEMORY_ONLY
+ testRecordcount(isBlockManagerBasedBlockHandler,
StorageLevel.MEMORY_ONLY,
+ IteratorBlock((ArrayBuffer.fill(100)(0)).iterator), blockManager,
Some(100))
+ // IteratorBlock-MEMORY_ONLY_SER
+ testRecordcount(isBlockManagerBasedBlockHandler,
StorageLevel.MEMORY_ONLY_SER,
+ IteratorBlock((ArrayBuffer.fill(100)(0)).iterator), blockManager,
Some(100))
+ // IteratorBlock-DISK_ONLY
+ testRecordcount(isBlockManagerBasedBlockHandler,
StorageLevel.DISK_ONLY,
+ IteratorBlock((ArrayBuffer.fill(125)(0)).iterator), blockManager,
Some(125))
+ // IteratorBlock-MEMORY_AND_DISK
+ testRecordcount(isBlockManagerBasedBlockHandler,
StorageLevel.MEMORY_AND_DISK,
+ IteratorBlock((ArrayBuffer.fill(150)(0)).iterator), blockManager,
Some(150))
+ }
+
+ private def createBlockManager(
+ maxMem: Long,
+ conf: SparkConf,
+ name: String = SparkContext.DRIVER_IDENTIFIER): BlockManager = {
+ val transfer = new NioBlockTransferService(conf, securityMgr)
+ val manager = new BlockManager(name, rpcEnv, blockManagerMaster,
serializer, maxMem, conf,
+ mapOutputTracker, shuffleManager, transfer, securityMgr, 0)
+ manager.initialize("app-id")
+ blockManagerBuffer += manager
+ manager
+ }
+
+ /**
+ * Test storing of data using different types of Handler, StorageLevle
and ReceivedBlocks
+ * and verify the correct record count
+ */
+ private def testRecordcount(isBlockManagedBasedBlockHandler: Boolean,
+ sLevel: StorageLevel,
+ receivedBlock: ReceivedBlock,
+ bManager: BlockManager,
+ expectedNumRecords: Option[Long]
+ ) {
+ blockManager = bManager
+ storageLevel = sLevel
+ var bId: StreamBlockId = null
+ try {
+ if (isBlockManagedBasedBlockHandler) {
+ // test received block with BlockManager based handler
+ withBlockManagerBasedBlockHandler { handler =>
+ val (blockId, blockStoreResult) = storeSingleBlock(handler,
receivedBlock)
+ bId = blockId
+ assert(blockStoreResult.numRecords === expectedNumRecords,
+ "Message count not matches for a " +
+ receivedBlock.getClass.getName +
+ " being inserted using BlockManagerBasedBlockHandler with " +
sLevel)
+ }
+ } else {
+ // test received block with WAL based handler
+ withWriteAheadLogBasedBlockHandler { handler =>
+ val (blockId, blockStoreResult) = storeSingleBlock(handler,
receivedBlock)
+ bId = blockId
+ assert(blockStoreResult.numRecords === expectedNumRecords,
+ "Message count not matches for a " +
+ receivedBlock.getClass.getName +
+ " being inserted using WriteAheadLogBasedBlockHandler with " +
sLevel)
+ }
+ }
+ } finally {
+ // Removing the Block Id to use same blockManager for next test
+ blockManager.removeBlock(bId, true)
--- End diff --
The cost of setup / teardown for the fixtures in this suite should be
pretty low, so it seems unnecessarily complicated to try to re-use them like
this vs. simply wrapping each case in `test(..) {} ` and letting
`BeforeAndAfter` handle it.
---
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 infrastruct...@apache.org or file a JIRA ticket
with INFRA.
---
---------------------------------------------------------------------
To unsubscribe, e-mail: reviews-unsubscr...@spark.apache.org
For additional commands, e-mail: reviews-h...@spark.apache.org
