I have written a program to process Google 1-Gram dataset in order to obtain the list of most frequent words. This is my second more or less working program in Haskell, so I really want to know if it complies with current Haskell best practices.
I want to hear any your thought about virtually anything: performance, coding style, naming, even grammar and spelling in comments.
-- This script processes Google N-Gram dataset to obtain the list of the most
-- frequent words sorted by their frequency descending.
--
-- Usage example:
-- $ zcat googlebooks-eng-all-1gram-20120701-{a..z}.gz |
-- ./ProcessGoogleDataset 1950 5000 "$(echo {a..z})" > result.txt
--
-- where 1950 means to filter out only records newer then year 1950,
-- 5000 means number of the most frequent words to collect,
-- {a..z} means the list of letters allowed in words.
--
-- Copyright (c) 2015, Pavel Kretov.
-- Provided under the terms of Apache 2.0 License.
import Control.Monad.IO.Class (liftIO)
import Data.Char (isSpace, isLower)
import Data.Function (on)
import Data.List (foldl', sortBy, groupBy, partition)
import Data.List.Ordered (mergeBy)
import Data.Maybe (fromJust)
import Data.Ord (comparing)
import Data.Tuple (swap)
import System.Environment (getArgs)
import System.IO (stdin, stdout)
-- Heap seems convenient for finding most frequent items.
import qualified Data.Heap as H
-- Use Conduit package for IO as it has proven to be faster
-- then built-in (lines . readFile) approach.
import Data.Conduit (($$), (=$))
import qualified Data.Conduit.Text as CT
import qualified Data.Conduit.Binary as CB
import qualified Data.Conduit.List as CL
import qualified Data.Conduit.Lazy as CZ
-- Conduit thinks Data.Text is better then plain strings.
import qualified Data.Text as T
import qualified Data.Text.Read as TR
-- Conduit really insists on that for lazy IO.
import Control.Monad.Trans.Resource (runResourceT)
-- Parses input 'str' into integer value. Produces an error if string
-- cannot be parsed successfully.
parseInt :: T.Text -> Int
parseInt str = case TR.decimal str of
Left errorText -> error errorText
Right (val, _) -> val
-- Parses a line of raw input into a tuple of word, year, and frequency.
-- Produces an error if line cannot be parsed successfully.
parseEntry :: T.Text -> (T.Text, Int, Int)
parseEntry line = (word, parseInt year, parseInt freq)
where [word, year, freq, _] = T.words line
-- Filters away entries older then 'year'.
filterByYear :: Int -> [(T.Text, Int, Int)] -> [(T.Text, Int, Int)]
filterByYear year = filter (\(w,y,f) -> y >= year)
-- Removes year component from the input tuple.
removeYears :: [(T.Text, Int, Int)] -> [(T.Text, Int)]
removeYears = map $ \(w,y,f) -> (w,f)
-- Original dataset contains additional qualifiers as underscore suffices,
-- something like "_NOUN" or "_ADJ". This function removes them.
removeSuffices :: [(T.Text, Int)] -> [(T.Text, Int)]
removeSuffices = map (\(word, f) -> (truncateUnderscore word, f))
where truncateUnderscore = T.takeWhile (\c -> c /= '_')
-- Filter away entries where word contains non-alphabetic characters. The list
-- of characters comprising the aplabet is given as argument. Note that
-- character case matters for comparison.
filterAlphabetic :: [Char] -> [(T.Text, Int)] -> [(T.Text, Int)]
filterAlphabetic alphabet =
filter (\(word, f) -> T.all (`elem` alphabet) word)
-- Summs frequencies for all consecutive tuples with same word.
aggregateByWord :: [(T.Text, Int)] -> [(T.Text, Int)]
aggregateByWord xs = id
$ map (foldl (\(w,f) (w',f') -> (w',f+f')) (T.empty,0))
$ groupBy ((==) `on` fst) xs
resortByFirstLetter :: [(T.Text, Int)] -> [(T.Text, Int)]
resortByFirstLetter = id
. concatMap (sortBy (comparing fst))
. groupBy ((==) `on` (T.head . fst))
-- Naive 'mostFreq' implementation. It is not actually used in this program
-- and left just for reference.
naiveMostFreq :: Int -> [(T.Text, Int)] -> [(T.Text, Int)]
naiveMostFreq n xs = take n $ sortBy (flip $ comparing snd) xs
-- Optimized implementation of 'naiveMostFreq'. This function takes most
-- frequent items from the input list.
-- http://stackoverflow.com/a/31348931/1447225
mostFreq :: Int -> [(T.Text, Int)] -> [(T.Text, Int)]
mostFreq n dataset = final
where
pairs = map swap dataset
(first, rest) = splitAt n pairs
start = H.fromList first :: H.MinHeap (Int, T.Text)
stop = foldl' step start rest
step heap pair = if H.viewHead heap < Just pair
then H.insert pair (fromJust $ H.viewTail heap)
else heap
final = map swap (H.toList stop)
swap ~(a,b) = (b,a)
formatOutputLine :: (T.Text, Int) -> T.Text
formatOutputLine (word, freq) =
T.concat [word, T.pack "\t", T.pack $ show freq, T.pack "\n"]
main :: IO ()
main = runResourceT $ do
-- Parse command line parameters. Note that all whitespace characters are
-- filtered away from 'aplphabet' parameter for user convenience.
args <- liftIO getArgs
let year = parseInt (T.pack $ args !! 0)
let limit = parseInt (T.pack $ args !! 1)
let alphabet = filter (not . isSpace)
(args !! 2)
-- Read lines from STDIN as a lazy list.
-- Note here that UTF-8 is hardcoded.
rawLines <- CZ.lazyConsume
$ CB.sourceHandle stdin
=$ CT.decodeUtf8
=$ CT.lines
-- Do actual processing.
let result = id
$ sortBy (flip $ comparing snd)
$ mostFreq limit
$ aggregateByWord
$ resortByFirstLetter
$ aggregateByWord
$ filterAlphabetic alphabet
$ removeSuffices
$ removeYears
$ filterByYear year
$ map parseEntry rawLines
-- Write result to STDOUT.
CL.sourceList (map formatOutputLine result)
$$ CT.encodeUtf8
=$ CB.sinkHandle stdout
You may download source dataset by the following command, but note that it's quite large (about 4.5 GB in compressed form).
$ wget http://storage.googleapis.com/books/ngrams/books/googlebooks-eng-all-1gram-20120701-{a..z}.gz
