<commit_before>package com.johnsnowlabs.nlp.datasets

import com.johnsnowlabs.nlp.util.io.ResourceHelper.spark
import com.johnsnowlabs.nlp.{Annotation, AnnotatorType}
import org.apache.spark.sql.{Column, DataFrame}
import org.apache.spark.sql.expressions.UserDefinedFunction
import org.apache.spark.sql.functions.{concat_ws, lit, split, udf}
import org.apache.spark.sql.types.MetadataBuilder

import scala.collection.mutable.ArrayBuffer

case class POS() {

  /*
  * Add Metadata annotationType to output DataFrame
  * NOTE: This should be replaced by an existing function when it's accessible in next release
  * */

  def wrapColumnMetadata(col: Column, annotatorType: String, outPutColName: String): Column = {
    val metadataBuilder: MetadataBuilder = new MetadataBuilder()
    metadataBuilder.putString("annotatorType", annotatorType)
    col.as(outPutColName, metadataBuilder.build)
  }

  /*
   * This section is to help users to convert text files in token|tag style into DataFrame
   * with POS Annotation for training PerceptronApproach
   * */

  private def annotateTokensTags: UserDefinedFunction = udf { (tokens: Seq[String], tags: Seq[String], text: String) =>
    lazy val strTokens = tokens.mkString("#")
    lazy val strPosTags = tags.mkString("#")

    require(tokens.length == tags.length, s"Cannot train from DataFrame since there" +
      s" is a row with different amount of tags and tokens:\n$strTokens\n$strPosTags")

    val tokenTagAnnotation: ArrayBuffer[Annotation] = ArrayBuffer()
    def annotatorType: String = AnnotatorType.POS
    var lastIndex = 0

    for ((e, i) <- tokens.zipWithIndex) {

      val beginOfToken = text.indexOfSlice(e, lastIndex)
      val endOfToken = (beginOfToken + e.length) - 1

      val fullPOSAnnotatorStruct = new Annotation(
        annotatorType = annotatorType,
        begin=beginOfToken,
        end=endOfToken,
        result=tags(i),
        metadata=Map("word" -> e)
      )
      tokenTagAnnotation += fullPOSAnnotatorStruct
      lastIndex = text.indexOfSlice(e, lastIndex)
    }
    tokenTagAnnotation
  }

  private def extractTokensAndTags: UserDefinedFunction = udf { (tokensTags: Seq[String], delimiter: String, condition: String) =>

    val tempArray: ArrayBuffer[String] = ArrayBuffer()

    for (e <- tokensTags.zipWithIndex) {
      val splittedTokenTag: Array[String] = e._1.split(delimiter.mkString)
      if(splittedTokenTag.length > 1){
        condition.mkString match {
          case "token" =>
            tempArray += splittedTokenTag(0)

          case "tag" =>
            tempArray += splittedTokenTag(1)
        }
      }
    }
    tempArray
  }

  def readDataset(path: String, delimiter: String, outPutColName: String): DataFrame = {
    import spark.implicits._
    def annotatorType: String = AnnotatorType.POS

    val tempDataFrame = spark.read.text(path).toDF
      .filter(row => !(row.mkString("").isEmpty && row.length>0))
      .withColumn("token_tags", split($"value", " "))
      .select("token_tags")
      .withColumn("tokens", extractTokensAndTags($"token_tags", lit(delimiter), lit("token")))
      .withColumn("tags", extractTokensAndTags($"token_tags", lit(delimiter), lit("tag")))
      .withColumn("text",  concat_ws(" ", $"tokens"))
      .withColumn(outPutColName, annotateTokensTags($"tokens", $"tags", $"text"))
      .select(outPutColName) // this will also generate ("text", "tokens", "tags")

    tempDataFrame.withColumn(
      outPutColName,
      wrapColumnMetadata(tempDataFrame(outPutColName), annotatorType, outPutColName)
    )
  }

}<commit_msg>Adding readDataframe function and refactoring<commit_after>package com.johnsnowlabs.nlp.datasets

import com.johnsnowlabs.nlp.util.io.ResourceHelper.spark
import com.johnsnowlabs.nlp.{Annotation, AnnotatorType}
import org.apache.spark.sql.{Column, DataFrame}
import org.apache.spark.sql.expressions.UserDefinedFunction
import org.apache.spark.sql.functions.{concat_ws, lit, split, udf, col}
import org.apache.spark.sql.types.MetadataBuilder

import scala.collection.mutable.ArrayBuffer

case class POS() {

  /*
  * Add Metadata annotationType to output DataFrame
  * NOTE: This should be replaced by an existing function when it's accessible in next release
  * */

  def wrapColumnMetadata(col: Column, annotatorType: String, outPutColName: String): Column = {
    val metadataBuilder: MetadataBuilder = new MetadataBuilder()
    metadataBuilder.putString("annotatorType", annotatorType)
    col.as(outPutColName, metadataBuilder.build)
  }

  /*
   * This section is to help users to convert text files in token|tag style into DataFrame
   * with POS Annotation for training PerceptronApproach
   * */

  private def annotateTokensTags: UserDefinedFunction = udf { (tokens: Seq[String], tags: Seq[String], text: String) =>
    lazy val strTokens = tokens.mkString("#")
    lazy val strPosTags = tags.mkString("#")

    require(tokens.length == tags.length, s"Cannot train from DataFrame since there" +
      s" is a row with different amount of tags and tokens:\n$strTokens\n$strPosTags")

    val tokenTagAnnotation: ArrayBuffer[Annotation] = ArrayBuffer()
    def annotatorType: String = AnnotatorType.POS
    var lastIndex = 0

    for ((e, i) <- tokens.zipWithIndex) {

      val beginOfToken = text.indexOfSlice(e, lastIndex)
      val endOfToken = (beginOfToken + e.length) - 1

      val fullPOSAnnotatorStruct = new Annotation(
        annotatorType = annotatorType,
        begin=beginOfToken,
        end=endOfToken,
        result=tags(i),
        metadata=Map("word" -> e)
      )
      tokenTagAnnotation += fullPOSAnnotatorStruct
      lastIndex = text.indexOfSlice(e, lastIndex)
    }
    tokenTagAnnotation
  }

  private def extractTokensAndTags: UserDefinedFunction = udf { (tokensTags: Seq[String], delimiter: String, condition: String) =>

    val tempArray: ArrayBuffer[String] = ArrayBuffer()

    for (e <- tokensTags.zipWithIndex) {
      val splittedTokenTag: Array[String] = e._1.split(delimiter.mkString)
      if(splittedTokenTag.length > 1){
        condition.mkString match {
          case "token" =>
            tempArray += splittedTokenTag(0)

          case "tag" =>
            tempArray += splittedTokenTag(1)
        }
      }
    }
    tempArray
  }

  def readDataset(path: String, delimiter: String, outPutPosColName: String = "tags", outPutDocColName: String = "text"): DataFrame = {
    import spark.implicits._
    def annotatorType: String = AnnotatorType.POS

    val tempDataFrame = spark.read.text(path).toDF
      .filter(row => !(row.mkString("").isEmpty && row.length>0))
      .withColumn("token_tags", split($"value", " "))
      .select("token_tags")
      .withColumn("tokens", extractTokensAndTags($"token_tags", lit(delimiter), lit("token")))
      .withColumn("tags", extractTokensAndTags($"token_tags", lit(delimiter), lit("tag")))
      .withColumn(outPutDocColName,  concat_ws(" ", $"tokens"))
      .withColumn(outPutPosColName, annotateTokensTags($"tokens", $"tags", col(outPutDocColName)))
      .drop("tokens", "token_tags")

    tempDataFrame.withColumn(
      outPutPosColName,
      wrapColumnMetadata(tempDataFrame(outPutPosColName), annotatorType, outPutPosColName)
    )
  }

  // For testing purposes when there is an array of tokens and an array of labels
  def readDatframe(posDataframe: DataFrame, tokensCol: String = "tokens", labelsCol: String = "labels",
                   outPutDocColName: String = "text", outPutPosColName: String = "tags"): DataFrame = {
    def annotatorType: String = AnnotatorType.POS

    val tempDataFrame = posDataframe
      .withColumn(outPutDocColName,  concat_ws(" ", col(tokensCol)))
      .withColumn(outPutPosColName, annotateTokensTags(col(tokensCol), col(labelsCol), col(outPutDocColName)))
      .drop(tokensCol, labelsCol)

    tempDataFrame.withColumn(
      outPutPosColName,
      wrapColumnMetadata(tempDataFrame(outPutPosColName), annotatorType, outPutPosColName)
    )
  }

}
{'fraction_non_alphanumeric': 0.11318626222110804, 'fraction_numerical': 0.0017548257708698922, 'mean_word_length': 3.7921921921921924, 'pattern_counts': {'":': 0, '<': 7, '<?xml version=': 0, '>': 19, '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': '1559455', 'n_tokens_mistral': 2482, 'n_tokens_neox': 2410, 'n_words': 551}