Merge pull request #1270 from datastax/SPARKC-606-2.5

SPARKC-606 honor ContinuousPaging properties

Author: jtgrabowski

connector/src/it/scala/com/datastax/spark/connector/cql/ContinuousPagingScannerSpec.scala

@@ -7,9 +7,13 @@
 package com.datastax.spark.connector.cql
 
 import com.datastax.bdp.spark.ContinuousPagingScanner
+import com.datastax.dse.driver.api.core.config.DseDriverOption
+import com.datastax.oss.driver.api.core.cql.Statement
 import com.datastax.spark.connector._
 import com.datastax.spark.connector.cluster.DefaultCluster
 import com.datastax.spark.connector.rdd.ReadConf
+import org.mockito.ArgumentCaptor
+import org.mockito.Mockito._
 import org.scalatest.concurrent.Eventually
 
 import scala.concurrent.Future
@@ -48,6 +52,23 @@ class ContinuousPagingScannerSpec extends SparkCassandraITFlatSpecBase with Defa
     }
   }
 
+  private def executeContinuousPagingScan(readConf: ReadConf): Statement[_] = {
+    // we don't want to use the session from CC as mockito is unable to spy on a Proxy
+    val cqlSession = conn.conf.connectionFactory.createSession(conn.conf)
+    try {
+      val sessionSpy = spy(cqlSession)
+      val scanner = ContinuousPagingScanner(readConf, conn.conf, IndexedSeq.empty, sessionSpy)
+      val stmt = sessionSpy.prepare(s"SELECT * FROM $ks.test1").bind()
+      val statementCaptor = ArgumentCaptor.forClass(classOf[Statement[_]])
+
+      scanner.scan(stmt)
+      verify(sessionSpy).executeContinuously(statementCaptor.capture())
+      statementCaptor.getValue
+    } finally {
+      cqlSession.close()
+    }
+  }
+
   "ContinuousPagingScanner" should "re-use a session in the same thread" in {
     val sessions = for (x <- 1 to 10) yield {
       val cps = ContinuousPagingScanner(ReadConf(), conn.conf, IndexedSeq.empty)
@@ -65,7 +86,9 @@ class ContinuousPagingScannerSpec extends SparkCassandraITFlatSpecBase with Defa
 
   it should "use a single CP session for all threads" in {
     CassandraConnector.evictCache()
-    eventually {CassandraConnector.sessionCache.cache.isEmpty}
+    eventually {
+      CassandraConnector.sessionCache.cache.isEmpty
+    }
     val rdd = sc.cassandraTable(ks, table).withReadConf(ReadConf(splitCount = Some(400)))
     rdd.partitions.length should be > 100 //Sanity check that we will have to reuse sessions
     rdd.count
@@ -76,6 +99,39 @@ class ContinuousPagingScannerSpec extends SparkCassandraITFlatSpecBase with Defa
 
     withClue(sessions.map(_.toString).mkString("\n"))(sessions.size should be(1))
   }
-}
 
+  it should "apply MB/s throughput limit" in skipIfNotDSE(conn) {
+    val readConf = ReadConf(throughputMiBPS = Some(32.0))
+    val executedStmt = executeContinuousPagingScan(readConf)
+
+    executedStmt.getExecutionProfile.getBoolean(DseDriverOption.CONTINUOUS_PAGING_PAGE_SIZE_BYTES) should be(true)
+    executedStmt.getExecutionProfile.getInt(DseDriverOption.CONTINUOUS_PAGING_MAX_PAGES_PER_SECOND) should be(1000)
+    executedStmt.getExecutionProfile.getInt(DseDriverOption.CONTINUOUS_PAGING_PAGE_SIZE) should be(33554) // 32MB/s
+  }
+
+  it should "apply reads/s throughput limit" in skipIfNotDSE(conn) {
+    val readConf = ReadConf(fetchSizeInRows = 999, readsPerSec = Some(5))
+    val executedStmt = executeContinuousPagingScan(readConf)
 
+    executedStmt.getExecutionProfile.getBoolean(DseDriverOption.CONTINUOUS_PAGING_PAGE_SIZE_BYTES) should be(false)
+    executedStmt.getExecutionProfile.getInt(DseDriverOption.CONTINUOUS_PAGING_MAX_PAGES_PER_SECOND) should be(5)
+    executedStmt.getExecutionProfile.getInt(DseDriverOption.CONTINUOUS_PAGING_PAGE_SIZE) should be(999)
+  }
+
+  it should "throw a meaningful exception when pages per second does not fall int (0, Int.MaxValue)" in skipIfNotDSE(conn) {
+    val readConfs = Seq(
+      ReadConf(throughputMiBPS = Some(1.0 + Int.MaxValue), readsPerSec = Some(1)),
+      ReadConf(throughputMiBPS = Some(-1)),
+      ReadConf(throughputMiBPS = Some(0)))
+
+    for (readConf <- readConfs) {
+      withClue(s"Expected IllegalArgumentException for invalid throughput argument: ${readConf.throughputMiBPS}.") {
+        val exc = intercept[IllegalArgumentException] {
+          executeContinuousPagingScan(readConf)
+        }
+        exc.getMessage should include(s"This number must be positive, non-zero and smaller than ${Int.MaxValue}")
+      }
+    }
+  }
+
+}

connector/src/main/scala/com/datastax/bdp/spark/ContinuousPagingScanner.scala

@@ -7,8 +7,8 @@ package com.datastax.bdp.spark
 
 import java.io.IOException
 
-import scala.collection.JavaConverters._
-import com.datastax.dse.driver.api.core.cql.continuous.{ContinuousResultSet, ContinuousSession}
+import com.datastax.dse.driver.api.core.config.DseDriverOption
+import com.datastax.dse.driver.api.core.cql.continuous.ContinuousResultSet
 import com.datastax.oss.driver.api.core.CqlSession
 import com.datastax.oss.driver.api.core.`type`.codec.TypeCodec
 import com.datastax.oss.driver.api.core.cql.{BoundStatement, Statement}
@@ -19,62 +19,52 @@ import com.datastax.spark.connector.rdd.ReadConf
 import com.datastax.spark.connector.util.DriverUtil.toName
 import com.datastax.spark.connector.util._
 
+import scala.collection.JavaConverters._
+
 case class ContinuousPagingScanner(
   readConf: ReadConf,
   connConf: CassandraConnectorConf,
-  columnNames: IndexedSeq[String]) extends Scanner with Logging {
+  columnNames: IndexedSeq[String],
+  cqlSession: CqlSession) extends Scanner with Logging {
 
   val TARGET_PAGE_SIZE_IN_BYTES: Int = 5000 * 50 // 5000 rows * 50 bytes per row
   val MIN_PAGES_PER_SECOND = 1000
 
-  //TODO This must be moved to session initilization? We can no longer pass options at execution time without deriving a new profile
-  // I think the right thing to do, to support old configurations as well as new is to create a new profile based on options as
-  // Set using derivied profiles, but this probably can't happen here
-  /**
-  private val cpOptions =  readConf.throughputMiBPS match {
+  private lazy val cpProfile = readConf.throughputMiBPS match {
     case Some(throughput) =>
-      val bytesPerSecond = (throughput * 1024 * 1024).toInt
+      val bytesPerSecond = (throughput * 1024 * 1024).toLong
       val fallBackPagesPerSecond = math.max(MIN_PAGES_PER_SECOND, bytesPerSecond / TARGET_PAGE_SIZE_IN_BYTES)
-      val pagesPerSecond: Int = readConf.readsPerSec.getOrElse(fallBackPagesPerSecond)
+      val pagesPerSecond = readConf.readsPerSec.map(_.toLong).getOrElse(fallBackPagesPerSecond)
       if (readConf.readsPerSec.isEmpty) {
         logInfo(s"Using a pages per second of $pagesPerSecond since " +
           s"${ReadConf.ReadsPerSecParam.name} is not set")
       }
-      val bytesPerPage = (bytesPerSecond / pagesPerSecond ).toInt
+      val bytesPerPage = bytesPerSecond / pagesPerSecond
 
-      if (bytesPerPage <= 0) {
+      if (bytesPerPage <= 0 || bytesPerPage > Int.MaxValue) {
         throw new IllegalArgumentException(
           s"""Read Throttling set to $throughput MBPS, but with the current
              | ${ReadConf.ReadsPerSecParam.name} value of $pagesPerSecond that equates to
-             | $bytesPerPage bytes per page. This number must be positive and non-zero.
+             | $bytesPerPage bytes per page.
+             | This number must be positive, non-zero and smaller than ${Int.MaxValue}.
            """.stripMargin)
       }
 
       logDebug(s"Read Throttling set to $throughput mbps. Pages of $bytesPerPage with ${readConf.readsPerSec} max" +
         s"pages per second. ${ReadConf.FetchSizeInRowsParam.name} will be ignored.")
-      ContinuousPagingOptions
-        .builder()
-        .withPageSize(bytesPerPage, ContinuousPagingOptions.PageUnit.BYTES)
-        .withMaxPagesPerSecond(pagesPerSecond)
-        .build()
+      cqlSession.getContext.getConfig.getDefaultProfile
+        .withBoolean(DseDriverOption.CONTINUOUS_PAGING_PAGE_SIZE_BYTES, true)
+        .withInt(DseDriverOption.CONTINUOUS_PAGING_PAGE_SIZE, bytesPerPage.toInt)
+        .withInt(DseDriverOption.CONTINUOUS_PAGING_MAX_PAGES_PER_SECOND, pagesPerSecond.toInt)
 
     case None =>
-      ContinuousPagingOptions
-        .builder()
-        .withPageSize(readConf.fetchSizeInRows, ContinuousPagingOptions.PageUnit.ROWS)
-        .withMaxPagesPerSecond(readConf.readsPerSec.getOrElse(Integer.MAX_VALUE))
-        .build()
+      cqlSession.getContext.getConfig.getDefaultProfile
+        .withBoolean(DseDriverOption.CONTINUOUS_PAGING_PAGE_SIZE_BYTES, false)
+        .withInt(DseDriverOption.CONTINUOUS_PAGING_PAGE_SIZE, readConf.fetchSizeInRows)
+        .withInt(DseDriverOption.CONTINUOUS_PAGING_MAX_PAGES_PER_SECOND, readConf.readsPerSec.getOrElse(0))
   }
-    **/
-
-  /**
-    * Attempts to get or create a session for this execution thread.
-    */
-  private val cpSession = new CassandraConnector(connConf)
-    .openSession()
-    .asInstanceOf[CqlSession with ContinuousSession]
 
-  private val codecRegistry = cpSession.getContext.getCodecRegistry
+  private val codecRegistry = cqlSession.getContext.getCodecRegistry
 
   private def asBoundStatement(statement: Statement[_]): Option[BoundStatement] = {
     statement match {
@@ -91,22 +81,22 @@ case class ContinuousPagingScanner(
     * Calls SessionProxy Close which issues a deferred close request on the session if no
     * references are requested to it in the next keep_alive ms
     */
-  override def close(): Unit = cpSession.close
+  override def close(): Unit = cqlSession.close()
 
-  override def getSession(): CqlSession = cpSession
+  override def getSession(): CqlSession = cqlSession
 
   override def scan[StatementT <: Statement[StatementT]](statement: StatementT): ScanResult = {
-     val authStatement = maybeExecutingAs(statement, readConf.executeAs)
+    val authStatement = maybeExecutingAs(statement, readConf.executeAs)
 
     if (isSolr(authStatement)) {
       logDebug("Continuous Paging doesn't work with Search, Falling back to default paging")
-      val regularResult = cpSession.execute(authStatement)
+      val regularResult = cqlSession.execute(authStatement)
       val regularIterator = regularResult.iterator().asScala
       ScanResult(regularIterator, CassandraRowMetadata.fromResultSet(columnNames, regularResult, codecRegistry))
 
     } else {
       try {
-        val cpResult = cpSession.executeContinuously(authStatement)
+        val cpResult = cqlSession.executeContinuously(authStatement.setExecutionProfile(cpProfile))
         val cpIterator = cpResult.iterator().asScala
         ScanResult(cpIterator, getMetaData(cpResult))
       } catch {
@@ -120,7 +110,7 @@ case class ContinuousPagingScanner(
     }
   }
 
-  private def getMetaData(result: ContinuousResultSet) = {
+  private def getMetaData(result: ContinuousResultSet): CassandraRowMetadata = {
     import scala.collection.JavaConverters._
     val columnDefs = result.getColumnDefinitions.asScala
     val rsColumnNames = columnDefs.map(c => toName(c.getName))
@@ -131,3 +121,16 @@ case class ContinuousPagingScanner(
     CassandraRowMetadata(columnNames, Some(rsColumnNames.toIndexedSeq), codecs.toIndexedSeq)
   }
 }
+
+object ContinuousPagingScanner {
+  def apply(
+      readConf: ReadConf,
+      connConf: CassandraConnectorConf,
+      columnNames: IndexedSeq[String]): ContinuousPagingScanner = {
+    /**
+      * Attempts to get or create a session for this execution thread.
+      */
+    val cqlSession = new CassandraConnector(connConf).openSession()
+    new ContinuousPagingScanner(readConf, connConf, columnNames, cqlSession)
+  }
+}

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

@@ -187,7 +187,7 @@ object DefaultConnectionFactory extends CassandraConnectionFactory {
 
     if (isContinuousPagingEnabled) {
       logger.debug("Using ContinousPagingScanner")
-      new ContinuousPagingScanner(readConf, connConf, columnNames)
+      ContinuousPagingScanner(readConf, connConf, columnNames)
     } else {
       logger.debug("Not Connected to DSE 5.1 or Greater Falling back to Non-Continuous Paging")
       new DefaultScanner(readConf, connConf, columnNames)

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

@@ -331,14 +331,12 @@ class CassandraTableScanRDD[R] private[connector](
     range: CqlTokenRange[_, _],
     inputMetricsUpdater: InputMetricsUpdater): Iterator[R] = {
 
-    val session = scanner.getSession()
-
     val (cql, values) = tokenRangeToCqlQuery(range)
     logDebug(
       s"Fetching data for range ${range.cql(partitionKeyStr)} " +
         s"with $cql " +
         s"with params ${values.mkString("[", ",", "]")}")
-    val stmt = createStatement(session, cql, values: _*)
+    val stmt = createStatement(scanner.getSession(), cql, values: _*)
       .setRoutingToken(range.range.startNativeToken())
 
     try {