Spark内存预测 sizeEstimate

/**
 * Estimate the current size of the collection in bytes. O(1) time.
 */
def estimateSize(): Long = {
  assert(samples.nonEmpty)
  val extrapolatedDelta = bytesPerUpdate * (numUpdates - samples.last.numUpdates)
  (samples.last.size + extrapolatedDelta).toLong
}

/**
 * Controls the base of the exponential which governs the rate of sampling.
 * E.g., a value of 2 would mean we sample at 1, 2, 4, 8, ... elements.
 */
private val SAMPLE_GROWTH_RATE = 1.1

/** Samples taken since last resetSamples(). Only the last two are kept for extrapolation. */
private val samples = new mutable.Queue[Sample]

/** The average number of bytes per update between our last two samples. */
private var bytesPerUpdate: Double = _

/** Total number of insertions and updates into the map since the last resetSamples(). */
private var numUpdates: Long = _

/** The value of 'numUpdates' at which we will take our next sample. */
private var nextSampleNum: Long = _

/**
 * Reset samples collected so far.
 * This should be called after the collection undergoes a dramatic change in size.
 */
protected def resetSamples(): Unit = {
  numUpdates = 1
  nextSampleNum = 1
  samples.clear()
  takeSample()
}

/**
 * Callback to be invoked after every update.
 */
protected def afterUpdate(): Unit = {
  numUpdates += 1
  if (nextSampleNum == numUpdates) {
    takeSample()
  }
}

/**
 * Take a new sample of the current collection's size.
 */
private def takeSample(): Unit = {
  samples.enqueue(Sample(SizeEstimator.estimate(this), numUpdates))
  // Only use the last two samples to extrapolate
  if (samples.size > 2) {
    samples.dequeue()
  }
  val bytesDelta = samples.toList.reverse match {
    case latest :: previous :: tail =>
      (latest.size - previous.size).toDouble / (latest.numUpdates - previous.numUpdates)
    // If fewer than 2 samples, assume no change
    case _ => 0
  }
  bytesPerUpdate = math.max(0, bytesDelta)
  nextSampleNum = math.ceil(numUpdates * SAMPLE_GROWTH_RATE).toLong
}

/**
 * Estimate the current size of the collection in bytes. O(1) time.
 */
def estimateSize(): Long = {
  assert(samples.nonEmpty)
  val extrapolatedDelta = bytesPerUpdate * (numUpdates - samples.last.numUpdates)
  (samples.last.size + extrapolatedDelta).toLong
}

private def estimate(obj: AnyRef, visited: IdentityHashMap[AnyRef, AnyRef]): Long = {
  val state = new SearchState(visited)
  state.enqueue(obj)
  while (!state.isFinished) {
    visitSingleObject(state.dequeue(), state)
  }
  state.size
}

private def visitSingleObject(obj: AnyRef, state: SearchState) {
  val cls = obj.getClass
  if (cls.isArray) {
    visitArray(obj, cls, state)
  } else if (cls.getName.startsWith("scala.reflect")) {
    // Many objects in the scala.reflect package reference global reflection objects which, in
    // turn, reference many other large global objects. Do nothing in this case.
  } else if (obj.isInstanceOf[ClassLoader] || obj.isInstanceOf[Class[_]]) {
    // Hadoop JobConfs created in the interpreter have a ClassLoader, which greatly confuses
    // the size estimator since it references the whole REPL. Do nothing in this case. In
    // general all ClassLoaders and Classes will be shared between objects anyway.
  } else {
    obj match {
      case s: KnownSizeEstimation =>
        state.size += s.estimatedSize
      case _ =>
        val classInfo = getClassInfo(cls)
        state.size += alignSize(classInfo.shellSize)
        for (field <- classInfo.pointerFields) {
          state.enqueue(field.get(obj))
        }
    }
  }
}

// Estimate the size of arrays larger than ARRAY_SIZE_FOR_SAMPLING by sampling.
private val ARRAY_SIZE_FOR_SAMPLING = 400
private val ARRAY_SAMPLE_SIZE = 100 // should be lower than ARRAY_SIZE_FOR_SAMPLING

private def visitArray(array: AnyRef, arrayClass: Class[_], state: SearchState) {
  val length = ScalaRunTime.array_length(array)
  val elementClass = arrayClass.getComponentType()

  // Arrays have object header and length field which is an integer
  var arrSize: Long = alignSize(objectSize + INT_SIZE)

  if (elementClass.isPrimitive) {
    arrSize += alignSize(length.toLong * primitiveSize(elementClass))
    state.size += arrSize
  } else {
    arrSize += alignSize(length.toLong * pointerSize)
    state.size += arrSize

    if (length <= ARRAY_SIZE_FOR_SAMPLING) {
      var arrayIndex = 0
      while (arrayIndex < length) {
        state.enqueue(ScalaRunTime.array_apply(array, arrayIndex).asInstanceOf[AnyRef])
        arrayIndex += 1
      }
    } else {
      // Estimate the size of a large array by sampling elements without replacement.
      // To exclude the shared objects that the array elements may link, sample twice
      // and use the min one to calculate array size.
      val rand = new Random(42)
      val drawn = new OpenHashSet[Int](2 * ARRAY_SAMPLE_SIZE)
      val s1 = sampleArray(array, state, rand, drawn, length)
      val s2 = sampleArray(array, state, rand, drawn, length)
      val size = math.min(s1, s2)
      state.size += math.max(s1, s2) +
        (size * ((length - ARRAY_SAMPLE_SIZE) / (ARRAY_SAMPLE_SIZE))).toLong
    }
  }
}