package mobvista.dmp.main

import org.apache.hadoop.io.{IntWritable, Text}
import org.apache.orc.TypeDescription
import org.apache.orc.mapred.OrcStruct
import org.apache.spark.sql.Row
import org.apache.spark.sql.types.StructType

/**
  * @package: mobvista.dmp.main
  * @author: wangjf
  * @date: 2020-04-23 17:21:58
  * @email: jinfeng.wang@mobvista.com
  * @phone: 152-1062-7698
  */
class CustomIterator(schema: StructType, iter: Iterator[Row], output: String, unMatchOutput: String) extends Iterator[(Text, OrcStruct)] {
  def hasNext: Boolean = {
    iter.hasNext
  }

  def next: (Text, OrcStruct) = {
    val row = iter.next

    val struct = OrcStruct.createValue(TypeDescription.fromString(schema.simpleString)).asInstanceOf[OrcStruct]
    for (elem <- schema.fieldNames) {
      if (row.getAs(elem).isInstanceOf[Int]) {
        struct.setFieldValue(elem, new IntWritable(row.getAs[Int](elem)))
      } else {
        struct.setFieldValue(elem, new Text(row.getAs[String](elem)))
      }
    }
    if (row.getAs[Boolean]("flag")) {
      (new Text(output), struct)
    } else {
      (new Text(unMatchOutput), struct)
    }
  }
}