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

import java.io.IOException

import com.datastax.driver.core.{Metadata, ConsistencyLevel, ProtocolVersion, Session}
import com.datastax.spark.connector._
import com.datastax.spark.connector.cql.{TableDef, CassandraConnector, Schema}
import com.datastax.spark.connector.rdd.reader._
import com.datastax.spark.connector.util.NameTools

import scala.reflect.ClassTag

/**
 * Used to get a RowReader of type [R] for transforming the rows of a particular Cassandra table into
 * scala objects. Performs necessary checking of the schema and output class to make sure they are
 * compatible.
 * @see [[CassandraTableScanRDD]]
 * @see [[CassandraJoinRDD]]
 */
trait CassandraTableRowReaderProvider[R] {

  protected def connector: CassandraConnector

  protected def keyspaceName: String

  protected def tableName: String

  protected def columnNames: ColumnSelector

  protected def readConf: ReadConf

  protected def fetchSize: Int = readConf.fetchSize

  protected def splitSize: Int = readConf.splitSize

  protected def consistencyLevel: ConsistencyLevel = readConf.consistencyLevel

  /** RowReaderFactory and ClassTag should be provided from implicit parameters in the constructor
    * of the class implementing this trait
    * @see CassandraTableScanRDD */
  protected val rowReaderFactory: RowReaderFactory[R]

  protected val classTag: ClassTag[R]

  protected lazy val rowReader: RowReader[R] =
    rowReaderFactory.rowReader(tableDef, RowReaderOptions(aliasToColumnName = aliasToColumnName))

  private lazy val aliasToColumnName = columnNames.aliases

  lazy val tableDef: TableDef = {
    Schema.fromCassandra(connector, Some(keyspaceName), Some(tableName)).tables.headOption match {
      case Some(t) => t
      case None =>
        val metadata: Metadata = connector.withClusterDo(_.getMetadata)
        val suggestions = NameTools.getSuggestions(metadata, keyspaceName, tableName)
        val errorMessage = NameTools.getErrorString(keyspaceName, tableName, suggestions)
        throw new IOException(errorMessage)
    }
  }

  protected def checkColumnsExistence(columns: Seq[SelectableColumnRef]): Seq[SelectableColumnRef] = {
    val allColumnNames = tableDef.allColumns.map(_.columnName).toSet
    val regularColumnNames = tableDef.regularColumns.map(_.columnName).toSet

    def checkSingleColumn(column: NamedColumnRef) = {
      if (!allColumnNames.contains(column.columnName))
        throw new IOException(s"Column $column not found in table $keyspaceName.$tableName")

      column match {
        case ColumnName(_, _) =>

        case TTL(columnName, _) =>
          if (!regularColumnNames.contains(columnName))
            throw new IOException(s"TTL can be obtained only for regular columns, " +
              s"but column $columnName is not a regular column in table $keyspaceName.$tableName.")

        case WriteTime(columnName, _) =>
          if (!regularColumnNames.contains(columnName))
            throw new IOException(s"TTL can be obtained only for regular columns, " +
              s"but column $columnName is not a regular column in table $keyspaceName.$tableName.")
      }

      column
    }

    columns.map {
      case namedColumnRef: NamedColumnRef => checkSingleColumn(namedColumnRef)
      case columnRef => columnRef
    }
  }

  /** Returns the names of columns to be selected from the table.*/
  lazy val selectedColumnRefs: Seq[SelectableColumnRef] = {
    val providedColumnRefs =<commit_msg>Update the WriteTime error exception for reader<commit_after>package com.datastax.spark.connector.rdd

import java.io.IOException

import com.datastax.driver.core.{Metadata, ConsistencyLevel, ProtocolVersion, Session}
import com.datastax.spark.connector._
import com.datastax.spark.connector.cql.{TableDef, CassandraConnector, Schema}
import com.datastax.spark.connector.rdd.reader._
import com.datastax.spark.connector.util.NameTools

import scala.reflect.ClassTag

/**
 * Used to get a RowReader of type [R] for transforming the rows of a particular Cassandra table into
 * scala objects. Performs necessary checking of the schema and output class to make sure they are
 * compatible.
 * @see [[CassandraTableScanRDD]]
 * @see [[CassandraJoinRDD]]
 */
trait CassandraTableRowReaderProvider[R] {

  protected def connector: CassandraConnector

  protected def keyspaceName: String

  protected def tableName: String

  protected def columnNames: ColumnSelector

  protected def readConf: ReadConf

  protected def fetchSize: Int = readConf.fetchSize

  protected def splitSize: Int = readConf.splitSize

  protected def consistencyLevel: ConsistencyLevel = readConf.consistencyLevel

  /** RowReaderFactory and ClassTag should be provided from implicit parameters in the constructor
    * of the class implementing this trait
    * @see CassandraTableScanRDD */
  protected val rowReaderFactory: RowReaderFactory[R]

  protected val classTag: ClassTag[R]

  protected lazy val rowReader: RowReader[R] =
    rowReaderFactory.rowReader(tableDef, RowReaderOptions(aliasToColumnName = aliasToColumnName))

  private lazy val aliasToColumnName = columnNames.aliases

  lazy val tableDef: TableDef = {
    Schema.fromCassandra(connector, Some(keyspaceName), Some(tableName)).tables.headOption match {
      case Some(t) => t
      case None =>
        val metadata: Metadata = connector.withClusterDo(_.getMetadata)
        val suggestions = NameTools.getSuggestions(metadata, keyspaceName, tableName)
        val errorMessage = NameTools.getErrorString(keyspaceName, tableName, suggestions)
        throw new IOException(errorMessage)
    }
  }

  protected def checkColumnsExistence(columns: Seq[SelectableColumnRef]): Seq[SelectableColumnRef] = {
    val allColumnNames = tableDef.allColumns.map(_.columnName).toSet
    val regularColumnNames = tableDef.regularColumns.map(_.columnName).toSet

    def checkSingleColumn(column: NamedColumnRef) = {
      if (!allColumnNames.contains(column.columnName))
        throw new IOException(s"Column $column not found in table $keyspaceName.$tableName")

      column match {
        case ColumnName(_, _) =>

        case TTL(columnName, _) =>
          if (!regularColumnNames.contains(columnName))
            throw new IOException(s"TTL can be obtained only for regular columns, " +
              s"but column $columnName is not a regular column in table $keyspaceName.$tableName.")

        case WriteTime(columnName, _) =>
          if (!regularColumnNames.contains(columnName))
            throw new IOException(s"Write time can be obtained only for regular columns, " +
              s"but column $columnName is not a regular column in table $keyspaceName.$tableName.")
      }

      column
    }

    columns.map {
      case namedColumnRef: NamedColumnRef => checkSingleColumn(namedColumnRef)
      case columnRef => columnRef
    }
  }

  /** Returns the names of columns to be selected from the table.*/
  lazy val selectedColumnRefs: Seq[SelectableColumnRef] = {
    val providedColumnRefs =
{'fraction_non_alphanumeric': 0.07887643637395375, 'fraction_numerical': 0.0, 'mean_word_length': 4.0071022727272725, 'pattern_counts': {'":': 0, '<': 3, '<?xml version=': 0, '>': 17, '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': '5223084', 'n_tokens_mistral': 1942, 'n_tokens_neox': 1881, 'n_words': 553}