0257838 · streamly · git.iko.soy

git.iko.soy/streamly

on master help sign in

Use the demux input to determine the fold

0257838Harendra Kumar5 years ago

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

compare with master

unifiedsplit

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

193

194{-# INLINE demuxWith #-}

195demuxWith :: (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)

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

198

199{-# INLINE demuxWithInt #-}

200demuxWithInt :: 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)

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

203

204{-# INLINE demuxWithHash #-}

205demuxWithHash :: (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)

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

208

209{-# INLINE demuxMutWith #-}

210demuxMutWith :: (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)

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

213

214{-# INLINE demuxMutWithHash #-}

215demuxMutWithHash :: (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)

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

218

219{-# INLINE classifyWith #-}

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

429 , bgroup "key-value"

430 [

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

432− $ demuxWith (getKey 64) getFold

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

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

434− $ demuxWithInt (getKey 64) getFold

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

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

436− $ demuxWithHash (getKey 64) getFold

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

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

438− $ demuxMutWith (getKey 64) getFold

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

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

440− $ demuxMutWithHash (getKey 64) getFold

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

441

442 -- classify: immutable

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

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

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

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

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

1484-- 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.

1485

1489

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

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

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

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

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

1497{-# INLINE demuxScanWith #-}

1501{-# INLINE demuxScanWith #-}

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

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

1499 (a -> Key f)

1503 (a -> Key f)

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

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

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

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

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

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

1503

1507

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

1521

1525

1522 step (Tuple' kv _) a = do

1526 step (Tuple' kv _) a = do

1523 let k = getKey a

1527 let k = getKey a

1524− fld <- getFold k

1525 case IsMap.mapLookup k kv of

1528 case IsMap.mapLookup k kv of

1529+ Nothing -> do

1530+ fld <- getFold a

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

1531+ runFold kv fld (k, a)

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

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

1528

1533

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

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

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

1544{-# INLINE demuxScanMutWith #-}

1549{-# INLINE demuxScanMutWith #-}

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

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

1546 (a -> Key f)

1551 (a -> Key f)

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

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

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

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

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

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

1550

1555

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

1586 let k = getKey a

1591 let k = getKey a

1587 case IsMap.mapLookup k kv of

1592 case IsMap.mapLookup k kv of

1588 Nothing -> do

1593 Nothing -> do

1589− f <- getFold k

1594+ f <- getFold a

1590 initFold kv f (k, a)

1595 initFold kv f (k, a)

1591 Just ref -> do

1596 Just ref -> do

1592 f <- liftIO $ readIORef ref

1597 f <- liftIO $ readIORef ref

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

1607--

1612--

1608{-# INLINE demuxWith #-}

1613{-# INLINE demuxWith #-}

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

1614demuxWith :: (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)

1611demuxWith getKey getFold =

1616demuxWith getKey getFold =

1612 let

1617 let

1613 classifier = demuxScanWith getKey getFold

1618 classifier = demuxScanWith getKey getFold

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

1625--

1630--

1626{-# INLINE demuxMutWith #-}

1631{-# INLINE demuxMutWith #-}

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

1632demuxMutWith :: (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)

1629demuxMutWith getKey getFold =

1634demuxMutWith getKey getFold =

1630 let

1635 let

1631 classifier = demuxScanMutWith getKey getFold

1636 classifier = demuxScanMutWith getKey getFold

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

1654{-# INLINE demux #-}

1659{-# INLINE demux #-}

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

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

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

1661 (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))

1658

1663

1659------------------------------------------------------------------------------

1664------------------------------------------------------------------------------

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

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

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

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

590

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

592 in Stream.fold

593− (F.demuxWith getKey getFold)

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

594 input

595 `shouldReturn`

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