SPARKC-619 restore PrefetchingResultSetIterator

The iterator tries to prefetch the next result page. Once the
current page is exhausted, the next page (hopefully
materialized at this point) becomes the current page and a
fetch request is sent.

fetchSize argument was removed, the supplied statement should
have pageSize parameter set by the caller.

Author: jtgrabowski

connector/src/it/scala/com/datastax/spark/connector/SparkCassandraITFlatSpecBase.scala

@@ -183,6 +183,10 @@ trait SparkCassandraITSpecBase
 
   implicit val ec = SparkCassandraITSpecBase.ec
 
+  def await[T](unit: Future[T]): T = {
+    Await.result(unit, Duration.Inf)
+  }
+
   def awaitAll[T](units: Future[T]*): Seq[T] = {
     Await.result(Future.sequence(units), Duration.Inf)
   }

connector/src/it/scala/com/datastax/spark/connector/rdd/reader/PrefetchingResultSetIteratorSpec.scala

@@ -0,0 +1,67 @@
+package com.datastax.spark.connector.rdd.reader
+
+import com.codahale.metrics.Timer
+import com.datastax.oss.driver.api.core.cql.SimpleStatement.newInstance
+import com.datastax.spark.connector.SparkCassandraITFlatSpecBase
+import com.datastax.spark.connector.cluster.DefaultCluster
+import com.datastax.spark.connector.cql.CassandraConnector
+import org.scalatest.concurrent.Eventually.{eventually, timeout}
+import org.scalatest.time.{Seconds, Span}
+
+class PrefetchingResultSetIteratorSpec extends SparkCassandraITFlatSpecBase with DefaultCluster {
+
+  private val table = "prefetching"
+  private val emptyTable = "empty_prefetching"
+  override lazy val conn = CassandraConnector(sparkConf)
+
+  override def beforeClass {
+    conn.withSessionDo { session =>
+      session.execute(
+        s"CREATE KEYSPACE IF NOT EXISTS $ks WITH REPLICATION = { 'class': 'SimpleStrategy', 'replication_factor': 1 }")
+
+      session.execute(
+        s"CREATE TABLE IF NOT EXISTS $ks.$table (key INT, x INT, PRIMARY KEY (key))")
+
+      session.execute(
+        s"CREATE TABLE IF NOT EXISTS $ks.$emptyTable (key INT, x INT, PRIMARY KEY (key))")
+
+      awaitAll(
+        for (i <- 1 to 999) yield {
+          executor.executeAsync(newInstance(s"INSERT INTO $ks.$table (key, x) values ($i, $i)"))
+        }
+      )
+    }
+  }
+
+  "PrefetchingResultSetIterator" should "return all rows regardless of the  page sizes" in {
+    val pageSizes = Seq(1, 2, 5, 111, 998, 999, 1000, 1001)
+    for (pageSize <- pageSizes) {
+      withClue(s"Prefetching iterator failed for the page size: $pageSize") {
+        val statement = newInstance(s"select * from $ks.$table").setPageSize(pageSize)
+        val result = executor.executeAsync(statement).map(new PrefetchingResultSetIterator(_))
+        await(result).toList should have size 999
+      }
+    }
+  }
+
+  it should "be empty for an empty table" in {
+    val statement = newInstance(s"select * from $ks.$emptyTable")
+    val result = executor.executeAsync(statement).map(new PrefetchingResultSetIterator(_))
+
+    await(result).hasNext should be(false)
+    intercept[NoSuchElementException] {
+      await(result).next()
+    }
+  }
+
+  it should "update the provided timer" in {
+    val statement = newInstance(s"select * from $ks.$table").setPageSize(200)
+    val timer = new Timer()
+    val result = executor.executeAsync(statement).map(rs => new PrefetchingResultSetIterator(rs, Option(timer)))
+    await(result).toList
+
+    eventually(timeout(Span(2, Seconds))) {
+      timer.getCount should be(4)
+    }
+  }
+}

connector/src/main/scala/com/datastax/bdp/util/ScalaJavaUtil.scala

@@ -11,7 +11,7 @@ import java.util.concurrent.{Callable, CompletionStage}
 import java.util.function
 import java.util.function.{BiConsumer, Consumer, Predicate, Supplier}
 
-import scala.concurrent.{ExecutionContext, Future, Promise}
+import scala.concurrent.{ExecutionContext, ExecutionContextExecutor, Future, Promise}
 import scala.concurrent.duration.{Duration => ScalaDuration}
 import scala.language.implicitConversions
 
@@ -48,7 +48,7 @@ object ScalaJavaUtil {
   def asScalaFunction[T, R](f: java.util.function.Function[T, R]): T => R = x => f(x)
 
   def asScalaFuture[T](completionStage: CompletionStage[T])
-                      (implicit context: ExecutionContext): Future[T] = {
+                      (implicit context: ExecutionContextExecutor): Future[T] = {
     val promise = Promise[T]()
     completionStage.whenCompleteAsync(new BiConsumer[T, java.lang.Throwable] {
       override def accept(t: T, throwable: Throwable): Unit = {
@@ -58,7 +58,7 @@ object ScalaJavaUtil {
           promise.failure(throwable)
 
       }
-    })
+    }, context)
     promise.future
   }
 }

connector/src/main/scala/com/datastax/spark/connector/cql/Scanner.scala

@@ -44,7 +44,7 @@ class DefaultScanner (
     val rs = session.executeAsync(maybeExecutingAs(statement, readConf.executeAs))
     val scanResult = asScalaFuture(rs).map { rs =>
       val columnMetaData = CassandraRowMetadata.fromResultSet(columnNames, rs, codecRegistry)
-      val prefetchingIterator = new PrefetchingResultSetIterator(rs, readConf.fetchSizeInRows)
+      val prefetchingIterator = new PrefetchingResultSetIterator(rs)
       val rateLimitingIterator = readConf.throughputMiBPS match {
         case Some(throughput) =>
           val rateLimiter = new RateLimiter((throughput * 1024 * 1024).toLong, 1024 * 1024)

connector/src/main/scala/com/datastax/spark/connector/rdd/CassandraCoGroupedRDD.scala

@@ -168,7 +168,7 @@ class CassandraCoGroupedRDD[T](
     val fetchResult = asScalaFuture(session.executeAsync(stmt)).map { rs =>
       val columnNames = fromRDD.selectedColumnRefs.map(_.selectedAs).toIndexedSeq ++ Seq(TokenColumn)
       val columnMetaData = CassandraRowMetadata.fromResultSet(columnNames, rs, session.getContext.getCodecRegistry)
-      val iterator = new PrefetchingResultSetIterator(rs, fromRDD.readConf.fetchSizeInRows)
+      val iterator = new PrefetchingResultSetIterator(rs)
       val iteratorWithMetrics = iterator.map(inputMetricsUpdater.updateMetrics)
       logDebug(s"Row iterator for range $range obtained successfully.")
       (columnMetaData, iteratorWithMetrics)

connector/src/main/scala/com/datastax/spark/connector/rdd/CassandraJoinRDD.scala

@@ -124,7 +124,6 @@ class CassandraJoinRDD[L, R] (
     metricsUpdater: InputMetricsUpdater
   ): Iterator[(L, R)] = {
 
-
     val queryExecutor = QueryExecutor(session, readConf.parallelismLevel, None, None)
 
     def pairWithRight(left: L): SettableFuture[Iterator[(L, R)]] = {
@@ -133,9 +132,12 @@ class CassandraJoinRDD[L, R] (
 
       import com.datastax.spark.connector.util.Threads.BlockingIOExecutionContext
 
-      queryExecutor.executeAsync(bsb.bind(left).executeAs(readConf.executeAs)).onComplete {
+      val stmt = bsb.bind(left)
+        .update(_.setPageSize(readConf.fetchSizeInRows))
+        .executeAs(readConf.executeAs)
+      queryExecutor.executeAsync(stmt).onComplete {
         case Success(rs) =>
-          val resultSet = new PrefetchingResultSetIterator(rs, fetchSize)
+          val resultSet = new PrefetchingResultSetIterator(rs)
           val iteratorWithMetrics = resultSet.map(metricsUpdater.updateMetrics)
           /* This is a much less than ideal place to actually rate limit, we are buffering
           these futures this means we will most likely exceed our threshold*/

connector/src/main/scala/com/datastax/spark/connector/rdd/CassandraLeftJoinRDD.scala

@@ -157,9 +157,12 @@ class CassandraLeftJoinRDD[L, R] (
       val resultFuture = SettableFuture.create[Iterator[(L, Option[R])]]
       val leftSide = Iterator.continually(left)
 
-      queryExecutor.executeAsync(bsb.bind(left).executeAs(readConf.executeAs)).onComplete {
+      val stmt = bsb.bind(left)
+        .update(_.setPageSize(readConf.fetchSizeInRows))
+        .executeAs(readConf.executeAs)
+      queryExecutor.executeAsync(stmt).onComplete {
         case Success(rs) =>
-          val resultSet = new PrefetchingResultSetIterator(rs, fetchSize)
+          val resultSet = new PrefetchingResultSetIterator(rs)
           val iteratorWithMetrics = resultSet.map(metricsUpdater.updateMetrics)
           /* This is a much less than ideal place to actually rate limit, we are buffering
           these futures this means we will most likely exceed our threshold*/

connector/src/main/scala/com/datastax/spark/connector/rdd/CassandraMergeJoinRDD.scala

@@ -160,7 +160,7 @@ class CassandraMergeJoinRDD[L,R](
     val fetchResult = asScalaFuture(session.executeAsync(stmt)).map { rs =>
       val columnNames = fromRDD.selectedColumnRefs.map(_.selectedAs).toIndexedSeq ++ Seq(TokenColumn)
       val columnMetaData = CassandraRowMetadata.fromResultSet(columnNames, rs, session)
-      val iterator = new PrefetchingResultSetIterator(rs, fromRDD.readConf.fetchSizeInRows)
+      val iterator = new PrefetchingResultSetIterator(rs)
       val iteratorWithMetrics = iterator.map(inputMetricsUpdater.updateMetrics)
       logDebug(s"Row iterator for range $range obtained successfully.")
       (columnMetaData, iteratorWithMetrics)

connector/src/main/scala/com/datastax/spark/connector/rdd/reader/PrefetchingResultSetIterator.scala

@@ -1,45 +1,55 @@
 package com.datastax.spark.connector.rdd.reader
 
+import java.util.concurrent.TimeUnit
+
 import com.codahale.metrics.Timer
+import com.datastax.bdp.util.ScalaJavaUtil
 import com.datastax.oss.driver.api.core.cql.{AsyncResultSet, Row}
-import com.datastax.oss.driver.internal.core.cql.ResultSets
+import com.datastax.spark.connector.util.Threads.BlockingIOExecutionContext
+
+import scala.concurrent.duration.Duration
+import scala.concurrent.{Await, Future}
 
 /** Allows to efficiently iterate over a large, paged ResultSet,
   * asynchronously prefetching the next page.
-  * 
+  *
+  * This iterator is NOT thread safe. Attempting to retrieve elements from many threads without synchronization
+  * may yield unspecified results.
+  *
   * @param resultSet result set obtained from the Java driver
-  * @param prefetchWindowSize if there are less than this rows available without blocking,
-  *                           initiates fetching the next page
-  * @param timer a Codahale timer to optionally gather the metrics of fetching time
+  * @param timer     a Codahale timer to optionally gather the metrics of fetching time
   */
-class PrefetchingResultSetIterator(resultSet: AsyncResultSet, prefetchWindowSize: Int, timer: Option[Timer] = None)
-  extends Iterator[Row] {
-
-  private val iterator = ResultSets.newInstance(resultSet).iterator() //TODO
-
-  override def hasNext = iterator.hasNext
-
-// TODO: implement async page fetching. Following implementation might call fetchMoreResults up to prefetchWindowSize
-//       times to fetch the same page. Is this behaviour still valid in the new driver?
-//       This class should take AsyncResultSet as constructor param (not ResultSet)
-
-//  private[this] def maybePrefetch(): Unit = {
-//    if (!resultSet.isFullyFetched && resultSet.getAvailableWithoutFetching < prefetchWindowSize) {
-//      val t0 = System.nanoTime()
-//      val future: ListenableFuture[ResultSet] = resultSet.fetchMoreResults()
-//      if (timer.isDefined)
-//        Futures.addCallback(future, new FutureCallback[ResultSet] {
-//          override def onSuccess(ignored: ResultSet): Unit = {
-//            timer.get.update(System.nanoTime() - t0, TimeUnit.NANOSECONDS)
-//          }
-//
-//          override def onFailure(ignored: Throwable): Unit = { }
-//        })
-//    }
-//  }
-
-  override def next() = {
-//    maybePrefetch()
-    iterator.next()
+class PrefetchingResultSetIterator(resultSet: AsyncResultSet, timer: Option[Timer] = None) extends Iterator[Row] {
+  private var currentIterator = resultSet.currentPage().iterator()
+  private var currentResultSet = resultSet
+  private var nextResultSet = fetchNextPage()
+
+  private def fetchNextPage(): Option[Future[AsyncResultSet]] = {
+    if (currentResultSet.hasMorePages) {
+      val t0 = System.nanoTime();
+      val next = ScalaJavaUtil.asScalaFuture(currentResultSet.fetchNextPage())
+      timer.foreach { t =>
+        next.foreach(_ => t.update(System.nanoTime() - t0, TimeUnit.NANOSECONDS))
+      }
+      Option(next)
+    } else
+      None
+  }
+
+  private def maybePrefetch(): Unit = {
+    if (!currentIterator.hasNext && currentResultSet.hasMorePages) {
+      currentResultSet = Await.result(nextResultSet.get, Duration.Inf)
+      currentIterator = currentResultSet.currentPage().iterator()
+      nextResultSet = fetchNextPage()
+    }
+  }
+
+  override def hasNext: Boolean =
+    currentIterator.hasNext || currentResultSet.hasMorePages
+
+  override def next(): Row = {
+    val row = currentIterator.next() // let's try to exhaust the current iterator first
+    maybePrefetch()
+    row
   }
 }

connector/src/main/scala/com/datastax/spark/connector/util/Threads.scala

@@ -4,14 +4,19 @@ import java.util.concurrent.{Executors, LinkedBlockingQueue, ThreadPoolExecutor,
 
 import com.google.common.util.concurrent.ThreadFactoryBuilder
 
-import scala.concurrent.ExecutionContext
+import scala.concurrent.{ExecutionContext, ExecutionContextExecutorService}
 
-object Threads {
+object Threads extends Logging {
 
-  implicit val BlockingIOExecutionContext = {
+  implicit val BlockingIOExecutionContext: ExecutionContextExecutorService = {
     val threadFactory = new ThreadFactoryBuilder()
       .setDaemon(true)
       .setNameFormat("spark-cassandra-connector-io" + "%d")
+      .setUncaughtExceptionHandler(new Thread.UncaughtExceptionHandler {
+        override def uncaughtException(t: Thread, e: Throwable): Unit = {
+          logWarning(s"Unhandled exception in thread ${t.getName}.", e)
+        }
+      })
       .build
     ExecutionContext.fromExecutorService(Executors.newCachedThreadPool(threadFactory))
   }