map-reduce-folds-0.1.2: src/Control/MapReduce/Engines/Streamly.hs
{-# LANGUAGE CPP #-}
{-# LANGUAGE DataKinds #-}
{-# LANGUAGE FlexibleContexts #-}
{-# LANGUAGE GADTs #-}
{-# LANGUAGE ScopedTypeVariables #-}
{-# LANGUAGE TypeOperators #-}
{-# LANGUAGE RankNTypes #-}
{-# LANGUAGE PolyKinds #-}
{-# LANGUAGE TypeFamilies #-}
{-# LANGUAGE ConstraintKinds #-}
{-# LANGUAGE MultiParamTypeClasses #-}
{-# LANGUAGE FlexibleInstances #-}
{-# LANGUAGE UndecidableInstances #-}
{-# LANGUAGE AllowAmbiguousTypes #-}
{-# LANGUAGE TypeApplications #-}
{-# LANGUAGE BangPatterns #-}
{-# OPTIONS_GHC -O2 -fwarn-incomplete-patterns #-}
{-|
Module : Control.MapReduce.Engines.Streams
Description : map-reduce-folds builders
Copyright : (c) Adam Conner-Sax 2019
License : BSD-3-Clause
Maintainer : adam_conner_sax@yahoo.com
Stability : experimental
map-reduce engine (fold builder) using @Streamly@ streams as its intermediate and return type.
Notes:
1. These are polymorphic in the return stream type. Thought the streams do have to be @serial@ when @groupBy@ is called
So you have to specify the stream type in the call or it has to be inferrable from the use of the result.
2. There is a concurrent engine here, one that uses Streamly's concurrency features to map over the stream. I've not
been able to verify that this is faster on an appropriate task with appropriate runtime settings.
-}
module Control.MapReduce.Engines.Streamly
(
-- * Engines
streamlyEngine
, streamlyEngineM
, concurrentStreamlyEngine
-- * Streamly Combinators
, toStreamlyFold
, toStreamlyFoldM
-- * Result Extraction
, resultToList
, concatStream
, concatStreamFold
, concatStreamFoldM
, concatConcurrentStreamFold
-- * @groupBy@ Functions
, groupByHashableKey
, groupByOrderedKey
#if MIN_VERSION_streamly(0,9,0)
, groupByOrderedKeyIO
#endif
, groupByHashableKeyST
, groupByDiscriminatedKey
-- * Re-Exports
#if MIN_VERSION_streamly(0,9,0)
, Stream
#else
, SerialT
, WSerialT
, AheadT
, AsyncT
, WAsyncT
, ParallelT
, MonadAsync
, IsStream
#endif
)
where
import qualified Control.MapReduce.Core as MRC
import qualified Control.MapReduce.Engines as MRE
import Control.Arrow ( second )
import qualified Control.Foldl as FL
import Control.Monad.ST as ST
import qualified Data.Discrimination.Grouping as DG
import qualified Data.Foldable as F
import Data.Functor.Identity ( Identity(runIdentity) )
import Data.Hashable ( Hashable )
import qualified Data.HashMap.Strict as HMS
import qualified Data.HashTable.Class as HT
import qualified Data.HashTable.ST.Cuckoo as HTC
import qualified Data.List.NonEmpty as LNE
import qualified Data.Maybe as Maybe
import qualified Data.Map.Strict as MS
import qualified Data.Sequence as Seq
#if MIN_VERSION_streamly(0,9,0)
import Streamly.Data.Stream (Stream)
import qualified Streamly.Data.Stream as S
import qualified Streamly.Data.Stream.Prelude as SP
import qualified Streamly.Data.StreamK as StreamK
import qualified Streamly.Data.Fold as SF
import qualified Streamly.Data.Unfold as SU
import qualified Streamly.Internal.Data.Fold as SF
import Control.Monad (join)
import Control.Monad.IO.Class (MonadIO)
#elif MIN_VERSION_streamly(0,8,0)
import qualified Streamly.Prelude as S
import qualified Streamly.Internal.Data.Fold as SF
import Streamly.Prelude ( SerialT
, WSerialT
, AheadT
, AsyncT
, WAsyncT
, ParallelT
, MonadAsync
, IsStream
)
#else
import qualified Streamly.Prelude as S
import qualified Streamly as S
import qualified Streamly.Internal.Data.Fold as SF
import Streamly ( SerialT
, WSerialT
, AheadT
, AsyncT
, WAsyncT
, ParallelT
, MonadAsync
, IsStream
)
--import qualified Streamly.Internal.Data.Parser.ParserK.Type as Streamly
#endif
#if MIN_VERSION_streamly(0,9,0)
-- | convert a Control.Foldl FoldM into a Streamly.Data.Fold fold
toStreamlyFoldM :: Functor m => FL.FoldM m a b -> SF.Fold m a b
toStreamlyFoldM (FL.FoldM step start done) = SF.Fold step' (SF.Partial <$> start) done where
step' s a = SF.Partial <$> step s a
-- | convert a Control.Foldl Fold into a Streamly.Data.Fold fold
toStreamlyFold :: Monad m => FL.Fold a b -> SF.Fold m a b
toStreamlyFold (FL.Fold step start done) = SF.Fold step' (pure $ SF.Partial $ start) (pure . done) where
step' s a = pure $ SF.Partial $ step s a
#elif MIN_VERSION_streamly(0,8,0)
fromEffect :: (Monad m, IsStream t) => m a -> t m a
fromEffect = S.fromEffect
{-# INLINE fromEffect #-}
-- | convert a Control.Foldl FoldM into a Streamly.Data.Fold fold
toStreamlyFoldM :: Functor m => FL.FoldM m a b -> SF.Fold m a b
toStreamlyFoldM (FL.FoldM step start done) = SF.mkFoldM step' (SF.Partial <$> start) done where
step' s a = SF.Partial <$> step s a
-- | convert a Control.Foldl Fold into a Streamly.Data.Fold fold
toStreamlyFold :: Monad m => FL.Fold a b -> SF.Fold m a b
toStreamlyFold (FL.Fold step start done) = SF.mkFold step' (SF.Partial start) done where
step' s a = SF.Partial $ step s a
#else
fromEffect :: (Monad m, IsStream t) => m a -> t m a
fromEffect = S.yieldM
{-# INLINE fromEffect #-}
-- | convert a Control.Foldl FoldM into a Streamly.Data.Fold fold
toStreamlyFoldM :: FL.FoldM m a b -> SF.Fold m a b
toStreamlyFoldM (FL.FoldM step start done) = SF.mkFold step start done
-- | convert a Control.Foldl Fold into a Streamly.Data.Fold fold
toStreamlyFold :: Monad m => FL.Fold a b -> SF.Fold m a b
toStreamlyFold (FL.Fold step start done) = SF.mkPure step start done
#endif
-- | unpack for streamly based map/reduce
#if MIN_VERSION_streamly(0,9,0)
unpackStreamK :: MRC.Unpack x y -> StreamK.StreamK Identity x -> S.Stream Identity y
unpackStreamK (MRC.Filter t) = S.filter t . StreamK.toStream
unpackStreamK (MRC.Unpack f) = S.concatMap (StreamK.toStream . StreamK.fromFoldable . f) . StreamK.toStream
-- S.concatMap (StreamK.toStream . StreamK.fromFoldable . f)
{-# INLINABLE unpackStreamK #-}
-- | effectful (monadic) unpack for streamly based map/reduce
unpackStreamKM :: (Monad m) => MRC.UnpackM m x y -> StreamK.StreamK m x -> S.Stream m y
unpackStreamKM (MRC.FilterM t) = S.filterM t . StreamK.toStream
unpackStreamKM (MRC.UnpackM f) = S.concatMapM (fmap (StreamK.toStream . StreamK.fromFoldable) . f) . StreamK.toStream
{-# INLINABLE unpackStreamKM #-}
-- | make a stream into an (effectful) @[]@
resultToList :: (Monad m) => S.Stream m a -> m [a]
resultToList = S.toList
{-# INLINEABLE resultToList #-}
-- | mappend all in a monoidal stream
concatStream :: (Monad m, Monoid a) => S.Stream m a -> m a
concatStream = S.fold (SF.foldl' (<>) mempty)
{-# INLINEABLE concatStream #-}
-- | mappend everything in a pure Streamly fold
concatStreamFold :: Monoid b => FL.Fold a (S.Stream Identity b) -> FL.Fold a b
concatStreamFold = fmap (runIdentity . concatStream)
{-# INLINEABLE concatStreamFold #-}
-- | mappend everything in an effectful Streamly fold.
concatStreamFoldM
:: (Monad m, Monoid b) => FL.FoldM m a (S.Stream m b) -> FL.FoldM m a b
concatStreamFoldM = MRC.postMapM concatStream
{-# INLINEABLE concatStreamFoldM #-}
-- | mappend everything in a concurrent Streamly fold.
concatConcurrentStreamFold
:: (Monad m, Monoid b) => FL.Fold a (S.Stream m b) -> FL.FoldM m a b
concatConcurrentStreamFold = concatStreamFoldM . FL.generalize
{-# INLINEABLE concatConcurrentStreamFold #-}
-- | map-reduce-fold builder returning a @SerialT Identity d@ result
streamlyEngine
:: (Foldable g, Functor g)
=> (forall z . S.Stream Identity (k, z) -> S.Stream Identity (k, g z))
-> MRE.MapReduceFold y k c (S.Stream Identity) x d
streamlyEngine groupByKey u (MRC.Assign a) r = FL.Fold
(flip StreamK.cons)
StreamK.nil
( fmap (\(k, lc) -> MRE.reduceFunction r k lc)
. groupByKey
. fmap a
. unpackStreamK u
)
{-# INLINABLE streamlyEngine #-}
-- | unpack for concurrent streamly based map/reduce
unpackConcurrentlyK
:: (SP.MonadAsync m) => MRC.Unpack x y -> StreamK.StreamK m x -> S.Stream m y
unpackConcurrentlyK (MRC.Filter t) = S.filter t . StreamK.toStream
unpackConcurrentlyK (MRC.Unpack f) = S.concatMap ((StreamK.toStream . StreamK.fromFoldable) . f) . StreamK.toStream
{-# INLINABLE unpackConcurrentlyK #-}
-- | possibly (depending on chosen stream types) concurrent map-reduce-fold builder returning an @(Istream t, MonadAsync m) => t m d@ result
concurrentStreamlyEngine
:: forall m g y k c x d
. (SP.MonadAsync m, Foldable g, Functor g)
=> (forall z . S.Stream m (k, z) -> S.Stream m (k, g z))
-> MRE.MapReduceFold y k c (S.Stream m) x d
concurrentStreamlyEngine groupByKey u (MRC.Assign a) r = FL.Fold
(\s a' -> (pure a') `StreamK.consM` s)
StreamK.nil
( S.mapM (\(k, lc) -> return $ MRE.reduceFunction r k lc)
. groupByKey
. S.mapM (return . a)
. unpackConcurrentlyK u
)
{-# INLINABLE concurrentStreamlyEngine #-}
-- | effectful map-reduce-fold engine returning a (Istream t => t m d) result
-- The "MonadAsync" constraint here more or less requires us to run in IO, or something IO like.
streamlyEngineM
:: (Monad m, SP.MonadAsync m, Traversable g)
=> (forall z . S.Stream m (k, z) -> S.Stream m (k, g z))
-> MRE.MapReduceFoldM m y k c (S.Stream m) x d
streamlyEngineM groupByKey u (MRC.AssignM a) r =
FL.generalize
$ FL.Fold
(flip StreamK.cons)
StreamK.nil
( S.mapM (\(k, lc) -> MRE.reduceFunctionM r k lc)
. groupByKey -- this requires a serial stream.
. S.mapM a
. unpackStreamKM u
)
{-# INLINABLE streamlyEngineM #-}
toHashMap :: (Monad m, Eq k, Monoid a, Hashable k) => SF.Fold m (k, a) (HMS.HashMap k a)
toHashMap = SF.foldl' (\hm (k, a) -> HMS.insertWith (<>) k a hm) HMS.empty
{-# INLINEABLE toHashMap #-}
streamMeta :: Monad m => SF.Fold m a b -> SU.Unfold m b c -> S.Stream m a -> S.Stream m c
streamMeta fld unfld = S.concatEffect . fmap (S.unfold unfld) . S.fold fld
{-# INLINEABLE streamMeta #-}
-- | Group streamly stream of @(k,c)@ by @hashable@ key.
groupByHashableKey
:: (Monad m, Hashable k, Eq k)
=> S.Stream m (k, c)
-> S.Stream m (k, Seq.Seq c)
groupByHashableKey = streamMeta toHashMap (SU.lmap HMS.toList $ SU.fromList) . fmap (second Seq.singleton)
{-# INLINABLE groupByHashableKey #-}
toMap :: (Monad m, Monoid a, Ord k) => SF.Fold m (k, a) (MS.Map k a)
toMap = SF.foldl' (\hm (k, a) -> MS.insertWith (<>) k a hm) MS.empty
{-# INLINEABLE toMap #-}
-- | Group streamly stream of @(k,c)@ by ordered key.
groupByOrderedKey
:: (Monad m, Ord k) => S.Stream m (k, c) -> S.Stream m (k, Seq.Seq c)
groupByOrderedKey = --streamMeta (SF.toMap fst (SF.lmap snd toSeq)) (SU.lmap MS.toList $ SU.fromList)
streamMeta toMap (SU.lmap MS.toList $ SU.fromList) . fmap (second Seq.singleton)
{-# INLINABLE groupByOrderedKey #-}
toSeq :: Monad m => SF.Fold m a (Seq.Seq a)
toSeq = SF.foldl' (\s a -> s <> Seq.singleton a) mempty
{-# INLINEABLE toSeq #-}
-- | Group streamly stream of @(k,c)@ by ordered key. Using toMapIO for mutable cells in the fold.
groupByOrderedKeyIO
:: (Monad m, MonadIO m, Ord k) => S.Stream m (k, c) -> S.Stream m (k, Seq.Seq c)
groupByOrderedKeyIO = streamMeta (SF.toMapIO fst (SF.lmap snd toSeq)) (SU.lmap MS.toList $ SU.fromList)
{-# INLINABLE groupByOrderedKeyIO #-}
{-
toHashMapSeqST :: (Monad m, Eq k, Hashable k) => SF.Fold m (k, a) (ST.ST s (HTC.HashTable s k (Seq.Seq a)))
toHashMapSeqST = SF.foldl' (\hm (k, a) -> traverse (mutate k (f a)) hm) HTC.new where
mutate k f hm = HTC.mutate hm k f
f :: forall s . a -> Maybe (Seq.Seq a) -> Maybe (Seq.Seq a, ((HTC.HashTable s k (Seq.Seq a))))
f a ht saM = case saM of
Nothing -> (Seq.singleton a, ())
Just sa -> (sa Seq.|> a, ())
{-# INLINEABLE toHashMapSeqST #-}
-}
-- | Group streamly stream of @(k,c)@ by @hashable@ key. Uses mutable hashtables running in the ST monad.
-- NB: this function uses the fact that @SerialT m@ is a monad
groupByHashableKeyST
:: forall m k c . (Monad m, Hashable k, Eq k)
=> S.Stream m (k, c)
-> S.Stream m (k, Seq.Seq c)
groupByHashableKeyST st = S.concatEffect $ fmap listToStream $ S.toList st
-- lkc <- fromEffect (S.toList st)
where
listToStream :: [(k, c)] -> S.Stream m (k, Seq.Seq c)
listToStream l = S.fromList $ ST.runST $ listToList l
listToList :: [(k, c)] -> ST.ST s [(k, Seq.Seq c)]
listToList = join . fmap HT.toList . toHT
toHT :: [(k, c)] -> ST.ST s (HTC.HashTable s k (Seq.Seq c))
toHT l = (MRE.fromListWithHT @HTC.HashTable) (<>)
$ fmap (second Seq.singleton) l
{-# INLINABLE groupByHashableKeyST #-}
-- | Group streamly stream of @(k,c)@ by key with instance of Grouping from <http://hackage.haskell.org/package/discrimination>.
-- NB: this function uses the fact that @SerialT m@ is a monad
groupByDiscriminatedKey
:: (Monad m, DG.Grouping k)
=> S.Stream m (k, c)
-> S.Stream m (k, Seq.Seq c)
groupByDiscriminatedKey s =
S.concatEffect
$ (StreamK.toStream . StreamK.fromFoldable . Maybe.catMaybes . fmap (fmap g . LNE.nonEmpty) . DG.groupWith fst)
<$> S.toList s
where
g :: LNE.NonEmpty (k, c) -> (k, Seq.Seq c)
g x = let k = fst (LNE.head x) in (k, F.fold $ fmap (Seq.singleton . snd) x)
{-# INLINABLE groupByDiscriminatedKey #-}
#else
unpackStream :: S.IsStream t => MRC.Unpack x y -> t Identity x -> t Identity y
unpackStream (MRC.Filter t) = S.filter t
unpackStream (MRC.Unpack f) = S.concatMap (S.fromFoldable . f)
{-# INLINABLE unpackStream #-}
-- | effectful (monadic) unpack for streamly based map/reduce
unpackStreamM :: (S.IsStream t, Monad m) => MRC.UnpackM m x y -> t m x -> t m y
unpackStreamM (MRC.FilterM t) = S.filterM t
unpackStreamM (MRC.UnpackM f) = S.concatMapM (fmap S.fromFoldable . f)
{-# INLINABLE unpackStreamM #-}
-- | make a stream into an (effectful) @[]@
resultToList :: (Monad m, S.IsStream t) => t m a -> m [a]
resultToList = S.toList . S.adapt
-- | mappend all in a monoidal stream
concatStream :: (Monad m, Monoid a) => S.SerialT m a -> m a
concatStream = S.foldl' (<>) mempty
-- | mappend everything in a pure Streamly fold
concatStreamFold :: Monoid b => FL.Fold a (S.SerialT Identity b) -> FL.Fold a b
concatStreamFold = fmap (runIdentity . concatStream)
-- | mappend everything in an effectful Streamly fold.
concatStreamFoldM
:: (Monad m, Monoid b, S.IsStream t) => FL.FoldM m a (t m b) -> FL.FoldM m a b
concatStreamFoldM = MRC.postMapM (concatStream . S.adapt)
-- | mappend everything in a concurrent Streamly fold.
concatConcurrentStreamFold
:: (Monad m, Monoid b, S.IsStream t) => FL.Fold a (t m b) -> FL.FoldM m a b
concatConcurrentStreamFold = concatStreamFoldM . FL.generalize
-- | map-reduce-fold builder returning a @SerialT Identity d@ result
streamlyEngine
:: (Foldable g, Functor g)
=> (forall z . S.SerialT Identity (k, z) -> S.SerialT Identity (k, g z))
-> MRE.MapReduceFold y k c (SerialT Identity) x d
streamlyEngine groupByKey u (MRC.Assign a) r = FL.Fold
(flip S.cons)
S.nil
( S.map (\(k, lc) -> MRE.reduceFunction r k lc)
. groupByKey
. S.map a
. unpackStream u
)
{-# INLINABLE streamlyEngine #-}
-- | unpack for concurrent streamly based map/reduce
unpackConcurrently
:: (S.MonadAsync m, S.IsStream t) => MRC.Unpack x y -> t m x -> t m y
unpackConcurrently (MRC.Filter t) = S.filter t
unpackConcurrently (MRC.Unpack f) = S.concatMap (S.fromFoldable . f)
{-# INLINABLE unpackConcurrently #-}
-- | possibly (depending on chosen stream types) concurrent map-reduce-fold builder returning an @(Istream t, MonadAsync m) => t m d@ result
concurrentStreamlyEngine
:: forall tIn tOut m g y k c x d
. (S.IsStream tIn, S.IsStream tOut, S.MonadAsync m, Foldable g, Functor g)
=> (forall z . S.SerialT m (k, z) -> S.SerialT m (k, g z))
-> MRE.MapReduceFold y k c (tOut m) x d
concurrentStreamlyEngine groupByKey u (MRC.Assign a) r = FL.Fold
(\s a' -> (return a') `S.consM` s)
S.nil
( S.mapM (\(k, lc) -> return $ MRE.reduceFunction r k lc)
. S.adapt @SerialT @tOut -- make it concurrent for reducing
. groupByKey
. S.adapt @tIn @SerialT-- make it serial for grouping
. S.mapM (return . a)
. (S.|$) (unpackConcurrently u)
)
{-# INLINABLE concurrentStreamlyEngine #-}
-- | effectful map-reduce-fold engine returning a (Istream t => t m d) result
-- The "MonadAsync" constraint here more or less requires us to run in IO, or something IO like.
streamlyEngineM
:: (S.IsStream t, Monad m, S.MonadAsync m, Traversable g)
=> (forall z . SerialT m (k, z) -> SerialT m (k, g z))
-> MRE.MapReduceFoldM m y k c (t m) x d
streamlyEngineM groupByKey u (MRC.AssignM a) r =
FL.generalize
$ fmap S.adapt
$ FL.Fold
(flip S.cons)
S.nil
( S.mapM (\(k, lc) -> MRE.reduceFunctionM r k lc)
. groupByKey -- this requires a serial stream.
. S.mapM a
. unpackStreamM u
)
{-# INLINABLE streamlyEngineM #-}
-- TODO: Try using Streamly folds and Map.insertWith instead of toList and fromListWith. Prolly the same.
-- | Group streamly stream of @(k,c)@ by @hashable@ key.
-- NB: this function uses the fact that @SerialT m@ is a monad
groupByHashableKey
:: (Monad m, Hashable k, Eq k)
=> S.SerialT m (k, c)
-> S.SerialT m (k, Seq.Seq c)
groupByHashableKey s = do
lkc <- fromEffect (S.toList s)
let hm = HMS.fromListWith (<>) $ fmap (second $ Seq.singleton) lkc
HMS.foldrWithKey (\k lc s' -> S.cons (k, lc) s') S.nil hm
{-# INLINABLE groupByHashableKey #-}
-- TODO: Try using Streamly folds and Map.insertWith instead of toList and fromListWith. Prolly the same.
-- | Group streamly stream of @(k,c)@ by ordered key.
-- NB: this function uses the fact that @SerialT m@ is a monad
groupByOrderedKey
:: (Monad m, Ord k) => S.SerialT m (k, c) -> S.SerialT m (k, Seq.Seq c)
groupByOrderedKey s = do
lkc <- fromEffect (S.toList s)
let hm = MS.fromListWith (<>) $ fmap (second $ Seq.singleton) lkc
MS.foldrWithKey (\k lc s' -> S.cons (k, lc) s') S.nil hm
{-# INLINABLE groupByOrderedKey #-}
-- | Group streamly stream of @(k,c)@ by @hashable@ key. Uses mutable hashtables running in the ST monad.
-- NB: this function uses the fact that @SerialT m@ is a monad
groupByHashableKeyST
:: (Monad m, Hashable k, Eq k)
=> S.SerialT m (k, c)
-> S.SerialT m (k, Seq.Seq c)
groupByHashableKeyST st = do
lkc <- fromEffect (S.toList st)
ST.runST $ do
hm <- (MRE.fromListWithHT @HTC.HashTable) (<>)
$ fmap (second Seq.singleton) lkc
HT.foldM (\s' (k, sc) -> return $ S.cons (k, sc) s') S.nil hm
{-# INLINABLE groupByHashableKeyST #-}
-- | Group streamly stream of @(k,c)@ by key with instance of Grouping from <http://hackage.haskell.org/package/discrimination>.
-- NB: this function uses the fact that @SerialT m@ is a monad
groupByDiscriminatedKey
:: (Monad m, DG.Grouping k)
=> S.SerialT m (k, c)
-> S.SerialT m (k, Seq.Seq c)
groupByDiscriminatedKey s = do
lkc <- fromEffect (S.toList s)
let g :: LNE.NonEmpty (k, c) -> (k, Seq.Seq c)
g x = let k = fst (LNE.head x) in (k, F.fold $ fmap (Seq.singleton . snd) x)
S.fromFoldable $ Maybe.catMaybes . fmap (fmap g . LNE.nonEmpty) $ DG.groupWith fst lkc
{-# INLINABLE groupByDiscriminatedKey #-}
#endif