blob: 44638e0c2d8f9be93af4f0904037db9921c4a473 (
plain) (
blame)
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
|
package spark.storage
import spark.serializer.Serializer
private[spark]
class ShuffleWriterGroup(val id: Int, val writers: Array[BlockObjectWriter])
private[spark]
trait ShuffleBlocks {
def acquireWriters(mapId: Int): ShuffleWriterGroup
def releaseWriters(group: ShuffleWriterGroup)
}
private[spark]
class ShuffleBlockManager(blockManager: BlockManager) {
def forShuffle(shuffleId: Int, numBuckets: Int, serializer: Serializer): ShuffleBlocks = {
new ShuffleBlocks {
// Get a group of writers for a map task.
override def acquireWriters(mapId: Int): ShuffleWriterGroup = {
val bufferSize = System.getProperty("spark.shuffle.file.buffer.kb", "100").toInt * 1024
val writers = Array.tabulate[BlockObjectWriter](numBuckets) { bucketId =>
val blockId = ShuffleBlockManager.blockId(shuffleId, bucketId, mapId)
blockManager.getDiskBlockWriter(blockId, serializer, bufferSize)
}
new ShuffleWriterGroup(mapId, writers)
}
override def releaseWriters(group: ShuffleWriterGroup) = {
// Nothing really to release here.
}
}
}
}
private[spark]
object ShuffleBlockManager {
// Returns the block id for a given shuffle block.
def blockId(shuffleId: Int, bucketId: Int, groupId: Int): String = {
"shuffle_" + shuffleId + "_" + groupId + "_" + bucketId
}
// Returns true if the block is a shuffle block.
def isShuffle(blockId: String): Boolean = blockId.startsWith("shuffle_")
}
|
