<commit_before>package com.datastax.spark.connector.rdd.partitioner

import java.util.concurrent.CompletableFuture

import com.datastax.oss.driver.api.core.metadata.token.Token
import com.datastax.spark.connector.cluster.DefaultCluster
import com.datastax.spark.connector.cql.{CassandraConnector, Schema}
import com.datastax.spark.connector.writer.RowWriterFactory
import com.datastax.spark.connector.{PartitionKeyColumns, SparkCassandraITFlatSpecBase}

import scala.collection.JavaConversions._
import scala.concurrent.Future

class TokenGeneratorSpec extends SparkCassandraITFlatSpecBase with DefaultCluster {

  case class SimpleKey(key: Int)
  case class ComplexKey(key1: Int, key2: Int, key3: String)

  override lazy val conn = CassandraConnector(defaultConf)

  val simpleKeys = Seq(1, -500, 2801, 4000000, -95500).map(SimpleKey(_))
  val complexKey = simpleKeys.map{ case SimpleKey(x) => ComplexKey(x, x, x.toString)}

  conn.withSessionDo { session =>
    createKeyspace(session)
    awaitAll(
      Future{session.execute(s"CREATE TABLE $ks.simple(key INT PRIMARY KEY)")},
      Future{session.execute(s"CREATE TABLE $ks.complex(key1 INT, key2 INT, key3 text, PRIMARY KEY ((key1, key2, key3)))")}
    )
    val simplePS = session.prepare(s"INSERT INTO $ks.simple (key) VALUES (?)")
    val complexPS = session.prepare(s"INSERT INTO $ks.complex (key1, key2, key3) VALUES (?, ?, ?)")
    val insertFutures =
      simpleKeys.map( simpleKey =>
        session.executeAsync(simplePS.bind(simpleKey.key: java.lang.Integer)).toCompletableFuture) ++
      complexKey.map( complexKey =>
        session.executeAsync(complexPS.bind(
          complexKey.key1: java.lang.Integer,
          complexKey.key2: java.lang.Integer,
          complexKey.key3: java.lang.String)).toCompletableFuture)
    CompletableFuture.allOf(insertFutures: _*)
  }

  val simpleTokenMap: Map[SimpleKey, Token] = conn.withSessionDo { session =>
    val resultSet = session.execute(s"SELECT key, TOKEN(key) FROM $ks.simple").all()
    resultSet.map(row => SimpleKey(row.getInt("key")) -> row.getToken("TOKEN(key)")).toMap
  }

  val complexTokenMap: Map[ComplexKey, Token] = conn.withSessionDo { session =>
    val resultSet = session
      .execute(s"SELECT key1, key2, key3, TOKEN(key1, key2, key3) FROM $ks.complex")
      .all()

    resultSet.map(row =>
      ComplexKey(row.getInt("key1"), row.getInt("key2"), row.getString("key3")) ->
        row.getToken("TOKEN(key1, key2, key3)")
    ).toMap
  }

  val simpleTableDef = Schema.fromCassandra(conn, Some(ks), Some("simple")).tables.head
  val complexTableDef = Schema.fromCassandra(conn, Some(ks), Some("complex")).tables.head

  val simpleRW = implicitly[RowWriterFactory[SimpleKey]]
    .rowWriter(simpleTableDef, PartitionKeyColumns.selectFrom(simpleTableDef))

  val complexRW = implicitly[RowWriterFactory[ComplexKey]]
    .rowWriter(complexTableDef, PartitionKeyColumns.selectFrom(complexTableDef))

  "TokenGenerators" should "be able to determine simple partition keys" in {

    val tokenGenerator = new TokenGenerator[SimpleKey](conn, simpleTableDef, simpleRW)
    for ((key, token) <- simpleTokenMap) {
      tokenGenerator.getTokenFor(key) should be (token)
    }
  }

  it should "be able to determine composite keys" in {
    val tokenGenerator = new TokenGenerator[ComplexKey](conn, complexTableDef, complexRW)
    for ((key, token) <- complexTokenMap) {
      tokenGenerator.getTokenFor(key) should be (token)
    }
  }<commit_msg>DSP-18857 fix TokenGeneratorSpec<commit_after>package com.datastax.spark.connector.rdd.partitioner

import java.util.concurrent.CompletableFuture

import com.datastax.oss.driver.api.core.metadata.token.Token
import com.datastax.spark.connector.cluster.DefaultCluster
import com.datastax.spark.connector.cql.{CassandraConnector, Schema}
import com.datastax.spark.connector.writer.RowWriterFactory
import com.datastax.spark.connector.{PartitionKeyColumns, SparkCassandraITFlatSpecBase}

import scala.collection.JavaConversions._
import scala.concurrent.Future

class TokenGeneratorSpec extends SparkCassandraITFlatSpecBase with DefaultCluster {

  case class SimpleKey(key: Int)
  case class ComplexKey(key1: Int, key2: Int, key3: String)

  override lazy val conn = CassandraConnector(defaultConf)

  val simpleKeys = Seq(1, -500, 2801, 4000000, -95500).map(SimpleKey(_))
  val complexKey = simpleKeys.map{ case SimpleKey(x) => ComplexKey(x, x, x.toString)}

  conn.withSessionDo { session =>
    createKeyspace(session)
    awaitAll(
      Future{session.execute(s"CREATE TABLE $ks.simple(key INT PRIMARY KEY)")},
      Future{session.execute(s"CREATE TABLE $ks.complex(key1 INT, key2 INT, key3 text, PRIMARY KEY ((key1, key2, key3)))")}
    )
    val simplePS = session.prepare(s"INSERT INTO $ks.simple (key) VALUES (?)")
    val complexPS = session.prepare(s"INSERT INTO $ks.complex (key1, key2, key3) VALUES (?, ?, ?)")
    val insertFutures =
      simpleKeys.map( simpleKey =>
        session.executeAsync(simplePS.bind(simpleKey.key: java.lang.Integer)).toCompletableFuture) ++
      complexKey.map( complexKey =>
        session.executeAsync(complexPS.bind(
          complexKey.key1: java.lang.Integer,
          complexKey.key2: java.lang.Integer,
          complexKey.key3: java.lang.String)).toCompletableFuture)
    CompletableFuture.allOf(insertFutures: _*)
  }

  val simpleTokenMap: Map[SimpleKey, Token] = conn.withSessionDo { session =>
    val resultSet = session.execute(s"SELECT key, TOKEN(key) FROM $ks.simple").all()
    resultSet.map(row => SimpleKey(row.getInt("key")) -> row.getToken("system.TOKEN(key)")).toMap
  }

  val complexTokenMap: Map[ComplexKey, Token] = conn.withSessionDo { session =>
    val resultSet = session
      .execute(s"SELECT key1, key2, key3, TOKEN(key1, key2, key3) FROM $ks.complex")
      .all()

    resultSet.map(row =>
      ComplexKey(row.getInt("key1"), row.getInt("key2"), row.getString("key3")) ->
        row.getToken("system.TOKEN(key1, key2, key3)")
    ).toMap
  }

  val simpleTableDef = Schema.fromCassandra(conn, Some(ks), Some("simple")).tables.head
  val complexTableDef = Schema.fromCassandra(conn, Some(ks), Some("complex")).tables.head

  val simpleRW = implicitly[RowWriterFactory[SimpleKey]]
    .rowWriter(simpleTableDef, PartitionKeyColumns.selectFrom(simpleTableDef))

  val complexRW = implicitly[RowWriterFactory[ComplexKey]]
    .rowWriter(complexTableDef, PartitionKeyColumns.selectFrom(complexTableDef))

  "TokenGenerators" should "be able to determine simple partition keys" in {

    val tokenGenerator = new TokenGenerator[SimpleKey](conn, simpleTableDef, simpleRW)
    for ((key, token) <- simpleTokenMap) {
      tokenGenerator.getTokenFor(key) should be (token)
    }
  }

  it should "be able to determine composite keys" in {
    val tokenGenerator = new TokenGenerator[ComplexKey](conn, complexTableDef, complexRW)
    for ((key, token) <- complexTokenMap) {
      tokenGenerator.getTokenFor(key) should be (token)
    }
  }
{'fraction_non_alphanumeric': 0.12248106331284836, 'fraction_numerical': 0.014148920966128341, 'mean_word_length': 4.580542264752791, 'pattern_counts': {'":': 0, '<': 7, '<?xml version=': 0, '>': 23, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '381313', 'n_tokens_mistral': 2268, 'n_tokens_neox': 2177, 'n_words': 448}