e078bae · streamly · git.iko.soy

git.iko.soy/streamly

on master help sign in

Allow different types of Map in classify

e078baeHarendra Kumar5 years ago

benchmark/Streamly/Benchmark/Data/Fold.hs+37−3

compare with master

unifiedsplit

@@ -6,12 +6,18 @@

6-- Maintainer : streamly@composewell.com

8{-# LANGUAGE FlexibleContexts #-}

9+{-# LANGUAGE RankNTypes #-}

10+{-# LANGUAGE ScopedTypeVariables #-}

10module Main (main) where

12module Main (main) where

12import Control.DeepSeq (NFData(..))

14import Control.DeepSeq (NFData(..))

15+import Control.Monad.IO.Class (MonadIO(..))

13import Data.Functor.Identity (Identity)

16import Data.Functor.Identity (Identity)

14import Data.Map.Strict (Map)

17import Data.Map.Strict (Map)

18+import Data.Hashable (Hashable)

19+import Data.HashMap.Strict (HashMap)

20+import Data.IntMap.Strict (IntMap)

15import Data.Monoid (Last(..), Sum(..))

21import Data.Monoid (Last(..), Sum(..))

16import System.Random (randomRIO)

22import System.Random (randomRIO)

@@ -211,8 +217,29 @@ classifyWith f = S.fold (FL.classifyWith f FL.sum)

211

217

212{-# INLINE classifyScanWith #-}

218{-# INLINE classifyScanWith #-}

213classifyScanWith ::

219classifyScanWith ::

214− (Monad m, Ord k, Num a) => (a -> k) -> SerialT m a -> m ()

220+ forall m k a. (Monad m, Ord k, Num a) => (a -> k) -> SerialT m a -> m ()

221+classifyScanWith f =

222+ S.drain . S.postscan (g f FL.sum)

223+

224+ where

225+ g =

215−classifyScanWith f = S.drain . S.postscan (FL.classifyScanWith f FL.sum)

226+ FL.classifyScanWith ::

227+ (a -> k) -> Fold m a b -> Fold m a (m (Map k b), Maybe (k, b))

228+

229+{-# INLINE classifyMutWith #-}

230+classifyMutWith ::

231+ (MonadIO m, Ord k, Num a) => (a -> k) -> SerialT m a -> m (Map k a)

232+classifyMutWith f = S.fold (FL.classifyMutWith f FL.sum)

233+

234+{-# INLINE classifyMutWithInt #-}

235+classifyMutWithInt ::

236+ (MonadIO m, Num a) => (a -> Int) -> SerialT m a -> m (IntMap a)

237+classifyMutWithInt f = S.fold (FL.classifyMutWith f FL.sum)

238+

239+{-# INLINE classifyMutWithHash #-}

240+classifyMutWithHash :: (MonadIO m, Ord k, Num a, Hashable k) =>

241+ (a -> k) -> SerialT m a -> m (HashMap k a)

242+classifyMutWithHash f = S.fold (FL.classifyMutWith f FL.sum)

216

243

217-------------------------------------------------------------------------------

244-------------------------------------------------------------------------------

218-- unzip

245-- unzip

@@ -374,10 +401,17 @@ o_1_space_serial_composition value =

374 , benchIOSink value "demuxDefaultWith [sum, length] sum"

401 , benchIOSink value "demuxDefaultWith [sum, length] sum"

375 $ demuxDefaultWith fn mp

402 $ demuxDefaultWith fn mp

376 , benchIOSink value "demuxWith [sum, length]" $ demuxWith fn mp

403 , benchIOSink value "demuxWith [sum, length]" $ demuxWith fn mp

377− , benchIOSink value "classifyWith sum" $ classifyWith (fst . fn)

404+ , benchIOSink value "classifyWith sum"

405+ $ classifyWith (fst . fn)

378 , benchIOSink value "classifyScanWith sum"

406 , benchIOSink value "classifyScanWith sum"

379 $ classifyScanWith (fst . fn)

407 $ classifyScanWith (fst . fn)

380 , bench "unfoldMany" $ unfoldMany value

408 , bench "unfoldMany" $ unfoldMany value

409+ , benchIOSink value "classifyMutWith sum"

410+ $ classifyMutWith (fst . fn)

411+ , benchIOSink value "classifyMutWithHash sum"

412+ $ classifyMutWithHash (fst . fn)

413+ , benchIOSink value "classifyMutWithInt sum"

414+ $ classifyMutWithInt (fst . fn)

381 ]

415 ]

382 ]

416 ]

383

417

benchmark/streamly-benchmarks.cabal+2

@@ -125,6 +125,8 @@ common bench-depends

125 , random >= 1.0 && < 2.0

126 , transformers >= 0.4 && < 0.7

127 , containers >= 0.5 && < 0.7

128+ , hashable >= 1.3 && < 1.5

129+ , unordered-containers >= 0.2 && < 0.3

128 , process >= 1.4 && < 1.7

130 , process >= 1.4 && < 1.7

129 , directory >= 1.2.2 && < 1.4

131 , directory >= 1.2.2 && < 1.4

130 , filepath >= 1.4.1 && < 1.5

132 , filepath >= 1.4.1 && < 1.5

core/src/Streamly/Internal/Data/Fold.hs+87−9

@@ -214,6 +214,7 @@ module Streamly.Internal.Data.Fold

214 -- in individual output buckets using the given fold.

215 , classify

216 , classifyWith

217+ , classifyMutWith

217 , classifyScanWith

218 , classifyScanWith

218 -- , classifyWithSel

219 -- , classifyWithSel

219 -- , classifyWithMin

220 -- , classifyWithMin

@@ -260,14 +261,17 @@ where

260

261

261#include "inline.hs"

262#include "inline.hs"

262

263

264+import Control.Monad.IO.Class (MonadIO(..))

263import Data.Bifunctor (first)

265import Data.Bifunctor (first)

264import Data.Either (isLeft, isRight, fromLeft, fromRight)

266import Data.Either (isLeft, isRight, fromLeft, fromRight)

265import Data.Functor.Identity (Identity(..))

267import Data.Functor.Identity (Identity(..))

266import Data.Int (Int64)

268import Data.Int (Int64)

269+import Data.IORef (newIORef, readIORef, writeIORef)

267import Data.Map.Strict (Map)

270import Data.Map.Strict (Map)

268import Data.Maybe (isJust, fromJust)

271import Data.Maybe (isJust, fromJust)

269import Streamly.Internal.Data.Either.Strict

272import Streamly.Internal.Data.Either.Strict

270 (Either'(..), fromLeft', fromRight', isLeft', isRight')

273 (Either'(..), fromLeft', fromRight', isLeft', isRight')

274+import Streamly.Internal.Data.IsMap (IsMap(..))

271import Streamly.Internal.Data.Pipe.Type (Pipe (..), PipeState(..))

275import Streamly.Internal.Data.Pipe.Type (Pipe (..), PipeState(..))

272import Streamly.Internal.Data.Unfold.Type (Unfold(..))

276import Streamly.Internal.Data.Unfold.Type (Unfold(..))

273import Streamly.Internal.Data.Tuple.Strict (Tuple'(..), Tuple3'(..))

277import Streamly.Internal.Data.Tuple.Strict (Tuple'(..), Tuple3'(..))

@@ -276,6 +280,7 @@ import Streamly.Internal.Data.Stream.Serial (SerialT(..))

276import qualified Data.IntSet as IntSet

280import qualified Data.IntSet as IntSet

277import qualified Data.Map.Strict as Map

281import qualified Data.Map.Strict as Map

278import qualified Data.Set as Set

282import qualified Data.Set as Set

283+import qualified Streamly.Internal.Data.IsMap as IsMap

279import qualified Streamly.Internal.Data.Pipe.Type as Pipe

284import qualified Streamly.Internal.Data.Pipe.Type as Pipe

280import qualified Streamly.Internal.Data.Stream.StreamD.Type as StreamD

285import qualified Streamly.Internal.Data.Stream.StreamD.Type as StreamD

281-- import qualified Streamly.Internal.Data.Stream.IsStream.Enumeration as Stream

286-- import qualified Streamly.Internal.Data.Stream.IsStream.Enumeration as Stream

@@ -1625,19 +1630,23 @@ demuxDefault :: (Monad m, Ord k)

1625 => Map k (Fold m a b) -> Fold m (k, a) b -> Fold m (k, a) (Map k b, b)

1630 => Map k (Fold m a b) -> Fold m (k, a) b -> Fold m (k, a) (Map k b, b)

1626demuxDefault = demuxDefaultWith id

1631demuxDefault = demuxDefaultWith id

1627

1632

1633+-- | Folds the values for each key using the supplied fold. When scanning, as

1634+-- soon as the fold is complete, its result is available in the second

1635+-- component of the tuple. The first component of the tuple is a snapshot of

1636+-- the in-progress folds.

1628{-# INLINE classifyScanWith #-}

1637{-# INLINE classifyScanWith #-}

1629−classifyScanWith :: (Monad m, Ord k) =>

1638+classifyScanWith :: (Monad m, IsMap f, Traversable f, Ord (Key f)) =>

1630 -- Note: we need to return the Map itself to display the in-progress values

1639 -- Note: we need to return the Map itself to display the in-progress values

1631 -- e.g. to implement top. We could possibly create a separate abstraction

1640 -- e.g. to implement top. We could possibly create a separate abstraction

1632 -- for that use case. We return an action because we want it to be lazy so

1641 -- for that use case. We return an action because we want it to be lazy so

1633 -- that the downstream consumers can choose to process or discard it.

1642 -- that the downstream consumers can choose to process or discard it.

1634− (a -> k) -> Fold m a b -> Fold m a (m (Map k b), Maybe (k, b))

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

1635classifyScanWith f (Fold step1 initial1 extract1) =

1644classifyScanWith f (Fold step1 initial1 extract1) =

1636 fmap extract $ foldlM' step initial

1645 fmap extract $ foldlM' step initial

1637

1646

1638 where

1647 where

1639

1648

1640− initial = return $ Tuple3' Map.empty Set.empty Nothing

1649+ initial = return $ Tuple3' IsMap.mapEmpty Set.empty Nothing

1641

1650

1642 {-# INLINE initFold #-}

1651 {-# INLINE initFold #-}

1643 initFold kv set k a = do

1652 initFold kv set k a = do

@@ -1648,14 +1657,14 @@ classifyScanWith f (Fold step1 initial1 extract1) =

1648 return

1657 return

1649 $ case r of

1658 $ case r of

1650 Partial s1 ->

1659 Partial s1 ->

1651− Tuple3' (Map.insert k s1 kv) set Nothing

1660+ Tuple3' (IsMap.mapInsert k s1 kv) set Nothing

1652 Done b ->

1661 Done b ->

1653 Tuple3' kv set (Just (k, b))

1662 Tuple3' kv set (Just (k, b))

1654 Done b -> return (Tuple3' kv (Set.insert k set) (Just (k, b)))

1663 Done b -> return (Tuple3' kv (Set.insert k set) (Just (k, b)))

1655

1664

1656 step (Tuple3' kv set _) a = do

1665 step (Tuple3' kv set _) a = do

1657 let k = f a

1666 let k = f a

1658− case Map.lookup k kv of

1667+ case IsMap.mapLookup k kv of

1659 Nothing -> do

1668 Nothing -> do

1660 if Set.member k set

1669 if Set.member k set

1661 then return (Tuple3' kv set Nothing)

1670 then return (Tuple3' kv set Nothing)

@@ -1665,16 +1674,69 @@ classifyScanWith f (Fold step1 initial1 extract1) =

1665 return

1674 return

1666 $ case r of

1675 $ case r of

1667 Partial s1 ->

1676 Partial s1 ->

1668− Tuple3' (Map.insert k s1 kv) set Nothing

1677+ Tuple3' (IsMap.mapInsert k s1 kv) set Nothing

1669 Done b ->

1678 Done b ->

1670− let kv1 = Map.delete k kv

1679+ let kv1 = IsMap.mapDelete k kv

1671 in Tuple3' kv1 (Set.insert k set) (Just (k, b))

1680 in Tuple3' kv1 (Set.insert k set) (Just (k, b))

1672

1681

1673 extract (Tuple3' kv _ x) = (Prelude.mapM extract1 kv, x)

1682 extract (Tuple3' kv _ x) = (Prelude.mapM extract1 kv, x)

1674

1683

1684+-- The code is almost the same as classifyScanWith except the IORef operations.

1685+--

1686+-- | Same as classifyScanWith except that it uses mutable IORef cells in the

1687+-- Map providing better performance. Be aware that if this is used as a scan,

1688+-- the values in the intermediate Maps would be mutable.

1689+{-# INLINE classifyScanMutWith #-}

1690+classifyScanMutWith :: (MonadIO m, IsMap f, Traversable f, Ord (Key f)) =>

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

1692+classifyScanMutWith f (Fold step1 initial1 extract1) =

1693+ fmap extract $ foldlM' step initial

1694+

1695+ where

1696+

1697+ initial = return $ Tuple3' IsMap.mapEmpty Set.empty Nothing

1698+

1699+ {-# INLINE initFold #-}

1700+ initFold kv set k a = do

1701+ x <- initial1

1702+ case x of

1703+ Partial s -> do

1704+ r <- step1 s a

1705+ case r of

1706+ Partial s1 -> do

1707+ ref <- liftIO $ newIORef s1

1708+ return $ Tuple3' (IsMap.mapInsert k ref kv) set Nothing

1709+ Done b ->

1710+ return $ Tuple3' kv set (Just (k, b))

1711+ Done b -> return (Tuple3' kv (Set.insert k set) (Just (k, b)))

1712+

1713+ step (Tuple3' kv set _) a = do

1714+ let k = f a

1715+ case IsMap.mapLookup k kv of

1716+ Nothing -> do

1717+ if Set.member k set

1718+ then return (Tuple3' kv set Nothing)

1719+ else initFold kv set k a

1720+ Just ref -> do

1721+ s <- liftIO $ readIORef ref

1722+ r <- step1 s a

1723+ case r of

1724+ Partial s1 -> do

1725+ liftIO $ writeIORef ref s1

1726+ return $ Tuple3' kv set Nothing

1727+ Done b ->

1728+ let kv1 = IsMap.mapDelete k kv

1729+ in return

1730+ $ Tuple3' kv1 (Set.insert k set) (Just (k, b))

1731+

1732+ extract (Tuple3' kv _ x) =

1733+ (Prelude.mapM (\ref -> liftIO (readIORef ref) >>= extract1) kv, x)

1734+

1735+-- | Fold a key value stream to a key-value container. If the same key appears

1736+-- multiple times, only the last value is retained.

1675{-# INLINE toMap #-}

1737{-# INLINE toMap #-}

1676−toMap :: (Monad m, Ord k) => Fold m (k, a) (Map k a)

1738+toMap :: (Monad m, IsMap f) => Fold m (Key f, a) (f a)

1677−toMap = foldl' (\kv (k, v) -> Map.insert k v kv) Map.empty

1739+toMap = foldl' (\kv (k, v) -> IsMap.mapInsert k v kv) IsMap.mapEmpty

1678

1740

1679-- | Split the input stream based on a key field and fold each split using the

1741-- | Split the input stream based on a key field and fold each split using the

1680-- given fold. Useful for map/reduce, bucketizing the input in different bins

1742-- given fold. Useful for map/reduce, bucketizing the input in different bins

@@ -1711,6 +1773,22 @@ classifyWith f fld =

1711 (lmap snd $ catMaybes toMap)

1773 (lmap snd $ catMaybes toMap)

1712 in postscan classifier aggregator

1774 in postscan classifier aggregator

1713

1775

1776+-- | Like classifyWith but maybe faster because it uses mutable cells as fold

1777+-- accumulators in the Map.

1778+{-# INLINE classifyMutWith #-}

1779+classifyMutWith :: (MonadIO m, IsMap f, Traversable f, Ord (Key f)) =>

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

1781+classifyMutWith f fld =

1782+ let

1783+ classifier = classifyScanMutWith f fld

1784+ getMap Nothing = pure IsMap.mapEmpty

1785+ getMap (Just action) = action

1786+ aggregator =

1787+ teeWith IsMap.mapUnion

1788+ (rmapM getMap $ lmap fst last)

1789+ (lmap snd $ catMaybes toMap)

1790+ in postscan classifier aggregator

1791+

1714-- | Given an input stream of key value pairs and a fold for values, fold all

1792-- | Given an input stream of key value pairs and a fold for values, fold all

1715-- the values belonging to each key. Useful for map/reduce, bucketizing the

1793-- the values belonging to each key. Useful for map/reduce, bucketizing the

1716-- input in different bins or for generating histograms.

1794-- input in different bins or for generating histograms.

core/src/Streamly/Internal/Data/IsMap.hsnew file+61

@@ -0,0 +1,61 @@

1+-- |

2+-- Module : Streamly.Internal.Data.IsMap

4+-- License : BSD-3-Clause

5+-- Maintainer : streamly@composewell.com

6+-- Stability : experimental

7+-- Portability : GHC

9+module Streamly.Internal.Data.IsMap (IsMap(..)) where

10+

11+import Data.Hashable (Hashable)

12+import Data.Kind (Type)

13+import Data.Map.Strict (Map)

14+

15+import qualified Data.HashMap.Strict as HashMap

16+import qualified Data.IntMap.Strict as IntMap

17+import qualified Data.Map.Strict as Map

18+

19+-- XXX Try unpacked-containers

20+

21+class IsMap f where

22+ type Key f :: Type

23+

24+ mapEmpty :: f a

25+ mapAlterF :: Functor g =>

26+ (Maybe a -> g (Maybe a)) -> Key f -> f a -> g (f a)

27+ -- These can be implemented in terms of alterF itself

28+ mapLookup :: Key f -> f a -> Maybe a

29+ mapInsert :: Key f -> a -> f a -> f a

30+ mapDelete :: Key f -> f a -> f a

31+ mapUnion :: f a -> f a -> f a

32+

33+instance Ord k => IsMap (Map k) where

34+ type Key (Map k) = k

35+

36+ mapEmpty = Map.empty

37+ mapAlterF = Map.alterF

38+ mapLookup = Map.lookup

39+ mapInsert = Map.insert

40+ mapDelete = Map.delete

41+ mapUnion = Map.union

42+

43+instance IsMap IntMap.IntMap where

44+ type Key IntMap.IntMap = Int

45+

46+ mapEmpty = IntMap.empty

47+ mapAlterF = IntMap.alterF

48+ mapLookup = IntMap.lookup

49+ mapInsert = IntMap.insert

50+ mapDelete = IntMap.delete

51+ mapUnion = IntMap.union

52+

53+instance (Hashable k, Eq k) => IsMap (HashMap.HashMap k) where

54+ type Key (HashMap.HashMap k) = k

55+

56+ mapEmpty = HashMap.empty

57+ mapAlterF = HashMap.alterF

58+ mapLookup = HashMap.lookup

59+ mapInsert = HashMap.insert

60+ mapDelete = HashMap.delete

61+ mapUnion = HashMap.union

core/streamly-core.cabal+3−1

@@ -221,6 +221,7 @@ library

221 , Streamly.Internal.Data.Cont

222 , Streamly.Internal.Foreign.Malloc

223 , Streamly.Internal.System.IO

224+ , Streamly.Internal.Data.IsMap

224

225

225 -- streamly-strict-data

226 -- streamly-strict-data

226 , Streamly.Internal.Data.Tuple.Strict

227 , Streamly.Internal.Data.Tuple.Strict

@@ -337,7 +338,8 @@ library

337 , primitive >= 0.5.4 && < 0.8

338 , primitive >= 0.5.4 && < 0.8

338 , heaps >= 0.3 && < 0.5

339 , heaps >= 0.3 && < 0.5

339 , deque >= 0.4 && < 0.5

340 , deque >= 0.4 && < 0.5

340−

341+ , hashable >= 1.3 && < 1.5

342+ , unordered-containers >= 0.2 && < 0.3

341 if flag(use-unliftio)

343 if flag(use-unliftio)

342 build-depends: unliftio-core >= 0.2 && < 0.3

344 build-depends: unliftio-core >= 0.2 && < 0.3

343 else

345 else

streamly.cabal+1

@@ -145,6 +145,7 @@ extra-source-files:

145 core/src/Streamly/Internal/Data/Cont.hs

146 core/src/Streamly/Internal/Foreign/Malloc.hs

147 core/src/Streamly/Internal/System/IO.hs

148+ core/src/Streamly/Internal/Data/IsMap.hs

148 core/src/Streamly/Internal/Data/Tuple/Strict.hs

149 core/src/Streamly/Internal/Data/Tuple/Strict.hs

149 core/src/Streamly/Internal/Data/Maybe/Strict.hs

150 core/src/Streamly/Internal/Data/Maybe/Strict.hs

150 core/src/Streamly/Internal/Data/Either/Strict.hs

151 core/src/Streamly/Internal/Data/Either/Strict.hs