Simple search engine using TF-IDF (200 LOC)

I wrote simple search engine based on TF-IDF statistic and cosine similarity. I am still learning Scala, so would be glad to hear comments (here or at GitHub) regarding my code from Scala professionals.

package search

package object types {
type Word = String
type DocumentTitle = String
type SearchResult = (DocumentTitle, Double)
}

import search.types.{DocumentTitle, SearchResult, Word}

import scala.collection.mutable.ListBuffer

/**
* Representation of a document as mapping from a word to number of instances in the document.
* When you create and ad-hoc vector, which doesn't have document title assigned you can use None
* as document title.
*/
case class WordVector(private val sentence: Traversable[Word],
documentTitle: Option[DocumentTitle]) extends Traversable[(Word, Int)] {
private val wordCount: Map[Word, Int] = sentence.groupBy(identity).mapValues(_.size)

val totalWordCount: Int = wordCount map (_._2) sum

def countOccurrencesOf(word: Word): Int = wordCount.getOrElse(word, 0)

override def foreach[U](f: ((Word, Int)) => U): Unit = wordCount.foreach(f)
}

case class TfidfVector(private val data: Map[Word, Double],
documentTitle: Option[DocumentTitle]) {
val length: Double = Math.sqrt(data.values.map(a => a * a).sum)

def apply(word: Word): Double = data.getOrElse(word, 0)

def commonWordsWith(other: TfidfVector) = data.keySet.intersect(other.data.keySet)
}

class IndexedDocuments(documentsSpace: Seq[WordVector],
docsContainingWord: Map[Word, Set[DocumentTitle]]
) {
val docCount = documentsSpace.size

/**
* Think about a corpus as a space where every document is represented as a vector. This space has
* M dimensions, where M is the number of unique words in all documents. Naturally, the name/label
* of the vector is the title of the document.
*
* We need named vectors because at the end of the calculations we need to show the titles
* of the documents to the user.
*/
val documentsAsTfidfSpace: Seq[TfidfVector] = documentsSpace map wordVectorToTfIdfVector toSeq

def wordVectorToTfIdfVector(wordVector: WordVector): TfidfVector = {
val data: Map[Word, Double] = wordVector map {
case (word, _) => (word, tfidf(word, wordVector))
} toMap

TfidfVector(data, wordVector.documentTitle)
}

def tfidf(word: Word, wordVector: WordVector): Double = {
if (documentsSpace.isEmpty) {
0
}
else {
val occurrencesInDoc: Double = wordVector.countOccurrencesOf(word).toDouble
val tf = occurrencesInDoc / wordVector.totalWordCount
val numDocsContainingWord = docsContainingWord.getOrElse(word, Seq.empty).size
if (numDocsContainingWord == 0) {
0
} else {
val idf = docCount / numDocsContainingWord.toDouble
val tfidf = tf * Math.log(idf)
tfidf
}
}
}

/**
* Calculates cosine similarity between two documents
*/
def compareWithQuery(vectorFromUser: TfidfVector)(vectorFromCorpus: TfidfVector): (DocumentTitle, Double) = {
val vec = vectorFromCorpus

val commonWords = vectorFromUser.commonWordsWith(vec)
val numerator = commonWords map (word => vectorFromUser(word) * vec(word)) sum
val denominator = vectorFromUser.length * vec.length
(vec.documentTitle.get, numerator / denominator)
}

/**
* @param sentence Has to be normalized (e.g. lowercased)
*/
def search(sentence: Traversable[Word], topN: Int): Seq[SearchResult] = {
require(topN > 0, s"Top N has to be greater than 0 but was $topN") val queryTfidfVector = wordVectorToTfIdfVector(WordVector(sentence, None)) val scoredDocuments: Seq[(DocumentTitle, Double)] = documentsAsTfidfSpace.map(compareWithQuery(queryTfidfVector)) scoredDocuments .sortWith(_._2 > _._2) // by score descending .filterNot(_._2.isNaN) .filterNot(_._2 < 0.000001) .take(topN) // .map(_._1) // keep only title } } object Indexer { def indexDocuments(documents: Iterator[(DocumentTitle, List[Word])]): IndexedDocuments = { import scala.collection.mutable val documentsSpace = new ListBuffer[WordVector] val docsContainingWord = new mutable.HashMap[Word, Set[DocumentTitle]] documents.foreach { case (title, words) => documentsSpace += WordVector(words, Some(title)) words.foreach { word => val currentDocsWithThisWord = docsContainingWord.getOrElse(word, Set.empty) docsContainingWord += word -> (currentDocsWithThisWord + title) } } new IndexedDocuments( documentsSpace.toVector, docsContainingWord.toMap ) } } package search import search.types.{DocumentTitle, Word} object SearchApp extends App { require(args.length > 0, "Usage: As an argument you should pass the path to the folder containing documents (program will search through subdirectories") val topNResults: Int = 10 val extractWords: (String) => Seq[String] = s => s.split("\\W+").map(_.toLowerCase).filterNot(_.trim.isEmpty) val documentLoader: Iterator[(DocumentTitle, List[Word])] = { import java.io.File def recursiveListFiles(f: File): Array[File] = { // src: http://stackoverflow.com/questions/2637643/how-do-i-list-all-files-in-a-subdirectory-in-scala val these = f.listFiles these ++ these.filter(_.isDirectory).flatMap(recursiveListFiles) } val files = recursiveListFiles(new File(args.head)).filter(_.isFile) println(s"Found${files.length} documents")

files zip Stream.from(1) map { case (f, i) =>
if (i % 200 == 0) println(s"Loaded \$i documents so far")
(f.getName, io.Source.fromFile(f).getLines().flatMap(extractWords).toList)
} toIterator

// test data
//    List(
//      ("doc1", List("welcome", "to", "scala", "labs")),
//      ("doc2", List("welcome", "to", "Toronto")),
//      ("doc3", List("introduce", "scala", "and", "enjoy", "scala")),
//      ("doc4", List("hello", "scala"))
//    )
}

println("Indexing... DONE")

while (true) {
println()
println("Enter a sentence or 'q' to quit")
if (input == "q")
System.exit(0)
println(s"Querying...")
val topMatches = indexedDocuments.search(extractWords(input), topNResults)
if (topMatches.isEmpty)
println("Couldn't find any relevant documents")
else {
println(s"Top results:")
topMatches.foreach(println)
}
}
}


I'm also just a beginner with Scala so I only have small chips for you.

val totalWordCount: Int = wordCount map (_._2) sum


I think this is more intuitive and shorter too:

val totalWordCount: Int = wordCount.values.sum


Also, I think wordCount is not a very good name, because it sounds like an Int, just a single number. wordCounts or wordCountMap might be better.

import search.types.{DocumentTitle, SearchResult, Word}


Why not simply:

import search.types.*

• Great suggestions, thanks! The last one (i.e. import) is a matter of taste, I think. I like to see where types come from, which is useful during code review because I cannot use "jump to definition" – Marcin Oct 28 '14 at 2:20
• You're right, it may be a matter of taste. Of course, if it was import java.util.{HashMap, Set} that would be good like that, I just brought this up because the import is from your own defined type with only a few classes. But yeah, although your way is longer to type, it's more accurate. – janos Oct 28 '14 at 7:08
• I use IntelliJ IDEA and it adds missing import automatically when I save the file, so I don't type it actually :) – Marcin Oct 29 '14 at 1:42

It's clear that you took some time to polish your code, good job on that. I like how used method names like countOccurrencesOf, the allow for very readable code.

My suggestions (note that none of the code was tested or compiled):

It's 'weird' to combine a case class and a Traversable. A case class gives you a few things for free:

• public values for it's arguments of which you have marked one private, so that's weird in it's own right
• a copy method, to easily create new instances
• a companion object containing an apply method for easy construction
• implemented hash and equal methods based on the parameters passed into the (first parameter list of the) constructor
• an implementation of the Product class

With a case class you say: this class is defined by it's values. With a Traversable you say: this class is defined by it's contents. To me that is a bit confusing, is it about the parameters passed in, or about the contents?

 case class WordVector(...) extends Traversable


In a sense it's strange that a WordVector contains a document title. If you look at the compareWithQuery method you see it's actually used in a (somewhat) arbitrary way. From the outside it's unclear which document title is returned. Another indication is that the document title seems not to be used in the indexer, it uses the given document title.

You might want to consider some members to be lazy. You are now doing work on construction while it might be much later that a value is actually used. Note that this depends on the usecase.

 lazy val totalWordCount = wordCount.values.sum


I would probably write WordVector like this:

class WordVector(sentence: Traversable[Word]) extends Traversable[Word] {
private lazy val wordCount = sentence.groupBy(identity).mapValues(_.size)

lazy val totalWordCount = wordCount.values.sum

def countOccurrencesOf(word:Word):Int = wordCount.getOrElse(word, 0)

// personal preference, I do not add override when implementating abstract
// members it allows me to easier see when I (or someone else) overrides an
// actual implemented method. In that case it's most likely a bad design.
def foreach[U](f: Word => U):Unit = wordCount.keySet.foreach(f)
}

object WordVector {
def apply(sentence:Traversable[Word]) = new WordVector(sentence)
}


I would also change TfidfVector to a normal class without title and a lazy length field.

About the IndexedDocuments and the Indexer, they seem to be the very related. It seems you can construct an IndexedDocuments instance from an Iterator[(DocumentTitle, List[Word])], this would make the indexDocuments an ideal candidate for the factory of IndexedDocuments. Such a factory is commonly at the companion object.

Note that I would personally introduce a case class for document.

case class Document(title:DocumentTitle, words:List[Word])

object IndexedDocuments {

def apply(documents:Iterator[Document]) = ...
}


To me it's unclear why the indexDocuments method contains the code it does. It seems the code could also be in the IndexedDocuments class. Or, the other way around, the code for analysis could be in the factory (or another class).

I would not use mutable structures to split the document space and documents containing a word, this is more sensitive to programming mistakes. On top of that it disconnects two pieces of code that are connected, forcing our brains to do more work. A simplified example using option:

val o:Option[String] = ???

val value1 =
if (o.isDefined) {
val x:String = ???
...
x + o.get
} else {
val y:String = ???
...
y + "empty"
}

val value2 = o
.map { o =>
val x:String = ???
...
x + o
}
.getOrElse {
val y:String = ???
...
y + "empty"
}


For value1 the if line and the o.get line are connected. o.get will throw an exception if it is not guarded by o.isDefined. This is a simple example. In more complex cases the maintainer of your code might move some pieces around accidentally breaking the code.

If we apply that to splitting the documents into the two datastructures it would look like this:

val emptyDocumentSpace = Vector.empty[WordVector]
val emptyDocsContainingWord = Map.empty[Word, Set[DocumentTitle]] withDefaultValue Set.empty
val input = (emptyDocumentSpace, emptyDocsContainingWord)

val (documentsSpace, docsContainingWord) =
documents.foldLeft(input) {
case ((documentSpace, docsContainingWord), Document(title, words) =>
val newDocumentSpace = documentSpace + WordVector(words)
val newDocsContainingWord =
words.foldLeft(docsContainingWord) { (docsContainingWord, word) =>
docsContainingWord.updated(word, docsContainingWord(word) + title)
}

(newDocumentSpace, newDocsContainingWord)
}

new IndexedDocuments(documentSpace, docsContainingWord)


The IndexedDocuments class contains quite a lot of public members that are not really suited for public use. It's a good practice to mark methods as private that are only for internal use. This has (at least) two upsides:

1. You don't bother users of your class (yourself) with the internal details
2. It allows you to easily identify which members can be safely refactored

It seems there is a clear relation between a WordVector and a TfidfVector, it seems reasonable to expect a factory method.

object TfidfVector {
def apply(wordVector:WordVector) = ...
}


I would probably write the factory method something like this:

object TfidfVector {

def apply(wordVector:WordVector, wordInDocs: Word => Int, documentCount:Int) =
new TfidfVector(extractDataFrom(wordVector, wordInDocs, documentCount))

def zero(wordVector:WordVector) =
new TfidfVector(wordVector.map(_ -> 0).toMap)

private def extractDataFrom(wordVector:WordVector, wordInDocs: Word => Int, documentCount:Int) = {
val totalWordCount = wordVector.totalWordCount
wordVector
.map { word =>
val occurrencesInVector = wordVector countOccurrencesOf word
word -> calculateTfidf(wordInDocs(word), occurrencesInVector, totalWordCount, documentCount)
}
.toMap
}

private calculateTfidf(numDocsContainingWord:Int, occurrencesInVector:Int, totalWordCount:Int, documentCount:Int) = {
val tf = occurrencesInVector / totalWordCount.toDouble
val idf = documentCount / numDocsContainingWord.toDouble

tf * Math.log(idf)
}
}


And then at the call site

// Note that a Map is also a partial function
val wordInDocs = docsContainingWord.mapValues(_.size) orElse { case _ => 0 }

if (documentSpace.isEmpty) TfidfVector.zero(wordVector)
else TfidfVector(wordVector, wordInDocs, documentCount)


The search method has a comment, requiring the user to do something before it can be used. Most programmers will not read a lot of documentation (apart from Ikea manuals) before using a library. In this case it might be best to help users of the search method with an explicit type.

// The constructor is private, forcing the user to use the factory
// method on the companion object
class Sentence private(val sentence:Traversable[Word])
object Sentence {
def apply(sentence:Traversable[Word]) = new Sentence(normalize(sentence))

private def normalize(sentence:Traversable[Word]) = ...
}

def search(sentence: Sentence, topN: Int): Seq[SearchResult] = ...


The SearchResult type is now modelled as a tuple. I would recommend a case class for it. That allows you have the code in search be more readable.

case class SearchResult(document:DocumentTitle, score:Double)
object SearchResult {
implicit val ordering = Ordering.fromLessThan[SearchResult](_.score > _.score)
}

val scoredDocuments = scoreDocumentsWith(sentence)
scoredDocuments
.filterNot(_.score.isNaN)
.filterNot(_.score < 0.000001)
.sorted
.take(topN)


Note that for the implementation of compareWithQuery the implementation of WordVector (without title) makes things a bit tricky. For that reason I would add a relation in the document space, either with a case class Document(title:DocumentTitle, sentences:Seq[WordVector]) or a pair. This would result in documentsAsTfidfSpace to be a map DocumentTitle -> TfidfVector.

As a result it would remove the .get call in compareWithQuery which could then be renamed to calculateCosineSimilarity allowing you to remove the comment. The two most common errors by Scala developers are the usage of .get on options and Await.result on futures. The only place these methods are allowed to be used is in test code. If have (to date) not found a valid usecase for either one in production code.

There is a small problem in your compareWithQuery method: denominator is an integer and thus could result in ArithmeticException: / by zero.

In the documentloader (and to some extend in the SearchApp) you have mixed println statements in your code. As a separation of concerns those should be moved out of the code performing the operations. The functional programming paradigm is really strict in this sense and some of their patterns can help you to do that. The book Functional programming in Scala contains a section about this exact problem.

Note that I would personally keep the println statements in the app constructor and only remove the ones from the document loader (which I would place in a different class). As with all things, always search for a balance between effort and 'ideal world perfection'.

A few general notes

I recommend you to reverse the general order of your code. Most important (or less detailed stuff) at the top. It's easier for readers to zoom in and see the important (public facing) bits. So if you have a program in a single file, put the App on top. Within the app, put the stuff that actually does something at the top of the file. I now have to scroll down and skip a lot of noise before I get to the actual stuff that is being done.

// Note: I skipped most of the println statements, it's just to give an idea
val indexedDocument = indexDocuments

withUserInput { input =>
val result = indexedDocuments.search(input)
if (result.nonEmpty) showResult(result)
else println("Could not find any relevant documents")
}

private def indexDocuments = ...

private def withUserInput(f:String => Unit) = ...

private def showResult(result:Seq[SearchResult]) = ...


I recommend the same structure for classes: public (and most important) members at the top. Then (in usage order) the other members. This is ofcourse personal preference.

In some cases it's not needed to specify a type. For public facing members it's good practice to write down the full type. For private (and function local) members it's not needed and might actually introduce noise that harms readability.

I noticed you are using ._1 and ._2 a lot. I try to not use them as they require the reader to backtrack to find the type of the tuple that contains them. There are a few techniques you can use to avoid their usage.

val (value1, value2) = value

someMap.map {
case (key, value) => ...
}

case class SomeName(value1, value2)

def test(t:(String, String) => String) = ...
test {
case (key, value) => ...
}


That's it all I have for now. I hope it helps you.