git.iko.soy/streamly
on master help sign in
Add countDistinct
74e6a24Harendra Kumar
core/src/Streamly/Internal/Data/Fold.hs+54
compare with master
unifiedsplit
@@ -55,6 +55,9 @@ module Streamly.Internal.Data.Fold
55 , drainBy
55 , drainBy
56 , last
56 , last
57 , length
57 , length
58+ , genericLength
59+ , countDistinct
60+ , countDistinctInt
58 , mean
61 , mean
59 , variance
62 , variance
60 , stdDev
63 , stdDev
@@ -268,6 +271,7 @@ import Streamly.Internal.Data.Pipe.Type (Pipe (..), PipeState(..))
268import Streamly.Internal.Data.Tuple.Strict (Tuple'(..), Tuple3'(..))
271import Streamly.Internal.Data.Tuple.Strict (Tuple'(..), Tuple3'(..))
269import Streamly.Internal.Data.Stream.Serial (SerialT(..))
272import Streamly.Internal.Data.Stream.Serial (SerialT(..))
270
273
274+import qualified Data.IntSet as IntSet
271import qualified Data.Map.Strict as Map
275import qualified Data.Map.Strict as Map
272import qualified Data.Set as Set
276import qualified Data.Set as Set
273import qualified Streamly.Internal.Data.Pipe.Type as Pipe
277import qualified Streamly.Internal.Data.Pipe.Type as Pipe
@@ -526,6 +530,56 @@ genericLength = foldl' (\n _ -> n + 1) 0
526length :: Monad m => Fold m a Int
530length :: Monad m => Fold m a Int
527length = genericLength
531length = genericLength
528
532
533+-- XXX Try Hash set
534+-- XXX Add a countDistinct window fold
535+-- XXX Add a bloom filter fold
536+
537+-- | Count non-duplicate elements in the stream.
538+--
539+-- The memory used is proportional to the number of distinct elements in the
540+-- stream, to guard against using too much memory use it as a scan and
541+-- terminate if the count reaches more than a threshold.
542+--
543+-- /Space/: \(\mathcal{O}(n)\)
544+--
545+-- /Pre-release/
546+--
547+{-# INLINE countDistinct #-}
548+countDistinct :: (Monad m, Ord a) => Fold m a Int
549+countDistinct = fmap (\(Tuple' _ n) -> n) $ foldl' step initial
550+
551+ where
552+
553+ initial = Tuple' Set.empty 0
554+
555+ step (Tuple' set n) x = do
556+ if Set.member x set
557+ then
558+ Tuple' set n
559+ else
560+ let cnt = n + 1
561+ in Tuple' (Set.insert x set) cnt
562+
563+-- | Like 'countDistinct' but specialized to a stream of 'Int', for better
564+-- performance.
565+--
566+-- /Pre-release/
567+{-# INLINE countDistinctInt #-}
568+countDistinctInt :: Monad m => Fold m Int Int
569+countDistinctInt = fmap (\(Tuple' _ n) -> n) $ foldl' step initial
570+
571+ where
572+
573+ initial = Tuple' IntSet.empty 0
574+
575+ step (Tuple' set n) x = do
576+ if IntSet.member x set
577+ then
578+ Tuple' set n
579+ else
580+ let cnt = n + 1
581+ in Tuple' (IntSet.insert x set) cnt
582+
529-- | Determine the sum of all elements of a stream of numbers. Returns additive
583-- | Determine the sum of all elements of a stream of numbers. Returns additive
530-- identity (@0@) when the stream is empty. Note that this is not numerically
584-- identity (@0@) when the stream is empty. Note that this is not numerically
531-- stable for floating point numbers.
585-- stable for floating point numbers.