Merge pull request #1302 from yohannrub/SPARKC-642-b2.5

SPARKC-642 repartitionByCassandraReplica relocates data to the local node only

Author: jtgrabowski

connector/src/it/scala/com/datastax/spark/connector/cluster/Fixtures.scala

@@ -187,6 +187,15 @@ trait TwoClustersWithOneNode extends Fixture {
   override def cluster(c: Int): Cluster = ClusterHolder.get(this)(c)
 }
 
+/** A fixture that bootstraps on single cluster with two nodes. */
+trait TwoNodeCluster extends SingleClusterFixture {
+
+  private[cluster] final override val configs: Seq[CcmConfig] = Seq(defaultConfig.copy(nodes = Seq(1, 2)))
+
+  private[cluster] override def connectionParameters(address: InetSocketAddress): Map[String, String] =
+    DefaultCluster.defaultConnectionParameters(address)
+}
+
 trait CETCluster extends DefaultCluster
 
 trait CSTCluster extends DefaultCluster

connector/src/it/scala/com/datastax/spark/connector/rdd/ReplicaRepartitionedCassandraRDDSpec.scala

@@ -0,0 +1,85 @@
+package com.datastax.spark.connector.rdd
+
+import java.lang.{Long => JLong}
+
+import com.datastax.spark.connector._
+import com.datastax.spark.connector.cluster.TwoNodeCluster
+import com.datastax.spark.connector.cql.CassandraConnector
+import com.datastax.spark.connector.rdd.partitioner.EndpointPartition
+
+import scala.concurrent.Future
+
+class ReplicaRepartitionedCassandraRDDSpec extends SparkCassandraITFlatSpecBase with TwoNodeCluster {
+
+  override lazy val conn = CassandraConnector(defaultConf)
+  val tableName = "key_value"
+  val keys = 0 to 200
+  val total = 0 to 10000
+
+  override def beforeClass {
+    conn.withSessionDo { session =>
+      createKeyspace(session)
+      val startTime = System.currentTimeMillis()
+
+      val executor = getExecutor(session)
+
+      awaitAll(
+        Future {
+          session.execute(
+            s"""
+               |CREATE TABLE $ks.$tableName (
+               |  key INT,
+               |  group BIGINT,
+               |  value TEXT,
+               |  PRIMARY KEY (key, group)
+               |)""".stripMargin)
+          val ps = session
+            .prepare(s"""INSERT INTO $ks.$tableName (key, group, value) VALUES (?, ?, ?)""")
+          awaitAll((for (value <- total) yield
+            executor.executeAsync(ps.bind(value: Integer, (value * 100).toLong: JLong, value.toString))): _*)
+        }
+      )
+      executor.waitForCurrentlyExecutingTasks()
+      println(s"Took ${(System.currentTimeMillis() - startTime) / 1000.0} Seconds to setup Suite Data")
+    }
+  }
+
+  def checkArrayCassandraRow[T](result: Array[(T, CassandraRow)]) = {
+    markup("Checking RightSide Join Results")
+    result.length should be(keys.length)
+    for (key <- keys) {
+      val sorted_result = result.map(_._2).sortBy(_.getInt(0))
+      sorted_result(key).getInt("key") should be(key)
+      sorted_result(key).getLong("group") should be(key * 100)
+      sorted_result(key).getString("value") should be(key.toString)
+    }
+  }
+
+  "A Tuple RDD specifying partition keys" should "be repartitionable" in {
+    val source = sc.parallelize(keys).map(Tuple1(_))
+    val repart = source.repartitionByCassandraReplica(ks, tableName, 10)
+    repart.partitions.length should be(conn.hosts.size * 10)
+    conn.hosts.size should be(2)
+    conn.hosts should be(cluster.addresses.toSet)
+    val someCass = repart.joinWithCassandraTable(ks, tableName)
+    someCass.partitions.foreach {
+      case e: EndpointPartition =>
+        conn.hostAddresses should contain(e.endpoints.head)
+      case _ =>
+        fail("Unable to get endpoints on repartitioned RDD, This means preferred locations will be broken")
+    }
+    val result = someCass.collect
+    checkArrayCassandraRow(result)
+  }
+
+  it should "be deterministically repartitionable" in {
+    val source = sc.parallelize(keys).map(Tuple1(_))
+    val repartRDDs = (1 to 10).map(_ =>
+      source
+        .repartitionByCassandraReplica(ks, tableName, 10)
+        .mapPartitionsWithIndex((index, it) => it.map((_, index))))
+    val first = repartRDDs(1).collect
+    repartRDDs.foreach(rdd => rdd.collect should be(first))
+  }
+
+}

connector/src/main/scala/com/datastax/spark/connector/RDDFunctions.scala

@@ -237,31 +237,6 @@ class RDDFunctions[T](rdd: RDD[T]) extends WritableToCassandra[T] with Serializa
     currentType: ClassTag[T],
     rwf: RowWriterFactory[T]): CassandraPartitionedRDD[T] = {
 
-    val replicaLocator = ReplicaLocator[T](connector, keyspaceName, tableName, partitionKeyMapper)
-    rdd.repartitionByCassandraReplica(
-      replicaLocator,
-      keyspaceName,
-      tableName,
-      partitionsPerHost,
-      partitionKeyMapper)
-  }
-
-
-  /**
-   * A Serializable version of repartitionByCassandraReplica which removes
-   * the implicit RowWriterFactory Dependency
-   */
-  private[connector] def repartitionByCassandraReplica(
-    replicaLocator: ReplicaLocator[T],
-    keyspaceName: String,
-    tableName: String,
-    partitionsPerHost: Int,
-    partitionKeyMapper: ColumnSelector)(
-  implicit
-    connector: CassandraConnector,
-    currentType: ClassTag[T],
-    rwf: RowWriterFactory[T]): CassandraPartitionedRDD[T] = {
-
     val partitioner = new ReplicaPartitioner[T](
       tableName,
       keyspaceName,

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

@@ -5,14 +5,13 @@ import java.net.{InetAddress, InetSocketAddress}
 
 import com.datastax.oss.driver.api.core.CqlSession
 import com.datastax.oss.driver.api.core.loadbalancing.NodeDistance
-import com.datastax.spark.connector.cql.CassandraConnectorConf.CassandraSSLConf
 import com.datastax.spark.connector.types.TypeAdapters.{ValueByNameAdapter, ValuesSeqAdapter}
 import com.datastax.spark.connector.types.{NullableTypeConverter, TypeConverter}
 import com.datastax.spark.connector.util.ConfigCheck.ConnectorConfigurationException
 import com.datastax.spark.connector.util.DriverUtil.toAddress
-import com.datastax.spark.connector.util.{DriverUtil, Logging, SerialShutdownHooks}
+import com.datastax.spark.connector.util.{Logging, SerialShutdownHooks}
 import org.apache.spark.sql.catalyst.expressions.GenericRowWithSchema
-import org.apache.spark.{SparkConf, SparkContext, SparkFiles}
+import org.apache.spark.{SparkConf, SparkContext}
 
 import scala.collection.JavaConverters._
 import scala.language.reflectiveCalls
@@ -65,14 +64,8 @@ class CassandraConnector(val conf: CassandraConnectorConf)
   def hosts: Set[InetSocketAddress] =
     // wrapped in a session, so we get full lists of hosts,
     // not only those explicitly passed in the conf
-    withSessionDo {
-      _.getMetadata
-        .getNodes
-        .values
-        .asScala
-        .filter(_.getDistance == NodeDistance.LOCAL)
-        .flatMap(DriverUtil.toAddress)
-        .toSet
+    withSessionDo { session =>
+      dataCenterNodes(session)
     }
 
   private[connector] def hostAddresses: Set[InetAddress] = hosts.map(_.getAddress)
@@ -187,14 +180,14 @@ object CassandraConnector extends Logging {
     logInfo(s"Disconnected from Cassandra cluster.")
   }
 
-  private def dataCenterNodes(conf: CassandraConnectorConf, ipConf: IpBasedContactInfo, session: CqlSession): Set[InetSocketAddress] = {
+  // LocalNodeFirstLoadBalancingPolicy assigns LOCAL or REMOTE (i.e. non-IGNORED) distance to local DC nodes
+  private def dataCenterNodes(session: CqlSession): Set[InetSocketAddress] = {
     val allNodes = session.getMetadata.getNodes.asScala.values.toSet
-    val dcToUse = conf.localDC.getOrElse(LocalNodeFirstLoadBalancingPolicy.determineDataCenter(ipConf.hosts, allNodes))
     val nodes = allNodes
-      .collect { case n if n.getDatacenter == dcToUse => toAddress(n) }
-      .flatten
+      .filter(_.getDistance != NodeDistance.IGNORED)
+      .flatMap(toAddress)
     if (nodes.isEmpty) {
-      throw new ConnectorConfigurationException(s"Could not determine suitable nodes for DC: $dcToUse and known nodes: " +
+      throw new ConnectorConfigurationException(s"Could not determine suitable nodes in local DC for known nodes: " +
         s"${allNodes.map(n => (n.getHostId, toAddress(n))).mkString(", ")}")
     }
     nodes
@@ -204,7 +197,7 @@ object CassandraConnector extends Logging {
   private def alternativeConnectionConfigs(conf: CassandraConnectorConf, session: CqlSession): Set[CassandraConnectorConf] = {
     conf.contactInfo match {
       case ipConf: IpBasedContactInfo =>
-        val nodes = dataCenterNodes(conf, ipConf, session)
+        val nodes = dataCenterNodes(session)
         nodes.map(n => conf.copy(contactInfo = ipConf.copy(hosts = Set(n)))) + conf.copy(contactInfo = ipConf.copy(hosts = nodes))
       case _ => Set.empty
     }