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
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
|
/*
* Licensed to the Apache Software Foundation (ASF) under one or more
* contributor license agreements. See the NOTICE file distributed with
* this work for additional information regarding copyright ownership.
* The ASF licenses this file to You under the Apache License, Version 2.0
* (the "License"); you may not use this file except in compliance with
* the License. You may obtain a copy of the License at
*
* http://www.apache.org/licenses/LICENSE-2.0
*
* Unless required by applicable law or agreed to in writing, software
* distributed under the License is distributed on an "AS IS" BASIS,
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
* See the License for the specific language governing permissions and
* limitations under the License.
*/
package org.apache.spark.storage
import scala.collection.Map
import scala.collection.mutable
import org.apache.spark.SparkContext
import org.apache.spark.util.Utils
private[spark]
class StorageStatus(
val blockManagerId: BlockManagerId,
val maxMem: Long,
val blocks: mutable.Map[BlockId, BlockStatus] = mutable.Map.empty) {
def memUsed() = blocks.values.map(_.memSize).reduceOption(_ + _).getOrElse(0L)
def memUsedByRDD(rddId: Int) =
rddBlocks.filterKeys(_.rddId == rddId).values.map(_.memSize).reduceOption(_ + _).getOrElse(0L)
def diskUsed() = blocks.values.map(_.diskSize).reduceOption(_ + _).getOrElse(0L)
def diskUsedByRDD(rddId: Int) =
rddBlocks.filterKeys(_.rddId == rddId).values.map(_.diskSize).reduceOption(_ + _).getOrElse(0L)
def memRemaining : Long = maxMem - memUsed()
def rddBlocks = blocks.flatMap {
case (rdd: RDDBlockId, status) => Some(rdd, status)
case _ => None
}
}
private[spark]
class RDDInfo(val id: Int, val name: String, val numPartitions: Int, val storageLevel: StorageLevel)
extends Ordered[RDDInfo] {
var numCachedPartitions = 0
var memSize = 0L
var diskSize = 0L
override def toString = {
("RDD \"%s\" (%d) Storage: %s; CachedPartitions: %d; TotalPartitions: %d; MemorySize: %s; " +
"DiskSize: %s").format(name, id, storageLevel.toString, numCachedPartitions,
numPartitions, Utils.bytesToString(memSize), Utils.bytesToString(diskSize))
}
override def compare(that: RDDInfo) = {
this.id - that.id
}
}
/* Helper methods for storage-related objects */
private[spark]
object StorageUtils {
/**
* Returns basic information of all RDDs persisted in the given SparkContext. This does not
* include storage information.
*/
def rddInfoFromSparkContext(sc: SparkContext): Array[RDDInfo] = {
sc.persistentRdds.values.map { rdd =>
val rddName = Option(rdd.name).getOrElse(rdd.id.toString)
val rddNumPartitions = rdd.partitions.size
val rddStorageLevel = rdd.getStorageLevel
val rddInfo = new RDDInfo(rdd.id, rddName, rddNumPartitions, rddStorageLevel)
rddInfo
}.toArray
}
/** Returns storage information of all RDDs persisted in the given SparkContext. */
def rddInfoFromStorageStatus(
storageStatuses: Seq[StorageStatus],
sc: SparkContext): Array[RDDInfo] = {
rddInfoFromStorageStatus(storageStatuses, rddInfoFromSparkContext(sc))
}
/** Returns storage information of all RDDs in the given list. */
def rddInfoFromStorageStatus(
storageStatuses: Seq[StorageStatus],
rddInfos: Seq[RDDInfo]): Array[RDDInfo] = {
// Mapping from RDD ID -> an array of associated BlockStatuses
val blockStatusMap = storageStatuses.flatMap(_.rddBlocks).toMap
.groupBy { case (k, _) => k.rddId }
.mapValues(_.values.toArray)
// Mapping from RDD ID -> the associated RDDInfo (with potentially outdated storage information)
val rddInfoMap = rddInfos.map { info => (info.id, info) }.toMap
val rddStorageInfos = blockStatusMap.flatMap { case (rddId, blocks) =>
// Add up memory and disk sizes
val persistedBlocks = blocks.filter { status => status.memSize + status.diskSize > 0 }
val memSize = persistedBlocks.map(_.memSize).reduceOption(_ + _).getOrElse(0L)
val diskSize = persistedBlocks.map(_.diskSize).reduceOption(_ + _).getOrElse(0L)
rddInfoMap.get(rddId).map { rddInfo =>
rddInfo.numCachedPartitions = persistedBlocks.length
rddInfo.memSize = memSize
rddInfo.diskSize = diskSize
rddInfo
}
}.toArray
scala.util.Sorting.quickSort(rddStorageInfos)
rddStorageInfos
}
/** Returns a mapping from BlockId to the locations of the associated block. */
def blockLocationsFromStorageStatus(
storageStatuses: Seq[StorageStatus]): Map[BlockId, Seq[String]] = {
val blockLocationPairs = storageStatuses.flatMap { storageStatus =>
storageStatus.blocks.map { case (bid, _) => (bid, storageStatus.blockManagerId.hostPort) }
}
blockLocationPairs.toMap
.groupBy { case (blockId, _) => blockId }
.mapValues(_.values.toSeq)
}
/** Filters the given list of StorageStatus by the given RDD ID. */
def filterStorageStatusByRDD(
storageStatuses: Seq[StorageStatus],
rddId: Int): Array[StorageStatus] = {
storageStatuses.map { status =>
val filteredBlocks = status.rddBlocks.filterKeys(_.rddId == rddId).toSeq
val filteredBlockMap = mutable.Map[BlockId, BlockStatus](filteredBlocks: _*)
new StorageStatus(status.blockManagerId, status.maxMem, filteredBlockMap)
}.toArray
}
}
|
