2-- Module : Streamly.Internal.Data.Fold.Container
3-- Copyright : (c) 2019 Composewell Technologies
5-- Maintainer : streamly@composewell.com
6-- Stability : experimental
10module Streamly.Internal.Data.Fold.Container
27 -- | Direct values in the input stream to different folds using an n-ary
28 -- fold selector. 'demux' is a generalization of 'classify' (and
29 -- 'partition') where each key of the classifier can use a different fold.
31 -- You need to see only 'demux' if you are looking to find the capabilities
32 -- of these combinators, all others are variants of that.
34 -- *** Output is a container
35 -- | The fold state snapshot returns the key-value container of in-progress
42 -- *** Input is explicit key-value tuple
43 -- | Like above but inputs are in explicit key-value pair form.
47 -- *** Scan of finished fold results
48 -- | Like above, but the resulting fold state snapshot contains the key
49 -- value container as well as the finished key result if a fold in the
50 -- container finished.
56 -- TODO: These can be implemented using the above operations
57 -- , demuxSel -- Stop when the fold for the specified key stops
58 -- , demuxMin -- Stop when any of the folds stop
59 -- , demuxAll -- Stop when all the folds stop (run once)
62 -- | In an input stream of key value pairs fold values for different keys
63 -- in individual output buckets using the given fold. 'classify' is a
64 -- special case of 'demux' where all the branches of the demultiplexer use
67 -- Different types of maps can be used with these combinators via the IsMap
68 -- type class. Hashmap performs better when there are more collisions, trie
69 -- Map performs better otherwise. Trie has an advantage of sorting the keys
70 -- at the same time. For example if we want to store a dictionary of words
71 -- and their meanings then trie Map would be better if we also want to
72 -- display them in sorted order.
91#include "ArrayMacros.h"
93import Control.Monad.IO.Class (MonadIO(..))
94import Data.IORef (newIORef, readIORef, writeIORef)
95import Data.Map.Strict (Map)
96import Data.IntSet (IntSet)
98import Streamly.Internal.Data.IsMap (IsMap(..))
99import Streamly.Internal.Data.Tuple.Strict (Tuple'(..), Tuple3'(..))
101import qualified Data.IntSet as IntSet
102import qualified Data.Set as Set
103import qualified Streamly.Internal.Data.IsMap as IsMap
105import Prelude hiding (Foldable(..))
106import Streamly.Internal.Data.Fold
110-- >>> :set -XFlexibleContexts
111-- >>> import qualified Data.Map as Map
112-- >>> import qualified Data.Set as Set
113-- >>> import qualified Data.IntSet as IntSet
114-- >>> import qualified Streamly.Data.Fold as Fold
115-- >>> import qualified Streamly.Data.Stream as Stream
116-- >>> import qualified Streamly.Internal.Data.Fold.Container as Fold
118-- | Fold the input to a set.
122-- >>> toSet = Fold.foldl' (flip Set.insert) Set.empty
125toSet :: (Monad m, Ord a) => Fold m a (Set a)
126toSet = foldl' (flip Set.insert) Set.empty
128-- | Fold the input to an int set. For integer inputs this performs better than
133-- >>> toIntSet = Fold.foldl' (flip IntSet.insert) IntSet.empty
135{-# INLINE toIntSet #-}
136toIntSet :: Monad m => Fold m Int IntSet
137toIntSet = foldl' (flip IntSet.insert) IntSet.empty
139-- XXX Name as nubOrd? Or write a nubGeneric
141-- | Used as a scan. Returns 'Just' for the first occurrence of an element,
142-- returns 'Nothing' for any other occurrences.
146-- >>> stream = Stream.fromList [1::Int,1,2,3,4,4,5,1,5,7]
147-- >>> Stream.fold Fold.toList $ Stream.scanMaybe Fold.nub stream
152nub :: (Monad m, Ord a) => Fold m a (Maybe a)
153nub = fmap (\(Tuple' _ x) -> x) $ foldl' step initial
157 initial = Tuple' Set.empty Nothing
159 step (Tuple' set _) x =
161 then Tuple' set Nothing
162 else Tuple' (Set.insert x set) (Just x)
164-- | Like 'nub' but specialized to a stream of 'Int', for better performance.
168nubInt :: Monad m => Fold m Int (Maybe Int)
169nubInt = fmap (\(Tuple' _ x) -> x) $ foldl' step initial
173 initial = Tuple' IntSet.empty Nothing
175 step (Tuple' set _) x =
176 if IntSet.member x set
177 then Tuple' set Nothing
178 else Tuple' (IntSet.insert x set) (Just x)
181-- XXX Add a countDistinct window fold
182-- XXX Add a bloom filter fold
184-- | Count non-duplicate elements in the stream.
188-- >>> countDistinct = fmap Set.size Fold.toSet
189-- >>> countDistinct = Fold.postscan Fold.nub $ Fold.catMaybes $ Fold.length
191-- The memory used is proportional to the number of distinct elements in the
192-- stream, to guard against using too much memory use it as a scan and
193-- terminate if the count reaches more than a threshold.
195-- /Space/: \(\mathcal{O}(n)\)
199{-# INLINE countDistinct #-}
200countDistinct :: (Monad m, Ord a) => Fold m a Int
201-- countDistinct = postscan nub $ catMaybes length
202countDistinct = fmap Set.size toSet
204countDistinct = fmap (\(Tuple' _ n) -> n) $ foldl' step initial
208 initial = Tuple' Set.empty 0
210 step (Tuple' set n) x = do
216 in Tuple' (Set.insert x set) cnt
219-- | Like 'countDistinct' but specialized to a stream of 'Int', for better
224-- >>> countDistinctInt = fmap IntSet.size Fold.toIntSet
225-- >>> countDistinctInt = Fold.postscan Fold.nubInt $ Fold.catMaybes $ Fold.length
228{-# INLINE countDistinctInt #-}
229countDistinctInt :: Monad m => Fold m Int Int
230-- countDistinctInt = postscan nubInt $ catMaybes length
231countDistinctInt = fmap IntSet.size toIntSet
233countDistinctInt = fmap (\(Tuple' _ n) -> n) $ foldl' step initial
237 initial = Tuple' IntSet.empty 0
239 step (Tuple' set n) x = do
240 if IntSet.member x set
245 in Tuple' (IntSet.insert x set) cnt
248------------------------------------------------------------------------------
249-- demux: in a key value stream fold each key sub-stream with a different fold
250------------------------------------------------------------------------------
252-- TODO Demultiplex an input element into a number of typed variants. We want
253-- to statically restrict the target values within a set of predefined types,
254-- an enumeration of a GADT.
256-- This is the consumer side dual of the producer side 'mux' operation (XXX to
259-- XXX If we use Refold in it, it can perhaps fuse/be more efficient. For
260-- example we can store just the result rather than storing the whole fold in
263-- Note: There are separate functions to determine Key and Fold from the input
264-- because key is to be determined on each input whereas fold is to be
265-- determined only once for a key.
267-- | This is the most general of all demux, classify operations.
269-- See 'demux' for documentation.
270{-# INLINE demuxGeneric #-}
271demuxGeneric :: (Monad m, IsMap f, Traversable f) =>
273 -> (a -> m (Fold m a b))
274 -> Fold m a (m (f b), Maybe (Key f, b))
275demuxGeneric getKey getFold = fmap extract $ foldlM' step initial
279 initial = return $ Tuple' IsMap.mapEmpty Nothing
281 {-# INLINE runFold #-}
282 runFold kv (Fold step1 initial1 extract1) (k, a) = do
290 let fld = Fold step1 (return res1) extract1
291 in Tuple' (IsMap.mapInsert k fld kv) Nothing
292 Done b -> Tuple' (IsMap.mapDelete k kv) (Just (k, b))
293 Done b -> return $ Tuple' kv (Just (k, b))
295 step (Tuple' kv _) a = do
297 case IsMap.mapLookup k kv of
300 runFold kv fld (k, a)
301 Just f -> runFold kv f (k, a)
303 extract (Tuple' kv x) = (Prelude.mapM f kv, x)
313-- | @demux getKey getFold@: In a key value stream, fold values corresponding
314-- to each key using a key specific fold. @getFold@ is invoked to generate a
315-- key specific fold when a key is encountered for the first time in the
318-- The first component of the output tuple is a key-value Map of in-progress
319-- folds. The fold returns the fold result as the second component of the
320-- output tuple whenever a fold terminates.
322-- If a fold terminates, another instance of the fold is started upon receiving
323-- an input with that key, @getFold@ is invoked again whenever the key is
326-- This can be used to scan a stream and collect the results from the scan
329-- Since the fold generator function is monadic we can add folds dynamically.
330-- For example, we can maintain a Map of keys to folds in an IORef and lookup
331-- the fold from that corresponding to a key. This Map can be changed
332-- dynamically, folds for new keys can be added or folds for old keys can be
333-- deleted or modified.
335-- Compare with 'classify', the fold in 'classify' is a static fold.
340demux :: (Monad m, Ord k) =>
342 -> (a -> m (Fold m a b))
343 -> Fold m a (m (Map k b), Maybe (k, b))
346-- | This is specialized version of 'demuxGeneric' that uses mutable IO cells
347-- as fold accumulators for better performance.
348{-# INLINE demuxGenericIO #-}
349demuxGenericIO :: (MonadIO m, IsMap f, Traversable f) =>
351 -> (a -> m (Fold m a b))
352 -> Fold m a (m (f b), Maybe (Key f, b))
353demuxGenericIO getKey getFold = fmap extract $ foldlM' step initial
357 initial = return $ Tuple' IsMap.mapEmpty Nothing
359 {-# INLINE initFold #-}
360 initFold kv (Fold step1 initial1 extract1) (k, a) = do
367 let fld = Fold step1 (return res1) extract1
368 ref <- liftIO $ newIORef fld
369 return $ Tuple' (IsMap.mapInsert k ref kv) Nothing
370 Done b -> return $ Tuple' kv (Just (k, b))
371 Done b -> return $ Tuple' kv (Just (k, b))
373 {-# INLINE runFold #-}
374 runFold kv ref (Fold step1 initial1 extract1) (k, a) = do
381 let fld = Fold step1 (return res1) extract1
382 liftIO $ writeIORef ref fld
383 return $ Tuple' kv Nothing
385 let kv1 = IsMap.mapDelete k kv
386 in return $ Tuple' kv1 (Just (k, b))
387 Done _ -> error "demuxGenericIO: unreachable"
389 step (Tuple' kv _) a = do
391 case IsMap.mapLookup k kv of
396 f <- liftIO $ readIORef ref
397 runFold kv ref f (k, a)
399 extract (Tuple' kv x) = (Prelude.mapM f kv, x)
404 (Fold _ i e) <- liftIO $ readIORef ref
410-- | This is specialized version of 'demux' that uses mutable IO cells as
411-- fold accumulators for better performance.
413-- Keep in mind that the values in the returned Map may be changed by the
414-- ongoing fold if you are using those concurrently in another thread.
416{-# INLINE demuxIO #-}
417demuxIO :: (MonadIO m, Ord k) =>
419 -> (a -> m (Fold m a b))
420 -> Fold m a (m (Map k b), Maybe (k, b))
421demuxIO = demuxGenericIO
423{-# INLINE demuxToContainer #-}
424demuxToContainer :: (Monad m, IsMap f, Traversable f) =>
425 (a -> Key f) -> (a -> m (Fold m a b)) -> Fold m a (f b)
426demuxToContainer getKey getFold =
428 classifier = demuxGeneric getKey getFold
429 getMap Nothing = pure IsMap.mapEmpty
430 getMap (Just action) = action
432 teeWith IsMap.mapUnion
433 (rmapM getMap $ lmap fst latest)
434 (lmap snd $ catMaybes kvToMapOverwriteGeneric)
435 in postscan classifier aggregator
437-- | This collects all the results of 'demux' in a Map.
439{-# INLINE demuxToMap #-}
440demuxToMap :: (Monad m, Ord k) =>
441 (a -> k) -> (a -> m (Fold m a b)) -> Fold m a (Map k b)
442demuxToMap = demuxToContainer
444{-# INLINE demuxToContainerIO #-}
445demuxToContainerIO :: (MonadIO m, IsMap f, Traversable f) =>
446 (a -> Key f) -> (a -> m (Fold m a b)) -> Fold m a (f b)
447demuxToContainerIO getKey getFold =
449 classifier = demuxGenericIO getKey getFold
450 getMap Nothing = pure IsMap.mapEmpty
451 getMap (Just action) = action
453 teeWith IsMap.mapUnion
454 (rmapM getMap $ lmap fst latest)
455 (lmap snd $ catMaybes kvToMapOverwriteGeneric)
456 in postscan classifier aggregator
458-- | Same as 'demuxToMap' but uses 'demuxIO' for better performance.
460{-# INLINE demuxToMapIO #-}
461demuxToMapIO :: (MonadIO m, Ord k) =>
462 (a -> k) -> (a -> m (Fold m a b)) -> Fold m a (Map k b)
463demuxToMapIO = demuxToContainerIO
465{-# INLINE demuxKvToContainer #-}
466demuxKvToContainer :: (Monad m, IsMap f, Traversable f) =>
467 (Key f -> m (Fold m a b)) -> Fold m (Key f, a) (f b)
468demuxKvToContainer f = demuxToContainer fst (\(k, _) -> fmap (lmap snd) (f k))
470-- | Fold a stream of key value pairs using a function that maps keys to folds.
474-- >>> demuxKvToMap f = Fold.demuxToContainer fst (Fold.lmap snd . f)
478-- >>> import Data.Map (Map)
480-- let f "SUM" = return Fold.sum
481-- f _ = return Fold.product
482-- input = Stream.fromList [("SUM",1),("PRODUCT",2),("SUM",3),("PRODUCT",4)]
483-- in Stream.fold (Fold.demuxKvToMap f) input :: IO (Map String Int)
485-- fromList [("PRODUCT",8),("SUM",4)]
488{-# INLINE demuxKvToMap #-}
489demuxKvToMap :: (Monad m, Ord k) =>
490 (k -> m (Fold m a b)) -> Fold m (k, a) (Map k b)
491demuxKvToMap = demuxKvToContainer
493------------------------------------------------------------------------------
494-- Classify: Like demux but uses the same fold for all keys.
495------------------------------------------------------------------------------
497-- XXX Change these to make the behavior similar to demux* variants. We can
498-- implement this using classifyScanManyWith. Maintain a set of done folds in
499-- the underlying monad, and when initial is called look it up, if the fold is
500-- done then initial would set a flag in the state to ignore the input or
503{-# INLINE classifyGeneric #-}
504classifyGeneric :: (Monad m, IsMap f, Traversable f, Ord (Key f)) =>
505 -- Note: we need to return the Map itself to display the in-progress values
506 -- e.g. to implement top. We could possibly create a separate abstraction
507 -- for that use case. We return an action because we want it to be lazy so
508 -- that the downstream consumers can choose to process or discard it.
509 (a -> Key f) -> Fold m a b -> Fold m a (m (f b), Maybe (Key f, b))
510classifyGeneric f (Fold step1 initial1 extract1) =
511 fmap extract $ foldlM' step initial
515 initial = return $ Tuple3' IsMap.mapEmpty Set.empty Nothing
517 {-# INLINE initFold #-}
518 initFold kv set k a = do
526 Tuple3' (IsMap.mapInsert k s1 kv) set Nothing
528 Tuple3' kv set (Just (k, b))
529 Done b -> return (Tuple3' kv (Set.insert k set) (Just (k, b)))
531 step (Tuple3' kv set _) a = do
533 case IsMap.mapLookup k kv of
536 then return (Tuple3' kv set Nothing)
537 else initFold kv set k a
543 Tuple3' (IsMap.mapInsert k s1 kv) set Nothing
545 let kv1 = IsMap.mapDelete k kv
546 in Tuple3' kv1 (Set.insert k set) (Just (k, b))
548 extract (Tuple3' kv _ x) = (Prelude.mapM extract1 kv, x)
550-- | Folds the values for each key using the supplied fold. When scanning, as
551-- soon as the fold is complete, its result is available in the second
552-- component of the tuple. The first component of the tuple is a snapshot of
553-- the in-progress folds.
555-- Once the fold for a key is done, any future values of the key are ignored.
559-- >>> classify f fld = Fold.demux f (const fld)
561{-# INLINE classify #-}
562classify :: (Monad m, Ord k) =>
563 (a -> k) -> Fold m a b -> Fold m a (m (Map k b), Maybe (k, b))
564classify = classifyGeneric
566-- XXX we can use a Prim IORef if we can constrain the state "s" to be Prim
568-- The code is almost the same as classifyGeneric except the IORef operations.
570{-# INLINE classifyGenericIO #-}
571classifyGenericIO :: (MonadIO m, IsMap f, Traversable f, Ord (Key f)) =>
572 (a -> Key f) -> Fold m a b -> Fold m a (m (f b), Maybe (Key f, b))
573classifyGenericIO f (Fold step1 initial1 extract1) =
574 fmap extract $ foldlM' step initial
578 initial = return $ Tuple3' IsMap.mapEmpty Set.empty Nothing
580 {-# INLINE initFold #-}
581 initFold kv set k a = do
588 ref <- liftIO $ newIORef s1
589 return $ Tuple3' (IsMap.mapInsert k ref kv) set Nothing
591 return $ Tuple3' kv set (Just (k, b))
592 Done b -> return (Tuple3' kv (Set.insert k set) (Just (k, b)))
594 step (Tuple3' kv set _) a = do
596 case IsMap.mapLookup k kv of
599 then return (Tuple3' kv set Nothing)
600 else initFold kv set k a
602 s <- liftIO $ readIORef ref
606 liftIO $ writeIORef ref s1
607 return $ Tuple3' kv set Nothing
609 let kv1 = IsMap.mapDelete k kv
611 $ Tuple3' kv1 (Set.insert k set) (Just (k, b))
613 extract (Tuple3' kv _ x) =
614 (Prelude.mapM (\ref -> liftIO (readIORef ref) >>= extract1) kv, x)
616-- | Same as classify except that it uses mutable IORef cells in the
617-- Map providing better performance. Be aware that if this is used as a scan,
618-- the values in the intermediate Maps would be mutable.
622-- >>> classifyIO f fld = Fold.demuxIO f (const fld)
624{-# INLINE classifyIO #-}
625classifyIO :: (MonadIO m, Ord k) =>
626 (a -> k) -> Fold m a b -> Fold m a (m (Map k b), Maybe (k, b))
627classifyIO = classifyGenericIO
629-- | Fold a key value stream to a key-value Map. If the same key appears
630-- multiple times, only the last value is retained.
631{-# INLINE kvToMapOverwriteGeneric #-}
632kvToMapOverwriteGeneric :: (Monad m, IsMap f) => Fold m (Key f, a) (f a)
633kvToMapOverwriteGeneric =
634 foldl' (\kv (k, v) -> IsMap.mapInsert k v kv) IsMap.mapEmpty
636{-# INLINE toContainer #-}
637toContainer :: (Monad m, IsMap f, Traversable f, Ord (Key f)) =>
638 (a -> Key f) -> Fold m a b -> Fold m a (f b)
641 classifier = classifyGeneric f fld
642 getMap Nothing = pure IsMap.mapEmpty
643 getMap (Just action) = action
645 teeWith IsMap.mapUnion
646 (rmapM getMap $ lmap fst latest)
647 (lmap snd $ catMaybes kvToMapOverwriteGeneric)
648 in postscan classifier aggregator
650-- | Split the input stream based on a key field and fold each split using the
651-- given fold. Useful for map/reduce, bucketizing the input in different bins
652-- or for generating histograms.
656-- >>> import Data.Map.Strict (Map)
658-- let input = Stream.fromList [("ONE",1),("ONE",1.1),("TWO",2), ("TWO",2.2)]
659-- classify = Fold.toMap fst (Fold.lmap snd Fold.toList)
660-- in Stream.fold classify input :: IO (Map String [Double])
662-- fromList [("ONE",[1.0,1.1]),("TWO",[2.0,2.2])]
664-- Once the classifier fold terminates for a particular key any further inputs
665-- in that bucket are ignored.
667-- Space used is proportional to the number of keys seen till now and
668-- monotonically increases because it stores whether a key has been seen or
671-- See 'demuxToMap' for a more powerful version where you can use a different
672-- fold for each key. A simpler version of 'toMap' retaining only the last
673-- value for a key can be written as:
675-- >>> toMap = Fold.foldl' (\kv (k, v) -> Map.insert k v kv) Map.empty
682toMap :: (Monad m, Ord k) =>
683 (a -> k) -> Fold m a b -> Fold m a (Map k b)
686{-# INLINE toContainerIO #-}
687toContainerIO :: (MonadIO m, IsMap f, Traversable f, Ord (Key f)) =>
688 (a -> Key f) -> Fold m a b -> Fold m a (f b)
691 classifier = classifyGenericIO f fld
692 getMap Nothing = pure IsMap.mapEmpty
693 getMap (Just action) = action
695 teeWith IsMap.mapUnion
696 (rmapM getMap $ lmap fst latest)
697 (lmap snd $ catMaybes kvToMapOverwriteGeneric)
698 in postscan classifier aggregator
700-- | Same as 'toMap' but maybe faster because it uses mutable cells as
701-- fold accumulators in the Map.
703{-# INLINE toMapIO #-}
704toMapIO :: (MonadIO m, Ord k) =>
705 (a -> k) -> Fold m a b -> Fold m a (Map k b)
706toMapIO = toContainerIO
708-- | Given an input stream of key value pairs and a fold for values, fold all
709-- the values belonging to each key. Useful for map/reduce, bucketizing the
710-- input in different bins or for generating histograms.
714-- >>> kvToMap = Fold.toMap fst . Fold.lmap snd
719-- let input = Stream.fromList [("ONE",1),("ONE",1.1),("TWO",2), ("TWO",2.2)]
720-- in Stream.fold (Fold.kvToMap Fold.toList) input
722-- fromList [("ONE",[1.0,1.1]),("TWO",[2.0,2.2])]
725{-# INLINE kvToMap #-}
726kvToMap :: (Monad m, Ord k) => Fold m a b -> Fold m (k, a) (Map k b)
727kvToMap = toMap fst . lmap snd
729-- | Determine the frequency of each element in the stream.
731-- You can just collect the keys of the resulting map to get the unique
732-- elements in the stream.
736-- >>> frequency = Fold.toMap id Fold.length
738{-# INLINE frequency #-}
739frequency :: (Monad m, Ord a) => Fold m a (Map a Int)
740frequency = toMap id length
keysesc to close
- j k
- move down, up
- enter
- open
- esc
- back, or up a level
- 1 2 3 4
- readme, files, commits, refs
- /
- filter
- y
- copy the clone command
- c
- compare the chosen commit with another (or shift-click both)
- ?
- show or hide this list