In the following example, we initialize a FastNLPProcessor object, annotate and tokenize the document into a list of Sentence, zip both the lemma and NER types, and finally return an array of recognized entities for each given sentence.

case class Entity(eType: String, eVal: String)

def processSentence(sentence: Sentence): List[Entity] = {
  val entities = sentence.lemmas.get
    .zip(sentence.entities.get)
    .map {
      case (eVal, eType) =>
        Entity(eType, eVal)
    }
}

def extractEntities(processor: Processor, corpus: String) = {
  val doc = processor.annotate(corpus)
  doc.sentences map processSentence
}

val t = "David Bowie was born in London"
val processor: Processor = new FastNLPProcessor()
val sentences = extractEntities(processor, t)

sentences foreach { sentence =>
  sentence foreach println
}
 
/*
Entity(David,PERSON)
Entity(Bowie,PERSON)
Entity(was,O)
Entity(born,O)
Entity(in,O) 
Entity(London,LOCATION) 
*/

Looking at the above output, you may notice that all the retrieved entities are not linked together, both David and Bowie being two distinct entities of a type [PERSON]. We recursively aggregate consecutive similar entities using the following methods.

def aggregate(entities: Array[Entity]) = {
  aggregateEntities(entities.head, entities.tail, List())
}

def aggregateEntity(e1: Entity, e2: Entity) = {
  Entity(e1.eType, e1.eVal + " " + e2.eVal)
}
 
def aggEntities(current: Entity, entities: Array[Entity], processed : List[Entity]): List[Entity] = {
  if(entities.isEmpty) {
    // End of recusion, no additional entity to process
    // Append our last un-processed entity to our list
    current :: processed
  } else {
    val entity = entities.head
    if(entity.eType == current.eType) {
      // Aggregate consecutive values only of a same entity type      val aggEntity = aggregateEntity(current, entity)
      // Process next record
      aggEntities(aggEntity, entities.tail, processed)
    } else {
      // Add current entity as a candidate for a next aggregation
      // Append our previous un-processed entity to our list      aggEntities(entity, entities.tail, current :: processed)
    }
  }
}
 
def processSentence(sentence: Sentence): List[Entity] = {
  val entities = sentence.lemmas.get
    .zip(sentence.entities.get)
    .map {
      case (eVal, eType) =>
        Entity(eType, eVal)
    }
  aggregate(entities)
}

Printing out the same content now gives us a much more consistent output.

/*
(PERSON,David Bowie)
(O,was born in)
(LOCATION,London) 
*/

We appreciate that working on an array of sentences (sentences being themselves an array of entities) might sound quite blurry. By experience, this will be much more confusing when running at scale, when several flatMap functions will be required for a simple transformation on a RDD. We wrap the results into a class Entities and expose the following methods:

case class Entities(sentences: Array[List[(String, String)]])
 {
 
  def getSentences = sentences

  def getEntities(entity: String) = {
    sentences flatMap { sentence =>
      sentence
    } filter { case (entityType, entityValue) =>
      entityType == entity
    } map { case (entityType, entityValue) =>
      entityValue
    } toSeq
  }

For all these reasons, embedding our code as-is within a map function would be terribly inefficient (and would probably blow all our available heap space). As per the below example, we leverage the mapPartitions pattern in order to optimize both the overhead time of loading and deserializing the models, as well as reducing the amount of memory used by our executors. Using mapPartitions forces the processing of the first record of each partition to evaluate the models inducing the model loading and deserializing process, and all subsequent calls on that executor will reuse those models within that partition, helping to limit the expensive model transfer and initialization costs to once per executor.

def extract(corpusRdd: RDD[String]): RDD[Entities] = {
  corpusRdd mapPartitions {
    case it=>
      val processor = new FastNLPProcessor()
      it map {
        corpus =>
          val entities = extractEntities(processor, corpus)
          new Entities(entities)
      }
    }
  }

The ultimate goal of this NLP scalability problem is to load the least possible number of models while processing as many records as possible. With one executor, we would load the models only once but would totally lose the point of parallel computing. With lots of executors, we will spend much more time deserializing models than actually processing our text content. This is discussed in the performance tuning section.

Because we designed our code locally before integrating it into Spark, we kept in mind writing things in the most convenient way. It is important because scalability is not only how fast you code works in a big data environment, but also how people feel about it, and how efficiently developers interact with your API. As a developer, if you need to chain nested flatMap functions in order to perform what should be a simple transformation, your code simply does not scale! Thanks to our data structure being totally abstracted inside of an Entities class, deriving the different RDDs from our NLP extraction can be done from a simple map function.

val entityRdd: RDD[Entities] = extract(corpusRdd)
entityRdd.persist(StorageLevel.DISK_ONLY)
entityRdd.count()
 
val perRdd = entityRdd.map(_.getEntities("PERSON"))
val locRdd = entityRdd.map(_.getEntities("LOCATION"))
val orgRdd = entityRdd.map(_.getEntities("ORGANIZATION"))

In order to bring this application to scale, you need to ask yourself the same key questions: Is this job I/O, memory, CPU, or network bound? NLP extraction is an expensive task, and loading a model is memory intensive. We may have to reduce the number of executors while allocating much more memory to each of them. To reflect these settings, we need to make sure our dataset will be evenly partitioned using at least as many partitions as the number of executors. We also need to enforce this repartitioning by caching our RDD and calling a simple count action that will evaluate all our previous transformations (including the partitioning itself).

val corpusRdd: RDD[String] = inputRdd.repartition(120)
corpusRdd.cache()
corpusRdd.count()

val entityRdd: RDD[Entities] = extract(corpusRdd)