git.iko.soy/streamly
on master help sign in
streamly/core/src/Streamly/Internal/Data/Stream/Container.hs
310 lines 11 kB history raw
1{-# LANGUAGE CPP #-}
2-- |
3-- Module : Streamly.Internal.Data.Stream.Container
4-- Copyright : (c) 2019 Composewell Technologies
5-- License : BSD-3-Clause
6-- Maintainer : streamly@composewell.com
7-- Stability : experimental
8-- Portability : GHC
9--
10-- Stream operations that require transformers or containers like Set or Map.
12module Streamly.Internal.Data.Stream.Container
13 (
14 -- * Deduplication
15 ordNub
17 -- * Joins
18 , leftJoin
19 , outerJoin
21 -- * Ord Joins
22 , innerOrdJoin
23 , leftOrdJoin
24 , outerOrdJoin
25 )
26where
28#include "inline.hs"
30import Control.Monad.IO.Class (MonadIO)
31import Control.Monad.Trans.State.Strict (get, put)
32import Data.Function ((&))
33import Data.Maybe (isJust)
34import Streamly.Internal.Data.Stream.Step (Step(..))
35import Streamly.Internal.Data.Stream.Type
36 (Stream(..), mkCross, unCross)
38import qualified Data.Map.Strict as Map
39import qualified Data.Set as Set
40import qualified Streamly.Data.Fold as Fold
41import qualified Streamly.Internal.Data.Array.Generic as Array
42import qualified Streamly.Internal.Data.MutArray.Type as MA
43import qualified Streamly.Internal.Data.Stream.Type as Stream
44import qualified Streamly.Internal.Data.Stream.Nesting as Stream
45import qualified Streamly.Internal.Data.Stream.Generate as Stream
46import qualified Streamly.Internal.Data.Stream.Transform as Stream
47import qualified Streamly.Internal.Data.Stream.Transformer as Stream
49#include "DocTestDataStream.hs"
51-- | @nub@ specialized to 'Ord' types for better performance. Returns a
52-- subsequence of the stream removing any duplicate elements.
53--
54-- The memory used is proportional to the number of unique elements in the
55-- stream. One way to limit the memory is to use @take@ on the resulting
56-- stream to limit the unique elements in the stream.
57{-# INLINE_NORMAL ordNub #-}
58ordNub :: (Monad m, Ord a) => Stream m a -> Stream m a
59ordNub (Stream step1 state1) = Stream step (Set.empty, state1)
61 where
63 step gst (set, st) = do
64 r <- step1 gst st
65 return
66 $ case r of
67 Yield x s ->
68 if Set.member x set
69 then Skip (set, s)
70 else Yield x (Set.insert x set, s)
71 Skip s -> Skip (set, s)
72 Stop -> Stop
74-- XXX Generate error if a duplicate insertion is attempted?
75toMap :: (Monad m, Ord k) => Stream m (k, v) -> m (Map.Map k v)
76toMap =
77 let f = Fold.foldl' (\kv (k, b) -> Map.insert k b kv) Map.empty
78 in Stream.fold f
80-- If the second stream is too big it can be partitioned based on hashes and
81-- then we can process one parition at a time.
82--
83-- XXX An IntMap may be faster when the keys are Int.
84-- XXX Use hashmap instead of map?
86-- | 'innerJoin' specialized to 'Ord' types for better performance.
87--
88-- If the input streams have duplicate keys, the behavior is undefined.
89--
90-- For space efficiency use the smaller stream as the second stream.
91--
92-- Space: O(n)
93--
94-- Time: O(m + n)
95--
96-- /Pre-release/
97{-# INLINE innerOrdJoin #-}
98innerOrdJoin :: (Monad m, Ord k) =>
99 Stream m (k, a) -> Stream m (k, b) -> Stream m (k, a, b)
100innerOrdJoin s1 s2 =
101 Stream.concatEffect $ do
102 km <- toMap s2
103 pure $ Stream.mapMaybe (joinAB km) s1
105 where
107 joinAB kvm (k, a) =
108 case k `Map.lookup` kvm of
109 Just b -> Just (k, a, b)
110 Nothing -> Nothing
112-- XXX We can do this concurrently.
113-- XXX Check performance of StreamD vs StreamK
114-- XXX If the second stream is sorted and passed as an Array or a seek capable
115-- stream then we could use binary search if we have an Ord instance or
116-- Ordering returning function. The time complexity would then become (m x log
117-- n).
119-- | Like 'innerJoin' but emits @(a, Just b)@ whenever a and b are equal, for
120-- those @a@'s that are not equal to any @b@ emits @(a, Nothing)@.
121--
122-- This is a generalization of 'innerJoin' to include all elements from the
123-- left stream and not just those which have an equal in the right stream. This
124-- is not a commutative operation, the order of the stream arguments matters.
125--
126-- All the caveats mentioned in 'innerJoin' apply here as well. Right join is
127-- not provided because it is just a flipped left join:
128--
129-- >>> rightJoin eq = flip (Stream.leftJoin eq)
130--
131-- Space: O(n) assuming the second stream is cached in memory.
132--
133-- Time: O(m x n)
134--
135-- /Unimplemented/
136{-# INLINE leftJoin #-}
137leftJoin :: Monad m =>
138 (a -> b -> Bool) -> Stream m a -> Stream m b -> Stream m (a, Maybe b)
139leftJoin eq s1 s2 = Stream.evalStateT (return False) $ unCross $ do
140 a <- mkCross (Stream.liftInner s1)
141 -- XXX should we use StreamD monad here?
142 -- XXX Is there a better way to perform some action at the end of a loop
143 -- iteration?
144 mkCross (Stream.fromEffect $ put False)
145 let final = Stream.concatEffect $ do
146 r <- get
147 if r
148 then pure Stream.nil
149 else pure (Stream.fromPure Nothing)
150 b <- mkCross (fmap Just (Stream.liftInner s2) `Stream.append` final)
151 case b of
152 Just b1 ->
153 if a `eq` b1
154 then do
155 mkCross (Stream.fromEffect $ put True)
156 return (a, Just b1)
157 else mkCross Stream.nil
158 Nothing -> return (a, Nothing)
160-- | 'leftJoin' specialized to 'Ord' types for better performance.
161--
162-- Space: O(n)
163--
164-- Time: O(m + n)
165--
166-- /Pre-release/
167{-# INLINE leftOrdJoin #-}
168leftOrdJoin :: (Ord k, Monad m) =>
169 Stream m (k, a) -> Stream m (k, b) -> Stream m (k, a, Maybe b)
170leftOrdJoin s1 s2 =
171 Stream.concatEffect $ do
172 km <- toMap s2
173 return $ fmap (joinAB km) s1
175 where
177 joinAB km (k, a) =
178 case k `Map.lookup` km of
179 Just b -> (k, a, Just b)
180 Nothing -> (k, a, Nothing)
182-- XXX We can do this concurrently.
183-- XXX Check performance of StreamD vs StreamK cross operation.
185-- | Like 'leftJoin' but emits a @(Just a, Just b)@. Like 'leftJoin', for those
186-- @a@'s that are not equal to any @b@ emit @(Just a, Nothing)@, but
187-- additionally, for those @b@'s that are not equal to any @a@ emit @(Nothing,
188-- Just b)@.
189--
190-- This is a generalization of left join to include all the elements from the
191-- right stream as well, in other words it is a combination of left and right
192-- joins. This is a commutative operation. The order of stream arguments can be
193-- changed without affecting results, except for the ordering of elements in
194-- the resulting tuple.
195--
196-- For space efficiency use the smaller stream as the second stream.
197--
198-- Space: O(n)
199--
200-- Time: O(m x n)
201--
202-- /Pre-release/
203{-# INLINE outerJoin #-}
204outerJoin :: MonadIO m =>
205 (a -> b -> Bool)
206 -> Stream m a
207 -> Stream m b
208 -> Stream m (Maybe a, Maybe b)
209outerJoin eq s1 s2 =
210 Stream.concatEffect $ do
211 inputArr <- Array.fromStream s2
212 let len = Array.length inputArr
213 foundArr <-
214 Stream.fold
215 (MA.writeN len)
216 (Stream.fromList (Prelude.replicate len False))
217 return $ go inputArr foundArr `Stream.append` leftOver inputArr foundArr
219 where
221 leftOver inputArr foundArr =
222 let stream1 = Array.read inputArr
223 stream2 = Stream.unfold MA.reader foundArr
224 in Stream.filter
225 isJust
226 ( Stream.zipWith (\x y ->
227 if y
228 then Nothing
229 else Just (Nothing, Just x)
230 ) stream1 stream2
231 ) & Stream.catMaybes
233 evalState = Stream.evalStateT (return False) . unCross
235 go inputArr foundArr = evalState $ do
236 a <- mkCross (Stream.liftInner s1)
237 -- XXX should we use StreamD monad here?
238 -- XXX Is there a better way to perform some action at the end of a loop
239 -- iteration?
240 mkCross (Stream.fromEffect $ put False)
241 let final = Stream.concatEffect $ do
242 r <- get
243 if r
244 then pure Stream.nil
245 else pure (Stream.fromPure Nothing)
246 (i, b) <-
247 let stream = Array.read inputArr
248 in mkCross
249 (Stream.indexed $ fmap Just (Stream.liftInner stream) `Stream.append` final)
251 case b of
252 Just b1 ->
253 if a `eq` b1
254 then do
255 mkCross (Stream.fromEffect $ put True)
256 MA.putIndex i foundArr True
257 return (Just a, Just b1)
258 else mkCross Stream.nil
259 Nothing -> return (Just a, Nothing)
261-- Put the b's that have been paired, in another hash or mutate the hash to set
262-- a flag. At the end go through @Stream m b@ and find those that are not in that
263-- hash to return (Nothing, b).
265-- | 'outerJoin' specialized to 'Ord' types for better performance.
266--
267-- Space: O(m + n)
268--
269-- Time: O(m + n)
270--
271-- /Pre-release/
272{-# INLINE outerOrdJoin #-}
273outerOrdJoin ::
274 (Ord k, MonadIO m) =>
275 Stream m (k, a) -> Stream m (k, b) -> Stream m (k, Maybe a, Maybe b)
276outerOrdJoin s1 s2 =
277 Stream.concatEffect $ do
278 km1 <- kvFold s1
279 km2 <- kvFold s2
281 -- XXX Not sure if toList/fromList would fuse optimally. We may have to
282 -- create a fused Map.toStream function.
283 let res1 = fmap (joinAB km2)
284 $ Stream.fromList $ Map.toList km1
285 where
286 joinAB km (k, a) =
287 case k `Map.lookup` km of
288 Just b -> (k, Just a, Just b)
289 Nothing -> (k, Just a, Nothing)
291 -- XXX We can take advantage of the lookups in the first pass above to
292 -- reduce the number of lookups in this pass. If we keep mutable cells
293 -- in the second Map, we can flag it in the first pass and not do any
294 -- lookup in the second pass if it is flagged.
295 let res2 = Stream.mapMaybe (joinAB km1)
296 $ Stream.fromList $ Map.toList km2
297 where
298 joinAB km (k, b) =
299 case k `Map.lookup` km of
300 Just _ -> Nothing
301 Nothing -> Just (k, Nothing, Just b)
303 return $ Stream.append res1 res2
305 where
307 -- XXX Generate error if a duplicate insertion is attempted?
308 kvFold =
309 let f = Fold.foldl' (\kv (k, b) -> Map.insert k b kv) Map.empty
310 in Stream.fold f