com.datastax.spark.connector.mapper.DataFrameColumnMapper.scala Maven / Gradle / Ivy

Go to download

Show more of this group Show more artifacts with this name
Show all versions of spark-cassandra-connector_2.11 Show documentation

A library that exposes YugaByte DB tables as Spark RDDs, writes Spark RDDs to YugaByte DB tables, and executes CQL queries in Spark applications using YugaByte DB's Cassandra-compatible API.

There is a newer version: 2.4-yb-4

Show newest version

package com.datastax.spark.connector.mapper

import com.datastax.driver.core.ProtocolVersion
import com.datastax.spark.connector.ColumnRef
import com.datastax.spark.connector.cql._
import com.datastax.spark.connector.types.ColumnType
import org.apache.spark.sql.types.StructType

class DataFrameColumnMapper[T](structType: StructType) extends ColumnMapper[T] {
  override def columnMapForWriting(struct: StructDef,
                                   selectedColumns: IndexedSeq[ColumnRef]): ColumnMapForWriting = ???

  override def columnMapForReading(struct: StructDef,
                                   selectedColumns: IndexedSeq[ColumnRef]): ColumnMapForReading = ???

  override def newTable(
    keyspaceName: String,
    tableName: String,
    protocolVersion: ProtocolVersion = ProtocolVersion.NEWEST_SUPPORTED): TableDef = {

    val columns = structType.zipWithIndex.map { case (field, i) => {
      val columnRole = if (i == 0) PartitionKeyColumn else RegularColumn
      ColumnDef(field.name, columnRole, ColumnType.fromSparkSqlType(field.dataType, protocolVersion))
    }}

    TableDef(keyspaceName, tableName, Seq(columns.head), Seq.empty, columns.tail)
  }
}