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
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.
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
39 -- *** Input is explicit key-value tuple
40 -- | Like above but inputs are in explicit key-value pair form.
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.
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)
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
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.
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)
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
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.
120-- >>> toSet = Fold.foldl' (flip Set.insert) Set.empty
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
131-- >>> toIntSet = Fold.foldl' (flip IntSet.insert) IntSet.empty
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.
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
150nub :: (Monad m, Ord a) => Fold m a (Maybe a)
151nub = fmap (\(Tuple' _ x) -> x) $ foldl' step initial
155 initial = Tuple' Set.empty Nothing
157 step (Tuple' set _) x =
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.
166nubInt :: Monad m => Fold m Int (Maybe Int)
167nubInt = fmap (\(Tuple' _ x) -> x) $ foldl' step initial
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)
179-- XXX Add a countDistinct window fold
180-- XXX Add a bloom filter fold
182-- | Count non-duplicate elements in the stream.
186-- >>> countDistinct = fmap Set.size Fold.toSet
187-- >>> countDistinct = Fold.postscan Fold.nub $ Fold.catMaybes $ Fold.length
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.
193-- /Space/: \(\mathcal{O}(n)\)
197{-# INLINE countDistinct #-}
198countDistinct :: (Monad m, Ord a) => Fold m a Int
199-- countDistinct = postscan nub $ catMaybes length
200countDistinct = fmap Set.size toSet
202countDistinct = fmap (\(Tuple' _ n) -> n) $ foldl' step initial
206 initial = Tuple' Set.empty 0
208 step (Tuple' set n) x = do
214 in Tuple' (Set.insert x set) cnt
217-- | Like 'countDistinct' but specialized to a stream of 'Int', for better
222-- >>> countDistinctInt = fmap IntSet.size Fold.toIntSet
223-- >>> countDistinctInt = Fold.postscan Fold.nubInt $ Fold.catMaybes $ Fold.length
226{-# INLINE countDistinctInt #-}
227countDistinctInt :: Monad m => Fold m Int Int
228-- countDistinctInt = postscan nubInt $ catMaybes length
229countDistinctInt = fmap IntSet.size toIntSet
231countDistinctInt = fmap (\(Tuple' _ n) -> n) $ foldl' step initial
235 initial = Tuple' IntSet.empty 0
237 step (Tuple' set n) x = do
238 if IntSet.member x set
243 in Tuple' (IntSet.insert x set) cnt
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.
254-- This is the consumer side dual of the producer side 'mux' operation (XXX to
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
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.
267-- See 'demux' for documentation.
268{-# INLINE demuxGeneric #-}
269demuxGeneric :: (Monad m, IsMap f, Traversable f) =>
271 -> (a -> m (Fold m a b))
272 -> Fold m a (m (f b), Maybe (Key f, b))
273demuxGeneric getKey getFold = fmap extract $ foldlM' step initial
277 initial = return $ Tuple' IsMap.mapEmpty Nothing
279 {-# INLINE runFold #-}
280 runFold kv (Fold step1 initial1 extract1) (k, a) = do
288 let fld = Fold step1 (return res1) extract1
289 in Tuple' (IsMap.mapInsert k fld kv) Nothing
290 Done b -> Tuple' (IsMap.mapDelete k kv) (Just (k, b))
291 Done b -> return $ Tuple' kv (Just (k, b))
293 step (Tuple' kv _) a = do
295 case IsMap.mapLookup k kv of
298 runFold kv fld (k, a)
299 Just f -> runFold kv f (k, a)
301 extract (Tuple' kv x) = (Prelude.mapM f kv, x)
311-- | @demux getKey getFold@: In a key value stream, fold values corresponding
312-- to each key using a key specific fold. @getFold@ is invoked to generate a
313-- key specific fold when a key is encountered for the first time in the
316-- The first component of the output tuple is a key-value Map of in-progress
317-- folds. The fold returns the fold result as the second component of the
318-- output tuple whenever a fold terminates.
320-- If a fold terminates, another instance of the fold is started upon receiving
321-- an input with that key, @getFold@ is invoked again whenever the key is
324-- This can be used to scan a stream and collect the results from the scan
327-- Since the fold generator function is monadic we can add folds dynamically.
328-- For example, we can maintain a Map of keys to folds in an IORef and lookup
329-- the fold from that corresponding to a key. This Map can be changed
330-- dynamically, folds for new keys can be added or folds for old keys can be
331-- deleted or modified.
333-- Compare with 'classify', the fold in 'classify' is a static fold.
338demux :: (Monad m, Ord k) =>
340 -> (a -> m (Fold m a b))
341 -> Fold m a (m (Map k b), Maybe (k, b))
344-- | This is specialized version of 'demuxGeneric' that uses mutable IO cells
345-- as fold accumulators for better performance.
346{-# INLINE demuxGenericIO #-}
347demuxGenericIO :: (MonadIO m, IsMap f, Traversable f) =>
349 -> (a -> m (Fold m a b))
350 -> Fold m a (m (f b), Maybe (Key f, b))
351demuxGenericIO getKey getFold = fmap extract $ foldlM' step initial
355 initial = return $ Tuple' IsMap.mapEmpty Nothing
357 {-# INLINE initFold #-}
358 initFold kv (Fold step1 initial1 extract1) (k, a) = do
365 let fld = Fold step1 (return res1) extract1
366 ref <- liftIO $ newIORef fld
367 return $ Tuple' (IsMap.mapInsert k ref kv) Nothing
368 Done b -> return $ Tuple' kv (Just (k, b))
369 Done b -> return $ Tuple' kv (Just (k, b))
371 {-# INLINE runFold #-}
372 runFold kv ref (Fold step1 initial1 extract1) (k, a) = do
379 let fld = Fold step1 (return res1) extract1
380 liftIO $ writeIORef ref fld
381 return $ Tuple' kv Nothing
383 let kv1 = IsMap.mapDelete k kv
384 in return $ Tuple' kv1 (Just (k, b))
385 Done _ -> error "demuxGenericIO: unreachable"
387 step (Tuple' kv _) a = do
389 case IsMap.mapLookup k kv of
394 f <- liftIO $ readIORef ref
395 runFold kv ref f (k, a)
397 extract (Tuple' kv x) = (Prelude.mapM f kv, x)
402 (Fold _ i e) <- liftIO $ readIORef ref
408-- | This is specialized version of 'demux' that uses mutable IO cells as
409-- fold accumulators for better performance.
411-- Keep in mind that the values in the returned Map may be changed by the
412-- ongoing fold if you are using those concurrently in another thread.
414{-# INLINE demuxIO #-}
415demuxIO :: (MonadIO m, Ord k) =>
417 -> (a -> m (Fold m a b))
418 -> Fold m a (m (Map k b), Maybe (k, b))
419demuxIO = demuxGenericIO
421{-# INLINE demuxToContainer #-}
422demuxToContainer :: (Monad m, IsMap f, Traversable f) =>
423 (a -> Key f) -> (a -> m (Fold m a b)) -> Fold m a (f b)
424demuxToContainer getKey getFold =
426 classifier = demuxGeneric getKey getFold
427 getMap Nothing = pure IsMap.mapEmpty
428 getMap (Just action) = action
430 teeWith IsMap.mapUnion
431 (rmapM getMap $ lmap fst latest)
432 (lmap snd $ catMaybes kvToMapOverwriteGeneric)
433 in postscan classifier aggregator
435-- | This collects all the results of 'demux' in a Map.
437{-# INLINE demuxToMap #-}
438demuxToMap :: (Monad m, Ord k) =>
439 (a -> k) -> (a -> m (Fold m a b)) -> Fold m a (Map k b)
440demuxToMap = demuxToContainer
442{-# INLINE demuxToContainerIO #-}
443demuxToContainerIO :: (MonadIO m, IsMap f, Traversable f) =>
444 (a -> Key f) -> (a -> m (Fold m a b)) -> Fold m a (f b)
445demuxToContainerIO getKey getFold =
447 classifier = demuxGenericIO getKey getFold
448 getMap Nothing = pure IsMap.mapEmpty
449 getMap (Just action) = action
451 teeWith IsMap.mapUnion
452 (rmapM getMap $ lmap fst latest)
453 (lmap snd $ catMaybes kvToMapOverwriteGeneric)
454 in postscan classifier aggregator
456-- | Same as 'demuxToMap' but uses 'demuxIO' for better performance.
458{-# INLINE demuxToMapIO #-}
459demuxToMapIO :: (MonadIO m, Ord k) =>
460 (a -> k) -> (a -> m (Fold m a b)) -> Fold m a (Map k b)
461demuxToMapIO = demuxToContainerIO
463{-# INLINE demuxKvToContainer #-}
464demuxKvToContainer :: (Monad m, IsMap f, Traversable f) =>
465 (Key f -> m (Fold m a b)) -> Fold m (Key f, a) (f b)
466demuxKvToContainer f = demuxToContainer fst (\(k, _) -> fmap (lmap snd) (f k))
468-- | Fold a stream of key value pairs using a function that maps keys to folds.
472-- >>> demuxKvToMap f = Fold.demuxToContainer fst (Fold.lmap snd . f)
476-- >>> import Data.Map (Map)
478-- let f "SUM" = return Fold.sum
479-- f _ = return Fold.product
480-- input = Stream.fromList [("SUM",1),("PRODUCT",2),("SUM",3),("PRODUCT",4)]
481-- in Stream.fold (Fold.demuxKvToMap f) input :: IO (Map String Int)
483-- fromList [("PRODUCT",8),("SUM",4)]
486{-# INLINE demuxKvToMap #-}
487demuxKvToMap :: (Monad m, Ord k) =>
488 (k -> m (Fold m a b)) -> Fold m (k, a) (Map k b)
489demuxKvToMap = demuxKvToContainer
491------------------------------------------------------------------------------
492-- Classify: Like demux but uses the same fold for all keys.
493------------------------------------------------------------------------------
495-- XXX Change these to make the behavior similar to demux* variants. We can
496-- implement this using classifyScanManyWith. Maintain a set of done folds in
497-- the underlying monad, and when initial is called look it up, if the fold is
498-- done then initial would set a flag in the state to ignore the input or
501{-# INLINE classifyGeneric #-}
502classifyGeneric :: (Monad m, IsMap f, Traversable f, Ord (Key f)) =>
503 -- Note: we need to return the Map itself to display the in-progress values
504 -- e.g. to implement top. We could possibly create a separate abstraction
505 -- for that use case. We return an action because we want it to be lazy so
506 -- that the downstream consumers can choose to process or discard it.
507 (a -> Key f) -> Fold m a b -> Fold m a (m (f b), Maybe (Key f, b))
508classifyGeneric f (Fold step1 initial1 extract1) =
509 fmap extract $ foldlM' step initial
513 initial = return $ Tuple3' IsMap.mapEmpty Set.empty Nothing
515 {-# INLINE initFold #-}
516 initFold kv set k a = do
524 Tuple3' (IsMap.mapInsert k s1 kv) set Nothing
526 Tuple3' kv set (Just (k, b))
527 Done b -> return (Tuple3' kv (Set.insert k set) (Just (k, b)))
529 step (Tuple3' kv set _) a = do
531 case IsMap.mapLookup k kv of
534 then return (Tuple3' kv set Nothing)
535 else initFold kv set k a
541 Tuple3' (IsMap.mapInsert k s1 kv) set Nothing
543 let kv1 = IsMap.mapDelete k kv
544 in Tuple3' kv1 (Set.insert k set) (Just (k, b))
546 extract (Tuple3' kv _ x) = (Prelude.mapM extract1 kv, x)
548-- | Folds the values for each key using the supplied fold. When scanning, as
549-- soon as the fold is complete, its result is available in the second
550-- component of the tuple. The first component of the tuple is a snapshot of
551-- the in-progress folds.
553-- Once the fold for a key is done, any future values of the key are ignored.
557-- >>> classify f fld = Fold.demux f (const fld)
559{-# INLINE classify #-}
560classify :: (Monad m, Ord k) =>
561 (a -> k) -> Fold m a b -> Fold m a (m (Map k b), Maybe (k, b))
562classify = classifyGeneric
564-- XXX we can use a Prim IORef if we can constrain the state "s" to be Prim
566-- The code is almost the same as classifyGeneric except the IORef operations.
568{-# INLINE classifyGenericIO #-}
569classifyGenericIO :: (MonadIO m, IsMap f, Traversable f, Ord (Key f)) =>
570 (a -> Key f) -> Fold m a b -> Fold m a (m (f b), Maybe (Key f, b))
571classifyGenericIO f (Fold step1 initial1 extract1) =
572 fmap extract $ foldlM' step initial
576 initial = return $ Tuple3' IsMap.mapEmpty Set.empty Nothing
578 {-# INLINE initFold #-}
579 initFold kv set k a = do
586 ref <- liftIO $ newIORef s1
587 return $ Tuple3' (IsMap.mapInsert k ref kv) set Nothing
589 return $ Tuple3' kv set (Just (k, b))
590 Done b -> return (Tuple3' kv (Set.insert k set) (Just (k, b)))
592 step (Tuple3' kv set _) a = do
594 case IsMap.mapLookup k kv of
597 then return (Tuple3' kv set Nothing)
598 else initFold kv set k a
600 s <- liftIO $ readIORef ref
604 liftIO $ writeIORef ref s1
605 return $ Tuple3' kv set Nothing
607 let kv1 = IsMap.mapDelete k kv
609 $ Tuple3' kv1 (Set.insert k set) (Just (k, b))
611 extract (Tuple3' kv _ x) =
612 (Prelude.mapM (\ref -> liftIO (readIORef ref) >>= extract1) kv, x)
614-- | Same as classify except that it uses mutable IORef cells in the
615-- Map providing better performance. Be aware that if this is used as a scan,
616-- the values in the intermediate Maps would be mutable.
620-- >>> classifyIO f fld = Fold.demuxIO f (const fld)
622{-# INLINE classifyIO #-}
623classifyIO :: (MonadIO m, Ord k) =>
624 (a -> k) -> Fold m a b -> Fold m a (m (Map k b), Maybe (k, b))
625classifyIO = classifyGenericIO
627-- | Fold a key value stream to a key-value Map. If the same key appears
628-- multiple times, only the last value is retained.
629{-# INLINE kvToMapOverwriteGeneric #-}
630kvToMapOverwriteGeneric :: (Monad m, IsMap f) => Fold m (Key f, a) (f a)
631kvToMapOverwriteGeneric =
632 foldl' (\kv (k, v) -> IsMap.mapInsert k v kv) IsMap.mapEmpty
634{-# INLINE toContainer #-}
635toContainer :: (Monad m, IsMap f, Traversable f, Ord (Key f)) =>
636 (a -> Key f) -> Fold m a b -> Fold m a (f b)
639 classifier = classifyGeneric f fld
640 getMap Nothing = pure IsMap.mapEmpty
641 getMap (Just action) = action
643 teeWith IsMap.mapUnion
644 (rmapM getMap $ lmap fst latest)
645 (lmap snd $ catMaybes kvToMapOverwriteGeneric)
646 in postscan classifier aggregator
648-- | Split the input stream based on a key field and fold each split using the
649-- given fold. Useful for map/reduce, bucketizing the input in different bins
650-- or for generating histograms.
654-- >>> import Data.Map.Strict (Map)
656-- let input = Stream.fromList [("ONE",1),("ONE",1.1),("TWO",2), ("TWO",2.2)]
657-- classify = Fold.toMap fst (Fold.lmap snd Fold.toList)
658-- in Stream.fold classify input :: IO (Map String [Double])
660-- fromList [("ONE",[1.0,1.1]),("TWO",[2.0,2.2])]
662-- Once the classifier fold terminates for a particular key any further inputs
663-- in that bucket are ignored.
665-- Space used is proportional to the number of keys seen till now and
666-- monotonically increases because it stores whether a key has been seen or
669-- See 'demuxToMap' for a more powerful version where you can use a different
670-- fold for each key. A simpler version of 'toMap' retaining only the last
671-- value for a key can be written as:
673-- >>> toMap = Fold.foldl' (\kv (k, v) -> Map.insert k v kv) Map.empty
680toMap :: (Monad m, Ord k) =>
681 (a -> k) -> Fold m a b -> Fold m a (Map k b)
684{-# INLINE toContainerIO #-}
685toContainerIO :: (MonadIO m, IsMap f, Traversable f, Ord (Key f)) =>
686 (a -> Key f) -> Fold m a b -> Fold m a (f b)
689 classifier = classifyGenericIO f fld
690 getMap Nothing = pure IsMap.mapEmpty
691 getMap (Just action) = action
693 teeWith IsMap.mapUnion
694 (rmapM getMap $ lmap fst latest)
695 (lmap snd $ catMaybes kvToMapOverwriteGeneric)
696 in postscan classifier aggregator
698-- | Same as 'toMap' but maybe faster because it uses mutable cells as
699-- fold accumulators in the Map.
701{-# INLINE toMapIO #-}
702toMapIO :: (MonadIO m, Ord k) =>
703 (a -> k) -> Fold m a b -> Fold m a (Map k b)
704toMapIO = toContainerIO
706-- | Given an input stream of key value pairs and a fold for values, fold all
707-- the values belonging to each key. Useful for map/reduce, bucketizing the
708-- input in different bins or for generating histograms.
712-- >>> kvToMap = Fold.toMap fst . Fold.lmap snd
717-- let input = Stream.fromList [("ONE",1),("ONE",1.1),("TWO",2), ("TWO",2.2)]
718-- in Stream.fold (Fold.kvToMap Fold.toList) input
720-- fromList [("ONE",[1.0,1.1]),("TWO",[2.0,2.2])]
723{-# INLINE kvToMap #-}
724kvToMap :: (Monad m, Ord k) => Fold m a b -> Fold m (k, a) (Map k b)
725kvToMap = toMap fst . lmap snd
727-- | Determine the frequency of each element in the stream.
729-- You can just collect the keys of the resulting map to get the unique
730-- elements in the stream.
734-- >>> frequency = Fold.toMap id Fold.length
736{-# INLINE frequency #-}
737frequency :: (Monad m, Ord a) => Fold m a (Map a Int)
738frequency = 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