core/src/main/scala/spark/util/CompletionIterator.scala


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

package spark.util

/**
 * Wrapper around an iterator which calls a completion method after it successfully iterates through all the elements
 */
abstract class CompletionIterator[+A, +I <: Iterator[A]](sub: I) extends Iterator[A]{
  def next = sub.next
  def hasNext = {
    val r = sub.hasNext
    if (!r) {
      completion
    }
    r
  }

  def completion()
}

object CompletionIterator {
  def apply[A, I <: Iterator[A]](sub: I, completionFunction: => Unit) : CompletionIterator[A,I] = {
    new CompletionIterator[A,I](sub) {
      def completion() = completionFunction
    }
  }
}