org.elasticsearch.spark.sql.DataFrameFieldExtractor.scala Maven / Gradle / Ivy

Go to download

Show more of this group Show more artifacts with this name
Show all versions of elasticsearch-spark-13_2.11 Show documentation

Elasticsearch Spark (for Spark 1.3-1.6)

There is a newer version: 7.0.0-beta1

Show newest version

package org.elasticsearch.spark.sql

import org.apache.spark.sql.Row
import org.apache.spark.sql.types.StructType
import org.elasticsearch.hadoop.serialization.field.FieldExtractor
import org.elasticsearch.spark.serialization.ScalaMapFieldExtractor

class DataFrameFieldExtractor extends ScalaMapFieldExtractor {

  override protected def extractField(target: AnyRef): AnyRef = {
    var obj = target
    for (in <- 0 until getFieldNames.size()) {
      val field = getFieldNames.get(in)
      obj = obj match {
        case (row: Row, struct: StructType) => {
          val index = struct.fieldNames.indexOf(field)
          if (index < 0) {
            FieldExtractor.NOT_FOUND
          } else {
            row(index).asInstanceOf[AnyRef]
          }
        }
        case _ => super.extractField(target)
      }
    }
    return obj
  }
}