PageRenderTime 54ms CodeModel.GetById 25ms RepoModel.GetById 1ms app.codeStats 0ms

/src/main/scala/org/riedelcastro/cmonnoun/clusterhub/EntityMentionAlignmentService.scala

https://github.com/riedelcastro/cmon-noun
Scala | 102 lines | 72 code | 27 blank | 3 comment | 4 complexity | 94653cd111e7c68e5c99067a13ee790d MD5 | raw file
  1. package org.riedelcastro.cmonnoun.clusterhub
  2. import org.bson.types.ObjectId
  3. import org.riedelcastro.cmonnoun.clusterhub.EntityService.ByIds
  4. import com.mongodb.casbah.commons.MongoDBObject
  5. import com.mongodb.casbah.Imports._
  6. import akka.actor.{ScalaActorRef, Actor}
  7. /**
  8. * @author sriedel
  9. */
  10. class EntityMentionAlignmentService(id: String) extends Actor with MongoSupport with StopWhenMailboxEmpty {
  11. import EntityMentionAlignmentService._
  12. lazy val coll = collFor("entityMentionAlign", id)
  13. coll.ensureIndex("entity")
  14. coll.ensureIndex("mention")
  15. def entityIdsFor(mentionId: Any): TraversableOnce[Any] = {
  16. for (dbo <- coll.find(MongoDBObject("mention" -> mentionId))) yield {
  17. dbo.as[String]("entity")
  18. }
  19. }
  20. def entities(): TraversableOnce[Any] = {
  21. for (dbo <- coll.find()) yield {
  22. dbo.as[String]("entity")
  23. }
  24. }
  25. def mentionIdsFor(entityId: Any): TraversableOnce[Any] = {
  26. for (dbo <- coll.find(MongoDBObject("entity" -> entityId))) yield {
  27. dbo.as[Any]("mention")
  28. }
  29. }
  30. def mentionsIdsFor(entityIds: Stream[Any]): TraversableOnce[Alignment] = {
  31. for (dbo <- coll.find(MongoDBObject("entity" -> MongoDBObject("$in" -> entityIds)))) yield {
  32. val m = dbo.as[Any]("mention")
  33. val e = dbo.as[Any]("entity")
  34. Alignment(m,e)
  35. }
  36. }
  37. def storeAlignment(mentionId: Any, entityId: Any) {
  38. val dbo = MongoDBObject(
  39. "mention" -> mentionId,
  40. "entity" -> entityId)
  41. coll += dbo
  42. }
  43. protected def receive = {
  44. stopWhenMailboxEmpty orElse {
  45. case StoreAlignment(alignment) =>
  46. storeAlignment(alignment.mentionId, alignment.entityId)
  47. case GetEntities(m, mentionId) =>
  48. val entityIds = entityIdsFor(mentionId)
  49. m.forward(EntityService.Query(ByIds(entityIds.toSeq)))
  50. case GetEntitiesWithMentions =>
  51. self.channel ! EntityIds(entities())
  52. case GetEntityIds(mentionId) =>
  53. self.channel ! EntityIds(entityIdsFor(mentionId))
  54. case GetAlignments(entityIds) =>
  55. self.channel ! Alignments(mentionsIdsFor(entityIds).map(a=>a.mentionId -> a.entityId).toMap)
  56. case GetMentions(m,entityIds) =>
  57. val mentionIds = mentionsIdsFor(entityIds)
  58. m.forward(EntityMentionService.Query(EntityMentionService.ByIds(mentionIds.toStream)))
  59. }
  60. }
  61. }
  62. object EntityMentionAlignmentService {
  63. type EntityId = Any
  64. type EntityMentionId = Any
  65. case class Alignment(mentionId: Any, entityId: Any)
  66. case class Alignments(alignments:Map[EntityMentionId,EntityId])
  67. case class StoreAlignment(alignment:Alignment)
  68. case class GetEntityIds(mentionId: Any)
  69. case class GetEntities(entityService: ScalaActorRef, mentionId: Any)
  70. case class GetAlignments(entityIds: Stream[Any])
  71. case class GetMentions(mentionService: ScalaActorRef, entityIds: Stream[Any])
  72. case object GetEntitiesWithMentions
  73. case class EntityIds(entityIds: TraversableOnce[Any])
  74. case class MentionIds(mentionIds:TraversableOnce[Any])
  75. }