Spark源码分析-topN源码

1. 介绍

TopN算子是取RDD的前N个元素。取TopN元素，我们就一定要对其进行严格排序吗？非也，也正是如此，加大了此算子的效率。

1. 源码理解

  

  def top(num: Int): JList[T] = {
  val comp = com.google.common.collect.Ordering.natural().asInstanceOf[Comparator[T]]
  top(num, comp)
}

我们需要传进去一个整形的参数，然后后面声明了一个默认的比较器comp，又调用了两个参数的top函数，然后可以看到又调用了RDD的takeOrdered函数：

def takeOrdered(num: Int)(implicit ord: Ordering[T]): Array[T] = withScope {
  if (num == 0) {
    Array.empty
  } else {
    val mapRDDs = mapPartitions { items =>
      // Priority keeps the largest elements, so let's reverse the ordering.
      val queue = new BoundedPriorityQueue[T](num)(ord.reverse)
      queue ++= collectionUtils.takeOrdered(items, num)(ord)
      Iterator.single(queue)
    }
    if (mapRDDs.partitions.length == 0) {
      Array.empty
    } else {
      mapRDDs.reduce { (queue1, queue2) =>
        queue1 ++= queue2
        queue1
      }.toArray.sorted(ord)
    }
  }
}

    RDD的takeOrdered函数里面使用的是mapPartitions算子，从这里就可以知道他首先需要在每个partition中求top信息，然后就在对所有的partition信息进行统一的排序。

    重点看collectionUtils.takeOrdered方法，里面调用了一个ordering.leastOf(input.asJava, num)方法，重点看这个方法里面的内容：

public <E extends T> List<E> leastOf(Iterator<E> elements, int k) {
  checkNotNull(elements);
  checkArgument(k >= 0, "k (%s) must be nonnegative", k);

  if (k == 0 || !elements.hasNext()) {
    return ImmutableList.of();
  } else if (k >= Integer.MAX_VALUE / 2) {
    // k is really large; just do a straightforward sorted-copy-and-sublist
    ArrayList<E> list = Lists.newArrayList(elements);
    Collections.sort(list, this);
    if (list.size() > k) {
      list.subList(k, list.size()).clear();
    }
    list.trimToSize();
    return Collections.unmodifiableList(list);
  }

  /*
   * Our goal is an O(n) algorithm using only one pass and O(k) additional
   * memory.
   *
   * We use the following algorithm: maintain a buffer of size 2*k. Every time
   * the buffer gets full, find the median and partition around it, keeping
   * only the lowest k elements.  This requires n/k find-median-and-partition
   * steps, each of which take O(k) time with a traditional quickselect.
   *
   * After sorting the output, the whole algorithm is O(n + k log k). It
   * degrades gracefully for worst-case input (descending order), performs
   * competitively or wins outright for randomly ordered input, and doesn't
   * require the whole collection to fit into memory.
   */
  int bufferCap = k * 2;
  @SuppressWarnings("unchecked") // we'll only put E's in
  E[] buffer = (E[]) new Object[bufferCap];
  E threshold = elements.next();
  buffer[0] = threshold;
  int bufferSize = 1;
  // threshold is the kth smallest element seen so far.  Once bufferSize >= k,
  // anything larger than threshold can be ignored immediately.

  while (bufferSize < k && elements.hasNext()) {
    E e = elements.next();
    buffer[bufferSize++] = e;
    threshold = max(threshold, e);
  }

  while (elements.hasNext()) {
    E e = elements.next();
    if (compare(e, threshold) >= 0) {
      continue;
    }

    buffer[bufferSize++] = e;
    if (bufferSize == bufferCap) {
      // We apply the quickselect algorithm to partition about the median,
      // and then ignore the last k elements.
      int left = 0;
      int right = bufferCap - 1;

      int minThresholdPosition = 0;
      // The leftmost position at which the greatest of the k lower elements
      // -- the new value of threshold -- might be found.

      while (left < right) {
        int pivotIndex = (left + right + 1) >>> 1;
        int pivotNewIndex = partition(buffer, left, right, pivotIndex);
        if (pivotNewIndex > k) {
          right = pivotNewIndex - 1;
        } else if (pivotNewIndex < k) {
          left = Math.max(pivotNewIndex, left + 1);
          minThresholdPosition = pivotNewIndex;
        } else {
          break;
        }
      }
      bufferSize = k;

      threshold = buffer[minThresholdPosition];
      for (int i = minThresholdPosition + 1; i < bufferSize; i++) {
        threshold = max(threshold, buffer[i]);
      }
    }
  }

  Arrays.sort(buffer, 0, bufferSize, this);

  bufferSize = Math.min(bufferSize, k);
  return Collections.unmodifiableList(
      Arrays.asList(ObjectArrays.arraysCopyOf(buffer, bufferSize)));
  // We can't use ImmutableList; we have to be null-friendly!
}

    从最上面开始，如果k过大，超过了Int类型最大数的一半，那么久直接用Collections.sort进行排序，底层的排序代码如下：

default void sort(Comparator<? super E> c) {
    Object[] a = this.toArray();
    Arrays.sort(a, (Comparator) c);
    ListIterator<E> i = this.listIterator();
    for (Object e : a) {
        i.next();
        i.set((E) e);
    }
}

    把原始的List放到一个数组中，然后用Arrays.sort进行排序，这个排序默认用的是Timsort方法排序（优化的归并排序-归并和插入排序集合），然后将排好序的数组重新放到list中，这也好解释，为什么上面有list.subList(k, list.size()).clear();代码，就是为了把所有多余的数据去掉。上面的是特殊情况，一般情况下没人取那么大的top值。

    下面接着看正常情况下的代码：

    首先声明了一个长度为k*2大小的数组buffer：

int bufferCap = k * 2;

E[] buffer = (E[]) new Object[bufferCap];

为什么是k*2的长度，是为了quickselect算法的实现：

在leastOf函数的第一个While循环中，直接将前K个元素复制到buffer数组中，并且记录前K个元素中的最大值，然后开始下面的while循环。总体思路如下：

开始遍历后面的元素，判断当前元素与前K元素中最大值比较，如果大于等于最大值，然后往buffer中放入值，直到长度为k*2的数组buffer元素填充满，然后开始取前K大值的操作：

int pivotIndex = (left + right + 1) >>> 1;

左右指针相加，再加一，然后无符号右移1位，取到buffer中位数的指针。然后后面的具体过程就是将所有的buffer中的元素前k大的元素移动到0到k-1之中，这个过程不需要排序，只需要移动动前k大的元素，并且需要记录此时前K大元素中最大值，然后再按照上面的过程再次将buffer填充满，再次移动找到前K大的值，直至遍历完所有的数。我们可以得到buffer中前K个元素就是该分区中最大的前K个值。

Arrays.sort(buffer, 0, bufferSize, this);

bufferSize = Math.min(bufferSize, k);
return Collections.unmodifiableList(
    Arrays.asList(ObjectArrays.arraysCopyOf(buffer, bufferSize)));

最后将数组排序，返回最前K大的值。