git.iko.soy/streamly
on 0529366 help sign in
streamly/src/Streamly/Internal/Data/Fold/Extra.hs
624 lines 21 kB history raw
1-- |
2-- Module : Streamly.Internal.Data.Fold.Extra
3-- Copyright : (c) 2019 Composewell Technologies
4-- License : BSD3
5-- Maintainer : streamly@composewell.com
6-- Stability : experimental
7-- Portability : GHC
8--
10module Streamly.Internal.Data.Fold.Extra
11 (
12 -- * Imports
13 -- $setup
15 countDistinct
16 , countDistinctInt
17 , frequency
18 , toMap
19 , nub
20 , nubInt
22 -- ** Demultiplexing
23 -- | Direct values in the input stream to different folds using an n-ary
24 -- fold selector. 'demux' is a generalization of 'classify' (and
25 -- 'partition') where each key of the classifier can use a different fold.
26 , demux
27 , demuxWith
28 , demuxScanWith
29 , demuxScanMutWith
30 , demuxMutWith
32 -- TODO: These can be implemented using the above operations
33 -- , demuxWithSel -- Stop when the fold for the specified key stops
34 -- , demuxWithMin -- Stop when any of the folds stop
35 -- , demuxWithAll -- Stop when all the folds stop (run once)
37 -- ** Classifying
38 -- | In an input stream of key value pairs fold values for different keys
39 -- in individual output buckets using the given fold. 'classify' is a
40 -- special case of 'demux' where all the branches of the demultiplexer use
41 -- the same fold.
42 --
43 -- Different types of maps can be used with these combinators via the IsMap
44 -- type class. Hashmap performs better when there are more collisions, trie
45 -- Map performs better otherwise. Trie has an advantage of sorting the keys
46 -- at the same time. For example if we want to store a dictionary of words
47 -- and their meanings then trie Map would be better if we also want to
48 -- display them in sorted order.
50 , classify
51 , classifyWith
52 , classifyMutWith
53 , classifyScanWith
54 -- , classifyWithSel
55 -- , classifyWithMin
56 )
57where
59#include "inline.hs"
60#include "ArrayMacros.h"
62import Control.Monad.IO.Class (MonadIO(..))
63import Data.IORef (newIORef, readIORef, writeIORef)
64import Data.Map.Strict (Map)
65import Streamly.Internal.Data.IsMap (IsMap(..))
66import Streamly.Internal.Data.Tuple.Strict (Tuple'(..), Tuple3'(..))
68import qualified Data.IntSet as IntSet
69import qualified Data.Set as Set
70import qualified Prelude
71import qualified Streamly.Internal.Data.IsMap as IsMap
73import Prelude hiding
74 ( filter, foldl1, drop, dropWhile, take, takeWhile, zipWith
75 , foldl, foldr, map, mapM_, sequence, all, any, sum, product, elem
76 , notElem, maximum, minimum, head, last, tail, length, null
77 , reverse, iterate, init, and, or, lookup, (!!)
78 , scanl, scanl1, replicate, concatMap, mconcat, foldMap, unzip
79 , span, splitAt, break, mapM, zip, maybe)
80import Streamly.Internal.Data.Fold
82-- $setup
83-- >>> :m
84-- >>> :set -XFlexibleContexts
85-- >>> import qualified Streamly.Data.Fold as Fold
86-- >>> import qualified Streamly.Data.Stream as Stream
87-- >>> import qualified Streamly.Internal.Data.Fold.Extra as Fold
89-- XXX Name as nubOrd? Or write a nubGeneric
91-- | Used as a scan. Returns 'Just' for the first occurrence of an element,
92-- returns 'Nothing' for any other occurrences.
93--
94-- Example:
95--
96-- >>> stream = Stream.fromList [1::Int,1,2,3,4,4,5,1,5,7]
97-- >>> Stream.fold Fold.toList $ Stream.scanMaybe Fold.nub stream
98-- [1,2,3,4,5,7]
99--
100-- /Pre-release/
101{-# INLINE nub #-}
102nub :: (Monad m, Ord a) => Fold m a (Maybe a)
103nub = fmap (\(Tuple' _ x) -> x) $ foldl' step initial
105 where
107 initial = Tuple' Set.empty Nothing
109 step (Tuple' set _) x =
110 if Set.member x set
111 then Tuple' set Nothing
112 else Tuple' (Set.insert x set) (Just x)
114-- | Like 'nub' but specialized to a stream of 'Int', for better performance.
115--
116-- Definition:
117--
118-- >>> nubInt = Fold.nub
119--
120-- /Pre-release/
121{-# INLINE nubInt #-}
122nubInt :: Monad m => Fold m Int (Maybe Int)
123nubInt = fmap (\(Tuple' _ x) -> x) $ foldl' step initial
125 where
127 initial = Tuple' IntSet.empty Nothing
129 step (Tuple' set _) x =
130 if IntSet.member x set
131 then Tuple' set Nothing
132 else Tuple' (IntSet.insert x set) (Just x)
134-- XXX Try Hash set
135-- XXX Add a countDistinct window fold
136-- XXX Add a bloom filter fold
138-- | Count non-duplicate elements in the stream.
139--
140-- Definition:
141--
142-- >>> countDistinct = Fold.postscan Fold.nub $ Fold.catMaybes $ Fold.length
143--
144-- The memory used is proportional to the number of distinct elements in the
145-- stream, to guard against using too much memory use it as a scan and
146-- terminate if the count reaches more than a threshold.
147--
148-- /Space/: \(\mathcal{O}(n)\)
149--
150-- /Pre-release/
151--
152{-# INLINE countDistinct #-}
153countDistinct :: (Monad m, Ord a) => Fold m a Int
154countDistinct = postscan nub $ catMaybes length
155{-
156countDistinct = fmap (\(Tuple' _ n) -> n) $ foldl' step initial
158 where
160 initial = Tuple' Set.empty 0
162 step (Tuple' set n) x = do
163 if Set.member x set
164 then
165 Tuple' set n
166 else
167 let cnt = n + 1
168 in Tuple' (Set.insert x set) cnt
169-}
171-- | Like 'countDistinct' but specialized to a stream of 'Int', for better
172-- performance.
173--
174-- Definition:
175--
176-- >>> countDistinctInt = Fold.postscan Fold.nubInt $ Fold.catMaybes $ Fold.length
177--
178-- /Pre-release/
179{-# INLINE countDistinctInt #-}
180countDistinctInt :: Monad m => Fold m Int Int
181countDistinctInt = postscan nubInt $ catMaybes length
182{-
183countDistinctInt = fmap (\(Tuple' _ n) -> n) $ foldl' step initial
185 where
187 initial = Tuple' IntSet.empty 0
189 step (Tuple' set n) x = do
190 if IntSet.member x set
191 then
192 Tuple' set n
193 else
194 let cnt = n + 1
195 in Tuple' (IntSet.insert x set) cnt
196 -}
198------------------------------------------------------------------------------
199-- demux: in a key value stream fold each key sub-stream with a different fold
200------------------------------------------------------------------------------
202-- TODO Demultiplex an input element into a number of typed variants. We want
203-- to statically restrict the target values within a set of predefined types,
204-- an enumeration of a GADT.
205--
206-- This is the consumer side dual of the producer side 'mux' operation (XXX to
207-- be implemented).
208--
209-- XXX If we use Refold in it, it can perhaps fuse/be more efficient. For
210-- example we can store just the result rather than storing the whole fold in
211-- the Map.
212--
213-- Note: There are separate functions to determine Key and Fold from the input
214-- because key is to be determined on each input whereas fold is to be
215-- determined only once for a key.
217-- | In a key value stream, fold values corresponding to each key with a key
218-- specific fold. The fold returns the fold result as the second component of
219-- the output tuple whenever a fold terminates. The first component of the
220-- tuple is a container of in-progress folds. If a fold terminates, another
221-- instance of the fold is started upon receiving an input with that key.
222--
223-- This can be used to scan a stream and collect the results from the scan
224-- output.
225--
226-- /Pre-release/
227--
228{-# INLINE demuxScanWith #-}
229demuxScanWith :: (Monad m, IsMap f, Traversable f) =>
230 (a -> Key f)
231 -> (a -> m (Fold m a b))
232 -> Fold m a (m (f b), Maybe (Key f, b))
233demuxScanWith getKey getFold = fmap extract $ foldlM' step initial
235 where
237 initial = return $ Tuple' IsMap.mapEmpty Nothing
239 {-# INLINE runFold #-}
240 runFold kv (Fold step1 initial1 extract1) (k, a) = do
241 res <- initial1
242 case res of
243 Partial s -> do
244 res1 <- step1 s a
245 return
246 $ case res1 of
247 Partial _ ->
248 let fld = Fold step1 (return res1) extract1
249 in Tuple' (IsMap.mapInsert k fld kv) Nothing
250 Done b -> Tuple' (IsMap.mapDelete k kv) (Just (k, b))
251 Done b -> return $ Tuple' kv (Just (k, b))
253 step (Tuple' kv _) a = do
254 let k = getKey a
255 case IsMap.mapLookup k kv of
256 Nothing -> do
257 fld <- getFold a
258 runFold kv fld (k, a)
259 Just f -> runFold kv f (k, a)
261 extract (Tuple' kv x) = (Prelude.mapM f kv, x)
263 where
265 f (Fold _ i e) = do
266 r <- i
267 case r of
268 Partial s -> e s
269 Done b -> return b
271-- | This is specialized version of 'demuxScanWith' that uses mutable cells as
272-- fold accumulators for better performance.
273--
274-- Definition:
275--
276-- >>> demuxScanMutWith = Fold.demuxScanWith
277--
278{-# INLINE demuxScanMutWith #-}
279demuxScanMutWith :: (MonadIO m, IsMap f, Traversable f) =>
280 (a -> Key f)
281 -> (a -> m (Fold m a b))
282 -> Fold m a (m (f b), Maybe (Key f, b))
283demuxScanMutWith getKey getFold = fmap extract $ foldlM' step initial
285 where
287 initial = return $ Tuple' IsMap.mapEmpty Nothing
289 {-# INLINE initFold #-}
290 initFold kv (Fold step1 initial1 extract1) (k, a) = do
291 res <- initial1
292 case res of
293 Partial s -> do
294 res1 <- step1 s a
295 case res1 of
296 Partial _ -> do
297 let fld = Fold step1 (return res1) extract1
298 ref <- liftIO $ newIORef fld
299 return $ Tuple' (IsMap.mapInsert k ref kv) Nothing
300 Done b -> return $ Tuple' kv (Just (k, b))
301 Done b -> return $ Tuple' kv (Just (k, b))
303 {-# INLINE runFold #-}
304 runFold kv ref (Fold step1 initial1 extract1) (k, a) = do
305 res <- initial1
306 case res of
307 Partial s -> do
308 res1 <- step1 s a
309 case res1 of
310 Partial _ -> do
311 let fld = Fold step1 (return res1) extract1
312 liftIO $ writeIORef ref fld
313 return $ Tuple' kv Nothing
314 Done b ->
315 let kv1 = IsMap.mapDelete k kv
316 in return $ Tuple' kv1 (Just (k, b))
317 Done _ -> error "demuxScanMutWith: unreachable"
319 step (Tuple' kv _) a = do
320 let k = getKey a
321 case IsMap.mapLookup k kv of
322 Nothing -> do
323 f <- getFold a
324 initFold kv f (k, a)
325 Just ref -> do
326 f <- liftIO $ readIORef ref
327 runFold kv ref f (k, a)
329 extract (Tuple' kv x) = (Prelude.mapM f kv, x)
331 where
333 f ref = do
334 (Fold _ i e) <- liftIO $ readIORef ref
335 r <- i
336 case r of
337 Partial s -> e s
338 Done b -> return b
340-- | This collects all the results of 'demuxScanWith' in a container.
341--
342{-# INLINE demuxWith #-}
343demuxWith :: (Monad m, IsMap f, Traversable f) =>
344 (a -> Key f) -> (a -> m (Fold m a b)) -> Fold m a (f b)
345demuxWith getKey getFold =
346 let
347 classifier = demuxScanWith getKey getFold
348 getMap Nothing = pure IsMap.mapEmpty
349 getMap (Just action) = action
350 aggregator =
351 teeWith IsMap.mapUnion
352 (rmapM getMap $ lmap fst latest)
353 (lmap snd $ catMaybes toMap)
354 in postscan classifier aggregator
356-- | Same as 'demuxWith' but uses 'demuxScanMutWith' for better performance.
357--
358-- Definition:
359--
360-- >>> demuxMutWith = Fold.demuxWith
361--
362{-# INLINE demuxMutWith #-}
363demuxMutWith :: (MonadIO m, IsMap f, Traversable f) =>
364 (a -> Key f) -> (a -> m (Fold m a b)) -> Fold m a (f b)
365demuxMutWith getKey getFold =
366 let
367 classifier = demuxScanMutWith getKey getFold
368 getMap Nothing = pure IsMap.mapEmpty
369 getMap (Just action) = action
370 aggregator =
371 teeWith IsMap.mapUnion
372 (rmapM getMap $ lmap fst latest)
373 (lmap snd $ catMaybes toMap)
374 in postscan classifier aggregator
376-- | Fold a stream of key value pairs using a function that maps keys to folds.
377--
378-- Definition:
379--
380-- >>> demux f = Fold.demuxWith fst (Fold.lmap snd . f)
381--
382-- Example:
383--
384-- >>> import Data.Map (Map)
385-- >>> :{
386-- let f "SUM" = return Fold.sum
387-- f _ = return Fold.product
388-- input = Stream.fromList [("SUM",1),("PRODUCT",2),("SUM",3),("PRODUCT",4)]
389-- in Stream.fold (Fold.demux f) input :: IO (Map String Int)
390-- :}
391-- fromList [("PRODUCT",8),("SUM",4)]
392--
393-- /Pre-release/
394{-# INLINE demux #-}
395demux :: (Monad m, IsMap f, Traversable f) =>
396 (Key f -> m (Fold m a b)) -> Fold m (Key f, a) (f b)
397demux f = demuxWith fst (\(k, _) -> fmap (lmap snd) (f k))
399------------------------------------------------------------------------------
400-- Classify: Like demux but uses the same fold for all keys.
401------------------------------------------------------------------------------
403-- XXX Change these to make the behavior similar to demux* variants. We can
404-- implement this using classifyScanManyWith. Maintain a set of done folds in
405-- the underlying monad, and when initial is called look it up, if the fold is
406-- done then initial would set a flag in the state to ignore the input or
407-- return an error.
409-- | Folds the values for each key using the supplied fold. When scanning, as
410-- soon as the fold is complete, its result is available in the second
411-- component of the tuple. The first component of the tuple is a snapshot of
412-- the in-progress folds.
413--
414-- Once the fold for a key is done, any future values of the key are ignored.
415--
416-- Definition:
417--
418-- >>> classifyScanWith f fld = Fold.demuxScanWith f (const fld)
419--
420{-# INLINE classifyScanWith #-}
421classifyScanWith :: (Monad m, IsMap f, Traversable f, Ord (Key f)) =>
422 -- Note: we need to return the Map itself to display the in-progress values
423 -- e.g. to implement top. We could possibly create a separate abstraction
424 -- for that use case. We return an action because we want it to be lazy so
425 -- that the downstream consumers can choose to process or discard it.
426 (a -> Key f) -> Fold m a b -> Fold m a (m (f b), Maybe (Key f, b))
427classifyScanWith f (Fold step1 initial1 extract1) =
428 fmap extract $ foldlM' step initial
430 where
432 initial = return $ Tuple3' IsMap.mapEmpty Set.empty Nothing
434 {-# INLINE initFold #-}
435 initFold kv set k a = do
436 x <- initial1
437 case x of
438 Partial s -> do
439 r <- step1 s a
440 return
441 $ case r of
442 Partial s1 ->
443 Tuple3' (IsMap.mapInsert k s1 kv) set Nothing
444 Done b ->
445 Tuple3' kv set (Just (k, b))
446 Done b -> return (Tuple3' kv (Set.insert k set) (Just (k, b)))
448 step (Tuple3' kv set _) a = do
449 let k = f a
450 case IsMap.mapLookup k kv of
451 Nothing -> do
452 if Set.member k set
453 then return (Tuple3' kv set Nothing)
454 else initFold kv set k a
455 Just s -> do
456 r <- step1 s a
457 return
458 $ case r of
459 Partial s1 ->
460 Tuple3' (IsMap.mapInsert k s1 kv) set Nothing
461 Done b ->
462 let kv1 = IsMap.mapDelete k kv
463 in Tuple3' kv1 (Set.insert k set) (Just (k, b))
465 extract (Tuple3' kv _ x) = (Prelude.mapM extract1 kv, x)
467-- XXX we can use a Prim IORef if we can constrain the state "s" to be Prim
468--
469-- The code is almost the same as classifyScanWith except the IORef operations.
471-- | Same as classifyScanWith except that it uses mutable IORef cells in the
472-- Map providing better performance. Be aware that if this is used as a scan,
473-- the values in the intermediate Maps would be mutable.
474--
475-- Definitions:
476--
477-- >>> classifyScanMutWith = Fold.classifyScanWith
478-- >>> classifyScanMutWith f fld = Fold.demuxScanMutWith f (const fld)
479--
480{-# INLINE classifyScanMutWith #-}
481classifyScanMutWith :: (MonadIO m, IsMap f, Traversable f, Ord (Key f)) =>
482 (a -> Key f) -> Fold m a b -> Fold m a (m (f b), Maybe (Key f, b))
483classifyScanMutWith f (Fold step1 initial1 extract1) =
484 fmap extract $ foldlM' step initial
486 where
488 initial = return $ Tuple3' IsMap.mapEmpty Set.empty Nothing
490 {-# INLINE initFold #-}
491 initFold kv set k a = do
492 x <- initial1
493 case x of
494 Partial s -> do
495 r <- step1 s a
496 case r of
497 Partial s1 -> do
498 ref <- liftIO $ newIORef s1
499 return $ Tuple3' (IsMap.mapInsert k ref kv) set Nothing
500 Done b ->
501 return $ Tuple3' kv set (Just (k, b))
502 Done b -> return (Tuple3' kv (Set.insert k set) (Just (k, b)))
504 step (Tuple3' kv set _) a = do
505 let k = f a
506 case IsMap.mapLookup k kv of
507 Nothing -> do
508 if Set.member k set
509 then return (Tuple3' kv set Nothing)
510 else initFold kv set k a
511 Just ref -> do
512 s <- liftIO $ readIORef ref
513 r <- step1 s a
514 case r of
515 Partial s1 -> do
516 liftIO $ writeIORef ref s1
517 return $ Tuple3' kv set Nothing
518 Done b ->
519 let kv1 = IsMap.mapDelete k kv
520 in return
521 $ Tuple3' kv1 (Set.insert k set) (Just (k, b))
523 extract (Tuple3' kv _ x) =
524 (Prelude.mapM (\ref -> liftIO (readIORef ref) >>= extract1) kv, x)
526-- | Fold a key value stream to a key-value container. If the same key appears
527-- multiple times, only the last value is retained.
528{-# INLINE toMap #-}
529toMap :: (Monad m, IsMap f) => Fold m (Key f, a) (f a)
530toMap = foldl' (\kv (k, v) -> IsMap.mapInsert k v kv) IsMap.mapEmpty
532-- | Split the input stream based on a key field and fold each split using the
533-- given fold. Useful for map/reduce, bucketizing the input in different bins
534-- or for generating histograms.
535--
536-- Example:
537--
538-- >>> import Data.Map.Strict (Map)
539-- >>> :{
540-- let input = Stream.fromList [("ONE",1),("ONE",1.1),("TWO",2), ("TWO",2.2)]
541-- classify = Fold.classifyWith fst (Fold.lmap snd Fold.toList)
542-- in Stream.fold classify input :: IO (Map String [Double])
543-- :}
544-- fromList [("ONE",[1.0,1.1]),("TWO",[2.0,2.2])]
545--
546-- Once the classifier fold terminates for a particular key any further inputs
547-- in that bucket are ignored.
548--
549-- Space used is proportional to the number of keys seen till now and
550-- monotonically increases because it stores whether a key has been seen or
551-- not.
552--
553-- /Stops: never/
554--
555-- /Pre-release/
556--
557{-# INLINE classifyWith #-}
558classifyWith :: (Monad m, IsMap f, Traversable f, Ord (Key f)) =>
559 (a -> Key f) -> Fold m a b -> Fold m a (f b)
560classifyWith f fld =
561 let
562 classifier = classifyScanWith f fld
563 getMap Nothing = pure IsMap.mapEmpty
564 getMap (Just action) = action
565 aggregator =
566 teeWith IsMap.mapUnion
567 (rmapM getMap $ lmap fst latest)
568 (lmap snd $ catMaybes toMap)
569 in postscan classifier aggregator
571-- | Same as 'classifyWith' but maybe faster because it uses mutable cells as
572-- fold accumulators in the Map.
573--
574-- Definition:
575--
576-- >>> classifyMutWith = Fold.classifyWith
577--
578{-# INLINE classifyMutWith #-}
579classifyMutWith :: (MonadIO m, IsMap f, Traversable f, Ord (Key f)) =>
580 (a -> Key f) -> Fold m a b -> Fold m a (f b)
581classifyMutWith f fld =
582 let
583 classifier = classifyScanMutWith f fld
584 getMap Nothing = pure IsMap.mapEmpty
585 getMap (Just action) = action
586 aggregator =
587 teeWith IsMap.mapUnion
588 (rmapM getMap $ lmap fst latest)
589 (lmap snd $ catMaybes toMap)
590 in postscan classifier aggregator
592-- | Given an input stream of key value pairs and a fold for values, fold all
593-- the values belonging to each key. Useful for map/reduce, bucketizing the
594-- input in different bins or for generating histograms.
595--
596-- Definition:
597--
598-- >>> classify = Fold.classifyWith fst . Fold.lmap snd
599--
600-- Example:
601--
602-- >>> :{
603-- let input = Stream.fromList [("ONE",1),("ONE",1.1),("TWO",2), ("TWO",2.2)]
604-- in Stream.fold (Fold.classify Fold.toList) input
605-- :}
606-- fromList [("ONE",[1.0,1.1]),("TWO",[2.0,2.2])]
607--
608-- /Pre-release/
609{-# INLINE classify #-}
610classify :: (Monad m, Ord k) => Fold m a b -> Fold m (k, a) (Map k b)
611classify = classifyWith fst . lmap snd
613-- | Determine the frequency of each element in the stream.
614--
615-- You can just collect the keys of the resulting map to get the unique
616-- elements in the stream.
617--
618-- Definition:
619--
620-- >>> frequency = Fold.classifyWith id Fold.length
621--
622{-# INLINE frequency #-}
623frequency :: (Monad m, Ord a) => Fold m a (Map a Int)
624frequency = classifyWith id length