Scala col类代码示例

//设置package包名称以及导入依赖的类
package me.danielpes.spark.datetime

import org.apache.spark.sql.Column
import org.apache.spark.sql.types.{DataType, DateType, TimestampType}
import org.apache.spark.sql.functions.{col, udf}

class DateTimeColumn(val col: Column, dataType: DataType = TimestampType) {

  def +(p: Period): Column = dataType match {
    case _: DateType => udf((d: java.sql.Date) => new RichDate(d) + p).apply(col)
    case _: TimestampType => udf((ts: java.sql.Timestamp) => new RichDate(ts) + p).apply(col)
  }

  def -(p: Period): Column = this.+(-p)

  override def toString: String = s"{column: ${col.toString}, type: ${dataType.toString}}"
}

object DateTimeColumn {

  def apply(col: Column, dataType: DataType = TimestampType) = new DateTimeColumn(col, dataType)
  def apply(col: Column, typeString: String) = new DateTimeColumn(col, typeFromString(typeString))
  def apply(cName: String) = new DateTimeColumn(col(cName), TimestampType)
  def apply(cName: String, dataType: DataType) = new DateTimeColumn(col(cName), dataType)
  def apply(cName: String, typeString: String) = new DateTimeColumn(col(cName), typeFromString(typeString))

  private def typeFromString(s: String): DataType = s match {
    case "date" => DateType
    case "timestamp" => TimestampType
  }
} 

//设置package包名称以及导入依赖的类
package com.koverse.example.spark

import org.apache.spark.rdd.RDD
import com.koverse.sdk.data.SimpleRecord
import org.apache.spark.sql.DataFrame
import org.apache.spark.sql.functions.col
import org.apache.spark.sql.functions.lower

class WordCounter(
    textFieldName: String,
    tokenizationString: String) extends java.io.Serializable {
  
  def count(inputRecordsRdd: RDD[SimpleRecord]): RDD[SimpleRecord] = {
    
    // for each Record, tokenize the specified text field and count each occurrence
    val wordCountRdd = inputRecordsRdd.flatMap { record => record.get(textFieldName).toString().split(tokenizationString) }
                           .map { token => token.toLowerCase().trim() }
                           .map { token => (token, 1) }
                           .reduceByKey { (a,b) => a + b }

    // wordCountRdd is an RDD[(String, Int)] so a (word,count) tuple.
    // turn each tuple into an output Record with a "word" and "count" fields
    val outputRdd = wordCountRdd.map { case(word, count) => {

      val record = new SimpleRecord()
      record.put("word", word)
      record.put("count", count)
      record
    }}
    
    outputRdd
  }
  
  def count(inputDataFrame: DataFrame): DataFrame = {
    
     // Take the column that contains the text and tokenize and count the words
    val wordDF = inputDataFrame.explode(textFieldName, "word") { (text: String) => text.split(tokenizationString) }
    wordDF.select(lower(col("word")).as("lowerWord"))
          .groupBy("lowerWord")
          .count()
  }
  
} 

//设置package包名称以及导入依赖的类
package spark.feature

import org.apache.spark.ml.Transformer
import org.apache.spark.ml.attribute.AttributeGroup
import org.apache.spark.ml.param.{ParamMap, _}
import org.apache.spark.ml.util._
import org.apache.spark.sql.functions.{col, udf}
import org.apache.spark.sql.types.StructType
import org.apache.spark.sql.{DataFrame, UserDefinedFunction}


  def setFunction(value: String=>Double) = set(function, value)
  def getFunction() =  $(function)

  override def transform(dataset: DataFrame): DataFrame = {
    val outputSchema = transformSchema(dataset.schema)
    val metadata = outputSchema($(outputCol)).metadata
    val dummy = udf { x: Any => $(expr) }
    var data = dataset.select(col("*"), dummy(col($(inputCols).head)).as("0"))
    val substitute: (String => ((String, Double) => String)) = name => (exp, elem) => exp.replace(name, elem.toString)
    def subst(v: String) = udf(substitute(v))
    $(inputCols).view.zipWithIndex foreach { case (v, i) => data = data.select(col("*"), subst(v)(data(i.toString), data(v)).as((i + 1).toString)).drop(i.toString) }
    val eval = udf($(function))
    data.select(col("*"), eval(data($(inputCols).length.toString)).as($(outputCol), metadata)).drop($(inputCols).length.toString)
  }


  override def transformSchema(schema: StructType): StructType = {
    // TODO: Assertions on inputCols
    val attrGroup = new AttributeGroup($(outputCol), $(numFeatures))
    val col = attrGroup.toStructField()
    require(!schema.fieldNames.contains(col.name), s"Column ${col.name} already exists.")
    StructType(schema.fields :+ col)
  }

  override def copy(extra: ParamMap): FeatureFuTransformer = defaultCopy(extra)
}