on da36a01 help sign in
629 lines 21 kB history raw
1-- |
2-- Module : Streamly.Internal.Data.Stream.Top
3-- Copyright : (c) 2020 Composewell Technologies
4-- License : BSD-3-Clause
5-- Maintainer : streamly@composewell.com
6-- Stability : experimental
7-- Portability : GHC
8--
9-- Top level module that can depend on all other lower level Stream modules.
11module Streamly.Internal.Data.Stream.Top
12 (
13 -- * Transformation
14 -- ** Sampling
15 -- | Value agnostic filtering.
16 sampleFromThen
18 -- ** Reordering
19 , sortBy
21 -- * Nesting
22 -- ** Set like operations
23 -- | These are not exactly set operations because streams are not
24 -- necessarily sets, they may have duplicated elements.
25 , intersectBy
26 , intersectBySorted
27 , differenceBy
28 , mergeDifferenceBy
29 , unionBy
30 , mergeUnionBy
32 -- ** Join operations
33 , crossJoin
34 , joinInner
35 , joinInnerMap
36 , joinInnerMerge
37 , joinLeft
38 , mergeLeftJoin
39 , joinLeftMap
40 , joinOuter
41 , mergeOuterJoin
42 , joinOuterMap
43 )
44where
46#include "inline.hs"
48import Control.Monad.Catch (MonadCatch)
49import Control.Monad.IO.Class (MonadIO(..))
50import Control.Monad.Trans.Class (lift)
51import Control.Monad.Trans.State.Strict (get, put)
52import Data.Function ((&))
53import Data.IORef (newIORef, readIORef, modifyIORef')
54import Data.Maybe (isJust)
55import Streamly.Internal.Data.Stream.Common ()
56import Streamly.Internal.Data.Stream.Type (Stream, fromStreamD, toStreamD)
58import qualified Data.List as List
59import qualified Data.Map.Strict as Map
60import qualified Streamly.Internal.Data.Array as Array
61import qualified Streamly.Internal.Data.Array.Unboxed.Mut.Type as MA
62import qualified Streamly.Internal.Data.Fold as Fold
63import qualified Streamly.Internal.Data.Parser as Parser
64import qualified Streamly.Internal.Data.Stream.Lift as Stream
65import qualified Streamly.Internal.Data.Stream.Eliminate as Stream
66import qualified Streamly.Internal.Data.Stream.Generate as Stream
67import qualified Streamly.Internal.Data.Stream.Expand as Stream
68import qualified Streamly.Internal.Data.Stream.Reduce as Stream
69import qualified Streamly.Internal.Data.Stream.Transform as Stream
70import qualified Streamly.Internal.Data.Stream.StreamD as StreamD
72import Prelude hiding (filter, zipWith, concatMap, concat)
74-- $setup
75-- >>> :m
76-- >>> import Prelude hiding (filter, zipWith, concatMap, concat)
77-- >>> import qualified Streamly.Internal.Data.Fold as Fold
78-- >>> import qualified Streamly.Internal.Data.Stream as Stream
80------------------------------------------------------------------------------
81-- Sampling
82------------------------------------------------------------------------------
84-- XXX We can implement this using addition instead of "mod" to make it more
85-- efficient.
86--
87-- | @sampleFromthen offset stride@ samples the element at @offset@ index and
88-- then every element at strides of @stride@.
89--
90-- >>> Stream.fold Fold.toList $ Stream.sampleFromThen 2 3 $ Stream.enumerateFromTo 0 10
91-- [2,5,8]
92--
93-- /Pre-release/
94--
95{-# INLINE sampleFromThen #-}
96sampleFromThen :: Monad m => Int -> Int -> Stream m a -> Stream m a
97sampleFromThen offset stride =
98 Stream.with Stream.indexed Stream.filter
99 (\(i, _) -> i >= offset && (i - offset) `mod` stride == 0)
101------------------------------------------------------------------------------
102-- Reordering
103------------------------------------------------------------------------------
104--
105-- We could possibly choose different algorithms depending on whether the
106-- input stream is almost sorted (ascending/descending) or random. We could
107-- serialize the stream to an array and use quicksort.
108--
109-- | Sort the input stream using a supplied comparison function.
110--
111-- /O(n) space/
112--
113-- Note: this is not the fastest possible implementation as of now.
114--
115-- /Pre-release/
116--
117{-# INLINE sortBy #-}
118sortBy :: MonadCatch m => (a -> a -> Ordering) -> Stream m a -> Stream m a
119-- sortBy f = Stream.concatPairsWith (Stream.mergeBy f) Stream.fromPure
120sortBy cmp =
121 let p =
122 Parser.groupByRollingEither
123 (\x -> (< GT) . cmp x)
124 Fold.toStreamRev
125 Fold.toStream
126 in Stream.concatPairsWith (Stream.mergeBy cmp) id
127 . Stream.parseMany (fmap (either id id) p)
129------------------------------------------------------------------------------
130-- SQL Joins
131------------------------------------------------------------------------------
132--
133-- Some references:
134-- * https://en.wikipedia.org/wiki/Relational_algebra
135-- * https://en.wikipedia.org/wiki/Join_(SQL)
137-- TODO: OrdSet/IntSet/hashmap based versions of these. With Eq only
138-- constraint, the best would be to use an Array with linear search. If the
139-- second stream is sorted we can also use a binary search, using Ord
140-- constraint or an ordering function.
141--
142-- For Storables we can cache the second stream into an unboxed array for
143-- possibly faster access/compact representation?
144--
145-- If we do not want to keep the stream in memory but always read it from the
146-- source (disk/network) every time we iterate through it then we can do that
147-- too by reading the stream every time, the stream must have immutable state
148-- in that case and the user is responsible for the behavior if the stream
149-- source changes during iterations. We can also use an Unfold instead of
150-- stream. We probably need a way to distinguish streams that can be read
151-- mutliple times without any interference (e.g. unfolding a stream using an
152-- immutable handle would work i.e. using pread/pwrite instead of maintianing
153-- an offset in the handle).
155-- XXX We can do this concurrently.
156--
157-- | This is the same as 'Streamly.Internal.Data.Unfold.outerProduct' but less
158-- efficient.
159--
160-- The second stream is evaluated multiple times. If the second stream is
161-- consume-once stream then it can be cached in an 'Data.Array.Array' before
162-- calling this function. Caching may also improve performance if the stream is
163-- expensive to evaluate.
164--
165-- Time: O(m x n)
166--
167-- /Pre-release/
168{-# INLINE crossJoin #-}
169crossJoin :: Monad m => Stream m a -> Stream m b -> Stream m (a, b)
170crossJoin s1 s2 = do
171 -- XXX use concatMap instead?
172 a <- s1
173 b <- s2
174 return (a, b)
176-- XXX We can do this concurrently.
177-- XXX If the second stream is sorted and passed as an Array we could use
178-- binary search if we have an Ord instance or Ordering returning function. The
179-- time complexity would then become (m x log n).
180--
181-- | For all elements in @Stream m a@, for all elements in @Stream m b@ if @a@ and @b@
182-- are equal by the given equality pedicate then return the tuple (a, b).
183--
184-- The second stream is evaluated multiple times. If the stream is a
185-- consume-once stream then the caller should cache it (e.g. in a
186-- 'Data.Array.Array') before calling this function. Caching may also improve
187-- performance if the stream is expensive to evaluate.
188--
189-- For space efficiency use the smaller stream as the second stream.
190--
191-- You should almost always use joinInnerMap instead of joinInner. joinInnerMap
192-- is an order of magnitude faster. joinInner may be used when the second
193-- stream is generated from a seed, therefore, need not be stored in memory and
194-- the amount of memory it takes is a concern.
195--
196-- Space: O(n) assuming the second stream is cached in memory.
197--
198-- Time: O(m x n)
199--
200-- /Pre-release/
201{-# INLINE joinInner #-}
202joinInner :: Monad m =>
203 (a -> b -> Bool) -> Stream m a -> Stream m b -> Stream m (a, b)
204joinInner eq s1 s2 = do
205 -- ConcatMap works faster than bind
206 Stream.concatMap (\a ->
207 Stream.concatMap (\b ->
208 if a `eq` b
209 then Stream.fromPure (a, b)
210 else Stream.nil
211 ) s2
212 ) s1
214-- XXX Generate error if a duplicate insertion is attempted?
215toMap :: (Monad m, Ord k) => Stream m (k, v) -> m (Map.Map k v)
216toMap =
217 let f = Fold.foldl' (\kv (k, b) -> Map.insert k b kv) Map.empty
218 in Stream.fold f
220-- If the second stream is too big it can be partitioned based on hashes and
221-- then we can process one parition at a time.
222--
223-- XXX An IntMap may be faster when the keys are Int.
224-- XXX Use hashmap instead of map?
225--
226-- | Like 'joinInner' but uses a 'Map' for efficiency.
227--
228-- If the input streams have duplicate keys, the behavior is undefined.
229--
230-- For space efficiency use the smaller stream as the second stream.
231--
232-- Space: O(n)
233--
234-- Time: O(m + n)
235--
236-- /Pre-release/
237{-# INLINE joinInnerMap #-}
238joinInnerMap :: (Monad m, Ord k) =>
239 Stream m (k, a) -> Stream m (k, b) -> Stream m (k, a, b)
240joinInnerMap s1 s2 =
241 Stream.concatM $ do
242 km <- toMap s2
243 pure $ Stream.mapMaybe (joinAB km) s1
245 where
247 joinAB kvm (k, a) =
248 case k `Map.lookup` kvm of
249 Just b -> Just (k, a, b)
250 Nothing -> Nothing
252-- | Like 'joinInner' but works only on sorted streams.
253--
254-- Space: O(1)
255--
256-- Time: O(m + n)
257--
258-- /Unimplemented/
259{-# INLINE joinInnerMerge #-}
260joinInnerMerge ::
261 (a -> b -> Ordering) -> Stream m a -> Stream m b -> Stream m (a, b)
262joinInnerMerge = undefined
264-- XXX We can do this concurrently.
265-- XXX If the second stream is sorted and passed as an Array or a seek capable
266-- stream then we could use binary search if we have an Ord instance or
267-- Ordering returning function. The time complexity would then become (m x log
268-- n).
269--
270-- | For all elements in @Stream m a@, for all elements in @Stream m b@ if @a@ and @b@
271-- are equal then return the tuple @(a, Just b)@. If @a@ is not present in @t
272-- m b@ then return @(a, Nothing)@.
273--
274-- The second stream is evaluated multiple times. If the stream is a
275-- consume-once stream then the caller should cache it in an 'Data.Array.Array'
276-- before calling this function. Caching may also improve performance if the
277-- stream is expensive to evaluate.
278--
279-- @
280-- rightJoin = flip joinLeft
281-- @
282--
283-- Space: O(n) assuming the second stream is cached in memory.
284--
285-- Time: O(m x n)
286--
287-- /Unimplemented/
288{-# INLINE joinLeft #-}
289joinLeft :: Monad m =>
290 (a -> b -> Bool) -> Stream m a -> Stream m b -> Stream m (a, Maybe b)
291joinLeft eq s1 s2 = Stream.evalStateT (return False) $ do
292 a <- Stream.liftInner s1
293 -- XXX should we use StreamD monad here?
294 -- XXX Is there a better way to perform some action at the end of a loop
295 -- iteration?
296 lift $ put False
297 let final = do
298 r <- lift get
299 if r
300 then Stream.nil
301 else Stream.fromPure Nothing
302 b <- fmap Just (Stream.liftInner s2) <> final
303 case b of
304 Just b1 ->
305 if a `eq` b1
306 then do
307 lift $ put True
308 return (a, Just b1)
309 else Stream.nil
310 Nothing -> return (a, Nothing)
312-- | Like 'joinLeft' but uses a hashmap for efficiency.
313--
314-- Space: O(n)
315--
316-- Time: O(m + n)
317--
318-- /Pre-release/
319{-# INLINE joinLeftMap #-}
320joinLeftMap :: (Ord k, Monad m) =>
321 Stream m (k, a) -> Stream m (k, b) -> Stream m (k, a, Maybe b)
322joinLeftMap s1 s2 =
323 Stream.concatM $ do
324 km <- toMap s2
325 return $ fmap (joinAB km) s1
327 where
329 joinAB km (k, a) =
330 case k `Map.lookup` km of
331 Just b -> (k, a, Just b)
332 Nothing -> (k, a, Nothing)
334-- | Like 'joinLeft' but works only on sorted streams.
335--
336-- Space: O(1)
337--
338-- Time: O(m + n)
339--
340-- /Unimplemented/
341{-# INLINE mergeLeftJoin #-}
342mergeLeftJoin :: -- Monad m =>
343 (a -> b -> Ordering) -> Stream m a -> Stream m b -> Stream m (a, Maybe b)
344mergeLeftJoin _eq _s1 _s2 = undefined
346-- XXX We can do this concurrently.
347--
348-- | For all elements in @t m a@, for all elements in @t m b@ if @a@ and @b@
349-- are equal by the given equality pedicate then return the tuple (Just a, Just
350-- b). If @a@ is not found in @t m b@ then return (a, Nothing), return
351-- (Nothing, b) for vice-versa.
352--
353-- For space efficiency use the smaller stream as the second stream.
354--
355-- Space: O(n)
356--
357-- Time: O(m x n)
358--
359-- /Pre-release/
360{-# INLINE joinOuter #-}
361joinOuter :: MonadIO m =>
362 (a -> b -> Bool)
363 -> Stream m a
364 -> Stream m b
365 -> Stream m (Maybe a, Maybe b)
366joinOuter eq s1 s =
367 Stream.concatM $ do
368 inputArr <- Array.fromStream s
369 let len = Array.length inputArr
370 foundArr <-
371 Stream.fold
372 (MA.writeN len)
373 (Stream.fromList (Prelude.replicate len False))
374 return $ go inputArr foundArr <> leftOver inputArr foundArr
376 where
378 leftOver inputArr foundArr =
379 let stream1 = Array.toStream inputArr
380 stream2 = Stream.unfold MA.read foundArr
381 in Stream.filter
382 isJust
383 ( Stream.zipWith (\x y ->
384 if y
385 then Nothing
386 else Just (Nothing, Just x)
387 ) stream1 stream2
388 ) & Stream.catMaybes
390 go inputArr foundArr = Stream.evalStateT (return False) $ do
391 a <- Stream.liftInner s1
392 -- XXX should we use StreamD monad here?
393 -- XXX Is there a better way to perform some action at the end of a loop
394 -- iteration?
395 lift $ put False
396 let final = do
397 r <- lift get
398 if r
399 then Stream.nil
400 else Stream.fromPure Nothing
401 (i, b) <-
402 let stream = Array.toStream inputArr
403 in Stream.indexed $ fmap Just (Stream.liftInner stream) <> final
405 case b of
406 Just b1 ->
407 if a `eq` b1
408 then do
409 lift $ put True
410 MA.putIndex i True foundArr
411 return (Just a, Just b1)
412 else Stream.nil
413 Nothing -> return (Just a, Nothing)
415-- Put the b's that have been paired, in another hash or mutate the hash to set
416-- a flag. At the end go through @Stream m b@ and find those that are not in that
417-- hash to return (Nothing, b).
418--
419-- | Like 'joinOuter' but uses a 'Map' for efficiency.
420--
421-- Space: O(m + n)
422--
423-- Time: O(m + n)
424--
425-- /Pre-release/
426{-# INLINE joinOuterMap #-}
427joinOuterMap ::
428 (Ord k, MonadIO m) =>
429 Stream m (k, a) -> Stream m (k, b) -> Stream m (k, Maybe a, Maybe b)
430joinOuterMap s1 s2 =
431 Stream.concatM $ do
432 km1 <- kvFold s1
433 km2 <- kvFold s2
435 -- XXX Not sure if toList/fromList would fuse optimally. We may have to
436 -- create a fused Map.toStream function.
437 let res1 = fmap (joinAB km2)
438 $ Stream.fromList $ Map.toList km1
439 where
440 joinAB km (k, a) =
441 case k `Map.lookup` km of
442 Just b -> (k, Just a, Just b)
443 Nothing -> (k, Just a, Nothing)
445 -- XXX We can take advantage of the lookups in the first pass above to
446 -- reduce the number of lookups in this pass. If we keep mutable cells
447 -- in the second Map, we can flag it in the first pass and not do any
448 -- lookup in the second pass if it is flagged.
449 let res2 = Stream.mapMaybe (joinAB km1)
450 $ Stream.fromList $ Map.toList km2
451 where
452 joinAB km (k, b) =
453 case k `Map.lookup` km of
454 Just _ -> Nothing
455 Nothing -> Just (k, Nothing, Just b)
457 return $ Stream.append res1 res2
459 where
461 -- XXX Generate error if a duplicate insertion is attempted?
462 kvFold =
463 let f = Fold.foldl' (\kv (k, b) -> Map.insert k b kv) Map.empty
464 in Stream.fold f
466-- | Like 'joinOuter' but works only on sorted streams.
467--
468-- Space: O(1)
469--
470-- Time: O(m + n)
471--
472-- /Unimplemented/
473{-# INLINE mergeOuterJoin #-}
474mergeOuterJoin :: -- Monad m =>
475 (a -> b -> Ordering)
476 -> Stream m a
477 -> Stream m b
478 -> Stream m (Maybe a, Maybe b)
479mergeOuterJoin _eq _s1 _s2 = undefined
481------------------------------------------------------------------------------
482-- Set operations (special joins)
483------------------------------------------------------------------------------
484--
485-- TODO: OrdSet/IntSet/hashmap based versions of these. With Eq only constraint
486-- the best would be to use an Array with linear search. If the second stream
487-- is sorted we can also use a binary search, using Ord constraint.
489-- | 'intersectBy' is essentially a filtering operation that retains only those
490-- elements in the first stream that are present in the second stream.
491--
492-- >>> Stream.fold Fold.toList $ Stream.intersectBy (==) (Stream.fromList [1,2,2,4]) (Stream.fromList [2,1,1,3])
493-- [1,2,2]
494--
495-- >>> Stream.fold Fold.toList $ Stream.intersectBy (==) (Stream.fromList [2,1,1,3]) (Stream.fromList [1,2,2,4])
496-- [2,1,1]
497--
498-- 'intersectBy' is similar to but not the same as 'joinInner':
499--
500-- >>> Stream.fold Fold.toList $ fmap fst $ Stream.joinInner (==) (Stream.fromList [1,2,2,4]) (Stream.fromList [2,1,1,3])
501-- [1,1,2,2]
502--
503-- Space: O(n) where @n@ is the number of elements in the second stream.
504--
505-- Time: O(m x n) where @m@ is the number of elements in the first stream and
506-- @n@ is the number of elements in the second stream.
507--
508-- /Pre-release/
509{-# INLINE intersectBy #-}
510intersectBy :: Monad m =>
511 (a -> a -> Bool) -> Stream m a -> Stream m a -> Stream m a
512intersectBy eq s1 s2 =
513 Stream.concatM
514 $ do
515 -- This may work well when s2 is small
516 xs <- Stream.fold Fold.toListRev $ Stream.uniqBy eq s2
517 return $ Stream.filter (\x -> List.any (eq x) xs) s1
519-- | Like 'intersectBy' but works only on streams sorted in ascending order.
520--
521-- Space: O(1)
522--
523-- Time: O(m+n)
524--
525-- /Pre-release/
526{-# INLINE intersectBySorted #-}
527intersectBySorted :: (Monad m) =>
528 (a -> a -> Ordering) -> Stream m a -> Stream m a -> Stream m a
529intersectBySorted eq s1 =
530 fromStreamD
531 . StreamD.intersectBySorted eq (toStreamD s1)
532 . toStreamD
534-- Roughly joinLeft s1 s2 = s1 `difference` s2 + s1 `intersection` s2
536-- | Delete first occurrences of those elements from the first stream that are
537-- present in the second stream. If an element occurs multiple times in the
538-- second stream as many occurrences of it are deleted from the first stream.
539--
540-- >>> Stream.fold Fold.toList $ Stream.differenceBy (==) (Stream.fromList [1,2,2]) (Stream.fromList [1,2,3])
541-- [2]
542--
543-- The following laws hold:
544--
545-- @
546-- (s1 `serial` s2) `differenceBy eq` s1 === s2
547-- (s1 `wSerial` s2) `differenceBy eq` s1 === s2
548-- @
549--
550-- Same as the list 'Data.List.//' operation.
551--
552-- Space: O(m) where @m@ is the number of elements in the first stream.
553--
554-- Time: O(m x n) where @m@ is the number of elements in the first stream and
555-- @n@ is the number of elements in the second stream.
556--
557-- /Pre-release/
558{-# INLINE differenceBy #-}
559differenceBy :: (Monad m) =>
560 (a -> a -> Bool) -> Stream m a -> Stream m a -> Stream m a
561differenceBy eq s1 s2 =
562 Stream.concatM
563 $ do
564 -- This may work well if s1 is small
565 -- If s1 is big we can go through s1, deleting elements from s2 and
566 -- not emitting an element if it was successfully deleted from s2.
567 -- we will need a deleteBy that can return whether the element was
568 -- deleted or not.
569 xs <- Stream.fold Fold.toList s1
570 let f = Fold.foldl' (flip (List.deleteBy eq)) xs
571 fmap Stream.fromList $ Stream.fold f s2
573-- | Like 'differenceBy' but works only on sorted streams.
574--
575-- Space: O(1)
576--
577-- /Unimplemented/
578{-# INLINE mergeDifferenceBy #-}
579mergeDifferenceBy :: -- (Monad m) =>
580 (a -> a -> Ordering) -> Stream m a -> Stream m a -> Stream m a
581mergeDifferenceBy _eq _s1 _s2 = undefined
583-- | This is essentially an append operation that appends all the extra
584-- occurrences of elements from the second stream that are not already present
585-- in the first stream.
586--
587-- >>> Stream.fold Fold.toList $ Stream.unionBy (==) (Stream.fromList [1,2,2,4]) (Stream.fromList [1,1,2,3])
588-- [1,2,2,4,3]
589--
590-- Equivalent to the following except that @s1@ is evaluated only once:
591--
592-- @
593-- unionBy eq s1 s2 = s1 \`serial` (s2 `differenceBy eq` s1)
594-- @
595--
596-- Similar to 'joinOuter' but not the same.
597--
598-- Space: O(n)
599--
600-- Time: O(m x n)
601--
602-- /Pre-release/
603{-# INLINE unionBy #-}
604unionBy :: MonadIO m =>
605 (a -> a -> Bool) -> Stream m a -> Stream m a -> Stream m a
606unionBy eq s1 s2 =
607 Stream.concatM
608 $ do
609 xs <- Stream.fold Fold.toList s2
610 -- XXX we can use postscanlMAfter' instead of IORef
611 ref <- liftIO $ newIORef $! List.nubBy eq xs
612 let f x = do
613 liftIO $ modifyIORef' ref (List.deleteBy eq x)
614 return x
615 s3 = Stream.concatM
616 $ do
617 xs1 <- liftIO $ readIORef ref
618 return $ Stream.fromList xs1
619 return $ Stream.mapM f s1 <> s3
621-- | Like 'unionBy' but works only on sorted streams.
622--
623-- Space: O(1)
624--
625-- /Unimplemented/
626{-# INLINE mergeUnionBy #-}
627mergeUnionBy :: -- (Monad m) =>
628 (a -> a -> Ordering) -> Stream m a -> Stream m a -> Stream m a
629mergeUnionBy _eq _s1 _s2 = undefined