0257838 · streamly · git.iko.soy

git.iko.soy/streamly

Use the demux input to determine the fold

We may want to use the entire value to determine the fold and not just the key. For example, we may have the key as an Int that uniquely determines the fold (for example a connection/request id). Once the fold starts we can always lookup the fold using this key. However, to determine what fold to use initially we may need more information than just the request-id (e.g. the type of the message).

commit02578386dab3479063e2f8e3e65efd5f1490a1d1 authorHarendra Kumar <harendra@composewell.com> dateMon 11 Apr 2022, 13:25 (5 years ago) parent6170d8d

3 files changed+24−19

compare with master

unifiedsplit

benchmark/Streamly/Benchmark/Data/Fold.hs+10−10 core/src/Streamly/Internal/Data/Fold.hs+13−8 test/Streamly/Test/Data/Fold.hs+1−1

benchmark/Streamly/Benchmark/Data/Fold.hs+10−10

@@ -193,27 +193,27 @@ partitionByMinM =

193193

194194{-# INLINE demuxWith #-}

195195demuxWith :: (Monad m, Ord k) =>

196− (a -> k) -> (k -> m (Fold m a b)) -> SerialT m a -> m (Map k b)

196+ (a -> k) -> (a -> m (Fold m a b)) -> SerialT m a -> m (Map k b)

197197demuxWith f g = S.fold (FL.demuxWith f g)

198198

199199{-# INLINE demuxWithInt #-}

200200demuxWithInt :: Monad m =>

201− (a -> Int) -> (Int -> m (Fold m a b)) -> SerialT m a -> m (IntMap b)

201+ (a -> Int) -> (a -> m (Fold m a b)) -> SerialT m a -> m (IntMap b)

202202demuxWithInt f g = S.fold (FL.demuxWith f g)

203203

204204{-# INLINE demuxWithHash #-}

205205demuxWithHash :: (Monad m, Ord k, Hashable k) =>

206− (a -> k) -> (k -> m (Fold m a b)) -> SerialT m a -> m (HashMap k b)

206+ (a -> k) -> (a -> m (Fold m a b)) -> SerialT m a -> m (HashMap k b)

207207demuxWithHash f g = S.fold (FL.demuxWith f g)

208208

209209{-# INLINE demuxMutWith #-}

210210demuxMutWith :: (MonadIO m, Ord k) =>

211− (a -> k) -> (k -> m (Fold m a b)) -> SerialT m a -> m (Map k b)

211+ (a -> k) -> (a -> m (Fold m a b)) -> SerialT m a -> m (Map k b)

212212demuxMutWith f g = S.fold (FL.demuxMutWith f g)

213213

214214{-# INLINE demuxMutWithHash #-}

215215demuxMutWithHash :: (MonadIO m, Ord k, Hashable k) =>

216− (a -> k) -> (k -> m (Fold m a b)) -> SerialT m a -> m (HashMap k b)

216+ (a -> k) -> (a -> m (Fold m a b)) -> SerialT m a -> m (HashMap k b)

217217demuxMutWithHash f g = S.fold (FL.demuxMutWith f g)

218218

219219{-# INLINE classifyWith #-}

@@ -429,15 +429,15 @@ o_n_heap_serial value =

429429 , bgroup "key-value"

430430 [

431431 benchIOSink value "demuxWith (64 buckets) [sum, length]"

432− $ demuxWith (getKey 64) getFold

432+ $ demuxWith (getKey 64) (getFold . getKey 64)

433433 , benchIOSink value "demuxWithInt (64 buckets) [sum, length]"

434− $ demuxWithInt (getKey 64) getFold

434+ $ demuxWithInt (getKey 64) (getFold . getKey 64)

435435 , benchIOSink value "demuxWithHash (64 buckets) [sum, length]"

436− $ demuxWithHash (getKey 64) getFold

436+ $ demuxWithHash (getKey 64) (getFold . getKey 64)

437437 , benchIOSink value "demuxMutWith (64 buckets) [sum, length]"

438− $ demuxMutWith (getKey 64) getFold

438+ $ demuxMutWith (getKey 64) (getFold . getKey 64)

439439 , benchIOSink value "demuxMutWithHash (64 buckets) [sum, length]"

440− $ demuxMutWithHash (getKey 64) getFold

440+ $ demuxMutWithHash (getKey 64) (getFold . getKey 64)

441441

442442 -- classify: immutable

443443 , benchIOSink value "classifyWith (64 buckets) sum"

core/src/Streamly/Internal/Data/Fold.hs+13−8

@@ -1482,6 +1482,10 @@ partition = partitionBy id

14821482-- XXX If we use Refold in it, it can perhaps fuse/be more efficient. For

14831483-- example we can store just the result rather than storing the whole fold in

14841484-- the Map.

1485+--

1486+-- Note: There are separate functions to determine Key and Fold from the input

1487+-- because key is to be determined on each input whereas fold is to be

1488+-- determined only once for a key.

14851489

14861490-- | In a key value stream, fold values corresponding to each key with a key

14871491-- specific fold. The fold returns the fold result as the second component of

@@ -1497,7 +1501,7 @@ partition = partitionBy id

14971501{-# INLINE demuxScanWith #-}

14981502demuxScanWith :: (Monad m, IsMap f, Traversable f) =>

14991503 (a -> Key f)

1500− -> (Key f -> m (Fold m a b))

1504+ -> (a -> m (Fold m a b))

15011505 -> Fold m a (m (f b), Maybe (Key f, b))

15021506demuxScanWith getKey getFold = fmap extract $ foldlM' step initial

15031507

@@ -1521,9 +1525,10 @@ demuxScanWith getKey getFold = fmap extract $ foldlM' step initial

15211525

15221526 step (Tuple' kv _) a = do

15231527 let k = getKey a

1524− fld <- getFold k

15251528 case IsMap.mapLookup k kv of

1526− Nothing -> runFold kv fld (k, a)

1529+ Nothing -> do

1530+ fld <- getFold a

1531+ runFold kv fld (k, a)

15271532 Just f -> runFold kv f (k, a)

15281533

15291534 extract (Tuple' kv x) = (Prelude.mapM f kv, x)

@@ -1544,7 +1549,7 @@ demuxScanWith getKey getFold = fmap extract $ foldlM' step initial

15441549{-# INLINE demuxScanMutWith #-}

15451550demuxScanMutWith :: (MonadIO m, IsMap f, Traversable f) =>

15461551 (a -> Key f)

1547− -> (Key f -> m (Fold m a b))

1552+ -> (a -> m (Fold m a b))

15481553 -> Fold m a (m (f b), Maybe (Key f, b))

15491554demuxScanMutWith getKey getFold = fmap extract $ foldlM' step initial

15501555

@@ -1586,7 +1591,7 @@ demuxScanMutWith getKey getFold = fmap extract $ foldlM' step initial

15861591 let k = getKey a

15871592 case IsMap.mapLookup k kv of

15881593 Nothing -> do

1589− f <- getFold k

1594+ f <- getFold a

15901595 initFold kv f (k, a)

15911596 Just ref -> do

15921597 f <- liftIO $ readIORef ref

@@ -1607,7 +1612,7 @@ demuxScanMutWith getKey getFold = fmap extract $ foldlM' step initial

16071612--

16081613{-# INLINE demuxWith #-}

16091614demuxWith :: (Monad m, IsMap f, Traversable f) =>

1610− (a -> Key f) -> (Key f -> m (Fold m a b)) -> Fold m a (f b)

1615+ (a -> Key f) -> (a -> m (Fold m a b)) -> Fold m a (f b)

16111616demuxWith getKey getFold =

16121617 let

16131618 classifier = demuxScanWith getKey getFold

@@ -1625,7 +1630,7 @@ demuxWith getKey getFold =

16251630--

16261631{-# INLINE demuxMutWith #-}

16271632demuxMutWith :: (MonadIO m, IsMap f, Traversable f) =>

1628− (a -> Key f) -> (Key f -> m (Fold m a b)) -> Fold m a (f b)

1633+ (a -> Key f) -> (a -> m (Fold m a b)) -> Fold m a (f b)

16291634demuxMutWith getKey getFold =

16301635 let

16311636 classifier = demuxScanMutWith getKey getFold

@@ -1654,7 +1659,7 @@ demuxMutWith getKey getFold =

16541659{-# INLINE demux #-}

16551660demux :: (Monad m, IsMap f, Traversable f) =>

16561661 (Key f -> m (Fold m a b)) -> Fold m (Key f, a) (f b)

1657−demux f = demuxWith fst (fmap (lmap snd) . f)

1662+demux f = demuxWith fst (\(k, _) -> fmap (lmap snd) (f k))

16581663

16591664------------------------------------------------------------------------------

16601665-- Classify: Like demux but uses the same fold for all keys.

test/Streamly/Test/Data/Fold.hs+1−1

@@ -590,7 +590,7 @@ demuxWith =

590590

591591 input = Stream.fromList [1, 2, 3, 4 :: Int]

592592 in Stream.fold

593− (F.demuxWith getKey getFold)

593+ (F.demuxWith getKey (getFold . getKey))

594594 input

595595 `shouldReturn`

596596 Data.Map.fromList [("PRODUCT",3),("SUM",6)]