com.exasol.spark.util.NextIterator.scala Maven / Gradle / Ivy

Go to download
Show more of this group Show more artifacts with this name
Show all versions of spark-connector_2.11 Show documentation
The Spark Exasol Connector
The newest version!
/*
 * 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 com.exasol.spark.util

/**
 * Provides a basic/boilerplate Iterator implementation.
 *
 * This is adapted from Spark codebase, [[org.apache.spark.util.NextIterator]].
 * It is private internally in Spark, copied here to be used in
 * [[com.exasol.spark.rdd.ExasolRDD]].
 */
private[spark] abstract class NextIterator[U] extends Iterator[U] {

  private var gotNext = false
  private var nextValue: U = _
  private var closed = false
  protected var finished = false

  /**
   * Method for subclasses to implement to provide the next element.
   *
   * If no next element is available, the subclass should set `finished` to
   * `true` and may return any value (it will be ignored).
   *
   * This convention is required because `null` may be a valid value, and using
   * `Option` seems like it might create unnecessary Some/None instances, given
   * some iterators might be called in a tight loop.
   *
   * @return U, or set 'finished' when done
   */
  protected def getNext(): U

  /**
   * Method for subclasses to implement when all elements have been successfully
   * iterated, and the iteration is done.
   *
   * Note: `NextIterator` cannot guarantee that `close` will be called
   * because it has no control over what happens when an exception happens in
   * the user code that is calling hasNext/next.
   *
   * Ideally you should have another try/catch, as in HadoopRDD, that ensures
   * any resources are closed should iteration fail.
   */
  protected def close(): Unit

  /**
   * Calls the subclass-defined close method, but only once.
   *
   * Usually calling `close` multiple times should be fine, but historically
   * there have been issues with some InputFormats throwing exceptions.
   */
  def closeIfNeeded(): Unit =
    if (!closed) {
      // Note: it's important that we set closed = true before calling close(),
      // since setting it afterwards would permit us to call close() multiple
      // times if close() threw an exception.
      closed = true
      close()
    }

  override def hasNext: Boolean = {
    if (!finished) {
      if (!gotNext) {
        nextValue = getNext()
        if (finished) {
          closeIfNeeded()
        }
        gotNext = true
      }
    }
    !finished
  }

  override def next(): U = {
    if (!hasNext) {
      throw new NoSuchElementException("End of stream")
    }
    gotNext = false
    nextValue
  }
}