git.iko.soy/streamly
on ed37c19 help sign in
streamly/core/src/Streamly/Internal/Data/Fold/Container.hs
783 lines 26 kB history raw
1-- |
2-- Module : Streamly.Internal.Data.Fold.Container
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.Container
11 (
12 -- * Set operations
13 toSet
14 , toIntSet
15 , countDistinct
16 , countDistinctInt
17 , nub
18 , nubInt
20 -- * Map operations
21 , frequency
23 -- ** Demultiplexing
24 -- | Direct values in the input stream to different folds using an n-ary
25 -- fold selector. 'demux' is a generalization of 'classify' (and
26 -- 'partition') where each key of the classifier can use a different fold.
27 --
28 -- You need to see only 'demux' if you are looking to find the capabilities
29 -- of these combinators, all others are variants of that.
31 -- *** Output is a container
32 -- | The fold state snapshot returns the key-value container of in-progress
33 -- folds.
34 , demuxToContainer
35 , demuxToContainerIO
36 , demuxToMap
37 , demuxToMapIO
39 -- *** Input is explicit key-value tuple
40 -- | Like above but inputs are in explicit key-value pair form.
41 , demuxKvToContainer
42 , demuxKvToMap
44 -- *** Scan of finished fold results
45 -- | Like above, but the resulting fold state snapshot contains the key
46 -- value container as well as the finished key result if a fold in the
47 -- container finished.
48 , demuxGeneric
49 , demux
50 , demuxGenericIO
51 , demuxIO
53 -- TODO: These can be implemented using the above operations
54 -- , demuxSel -- Stop when the fold for the specified key stops
55 -- , demuxMin -- Stop when any of the folds stop
56 -- , demuxAll -- Stop when all the folds stop (run once)
58 -- ** Classifying
59 -- | In an input stream of key value pairs fold values for different keys
60 -- in individual output buckets using the given fold. 'classify' is a
61 -- special case of 'demux' where all the branches of the demultiplexer use
62 -- the same fold.
63 --
64 -- Different types of maps can be used with these combinators via the IsMap
65 -- type class. Hashmap performs better when there are more collisions, trie
66 -- Map performs better otherwise. Trie has an advantage of sorting the keys
67 -- at the same time. For example if we want to store a dictionary of words
68 -- and their meanings then trie Map would be better if we also want to
69 -- display them in sorted order.
71 , kvToMap
73 , toContainer
74 , toContainerIO
75 , toMap
76 , toMapIO
78 , classifyGeneric
79 , classify
80 , classifyGenericIO
81 , classifyIO
82 -- , toContainerSel
83 -- , toContainerMin
84 )
85where
87#include "inline.hs"
88#include "ArrayMacros.h"
90import Control.Monad.IO.Class (MonadIO(..))
91import Data.IORef (newIORef, readIORef, writeIORef)
92import Data.Map.Strict (Map)
93import Data.IntSet (IntSet)
94import Data.Set (Set)
95import Streamly.Internal.Data.IsMap (IsMap(..))
96import Streamly.Internal.Data.Tuple.Strict (Tuple'(..), Tuple3'(..))
98import qualified Data.IntSet as IntSet
99import qualified Data.Set as Set
100import qualified Streamly.Internal.Data.IsMap as IsMap
102import Prelude hiding (Foldable(..))
103import Streamly.Internal.Data.Fold.Type
104import Streamly.Internal.Data.Fold.Combinators
106-- $setup
107-- >>> :m
108-- >>> :set -XFlexibleContexts
109-- >>> import qualified Data.Map as Map
110-- >>> import qualified Data.Set as Set
111-- >>> import qualified Data.IntSet as IntSet
112-- >>> import qualified Streamly.Data.Fold as Fold
113-- >>> import qualified Streamly.Data.Stream as Stream
114-- >>> import qualified Streamly.Internal.Data.Fold.Container as Fold
116-- | Fold the input to a set.
117--
118-- Definition:
119--
120-- >>> toSet = Fold.foldl' (flip Set.insert) Set.empty
121--
122{-# INLINE toSet #-}
123toSet :: (Monad m, Ord a) => Fold m a (Set a)
124toSet = foldl' (flip Set.insert) Set.empty
126-- | Fold the input to an int set. For integer inputs this performs better than
127-- 'toSet'.
128--
129-- Definition:
130--
131-- >>> toIntSet = Fold.foldl' (flip IntSet.insert) IntSet.empty
132--
133{-# INLINE toIntSet #-}
134toIntSet :: Monad m => Fold m Int IntSet
135toIntSet = foldl' (flip IntSet.insert) IntSet.empty
137-- XXX Name as nubOrd? Or write a nubGeneric
139-- | Used as a scan. Returns 'Just' for the first occurrence of an element,
140-- returns 'Nothing' for any other occurrences.
141--
142-- Example:
143--
144-- >>> stream = Stream.fromList [1::Int,1,2,3,4,4,5,1,5,7]
145-- >>> Stream.fold Fold.toList $ Stream.scanMaybe Fold.nub stream
146-- [1,2,3,4,5,7]
147--
148-- /Pre-release/
149{-# INLINE nub #-}
150nub :: (Monad m, Ord a) => Fold m a (Maybe a)
151nub = fmap (\(Tuple' _ x) -> x) $ foldl' step initial
153 where
155 initial = Tuple' Set.empty Nothing
157 step (Tuple' set _) x =
158 if Set.member x set
159 then Tuple' set Nothing
160 else Tuple' (Set.insert x set) (Just x)
162-- | Like 'nub' but specialized to a stream of 'Int', for better performance.
163--
164-- /Pre-release/
165{-# INLINE nubInt #-}
166nubInt :: Monad m => Fold m Int (Maybe Int)
167nubInt = fmap (\(Tuple' _ x) -> x) $ foldl' step initial
169 where
171 initial = Tuple' IntSet.empty Nothing
173 step (Tuple' set _) x =
174 if IntSet.member x set
175 then Tuple' set Nothing
176 else Tuple' (IntSet.insert x set) (Just x)
178-- XXX Try Hash set
179-- XXX Add a countDistinct window fold
180-- XXX Add a bloom filter fold
182-- | Count non-duplicate elements in the stream.
183--
184-- Definition:
185--
186-- >>> countDistinct = fmap Set.size Fold.toSet
187-- >>> countDistinct = Fold.postscan Fold.nub $ Fold.catMaybes $ Fold.length
188--
189-- The memory used is proportional to the number of distinct elements in the
190-- stream, to guard against using too much memory use it as a scan and
191-- terminate if the count reaches more than a threshold.
192--
193-- /Space/: \(\mathcal{O}(n)\)
194--
195-- /Pre-release/
196--
197{-# INLINE countDistinct #-}
198countDistinct :: (Monad m, Ord a) => Fold m a Int
199-- countDistinct = postscan nub $ catMaybes length
200countDistinct = fmap Set.size toSet
201{-
202countDistinct = fmap (\(Tuple' _ n) -> n) $ foldl' step initial
204 where
206 initial = Tuple' Set.empty 0
208 step (Tuple' set n) x = do
209 if Set.member x set
210 then
211 Tuple' set n
212 else
213 let cnt = n + 1
214 in Tuple' (Set.insert x set) cnt
215-}
217-- | Like 'countDistinct' but specialized to a stream of 'Int', for better
218-- performance.
219--
220-- Definition:
221--
222-- >>> countDistinctInt = fmap IntSet.size Fold.toIntSet
223-- >>> countDistinctInt = Fold.postscan Fold.nubInt $ Fold.catMaybes $ Fold.length
224--
225-- /Pre-release/
226{-# INLINE countDistinctInt #-}
227countDistinctInt :: Monad m => Fold m Int Int
228-- countDistinctInt = postscan nubInt $ catMaybes length
229countDistinctInt = fmap IntSet.size toIntSet
230{-
231countDistinctInt = fmap (\(Tuple' _ n) -> n) $ foldl' step initial
233 where
235 initial = Tuple' IntSet.empty 0
237 step (Tuple' set n) x = do
238 if IntSet.member x set
239 then
240 Tuple' set n
241 else
242 let cnt = n + 1
243 in Tuple' (IntSet.insert x set) cnt
244 -}
246------------------------------------------------------------------------------
247-- demux: in a key value stream fold each key sub-stream with a different fold
248------------------------------------------------------------------------------
250-- TODO Demultiplex an input element into a number of typed variants. We want
251-- to statically restrict the target values within a set of predefined types,
252-- an enumeration of a GADT.
253--
254-- This is the consumer side dual of the producer side 'mux' operation (XXX to
255-- be implemented).
256--
257-- XXX If we use Refold in it, it can perhaps fuse/be more efficient. For
258-- example we can store just the result rather than storing the whole fold in
259-- the Map.
260--
261-- Note: There are separate functions to determine Key and Fold from the input
262-- because key is to be determined on each input whereas fold is to be
263-- determined only once for a key.
265-- | This is the most general of all demux, classify operations.
266--
267-- See 'demux' for documentation.
268{-# INLINE demuxGeneric #-}
269demuxGeneric :: (Monad m, IsMap f, Traversable f) =>
270 (a -> Key f)
271 -> (a -> m (Fold m a b))
272 -> Fold m a (m (f b), Maybe (Key f, b))
273demuxGeneric getKey getFold =
274 Fold (\s a -> Partial <$> step s a) (Partial <$> initial) extract final
276 where
278 initial = return $ Tuple' IsMap.mapEmpty Nothing
280 {-# INLINE runFold #-}
281 runFold kv (Fold step1 initial1 extract1 final1) (k, a) = do
282 res <- initial1
283 case res of
284 Partial s -> do
285 res1 <- step1 s a
286 return
287 $ case res1 of
288 Partial _ ->
289 let fld = Fold step1 (return res1) extract1 final1
290 in Tuple' (IsMap.mapInsert k fld kv) Nothing
291 Done b -> Tuple' (IsMap.mapDelete k kv) (Just (k, b))
292 Done b -> return $ Tuple' kv (Just (k, b))
294 step (Tuple' kv _) a = do
295 let k = getKey a
296 case IsMap.mapLookup k kv of
297 Nothing -> do
298 fld <- getFold a
299 runFold kv fld (k, a)
300 Just f -> runFold kv f (k, a)
302 extract (Tuple' kv x) = return (Prelude.mapM f kv, x)
304 where
306 f (Fold _ i e _) = do
307 r <- i
308 case r of
309 Partial s -> e s
310 _ -> error "demuxGeneric: unreachable code"
312 final (Tuple' kv x) = return (Prelude.mapM f kv, x)
314 where
316 f (Fold _ i _ fin) = do
317 r <- i
318 case r of
319 Partial s -> fin s
320 _ -> error "demuxGeneric: unreachable code"
322-- | @demux getKey getFold@: In a key value stream, fold values corresponding
323-- to each key using a key specific fold. @getFold@ is invoked to generate a
324-- key specific fold when a key is encountered for the first time in the
325-- stream.
326--
327-- The first component of the output tuple is a key-value Map of in-progress
328-- folds. The fold returns the fold result as the second component of the
329-- output tuple whenever a fold terminates.
330--
331-- If a fold terminates, another instance of the fold is started upon receiving
332-- an input with that key, @getFold@ is invoked again whenever the key is
333-- encountered again.
334--
335-- This can be used to scan a stream and collect the results from the scan
336-- output.
337--
338-- Since the fold generator function is monadic we can add folds dynamically.
339-- For example, we can maintain a Map of keys to folds in an IORef and lookup
340-- the fold from that corresponding to a key. This Map can be changed
341-- dynamically, folds for new keys can be added or folds for old keys can be
342-- deleted or modified.
343--
344-- Compare with 'classify', the fold in 'classify' is a static fold.
345--
346-- /Pre-release/
347--
348{-# INLINE demux #-}
349demux :: (Monad m, Ord k) =>
350 (a -> k)
351 -> (a -> m (Fold m a b))
352 -> Fold m a (m (Map k b), Maybe (k, b))
353demux = demuxGeneric
355-- | This is specialized version of 'demuxGeneric' that uses mutable IO cells
356-- as fold accumulators for better performance.
357{-# INLINE demuxGenericIO #-}
358demuxGenericIO :: (MonadIO m, IsMap f, Traversable f) =>
359 (a -> Key f)
360 -> (a -> m (Fold m a b))
361 -> Fold m a (m (f b), Maybe (Key f, b))
362demuxGenericIO getKey getFold =
363 Fold (\s a -> Partial <$> step s a) (Partial <$> initial) extract final
365 where
367 initial = return $ Tuple' IsMap.mapEmpty Nothing
369 {-# INLINE initFold #-}
370 initFold kv (Fold step1 initial1 extract1 final1) (k, a) = do
371 res <- initial1
372 case res of
373 Partial s -> do
374 res1 <- step1 s a
375 case res1 of
376 Partial _ -> do
377 let fld = Fold step1 (return res1) extract1 final1
378 ref <- liftIO $ newIORef fld
379 return $ Tuple' (IsMap.mapInsert k ref kv) Nothing
380 Done b -> return $ Tuple' kv (Just (k, b))
381 Done b -> return $ Tuple' kv (Just (k, b))
383 {-# INLINE runFold #-}
384 runFold kv ref (Fold step1 initial1 extract1 final1) (k, a) = do
385 res <- initial1
386 case res of
387 Partial s -> do
388 res1 <- step1 s a
389 case res1 of
390 Partial _ -> do
391 let fld = Fold step1 (return res1) extract1 final1
392 liftIO $ writeIORef ref fld
393 return $ Tuple' kv Nothing
394 Done b ->
395 let kv1 = IsMap.mapDelete k kv
396 in return $ Tuple' kv1 (Just (k, b))
397 Done _ -> error "demuxGenericIO: unreachable"
399 step (Tuple' kv _) a = do
400 let k = getKey a
401 case IsMap.mapLookup k kv of
402 Nothing -> do
403 f <- getFold a
404 initFold kv f (k, a)
405 Just ref -> do
406 f <- liftIO $ readIORef ref
407 runFold kv ref f (k, a)
409 extract (Tuple' kv x) = return (Prelude.mapM f kv, x)
411 where
413 f ref = do
414 Fold _ i e _ <- liftIO $ readIORef ref
415 r <- i
416 case r of
417 Partial s -> e s
418 _ -> error "demuxGenericIO: unreachable code"
420 final (Tuple' kv x) = return (Prelude.mapM f kv, x)
422 where
424 f ref = do
425 Fold _ i _ fin <- liftIO $ readIORef ref
426 r <- i
427 case r of
428 Partial s -> fin s
429 _ -> error "demuxGenericIO: unreachable code"
431-- | This is specialized version of 'demux' that uses mutable IO cells as
432-- fold accumulators for better performance.
433--
434-- Keep in mind that the values in the returned Map may be changed by the
435-- ongoing fold if you are using those concurrently in another thread.
436--
437{-# INLINE demuxIO #-}
438demuxIO :: (MonadIO m, Ord k) =>
439 (a -> k)
440 -> (a -> m (Fold m a b))
441 -> Fold m a (m (Map k b), Maybe (k, b))
442demuxIO = demuxGenericIO
444-- | Fold a key value stream to a key-value Map. If the same key appears
445-- multiple times, only the last value is retained.
446{-# INLINE kvToMapOverwriteGeneric #-}
447kvToMapOverwriteGeneric :: (Monad m, IsMap f) => Fold m (Key f, a) (f a)
448kvToMapOverwriteGeneric =
449 foldl' (\kv (k, v) -> IsMap.mapInsert k v kv) IsMap.mapEmpty
451{-# INLINE demuxToContainer #-}
452demuxToContainer :: (Monad m, IsMap f, Traversable f) =>
453 (a -> Key f) -> (a -> m (Fold m a b)) -> Fold m a (f b)
454demuxToContainer getKey getFold =
455 let
456 classifier = demuxGeneric getKey getFold
457 getMap Nothing = pure IsMap.mapEmpty
458 getMap (Just action) = action
459 aggregator =
460 teeWith IsMap.mapUnion
461 (rmapM getMap $ lmap fst latest)
462 (lmap snd $ catMaybes kvToMapOverwriteGeneric)
463 in postscan classifier aggregator
465-- | This collects all the results of 'demux' in a Map.
466--
467{-# INLINE demuxToMap #-}
468demuxToMap :: (Monad m, Ord k) =>
469 (a -> k) -> (a -> m (Fold m a b)) -> Fold m a (Map k b)
470demuxToMap = demuxToContainer
472{-# INLINE demuxToContainerIO #-}
473demuxToContainerIO :: (MonadIO m, IsMap f, Traversable f) =>
474 (a -> Key f) -> (a -> m (Fold m a b)) -> Fold m a (f b)
475demuxToContainerIO getKey getFold =
476 let
477 classifier = demuxGenericIO getKey getFold
478 getMap Nothing = pure IsMap.mapEmpty
479 getMap (Just action) = action
480 aggregator =
481 teeWith IsMap.mapUnion
482 (rmapM getMap $ lmap fst latest)
483 (lmap snd $ catMaybes kvToMapOverwriteGeneric)
484 in postscan classifier aggregator
486-- | Same as 'demuxToMap' but uses 'demuxIO' for better performance.
487--
488{-# INLINE demuxToMapIO #-}
489demuxToMapIO :: (MonadIO m, Ord k) =>
490 (a -> k) -> (a -> m (Fold m a b)) -> Fold m a (Map k b)
491demuxToMapIO = demuxToContainerIO
493{-# INLINE demuxKvToContainer #-}
494demuxKvToContainer :: (Monad m, IsMap f, Traversable f) =>
495 (Key f -> m (Fold m a b)) -> Fold m (Key f, a) (f b)
496demuxKvToContainer f = demuxToContainer fst (\(k, _) -> fmap (lmap snd) (f k))
498-- | Fold a stream of key value pairs using a function that maps keys to folds.
499--
500-- Definition:
501--
502-- >>> demuxKvToMap f = Fold.demuxToContainer fst (Fold.lmap snd . f)
503--
504-- Example:
505--
506-- >>> import Data.Map (Map)
507-- >>> :{
508-- let f "SUM" = return Fold.sum
509-- f _ = return Fold.product
510-- input = Stream.fromList [("SUM",1),("PRODUCT",2),("SUM",3),("PRODUCT",4)]
511-- in Stream.fold (Fold.demuxKvToMap f) input :: IO (Map String Int)
512-- :}
513-- fromList [("PRODUCT",8),("SUM",4)]
514--
515-- /Pre-release/
516{-# INLINE demuxKvToMap #-}
517demuxKvToMap :: (Monad m, Ord k) =>
518 (k -> m (Fold m a b)) -> Fold m (k, a) (Map k b)
519demuxKvToMap = demuxKvToContainer
521------------------------------------------------------------------------------
522-- Classify: Like demux but uses the same fold for all keys.
523------------------------------------------------------------------------------
525-- XXX Change these to make the behavior similar to demux* variants. We can
526-- implement this using classifyScanManyWith. Maintain a set of done folds in
527-- the underlying monad, and when initial is called look it up, if the fold is
528-- done then initial would set a flag in the state to ignore the input or
529-- return an error.
531{-# INLINE classifyGeneric #-}
532classifyGeneric :: (Monad m, IsMap f, Traversable f, Ord (Key f)) =>
533 -- Note: we need to return the Map itself to display the in-progress values
534 -- e.g. to implement top. We could possibly create a separate abstraction
535 -- for that use case. We return an action because we want it to be lazy so
536 -- that the downstream consumers can choose to process or discard it.
537 (a -> Key f) -> Fold m a b -> Fold m a (m (f b), Maybe (Key f, b))
538classifyGeneric f (Fold step1 initial1 extract1 final1) =
539 Fold (\s a -> Partial <$> step s a) (Partial <$> initial) extract final
541 where
543 initial = return $ Tuple3' IsMap.mapEmpty Set.empty Nothing
545 {-# INLINE initFold #-}
546 initFold kv set k a = do
547 x <- initial1
548 case x of
549 Partial s -> do
550 r <- step1 s a
551 return
552 $ case r of
553 Partial s1 ->
554 Tuple3' (IsMap.mapInsert k s1 kv) set Nothing
555 Done b ->
556 Tuple3' kv set (Just (k, b))
557 Done b -> return (Tuple3' kv (Set.insert k set) (Just (k, b)))
559 step (Tuple3' kv set _) a = do
560 let k = f a
561 case IsMap.mapLookup k kv of
562 Nothing -> do
563 if Set.member k set
564 then return (Tuple3' kv set Nothing)
565 else initFold kv set k a
566 Just s -> do
567 r <- step1 s a
568 return
569 $ case r of
570 Partial s1 ->
571 Tuple3' (IsMap.mapInsert k s1 kv) set Nothing
572 Done b ->
573 let kv1 = IsMap.mapDelete k kv
574 in Tuple3' kv1 (Set.insert k set) (Just (k, b))
576 extract (Tuple3' kv _ x) = return (Prelude.mapM extract1 kv, x)
578 final (Tuple3' kv set x) = return (IsMap.mapTraverseWithKey f1 kv, x)
580 where
582 f1 k s = do
583 if Set.member k set
584 then extract1 s
585 else final1 s
587-- | Folds the values for each key using the supplied fold. When scanning, as
588-- soon as the fold is complete, its result is available in the second
589-- component of the tuple. The first component of the tuple is a snapshot of
590-- the in-progress folds.
591--
592-- Once the fold for a key is done, any future values of the key are ignored.
593--
594-- Definition:
595--
596-- >>> classify f fld = Fold.demux f (const fld)
597--
598{-# INLINE classify #-}
599classify :: (Monad m, Ord k) =>
600 (a -> k) -> Fold m a b -> Fold m a (m (Map k b), Maybe (k, b))
601classify = classifyGeneric
603-- XXX we can use a Prim IORef if we can constrain the state "s" to be Prim
604--
605-- The code is almost the same as classifyGeneric except the IORef operations.
607{-# INLINE classifyGenericIO #-}
608classifyGenericIO :: (MonadIO m, IsMap f, Traversable f, Ord (Key f)) =>
609 (a -> Key f) -> Fold m a b -> Fold m a (m (f b), Maybe (Key f, b))
610classifyGenericIO f (Fold step1 initial1 extract1 final1) =
611 Fold (\s a -> Partial <$> step s a) (Partial <$> initial) extract final
613 where
615 initial = return $ Tuple3' IsMap.mapEmpty Set.empty Nothing
617 {-# INLINE initFold #-}
618 initFold kv set k a = do
619 x <- initial1
620 case x of
621 Partial s -> do
622 r <- step1 s a
623 case r of
624 Partial s1 -> do
625 ref <- liftIO $ newIORef s1
626 return $ Tuple3' (IsMap.mapInsert k ref kv) set Nothing
627 Done b ->
628 return $ Tuple3' kv set (Just (k, b))
629 Done b -> return (Tuple3' kv (Set.insert k set) (Just (k, b)))
631 step (Tuple3' kv set _) a = do
632 let k = f a
633 case IsMap.mapLookup k kv of
634 Nothing -> do
635 if Set.member k set
636 then return (Tuple3' kv set Nothing)
637 else initFold kv set k a
638 Just ref -> do
639 s <- liftIO $ readIORef ref
640 r <- step1 s a
641 case r of
642 Partial s1 -> do
643 liftIO $ writeIORef ref s1
644 return $ Tuple3' kv set Nothing
645 Done b ->
646 let kv1 = IsMap.mapDelete k kv
647 in return
648 $ Tuple3' kv1 (Set.insert k set) (Just (k, b))
650 extract (Tuple3' kv _ x) = return (Prelude.mapM g kv, x)
652 where
654 g ref = liftIO (readIORef ref) >>= extract1
656 final (Tuple3' kv set x) = return (IsMap.mapTraverseWithKey g kv, x)
658 where
660 g k ref = do
661 s <- liftIO $ readIORef ref
662 if Set.member k set
663 then extract1 s
664 else final1 s
666-- | Same as classify except that it uses mutable IORef cells in the
667-- Map providing better performance. Be aware that if this is used as a scan,
668-- the values in the intermediate Maps would be mutable.
669--
670-- Definitions:
671--
672-- >>> classifyIO f fld = Fold.demuxIO f (const fld)
673--
674{-# INLINE classifyIO #-}
675classifyIO :: (MonadIO m, Ord k) =>
676 (a -> k) -> Fold m a b -> Fold m a (m (Map k b), Maybe (k, b))
677classifyIO = classifyGenericIO
679{-# INLINE toContainer #-}
680toContainer :: (Monad m, IsMap f, Traversable f, Ord (Key f)) =>
681 (a -> Key f) -> Fold m a b -> Fold m a (f b)
682toContainer f fld =
683 let
684 classifier = classifyGeneric f fld
685 getMap Nothing = pure IsMap.mapEmpty
686 getMap (Just action) = action
687 aggregator =
688 teeWith IsMap.mapUnion
689 (rmapM getMap $ lmap fst latest)
690 (lmap snd $ catMaybes kvToMapOverwriteGeneric)
691 in postscan classifier aggregator
693-- | Split the input stream based on a key field and fold each split using the
694-- given fold. Useful for map/reduce, bucketizing the input in different bins
695-- or for generating histograms.
696--
697-- Example:
698--
699-- >>> import Data.Map.Strict (Map)
700-- >>> :{
701-- let input = Stream.fromList [("ONE",1),("ONE",1.1),("TWO",2), ("TWO",2.2)]
702-- classify = Fold.toMap fst (Fold.lmap snd Fold.toList)
703-- in Stream.fold classify input :: IO (Map String [Double])
704-- :}
705-- fromList [("ONE",[1.0,1.1]),("TWO",[2.0,2.2])]
706--
707-- Once the classifier fold terminates for a particular key any further inputs
708-- in that bucket are ignored.
709--
710-- Space used is proportional to the number of keys seen till now and
711-- monotonically increases because it stores whether a key has been seen or
712-- not.
713--
714-- See 'demuxToMap' for a more powerful version where you can use a different
715-- fold for each key. A simpler version of 'toMap' retaining only the last
716-- value for a key can be written as:
717--
718-- >>> toMap = Fold.foldl' (\kv (k, v) -> Map.insert k v kv) Map.empty
719--
720-- /Stops: never/
721--
722-- /Pre-release/
723--
724{-# INLINE toMap #-}
725toMap :: (Monad m, Ord k) =>
726 (a -> k) -> Fold m a b -> Fold m a (Map k b)
727toMap = toContainer
729{-# INLINE toContainerIO #-}
730toContainerIO :: (MonadIO m, IsMap f, Traversable f, Ord (Key f)) =>
731 (a -> Key f) -> Fold m a b -> Fold m a (f b)
732toContainerIO f fld =
733 let
734 classifier = classifyGenericIO f fld
735 getMap Nothing = pure IsMap.mapEmpty
736 getMap (Just action) = action
737 aggregator =
738 teeWith IsMap.mapUnion
739 (rmapM getMap $ lmap fst latest)
740 (lmap snd $ catMaybes kvToMapOverwriteGeneric)
741 in postscan classifier aggregator
743-- | Same as 'toMap' but maybe faster because it uses mutable cells as
744-- fold accumulators in the Map.
745--
746{-# INLINE toMapIO #-}
747toMapIO :: (MonadIO m, Ord k) =>
748 (a -> k) -> Fold m a b -> Fold m a (Map k b)
749toMapIO = toContainerIO
751-- | Given an input stream of key value pairs and a fold for values, fold all
752-- the values belonging to each key. Useful for map/reduce, bucketizing the
753-- input in different bins or for generating histograms.
754--
755-- Definition:
756--
757-- >>> kvToMap = Fold.toMap fst . Fold.lmap snd
758--
759-- Example:
760--
761-- >>> :{
762-- let input = Stream.fromList [("ONE",1),("ONE",1.1),("TWO",2), ("TWO",2.2)]
763-- in Stream.fold (Fold.kvToMap Fold.toList) input
764-- :}
765-- fromList [("ONE",[1.0,1.1]),("TWO",[2.0,2.2])]
766--
767-- /Pre-release/
768{-# INLINE kvToMap #-}
769kvToMap :: (Monad m, Ord k) => Fold m a b -> Fold m (k, a) (Map k b)
770kvToMap = toMap fst . lmap snd
772-- | Determine the frequency of each element in the stream.
773--
774-- You can just collect the keys of the resulting map to get the unique
775-- elements in the stream.
776--
777-- Definition:
778--
779-- >>> frequency = Fold.toMap id Fold.length
780--
781{-# INLINE frequency #-}
782frequency :: (Monad m, Ord a) => Fold m a (Map a Int)
783frequency = toMap id length